Discard unused body of Unary and ClientStream methods in background

MODULE.bazel

@@ -54,9 +54,11 @@ use_repo(
     go_deps,
     "com_github_antihax_optional",
     "com_github_bazelbuild_buildtools_v7",
+    "com_github_coder_websocket",
     "com_github_golang_protobuf",
     "com_github_google_go_cmp",
     "com_github_rogpeppe_fastuuid",
+    "com_github_tmc_grpc_websocket_proxy",
     "in_gopkg_yaml_v3",
     "org_golang_google_genproto_googleapis_api",
     "org_golang_google_genproto_googleapis_rpc",

examples/internal/gateway/BUILD.bazel

@@ -14,6 +14,7 @@ go_library(
         "//examples/internal/proto/examplepb",
         "//examples/internal/proto/standalone",
         "//runtime",
+        "@com_github_tmc_grpc_websocket_proxy//wsproxy",
         "@org_golang_google_genproto_googleapis_rpc//errdetails",
         "@org_golang_google_grpc//:grpc",
         "@org_golang_google_grpc//connectivity",

examples/internal/gateway/handlers.go

@@ -8,6 +8,7 @@ import (
 	"path"
 	"strings"
 
+	"github.com/tmc/grpc-websocket-proxy/wsproxy"
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/connectivity"
 	"google.golang.org/grpc/grpclog"
@@ -112,3 +113,13 @@ func logRequestBody(h http.Handler) http.Handler {
 		}
 	})
 }
+
+// websocketGateway wraps handler in a wsproxy.WebsocketProxy, adding method
+// query parameter as method selector.
+func websocketGateway(h http.Handler) http.Handler {
+	proxy := wsproxy.WebsocketProxy(
+		h,
+		wsproxy.WithMethodParamOverride("method"),
+	)
+	return proxy
+}

examples/internal/gateway/main.go

