add header forwarding for grpc

go.mod

@@ -21,7 +21,8 @@ require (
 	github.com/tidwall/gjson v1.14.2
 	github.com/tidwall/sjson v1.2.5
 	github.com/vmihailenco/msgpack/v5 v5.3.5
-	golang.org/x/sync v0.0.0-20210220032951-036812b2e83c
+	golang.org/x/sync v0.4.0
+	google.golang.org/grpc v1.60.1
 )
 
 require (
@@ -34,7 +35,7 @@ require (
 	github.com/dgraph-io/ristretto v0.0.3 // indirect
 	github.com/dgryski/go-farm v0.0.0-20200201041132-a6ae2369ad13 // indirect
 	github.com/dustin/go-humanize v1.0.0 // indirect
-	github.com/golang/protobuf v1.5.2 // indirect
+	github.com/golang/protobuf v1.5.3 // indirect
 	github.com/golang/snappy v0.0.4 // indirect
 	github.com/klauspost/compress v1.12.3 // indirect
 	github.com/mattn/go-colorable v0.1.9 // indirect
@@ -46,9 +47,11 @@ require (
 	github.com/tidwall/match v1.1.1 // indirect
 	github.com/tidwall/pretty v1.2.0 // indirect
 	github.com/vmihailenco/tagparser/v2 v2.0.0 // indirect
-	golang.org/x/crypto v0.0.0-20210921155107-089bfa567519 // indirect
-	golang.org/x/net v0.0.0-20220607020251-c690dde0001d // indirect
-	golang.org/x/sys v0.0.0-20220520151302-bc2c85ada10a // indirect
-	google.golang.org/protobuf v1.26.0 // indirect
+	golang.org/x/crypto v0.14.0 // indirect
+	golang.org/x/net v0.16.0 // indirect
+	golang.org/x/sys v0.13.0 // indirect
+	golang.org/x/text v0.13.0 // indirect
+	google.golang.org/genproto/googleapis/rpc v0.0.0-20231002182017-d307bd883b97 // indirect
+	google.golang.org/protobuf v1.31.0 // indirect
 	gopkg.in/yaml.v3 v3.0.1 // indirect
 )

go.sum

@@ -63,13 +63,13 @@ github.com/fsnotify/fsnotify v1.4.7/go.mod h1:jwhsz4b93w/PPRr/qN1Yymfu8t87LnFCMo
 github.com/golang/protobuf v1.2.0/go.mod h1:6lQm79b+lXiMfvg/cZm0SGofjICqVBUtrP5yJMmIC1U=
 github.com/golang/protobuf v1.3.1/go.mod h1:6lQm79b+lXiMfvg/cZm0SGofjICqVBUtrP5yJMmIC1U=
 github.com/golang/protobuf v1.5.0/go.mod h1:FsONVRAS9T7sI+LIUmWTfcYkHO4aIWwzhcaSAoJOfIk=
-github.com/golang/protobuf v1.5.2 h1:ROPKBNFfQgOUMifHyP+KYbvpjbdoFNs+aK7DXlji0Tw=
-github.com/golang/protobuf v1.5.2/go.mod h1:XVQd3VNwM+JqD3oG2Ue2ip4fOMUkwXdXDdiuN0vRsmY=
+github.com/golang/protobuf v1.5.3 h1:KhyjKVUg7Usr/dYsdSqoFveMYd5ko72D+zANwlG1mmg=
+github.com/golang/protobuf v1.5.3/go.mod h1:XVQd3VNwM+JqD3oG2Ue2ip4fOMUkwXdXDdiuN0vRsmY=
 github.com/golang/snappy v0.0.3/go.mod h1:/XxbfmMg8lxefKM7IXC3fBNl/7bRcc72aCRzEWrmP2Q=
 github.com/golang/snappy v0.0.4 h1:yAGX7huGHXlcLOEtBnF4w7FQwA26wojNCwOYAEhLjQM=
 github.com/golang/snappy v0.0.4/go.mod h1:/XxbfmMg8lxefKM7IXC3fBNl/7bRcc72aCRzEWrmP2Q=
-github.com/google/go-cmp v0.5.5 h1:Khx7svrCpmxxtHBq5j2mp/xVjsi8hQMfNLvJFAlrGgU=
 github.com/google/go-cmp v0.5.5/go.mod h1:v8dTdLbMG2kIc/vJvl+f65V22dbkXbowE6jgT/gNBxE=
+github.com/google/go-cmp v0.5.9 h1:O2Tfq5qg4qc4AmwVlvv0oLiVAGB7enBSJ2x2DqQFi38=
 github.com/google/uuid v1.1.1/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo=
 github.com/google/uuid v1.6.0 h1:NIvaJDMOsjHA8n1jAhLSgzrAzy1Hgr+hNrb57e+94F0=
 github.com/google/uuid v1.6.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo=
@@ -157,18 +157,18 @@ golang.org/x/crypto v0.0.0-20190308221718-c2843e01d9a2/go.mod h1:djNgcEr1/C05ACk
 golang.org/x/crypto v0.0.0-20190426145343-a29dc8fdc734/go.mod h1:yigFU9vqHzYiE8UmvKecakEJjdnWj3jj499lnFckfCI=
 golang.org/x/crypto v0.0.0-20200115085410-6d4e4cb37c7d/go.mod h1:LzIPMQfyMNhhGPhUkYOs5KpL4U8rLKemX1yGLhDgUto=
 golang.org/x/crypto v0.0.0-20210322153248-0c34fe9e7dc2/go.mod h1:T9bdIzuCu7OtxOm1hfPfRQxPLYneinmdGuTeoZ9dtd4=
-golang.org/x/crypto v0.0.0-20210921155107-089bfa567519 h1:7I4JAnoQBe7ZtJcBaYHi5UtiO8tQHbUSXxL+pnGRANg=
-golang.org/x/crypto v0.0.0-20210921155107-089bfa567519/go.mod h1:GvvjBRRGRdwPK5ydBHafDWAxML/pGHZbMvKqRZ5+Abc=
+golang.org/x/crypto v0.14.0 h1:wBqGXzWJW6m1XrIKlAH0Hs1JJ7+9KBwnIO8v66Q9cHc=
+golang.org/x/crypto v0.14.0/go.mod h1:MVFd36DqK4CsrnJYDkBA3VC4m2GkXAM0PvzMCn4JQf4=
 golang.org/x/net v0.0.0-20180906233101-161cd47e91fd/go.mod h1:mL1N/T3taQHkDXs73rZJwtUhF3w3ftmwwsq0BUmARs4=
 golang.org/x/net v0.0.0-20190404232315-eb5bcb51f2a3/go.mod h1:t9HGtf8HONx5eT2rtn7q6eTqICYqUVnKs3thJo3Qplg=
 golang.org/x/net v0.0.0-20190620200207-3b0461eec859/go.mod h1:z5CRVTTTmAJ677TzLLGU+0bjPO0LkuOLi4/5GtJWs/s=
 golang.org/x/net v0.0.0-20210226172049-e18ecbb05110/go.mod h1:m0MpNAwzfU5UDzcl9v0D8zg8gWTRqZa9RBIspLL5mdg=
-golang.org/x/net v0.0.0-20220607020251-c690dde0001d h1:4SFsTMi4UahlKoloni7L4eYzhFRifURQLw+yv0QDCx8=
-golang.org/x/net v0.0.0-20220607020251-c690dde0001d/go.mod h1:XRhObCWvk6IyKnWLug+ECip1KBveYUHfp+8e9klMJ9c=
+golang.org/x/net v0.16.0 h1:7eBu7KsSvFDtSXUIDbh3aqlK4DPsZ1rByC8PFfBThos=
+golang.org/x/net v0.16.0/go.mod h1:NxSsAGuq816PNPmqtQdLE42eU2Fs7NoRIZrHJAlaCOE=
 golang.org/x/sync v0.0.0-20180314180146-1d60e4601c6f/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
 golang.org/x/sync v0.0.0-20200625203802-6e8e738ad208/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
-golang.org/x/sync v0.0.0-20210220032951-036812b2e83c h1:5KslGYwFpkhGh+Q16bwMP3cOontH8FOep7tGV86Y7SQ=
-golang.org/x/sync v0.0.0-20210220032951-036812b2e83c/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
+golang.org/x/sync v0.4.0 h1:zxkM55ReGkDlKSM+Fu41A+zmbZuaPVbGMzvvdUPznYQ=
+golang.org/x/sync v0.4.0/go.mod h1:FU7BRWz2tNW+3quACPkgCx/L+uEAv1htQ0V83Z9Rj+Y=
 golang.org/x/sys v0.0.0-20180909124046-d0be0721c37e/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY=
 golang.org/x/sys v0.0.0-20181205085412-a5c9d58dba9a/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY=
 golang.org/x/sys v0.0.0-20190215142949-d0b11bdaac8a/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY=
@@ -179,17 +179,23 @@ golang.org/x/sys v0.0.0-20200223170610-d5e6a3e2c0ae/go.mod h1:h1NjWce9XRLGQEsW7w
 golang.org/x/sys v0.0.0-20201119102817-f84b799fce68/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
 golang.org/x/sys v0.0.0-20210420205809-ac73e9fd8988/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
 golang.org/x/sys v0.0.0-20210630005230-0f9fa26af87c/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
-golang.org/x/sys v0.0.0-20220520151302-bc2c85ada10a h1:dGzPydgVsqGcTRVwiLJ1jVbufYwmzD3LfVPLKsKg+0k=
-golang.org/x/sys v0.0.0-20220520151302-bc2c85ada10a/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
+golang.org/x/sys v0.13.0 h1:Af8nKPmuFypiUBjVoU9V20FiaFXOcuZI21p0ycVYYGE=
+golang.org/x/sys v0.13.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
 golang.org/x/term v0.0.0-20201126162022-7de9c90e9dd1/go.mod h1:bj7SfCRtBDWHUb9snDiAeCFNEtKQo2Wmx5Cou7ajbmo=
 golang.org/x/text v0.3.0/go.mod h1:NqM8EUOU14njkJ3fqMW+pc6Ldnwhi/IjpwHt7yyuwOQ=
 golang.org/x/text v0.3.3/go.mod h1:5Zoc/QRtKVWzQhOtBMvqHzDpF6irO9z98xDceosuGiQ=
+golang.org/x/text v0.13.0 h1:ablQoSUd0tRdKxZewP80B+BaqeKJuVhuRxj/dkrun3k=
+golang.org/x/text v0.13.0/go.mod h1:TvPlkZtksWOMsz7fbANvkp4WM8x/WCo/om8BMLbz+aE=
 golang.org/x/tools v0.0.0-20180917221912-90fa682c2a6e/go.mod h1:n7NCudcB/nEzxVGmLbDWY5pfWTLqBcC2KZ6jyYvM4mQ=
 golang.org/x/xerrors v0.0.0-20191204190536-9bdfabe68543/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=
-golang.org/x/xerrors v0.0.0-20220517211312-f3a8303e98df h1:5Pf6pFKu98ODmgnpvkJ3kFUOQGGLIzLIkbzUHp47618=
+google.golang.org/genproto/googleapis/rpc v0.0.0-20231002182017-d307bd883b97 h1:6GQBEOdGkX6MMTLT9V+TjtIRZCw9VPD5Z+yHY9wMgS0=
+google.golang.org/genproto/googleapis/rpc v0.0.0-20231002182017-d307bd883b97/go.mod h1:v7nGkzlmW8P3n/bKmWBn2WpBjpOEx8Q6gMueudAmKfY=
+google.golang.org/grpc v1.60.1 h1:26+wFr+cNqSGFcOXcabYC0lUVJVRa2Sb2ortSK7VrEU=
+google.golang.org/grpc v1.60.1/go.mod h1:OlCHIeLYqSSsLi6i49B5QGdzaMZK9+M7LXN2FKz4eGM=
 google.golang.org/protobuf v1.26.0-rc.1/go.mod h1:jlhhOSvTdKEhbULTjvd4ARK9grFBp09yW+WbY/TyQbw=
-google.golang.org/protobuf v1.26.0 h1:bxAC2xTBsZGibn2RTntX0oH50xLsqy1OxA9tTL3p/lk=
 google.golang.org/protobuf v1.26.0/go.mod h1:9q0QmTI4eRPtz6boOQmLYwt+qCgq0jsYwAQnmE0givc=
+google.golang.org/protobuf v1.31.0 h1:g0LDEJHgrBl9N9r17Ru3sqWhkIx2NB67okBHPwC7hs8=
+google.golang.org/protobuf v1.31.0/go.mod h1:HV8QOd/L58Z+nl8r43ehVNZIU/HEI6OcFqwMG9pJV4I=
 gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
 gopkg.in/check.v1 v1.0.0-20190902080502-41f04d3bba15/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
 gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c h1:Hei/4ADfdWqJk1ZMxUNpqntNwaWcugrBjAiHlqqRiVk=

headerforwarder/context_headers.go

@@ -30,7 +30,12 @@ func ContextWithRosettaID(ctx context.Context) context.Context {
 }
 
 func RosettaIDFromContext(ctx context.Context) string {
-	return ctx.Value(requestIDKey).(string)
+	switch val := ctx.Value(requestIDKey).(type) {
+	case string:
+		return val
+	default:
+		return ""
+	}
 }
 
 func RosettaIDFromRequest(r *http.Request) string {

headerforwarder/forwarder.go

@@ -15,7 +15,12 @@
 package headerforwarder
 
 import (
+	"context"
+	"fmt"
 	"net/http"
+
+	"google.golang.org/grpc"
+	"google.golang.org/grpc/metadata"
 )
 
 // HeaderExtractingTransport is a utility to help a rosetta server forward headers to and from
@@ -34,27 +39,36 @@ type HeaderForwarder struct {
 	actualTransport    http.RoundTripper
 }
 
-func NewHeaderForwarder(interestingHeaders []string, transport http.RoundTripper) *HeaderForwarder {
+// TODO: make transport an optional parameter, add "WithTransport" style functions to make it easier
+// to add the actual RPC clients to this struct
+func NewHeaderForwarder(interestingHeaders []string, transport http.RoundTripper) (*HeaderForwarder, error) {
+	if len(interestingHeaders) == 0 {
+		return nil, fmt.Errorf("must provide at least one interesting header")
+	}
+
 	return &HeaderForwarder{
 		requestHeaders:     make(map[string]http.Header),
 		interestingHeaders: interestingHeaders,
 		actualTransport:    transport,
-	}
+	}, nil
 }
 
-// RoundTrip implements http.RoundTripper and will be used to construct an http Client which
-// saves the native node response headers if necessary.
-func (hf *HeaderForwarder) RoundTrip(req *http.Request) (*http.Response, error) {
-	resp, err := hf.actualTransport.RoundTrip(req)
+func (hf *HeaderForwarder) WithTransport(transport http.RoundTripper) *HeaderForwarder {
+	hf.actualTransport = transport
+	return hf
+}
 
-	if err == nil && hf.shouldRememberHeaders(req, resp) {
-		hf.rememberHeaders(req, resp)
-	}
+// RequestWithRequestID adds a unique ID to the request context. A new request is returned that contains the
+// new context
+func (hf *HeaderForwarder) RequestWithRequestID(req *http.Request) *http.Request {
+	ctx := req.Context()
+	ctxWithID := ContextWithRosettaID(ctx)
+	requestWithID := req.WithContext(ctxWithID)
 
-	return resp, err
+	return requestWithID
 }
 
-// shouldRememberHeaders is called to determine if response headers should be remembered for a
+// shouldRememberHeaders reports whether response headers should be remembered for a
 // given request. Response headers will only be remembered if the request does not contain all of
 // the interesting headers and the response contains at least one of the interesting headers.
 //
@@ -82,12 +96,11 @@ func (hf *HeaderForwarder) shouldRememberHeaders(req *http.Request, resp *http.R
 	return !requestHasAllHeaders && responseHasSomeHeaders
 }
 
-// rememberHeaders is called to save the native node response headers. The request object
-// here is a native node request (constructed by go-ethereum for geth-based rosetta implementations).
+// rememberHeaders saves the native node response headers. The request object here is a
+// native node request (e.g. one constructed by go-ethereum for geth-based rosetta implementations).
 // The response object is a native node response.
 func (hf *HeaderForwarder) rememberHeaders(req *http.Request, resp *http.Response) {
 	ctx := req.Context()
-	// rosettaRequestID := services.osettaIdFromContext(ctx)
 	rosettaRequestID := RosettaIDFromContext(ctx)
 
 	// Only remember interesting headers
@@ -99,7 +112,41 @@ func (hf *HeaderForwarder) rememberHeaders(req *http.Request, resp *http.Respons
 	hf.requestHeaders[rosettaRequestID] = headersToRemember
 }
 
-// GetResponseHeaders returns any native node response headers that were recorded for a request ID.
+// shouldRememberMetadata reports whether response metadata should be remembered for a grpc unary
+// RPC call. Response metadata will only be remembered if it contains any of the interesting headers.
+func (hf *HeaderForwarder) shouldRememberMetadata(ctx context.Context, req any, resp metadata.MD) bool {
+	rosettaID := RosettaIDFromContext(ctx)
+	if rosettaID == "" {
+		return false
+	}
+
+	// If any of the interesting headers are in the response metadata, remember it
+	for _, interestingHeader := range hf.interestingHeaders {
+		if _, responseHasHeader := resp[http.CanonicalHeaderKey(interestingHeader)]; responseHasHeader {
+			return true
+		}
+	}
+
+	return false
+}
+
+// rememberMetadata saves the native node response metadata. The response object is metadata retrieved
+// from a native node GRPC unary RPC call.
+func (hf *HeaderForwarder) rememberMetadata(ctx context.Context, req any, resp metadata.MD) {
+	rosettaID := RosettaIDFromContext(ctx)
+
+	headersToRemember := make(http.Header)
+	for _, interestingHeader := range hf.interestingHeaders {
+		for _, value := range resp.Get(interestingHeader) {
+			headersToRemember.Set(interestingHeader, value)
+		}
+	}
+
+	hf.requestHeaders[rosettaID] = headersToRemember
+}
+
+// GetResponseHeaders returns any headers that should be returned to a rosetta response. These
+// consist of native node response headers/metadata that were remembered for a request ID.
 func (hf *HeaderForwarder) getResponseHeaders(rosettaRequestID string) (http.Header, bool) {
 	headers, ok := hf.requestHeaders[rosettaRequestID]
 
@@ -115,13 +162,14 @@ func (hf *HeaderForwarder) getResponseHeaders(rosettaRequestID string) (http.Hea
 // those headers on the response
 func (hf *HeaderForwarder) HeaderForwarderHandler(next http.Handler) http.Handler {
 	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		fmt.Println("HeaderForwarder Handler")
 		// add a unique ID to the request context, and make a new request for it
-		requestWithID := hf.WithRequestID(r)
+		requestWithID := hf.RequestWithRequestID(r)
 
 		// Serve the request
-		// NOTE: ResponseWriter::WriteHeader() WILL be called here, so we can't set headers after this happens
-		// We include a wrapper around the response writer that allows us to set headers just before
-		// WriteHeader is called
+		// NOTE: for servers using github.com/coinbase/mesh-geth-sdk, ResponseWriter::WriteHeader() WILL
+		// be called here, so we can't set headers after this happens. We include a wrapper around the
+		// response writer that allows us to set headers just before WriteHeader is called
 		wrappedResponseWriter := NewResponseWriter(
 			w,
 			RosettaIDFromRequest(requestWithID),
@@ -131,12 +179,38 @@ func (hf *HeaderForwarder) HeaderForwarderHandler(next http.Handler) http.Handle
 	})
 }
 
-// WithRequestID adds a unique ID to the request context. A new request is returned that contains the
-// new context
-func (hf *HeaderForwarder) WithRequestID(req *http.Request) *http.Request {
-	ctx := req.Context()
-	ctxWithID := ContextWithRosettaID(ctx)
-	requestWithID := req.WithContext(ctxWithID)
+// RoundTrip implements http.RoundTripper and will be used to construct an http Client which
+// saves the native node response headers if necessary.
+func (hf *HeaderForwarder) RoundTrip(req *http.Request) (*http.Response, error) {
+	fmt.Println("HeaderForwarder RoundTrip")
+	resp, err := hf.actualTransport.RoundTrip(req)
 
-	return requestWithID
+	fmt.Println("HeaderForwarder RoundTrip: response headers", resp.Header)
+
+	if err == nil && hf.shouldRememberHeaders(req, resp) {
+		hf.rememberHeaders(req, resp)
+	}
+
+	return resp, err
+}
+
+func (hf *HeaderForwarder) UnaryClientInterceptor(ctx context.Context, method string, req, reply any, cc *grpc.ClientConn, invoker grpc.UnaryInvoker, opts ...grpc.CallOption) error {
+	fmt.Println("HeaderForwarder grpc interceptor")
+
+	fmt.Println("request id: ", RosettaIDFromContext(ctx))
+
+	// append a header DialOption to the request
+	var responseMD metadata.MD
+	opts = append(opts, grpc.Header(&responseMD))
+
+	err := invoker(ctx, method, req, reply, cc, opts...)
+
+	if hf.shouldRememberMetadata(ctx, req, responseMD) {
+		hf.rememberMetadata(ctx, req, responseMD)
+	}
+
+	// get headers from response
+	fmt.Println("HeaderForwarder grpc interceptor: headers from response", responseMD)
+
+	return err
 }