Move HTTP2 routing information into headers (#4001)

Signed-off-by: Joshua Kim <[email protected]>

Author: joshua-kim

api/grpcclient/client.go

@@ -8,16 +8,17 @@ import (
 	"fmt"
 
 	"google.golang.org/grpc"
+	"google.golang.org/grpc/metadata"
 
 	"github.com/ava-labs/avalanchego/ids"
 )
 
-// NewChainClient returns grpc.ClientConn that prefixes method calls with the
-// provided chainID prefix
+// NewChainClient returns a grpc.ClientConn that sets the chain-id header for
+// all requests
 func NewChainClient(uri string, chainID ids.ID, opts ...grpc.DialOption) (*grpc.ClientConn, error) {
 	dialOpts := []grpc.DialOption{
-		grpc.WithUnaryInterceptor(PrefixChainIDUnaryClientInterceptor(chainID)),
-		grpc.WithStreamInterceptor(PrefixChainIDStreamClientInterceptor(chainID)),
+		grpc.WithUnaryInterceptor(SetChainIDHeaderUnaryClientInterceptor(chainID)),
+		grpc.WithStreamInterceptor(SetChainIDHeaderStreamClientInterceptor(chainID)),
 	}
 
 	dialOpts = append(dialOpts, opts...)
@@ -30,9 +31,9 @@ func NewChainClient(uri string, chainID ids.ID, opts ...grpc.DialOption) (*grpc.
 	return conn, nil
 }
 
-// PrefixChainIDUnaryClientInterceptor prefixes unary grpc calls with the
-// provided chainID prefix
-func PrefixChainIDUnaryClientInterceptor(chainID ids.ID) grpc.UnaryClientInterceptor {
+// SetChainIDHeaderUnaryClientInterceptor sets the chain-id header for unary
+// requests
+func SetChainIDHeaderUnaryClientInterceptor(chainID ids.ID) grpc.UnaryClientInterceptor {
 	return func(
 		ctx context.Context,
 		method string,
@@ -42,13 +43,14 @@ func PrefixChainIDUnaryClientInterceptor(chainID ids.ID) grpc.UnaryClientInterce
 		invoker grpc.UnaryInvoker,
 		opts ...grpc.CallOption,
 	) error {
-		return invoker(ctx, prefix(chainID, method), req, reply, cc, opts...)
+		ctx = newContextWithChainIDHeader(ctx, chainID)
+		return invoker(ctx, method, req, reply, cc, opts...)
 	}
 }
 
-// PrefixChainIDStreamClientInterceptor prefixes streaming grpc calls with the
-// provided chainID prefix
-func PrefixChainIDStreamClientInterceptor(chainID ids.ID) grpc.StreamClientInterceptor {
+// SetChainIDHeaderStreamClientInterceptor sets the chain-id header for
+// streaming requests
+func SetChainIDHeaderStreamClientInterceptor(chainID ids.ID) grpc.StreamClientInterceptor {
 	return func(
 		ctx context.Context,
 		desc *grpc.StreamDesc,
@@ -57,11 +59,13 @@ func PrefixChainIDStreamClientInterceptor(chainID ids.ID) grpc.StreamClientInter
 		streamer grpc.Streamer,
 		opts ...grpc.CallOption,
 	) (grpc.ClientStream, error) {
-		return streamer(ctx, desc, cc, prefix(chainID, method), opts...)
+		ctx = newContextWithChainIDHeader(ctx, chainID)
+		return streamer(ctx, desc, cc, method, opts...)
 	}
 }
 
-// http/2 :path takes the form of /ChainID/Service/Method
-func prefix(chainID ids.ID, method string) string {
-	return "/" + chainID.String() + method
+// newContextWithChainHeader sets the chain-id header which the server uses
+// to route the client grpc request
+func newContextWithChainIDHeader(ctx context.Context, chainID ids.ID) context.Context {
+	return metadata.AppendToOutgoingContext(ctx, "chain-id", chainID.String())
 }

api/server/http2_router.go

@@ -5,7 +5,6 @@ package server
 
 import (
 	"net/http"
-	"strings"
 	"sync"
 
 	"github.com/ava-labs/avalanchego/ids"
@@ -25,15 +24,14 @@ func newHTTP2Router() *http2Router {
 }
 
 func (h *http2Router) ServeHTTP(w http.ResponseWriter, r *http.Request) {
-	// The :path pseudo-header takes the form of /Prefix/Path
-	parsed := strings.Split(r.URL.Path, "/")
-	if len(parsed) < 2 {
+	// the chain-id header must be set to route the request to the correct chain
+	// http2 handler
+	chainID := r.Header.Get("chain-id")
+	if len(chainID) == 0 {
 		w.WriteHeader(http.StatusBadRequest)
 		return
 	}
 
-	chainID := parsed[1]
-
 	h.lock.RLock()
 	handler, ok := h.handlers[chainID]
 	h.lock.RUnlock()
@@ -42,15 +40,7 @@ func (h *http2Router) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 		return
 	}
 
-	// Deep copy the request to avoid weird behavior from modifying r
-	requestDeepCopy := r.Clone(r.Context())
-	// Route this request to the http2 handler using the chain prefix
-	requestDeepCopy.URL.Path = strings.TrimPrefix(
-		requestDeepCopy.URL.Path,
-		"/"+chainID,
-	)
-
-	handler.ServeHTTP(w, requestDeepCopy)
+	handler.ServeHTTP(w, r)
 }
 
 func (h *http2Router) Add(chainID ids.ID, handler http.Handler) bool {

api/server/http2_router_test.go

@@ -4,10 +4,8 @@
 package server
 
 import (
-	"fmt"
 	"net/http"
 	"net/http/httptest"
-	"net/url"
 	"testing"
 
 	"github.com/stretchr/testify/require"
@@ -28,23 +26,26 @@ func TestHTTP2RouterServeHTTP(t *testing.T) {
 	tests := []struct {
 		name     string
 		chainIDs []ids.ID
-		path     string
+		header   http.Header
 		wantCode int
 	}{
 		{
-			name:     "invalid request",
-			path:     "foo",
+			name:     "missing chain-id header",
 			wantCode: http.StatusBadRequest,
 		},
 		{
-			name:     "invalid handler",
-			path:     "/foo/bar/method",
+			name: "unknown referenced chain-id",
+			header: http.Header{
+				http.CanonicalHeaderKey("chain-id"): []string{ids.Empty.String()},
+			},
 			wantCode: http.StatusNotFound,
 		},
 		{
 			name:     "valid handler",
-			chainIDs: []ids.ID{{'f', 'o', 'o'}},
-			path:     fmt.Sprintf("/%s/bar/method", ids.ID{'f', 'o', 'o'}.String()),
+			chainIDs: []ids.ID{ids.Empty},
+			header: http.Header{
+				http.CanonicalHeaderKey("chain-id"): []string{ids.Empty.String()},
+			},
 			wantCode: http.StatusOK,
 		},
 	}
@@ -57,9 +58,7 @@ func TestHTTP2RouterServeHTTP(t *testing.T) {
 			handler := http.HandlerFunc(func(http.ResponseWriter, *http.Request) {})
 			writer := httptest.NewRecorder()
 			request := httptest.NewRequest(http.MethodPost, "/", nil)
-			request.URL = &url.URL{
-				Path: tt.path,
-			}
+			request.Header = tt.header
 
 			for _, chainID := range tt.chainIDs {
 				require.True(h.Add(chainID, handler))