Ensure HTTP headers are propagated through the rpcchainvm (#3917)

Signed-off-by: Joshua Kim <20001595+joshua-kim@users.noreply.github.com>
Co-authored-by: Stephen Buttolph <stephen@avalabs.org>

Author: joshua-kim

proto/http/http.proto

@@ -2,14 +2,12 @@ syntax = "proto3";
 
 package http;
 
-import "google/protobuf/empty.proto";
-
 option go_package = "github.com/ava-labs/avalanchego/proto/pb/http";
 
 service HTTP {
   // Handle wraps http1 over http2 and provides support for websockets by implementing
   // net conn and responsewriter in http2.
-  rpc Handle(HTTPRequest) returns (google.protobuf.Empty);
+  rpc Handle(HTTPRequest) returns (HTTPResponse);
   // HandleSimple wraps http1 requests over http2 similar to Handle but only passes headers
   // and body bytes. Because the request and response are single protos with no inline
   // gRPC servers the CPU cost as well as file descriptor overhead is less
@@ -147,6 +145,11 @@ message HTTPRequest {
   Request request = 2;
 }
 
+message HTTPResponse {
+  // header is the http headers for the response
+  repeated Element header = 1;
+}
+
 message HandleSimpleHTTPRequest {
   // method specifies the HTTP method (GET, POST, PUT, etc.)
   string method = 1;

proto/pb/http/http.pb.go

@@ -156,10 +156,13 @@ func (c *Client) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 		}
 	}
 
-	_, err = c.client.Handle(r.Context(), req)
+	reply, err := c.client.Handle(r.Context(), req)
 	if err != nil {
 		http.Error(w, err.Error(), http.StatusInternalServerError)
+		return
 	}
+
+	grpcutils.SetHeaders(w.Header(), reply.Header)
 }
 
 // serveHTTPSimple converts an http request to a gRPC HTTPRequest and returns the
