client: allow overriding grpc-accept-encoding header

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

Author: iblancasa

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)