@@ -59,8 +59,12 @@ func Run(ctx context.Context, opts Options) error {
 	// Do not use logRequestBody for NoBodyPostServer because it will perform
 	// io.ReadAll and mask the issue:
 	// https://github.com/grpc-ecosystem/grpc-gateway/issues/5236
+	//
+	// Use WebSocket gateway for NoBodyPostServer to test whether we break
+	// streaming HTTP requests:
+	// https://github.com/grpc-ecosystem/grpc-gateway/issues/5326
 	hmux := http.NewServeMux()
-	hmux.Handle("/rpc/no-body/", allowCORS(mux))
+	hmux.Handle("/rpc/no-body/", websocketGateway(allowCORS(mux)))
 	hmux.Handle("/", logRequestBody(allowCORS(mux)))
 
 	s := &http.Server{

examples/internal/integration/BUILD.bazel

@@ -17,6 +17,7 @@ go_test(
         "//examples/internal/proto/sub",
         "//examples/internal/server",
         "//runtime",
+        "@com_github_coder_websocket//:websocket",
         "@com_github_google_go_cmp//cmp",
         "@com_github_rogpeppe_fastuuid//:fastuuid",
         "@org_golang_google_genproto_googleapis_rpc//status",

examples/internal/integration/integration_test.go

@@ -17,6 +17,7 @@ import (
 	"testing"
 	"time"
 
+	"github.com/coder/websocket"
 	"github.com/google/go-cmp/cmp"
 	"github.com/grpc-ecosystem/grpc-gateway/v2/examples/internal/proto/examplepb"
 	"github.com/grpc-ecosystem/grpc-gateway/v2/examples/internal/proto/pathenum"
@@ -2683,3 +2684,155 @@ func testNoBodyPostStream(t *testing.T, port int) {
 		t.Errorf("server context not done")
 	}
 }
+
+func TestNoBodyPostWebSocket(t *testing.T) {
+	if testing.Short() {
+		t.Skip()
+		return
+	}
+
+	testNoBodyPostWebSocketRPC(t, 8088)
+	testNoBodyPostWebSocketStream(t, 8088)
+}
+
+func testNoBodyPostWebSocketRPC(t *testing.T, port int) {
+	apiURL := fmt.Sprintf("ws://localhost:%d/rpc/no-body/rpc?method=POST", port)
+
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
+	conn, _, err := websocket.Dial(ctx, apiURL, nil)
+	if err != nil {
+		t.Fatalf("websocket.Dial(%q) failed with %v; want success", apiURL, err)
+	}
+	defer conn.Close(websocket.StatusNormalClosure, "")
+
+	// Write request even though server doesn't expect it
+	err = conn.Write(ctx, websocket.MessageText, []byte("{}"))
+	if err != nil {
+		t.Fatalf("conn.WriteMessage() failed with %v; want success", err)
+	}
+
+	// Wait for the server to start processing the request.
+	ctxServer := server.NoBodyPostServer_RetrieveContextRPC()
+	conn.Close(websocket.StatusNormalClosure, "")
+
+	// Wait for server context to be done
+	select {
+	case <-ctxServer.Done():
+	case <-time.After(time.Second):
+		t.Fatalf("server context not done")
+	}
+}
+
+func testNoBodyPostWebSocketStream(t *testing.T, port int) {
+	apiURL := fmt.Sprintf("ws://localhost:%d/rpc/no-body/stream?method=POST", port)
+
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
+	conn, _, err := websocket.Dial(ctx, apiURL, nil)
+	if err != nil {
+		t.Fatalf("websocket.DefaultDialer.DialContext(%q) failed with %v; want success", apiURL, err)
+	}
+	defer conn.Close(websocket.StatusNormalClosure, "")
+
+	// Write request even though server doesn't expect it
+	err = conn.Write(ctx, websocket.MessageText, []byte("{}"))
+	if err != nil {
+		t.Fatalf("conn.WriteMessage() failed with %v; want success", err)
+	}
+
+	// Wait for the server to start processing the request.
+	ctxServer := server.NoBodyPostServer_RetrieveContextStream()
+	conn.Close(websocket.StatusNormalClosure, "")
+
+	// Wait for server context to be done
+	select {
+	case <-ctxServer.Done():
+	case <-time.After(time.Second):
+		t.Fatalf("server context not done")
+	}
+}
+
+func TestWithBodyPostWebSocket(t *testing.T) {
+	if testing.Short() {
+		t.Skip()
+		return
+	}
+
+	testWithBodyPostWebSocketRPC(t, 8088)
+	testWithBodyPostWebSocketStream(t, 8088)
+}
+
+func testWithBodyPostWebSocketRPC(t *testing.T, port int) {
+	apiURL := fmt.Sprintf("ws://localhost:%d/rpc/no-body/rpc-with-body?method=POST", port)
+
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
+	conn, _, err := websocket.Dial(ctx, apiURL, nil)
+	if err != nil {
+		t.Fatalf("websocket.Dial(%q) failed with %v; want success", apiURL, err)
+	}
+	defer conn.Close(websocket.StatusNormalClosure, "")
+
+	// Write request because server expects it
+	err = conn.Write(ctx, websocket.MessageText, []byte("{}"))
+	if err != nil {
+		t.Fatalf("conn.WriteMessage() failed with %v; want success", err)
+	}
+
+	// Write another request even though server doesn't expect it
+	err = conn.Write(ctx, websocket.MessageText, []byte("{}"))
+	if err != nil {
+		t.Fatalf("conn.WriteMessage() failed with %v; want success", err)
+	}
+
+	// Wait for the server to start processing the request.
+	ctxServer := server.NoBodyPostServer_RetrieveContextRPC()
+	conn.Close(websocket.StatusNormalClosure, "")
+
+	// Wait for server context to be done
+	select {
+	case <-ctxServer.Done():
+	case <-time.After(time.Second):
+		t.Fatalf("server context not done")
+	}
+}
+
+func testWithBodyPostWebSocketStream(t *testing.T, port int) {
+	apiURL := fmt.Sprintf("ws://localhost:%d/rpc/no-body/stream-with-body?method=POST", port)
+
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
+	conn, _, err := websocket.Dial(ctx, apiURL, nil)
+	if err != nil {
+		t.Fatalf("websocket.DefaultDialer.DialContext(%q) failed with %v; want success", apiURL, err)
+	}
+	defer conn.Close(websocket.StatusNormalClosure, "")
+
+	// Write request because server expects it
+	err = conn.Write(ctx, websocket.MessageText, []byte("{}"))
+	if err != nil {
+		t.Fatalf("conn.WriteMessage() failed with %v; want success", err)
+	}
+
+	// Write another request even though server doesn't expect it
+	err = conn.Write(ctx, websocket.MessageText, []byte("{}"))
+	if err != nil {
+		t.Fatalf("conn.WriteMessage() failed with %v; want success", err)
+	}
+
+	// Wait for the server to start processing the request.
+	ctxServer := server.NoBodyPostServer_RetrieveContextStream()
+	conn.Close(websocket.StatusNormalClosure, "")
+
+	// Wait for server context to be done
+	select {
+	case <-ctxServer.Done():
+	case <-time.After(time.Second):
+		t.Fatalf("server context not done")
+	}
+}