add x-forward-for message in log

unclezoro · unclezoro · commit 0e797c53b7d2 · 2021-01-11T10:49:34.000+08:00
diff --git a/rpc/client.go b/rpc/client.go
@@ -552,9 +552,9 @@ func (c *Client) dispatch(codec ServerCodec) {
 		// Read path:
 		case op := <-c.readOp:
 			if op.batch {
-				conn.handler.handleBatch(op.msgs)
+				conn.handler.handleBatch(context.Background(), op.msgs)
 			} else {
-				conn.handler.handleMsg(op.msgs[0])
+				conn.handler.handleMsg(context.Background(), op.msgs[0])
 			}
 
 		case err := <-c.readErr:
diff --git a/rpc/handler.go b/rpc/handler.go
@@ -93,7 +93,7 @@ func newHandler(connCtx context.Context, conn jsonWriter, idgen func() ID, reg *
 }
 
 // handleBatch executes all messages in a batch and returns the responses.
-func (h *handler) handleBatch(msgs []*jsonrpcMessage) {
+func (h *handler) handleBatch(ctx context.Context, msgs []*jsonrpcMessage) {
 	// Emit error response for empty batches:
 	if len(msgs) == 0 {
 		h.startCallProc(func(cp *callProc) {
@@ -116,7 +116,7 @@ func (h *handler) handleBatch(msgs []*jsonrpcMessage) {
 	h.startCallProc(func(cp *callProc) {
 		answers := make([]*jsonrpcMessage, 0, len(msgs))
 		for _, msg := range calls {
-			if answer := h.handleCallMsg(cp, msg); answer != nil {
+			if answer := h.handleCallMsg(cp, ctx, msg); answer != nil {
 				answers = append(answers, answer)
 			}
 		}
@@ -131,12 +131,12 @@ func (h *handler) handleBatch(msgs []*jsonrpcMessage) {
 }
 
 // handleMsg handles a single message.
-func (h *handler) handleMsg(msg *jsonrpcMessage) {
+func (h *handler) handleMsg(ctx context.Context, msg *jsonrpcMessage) {
 	if ok := h.handleImmediate(msg); ok {
 		return
 	}
 	h.startCallProc(func(cp *callProc) {
-		answer := h.handleCallMsg(cp, msg)
+		answer := h.handleCallMsg(cp, ctx, msg)
 		h.addSubscriptions(cp.notifiers)
 		if answer != nil {
 			h.conn.writeJSON(cp.ctx, answer)
@@ -287,7 +287,7 @@ func (h *handler) handleResponse(msg *jsonrpcMessage) {
 }
 
 // handleCallMsg executes a call message and returns the answer.
-func (h *handler) handleCallMsg(ctx *callProc, msg *jsonrpcMessage) *jsonrpcMessage {
+func (h *handler) handleCallMsg(ctx *callProc, reqCtx context.Context, msg *jsonrpcMessage) *jsonrpcMessage {
 	start := time.Now()
 	switch {
 	case msg.isNotification():
@@ -297,7 +297,8 @@ func (h *handler) handleCallMsg(ctx *callProc, msg *jsonrpcMessage) *jsonrpcMess
 	case msg.isCall():
 		resp := h.handleCall(ctx, msg)
 		if resp.Error != nil {
-			h.log.Warn("Served "+msg.Method, "reqid", idForLog{msg.ID}, "t", time.Since(start), "err", resp.Error.Message)
+			xForward := reqCtx.Value("X-Forwarded-For")
+			h.log.Warn("Served "+msg.Method, "reqid", idForLog{msg.ID}, "t", time.Since(start), "err", resp.Error.Message, "X-Forwarded-For", xForward)
 		} else {
 			h.log.Debug("Served "+msg.Method, "reqid", idForLog{msg.ID}, "t", time.Since(start))
 		}
diff --git a/rpc/http.go b/rpc/http.go
@@ -227,6 +227,9 @@ func (s *Server) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 	if origin := r.Header.Get("Origin"); origin != "" {
 		ctx = context.WithValue(ctx, "Origin", origin)
 	}
+	if xForward := r.Header.Get("X-Forwarded-For"); xForward != "" {
+		ctx = context.WithValue(ctx, "X-Forwarded-For", xForward)
+	}
 
 	w.Header().Set("content-type", contentType)
 	codec := newHTTPServerConn(r, w)
diff --git a/rpc/server.go b/rpc/server.go
@@ -109,9 +109,9 @@ func (s *Server) serveSingleRequest(ctx context.Context, codec ServerCodec) {
 		return
 	}
 	if batch {
-		h.handleBatch(reqs)
+		h.handleBatch(ctx, reqs)
 	} else {
-		h.handleMsg(reqs[0])
+		h.handleMsg(ctx, reqs[0])
 	}
 }
 

Original file line number	Diff line number	Diff line change
`@@ -227,6 +227,9 @@ func (s Server) ServeHTTP(w http.ResponseWriter, r http.Request) {`
`227`	`227`	`if origin := r.Header.Get("Origin"); origin != "" {`
`228`	`228`	`ctx = context.WithValue(ctx, "Origin", origin)`
`229`	`229`	`}`
	`230`	`+ if xForward := r.Header.Get("X-Forwarded-For"); xForward != "" {`
	`231`	`+ ctx = context.WithValue(ctx, "X-Forwarded-For", xForward)`
	`232`	`+ }`
`230`	`233`
`231`	`234`	`w.Header().Set("content-type", contentType)`
`232`	`235`	`codec := newHTTPServerConn(r, w)`
Original file line number	Diff line number	Diff line change
`@@ -109,9 +109,9 @@ func (s *Server) serveSingleRequest(ctx context.Context, codec ServerCodec) {`
`109`	`109`	`return`
`110`	`110`	`}`
`111`	`111`	`if batch {`
`112`		`- h.handleBatch(reqs)`
	`112`	`+ h.handleBatch(ctx, reqs)`
`113`	`113`	`} else {`
`114`		`- h.handleMsg(reqs[0])`
	`114`	`+ h.handleMsg(ctx, reqs[0])`
`115`	`115`	`}`
`116`	`116`	`}`
`117`	`117`