@@ -205,7 +208,7 @@ func getHTTPSimpleRequest(r *http.Request) (*httppb.HandleSimpleHTTPRequest, err
 
 // convertWriteResponse converts a gRPC HandleSimpleHTTPResponse to an HTTP response.
 func convertWriteResponse(w http.ResponseWriter, resp *httppb.HandleSimpleHTTPResponse) error {
-	grpcutils.MergeHTTPHeader(resp.Headers, w.Header())
+	grpcutils.SetHeaders(w.Header(), resp.Headers)
 	w.WriteHeader(grpcutils.EnsureValidResponseCode(int(resp.Code)))
 	_, err := w.Write(resp.Body)
 	return err

proto/pb/http/http_grpc.pb.go

@@ -8,11 +8,10 @@ import (
 	"context"
 	"crypto/tls"
 	"crypto/x509"
+	"fmt"
 	"net/http"
 	"net/url"
 
-	"google.golang.org/protobuf/types/known/emptypb"
-
 	"github.com/ava-labs/avalanchego/proto/pb/io/reader"
 	"github.com/ava-labs/avalanchego/vms/rpcchainvm/ghttp/greader"
 	"github.com/ava-labs/avalanchego/vms/rpcchainvm/ghttp/gresponsewriter"
@@ -40,7 +39,7 @@ func NewServer(handler http.Handler) *Server {
 	}
 }
 
-func (s *Server) Handle(ctx context.Context, req *httppb.HTTPRequest) (*emptypb.Empty, error) {
+func (s *Server) Handle(ctx context.Context, req *httppb.HTTPRequest) (*httppb.HTTPResponse, error) {
 	clientConn, err := grpcutils.Dial(req.ResponseWriter.ServerAddr)
 	if err != nil {
 		return nil, err
@@ -142,7 +141,13 @@ func (s *Server) Handle(ctx context.Context, req *httppb.HTTPRequest) (*emptypb.
 
 	s.handler.ServeHTTP(writer, request)
 
-	return &emptypb.Empty{}, clientConn.Close()
+	if err := clientConn.Close(); err != nil {
+		return nil, fmt.Errorf("failed to close client conn: %w", err)
+	}
+
+	return &httppb.HTTPResponse{
+		Header: grpcutils.GetHTTPHeader(writerHeaders),
+	}, nil
 }
 
 // HandleSimple handles http requests over http2 using a simple request response model.
@@ -153,7 +158,7 @@ func (s *Server) HandleSimple(ctx context.Context, r *httppb.HandleSimpleHTTPReq
 		return nil, err
 	}
 
-	grpcutils.MergeHTTPHeader(r.Headers, req.Header)
+	grpcutils.SetHeaders(req.Header, r.Headers)
 
 	req = req.WithContext(ctx)
 	req.RequestURI = r.Url

vms/rpcchainvm/ghttp/http_client.go

@@ -4,16 +4,20 @@
 package ghttp
 
 import (
+	"bytes"
 	"io"
 	"net/http"
 	"net/http/httptest"
+	"strings"
 	"testing"
 
 	"github.com/stretchr/testify/require"
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/credentials/insecure"
 	"google.golang.org/grpc/test/bufconn"
 
+	"github.com/ava-labs/avalanchego/vms/rpcchainvm/grpcutils"
+
 	httppb "github.com/ava-labs/avalanchego/proto/pb/http"
 )
 
@@ -84,6 +88,77 @@ func TestRequestClientArbitrarilyLongBody(t *testing.T) {
 	client.ServeHTTP(w, r) // Shouldn't block forever reading the body
 }
 
+// Tests that writes to the http response in the server are propagated to the
+// client
+func TestHttpResponse(t *testing.T) {
+	tests := []struct {
+		name   string
+		header http.Header
+	}{
+		{
+			// Requests with an upgrade header do not use the "Simple*" http response
+			// apis and must be separately tested
+			name: "upgrade header specified",
+			header: http.Header{
+				"Upgrade": {"upgrade"},
+				"foo":     {"foo"},
+			},
+		},
+		{
+			name: "arbitrary headers",
+			header: http.Header{
+				"foo": {"foo"},
+			},
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			require := require.New(t)
+
+			handler := http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
+				w.Header()["Bar"] = []string{"bar"}
+				_, err := w.Write([]byte("baz"))
+				require.NoError(err)
+			})
+
+			listener, err := grpcutils.NewListener()
+			require.NoError(err)
+			server := grpc.NewServer()
+			httppb.RegisterHTTPServer(server, NewServer(handler))
+
+			go func() {
+				require.NoError(server.Serve(listener))
+			}()
+
+			conn, err := grpc.NewClient(
+				listener.Addr().String(),
+				grpc.WithTransportCredentials(insecure.NewCredentials()),
+			)
+			require.NoError(err)
+
+			recorder := &httptest.ResponseRecorder{
+				Body: bytes.NewBuffer(nil),
+			}
+
+			client := NewClient(httppb.NewHTTPClient(conn))
+			client.ServeHTTP(recorder, &http.Request{
+				Body:   io.NopCloser(strings.NewReader("foo")),
+				Header: tt.header,
+			})
+
+			require.Equal(http.StatusOK, recorder.Code)
+			require.Equal(
+				http.Header{
+					"Bar": []string{"bar"},
+				},
+				recorder.Header(),
+			)
+			require.Equal("baz", recorder.Body.String())
+		})
+	}
+}
+
 type infiniteStream struct{}
 
 func (infiniteStream) Read(p []byte) (n int, err error) {

vms/rpcchainvm/ghttp/http_server.go

@@ -64,10 +64,12 @@ func GetHTTPHeader(hs http.Header) []*httppb.Element {
 	return result
 }
 
-// MergeHTTPHeader takes a slice of Header and merges with http.Header map.
-func MergeHTTPHeader(hs []*httppb.Element, header http.Header) {
-	for _, h := range hs {
-		header[h.Key] = h.Values
+// SetHeaders sets headers to next
+func SetHeaders(headers http.Header, next []*httppb.Element) {
+	clear(headers)
+
+	for _, h := range next {
+		headers[h.Key] = h.Values
 	}
 }