client: allow overriding grpc-accept-encoding header

Signed-off-by: Israel Blancas <[email protected]>

Author: iblancasa

default_dial_option_server_option_test.go

@@ -176,3 +176,81 @@ func (s) TestHeaderListSizeDialOptionServerOption(t *testing.T) {
 		t.Fatalf("Unexpected s.opts.MaxHeaderListSizeDialOption.MaxHeaderListSize: %d != %d", serverHeaderListSize, maxHeaderListSize)
 	}
 }
+
+func (s) TestFormatAcceptedCompressors(t *testing.T) {
+	tests := []struct {
+		name  string
+		input []string
+		want  string
+	}{
+		{
+			name:  "empty input",
+			input: []string{},
+			want:  "",
+		},
+		{
+			name:  "nil input",
+			input: nil,
+			want:  "",
+		},
+		{
+			name:  "single valid compressor",
+			input: []string{"gzip"},
+			want:  "gzip",
+		},
+		{
+			name:  "multiple valid compressors",
+			input: []string{"gzip", "gzip"}, // Note: need to test with actually different compressors when available
+			want:  "gzip",                   // deduplicates
+		},
+		{
+			name:  "compressor with whitespace",
+			input: []string{" gzip ", " gzip"},
+			want:  "gzip",
+		},
+		{
+			name:  "empty strings",
+			input: []string{"", "  ", ""},
+			want:  "",
+		},
+		{
+			name:  "unregistered compressor",
+			input: []string{"unknown"},
+			want:  "",
+		},
+		{
+			name:  "mix of valid and invalid",
+			input: []string{"unknown", "gzip", "notreal"},
+			want:  "gzip",
+		},
+		{
+			name:  "duplicates with whitespace",
+			input: []string{"gzip", " gzip", "gzip "},
+			want:  "gzip",
+		},
+		{
+			name:  "empty string mixed with valid",
+			input: []string{"", "gzip", "  ", "gzip"},
+			want:  "gzip",
+		},
+		{
+			name:  "whitespace only",
+			input: []string{"   ", "\t", " "},
+			want:  "",
+		},
+		{
+			name:  "preserves order after deduplication",
+			input: []string{"gzip", "gzip"},
+			want:  "gzip",
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			got := formatAcceptedCompressors(tt.input)
+			if got != tt.want {
+				t.Errorf("formatAcceptedCompressors(%v) = %q, want %q", tt.input, got, tt.want)
+			}
+		})
+	}
+}

dialoptions.go

@@ -22,6 +22,7 @@ import (
 	"context"
 	"net"
 	"net/url"
+	"strings"
 	"time"
 
 	"google.golang.org/grpc/backoff"
@@ -31,6 +32,7 @@ import (
 	"google.golang.org/grpc/internal"
 	internalbackoff "google.golang.org/grpc/internal/backoff"
 	"google.golang.org/grpc/internal/binarylog"
+	"google.golang.org/grpc/internal/grpcutil"
 	"google.golang.org/grpc/internal/transport"
 	"google.golang.org/grpc/keepalive"
 	"google.golang.org/grpc/mem"
@@ -183,6 +185,16 @@ func WithSharedWriteBuffer(val bool) DialOption {
 	})
 }
 
+// WithAcceptedCompressionNames limits the values advertised in the
+// grpc-accept-encoding header for this ClientConn. Names must correspond to
+// compressors registered via encoding.RegisterCompressor. When empty, the
+// default set of registered compressors is used.
+func WithAcceptedCompressionNames(names ...string) DialOption {
+	return newFuncDialOption(func(o *dialOptions) {
+		o.copts.AcceptedCompressors = formatAcceptedCompressors(names)
+	})
+}
+
 // WithWriteBufferSize determines how much data can be batched before doing a
 // write on the wire. The default value for this buffer is 32KB.
 //
@@ -207,6 +219,32 @@ func WithReadBufferSize(s int) DialOption {
 	})
 }
 
+func formatAcceptedCompressors(names []string) string {
+	if len(names) == 0 {
+		return ""
+	}
+	sch := make([]string, 0, len(names))
+	seen := make(map[string]struct{}, len(names))
+	for _, name := range names {
+		name = strings.TrimSpace(name)
+		if name == "" {
+			continue
+		}
+		if !grpcutil.IsCompressorNameRegistered(name) {
+			continue
+		}
+		if _, ok := seen[name]; ok {
+			continue
+		}
+		seen[name] = struct{}{}
+		sch = append(sch, name)
+	}
+	if len(sch) == 0 {
+		return ""
+	}
+	return strings.Join(sch, ",")
+}
+
 // WithInitialWindowSize returns a DialOption which sets the value for initial
 // window size on a stream. The lower bound for window size is 64K and any value
 // smaller than that will be ignored.

internal/transport/http2_client.go

@@ -321,12 +321,17 @@ func NewHTTP2Client(connectCtx, ctx context.Context, addr resolver.Address, opts
 		maxHeaderListSize = *opts.MaxHeaderListSize
 	}
 
+	registeredCompressors := opts.AcceptedCompressors
+	if registeredCompressors == "" {
+		registeredCompressors = grpcutil.RegisteredCompressors()
+	}
+
 	t := &http2Client{
 		ctx:                   ctx,
 		ctxDone:               ctx.Done(), // Cache Done chan.
 		cancel:                cancel,
 		userAgent:             opts.UserAgent,
-		registeredCompressors: grpcutil.RegisteredCompressors(),
+		registeredCompressors: registeredCompressors,
 		address:               addr,
 		conn:                  conn,
 		remoteAddr:            conn.RemoteAddr(),

internal/transport/transport.go

@@ -531,6 +531,9 @@ type ConnectOptions struct {
 	BufferPool mem.BufferPool
 	// StaticWindowSize controls whether dynamic window sizing is enabled.
 	StaticWindowSize bool
+	// AcceptedCompressors overrides the grpc-accept-encoding header for this
+	// connection. When empty, all registered compressors are advertised.
+	AcceptedCompressors string
 }
 
 // WriteOptions provides additional hints and information for message

test/compressor_test.go

@@ -533,6 +533,30 @@ func (s) TestClientSupportedCompressors(t *testing.T) {
 	}
 }
 
+func (s) TestAcceptedCompressionNamesDialOption(t *testing.T) {
+	const want = "gzip"
+	ss := &stubserver.StubServer{
+		EmptyCallF: func(ctx context.Context, _ *testpb.Empty) (*testpb.Empty, error) {
+			md, _ := metadata.FromIncomingContext(ctx)
+			if got := md.Get("grpc-accept-encoding"); len(got) != 1 || got[0] != want {
+				t.Fatalf("unexpected grpc-accept-encoding header: %v", got)
+			}
+			return &testpb.Empty{}, nil
+		},
+	}
+	if err := ss.Start(nil, grpc.WithAcceptedCompressionNames(want)); err != nil {
+		t.Fatalf("failed to start server: %v", err)
+	}
+	defer ss.Stop()
+
+	ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
+	defer cancel()
+
+	if _, err := ss.Client.EmptyCall(ctx, &testpb.Empty{}); err != nil {
+		t.Fatalf("EmptyCall failed: %v", err)
+	}
+}
+
 func (s) TestCompressorRegister(t *testing.T) {
 	for _, e := range listTestEnv() {
 		testCompressorRegister(t, e)