From 6c6177c80a63d8d544bb047b8845bba9aa831742 Mon Sep 17 00:00:00 2001
From: Yuki Ito <mrno110y@gmail.com>
Date: Thu, 10 Nov 2022 15:58:43 +0900
Subject: [PATCH] client: add a call option for modifying the :authority header

---
 internal/transport/http2_client.go |  2 +-
 internal/transport/transport.go    |  3 +++
 rpc_util.go                        | 24 ++++++++++++++++++++++++
 stream.go                          | 13 ++++++++++++-
 4 files changed, 40 insertions(+), 2 deletions(-)

diff --git a/internal/transport/http2_client.go b/internal/transport/http2_client.go
index d518b07e16f7..394cfc2b4f7e 100644
--- a/internal/transport/http2_client.go
+++ b/internal/transport/http2_client.go
@@ -522,7 +522,7 @@ func (t *http2Client) createHeaderFields(ctx context.Context, callHdr *CallHdr)
 	headerFields = append(headerFields, hpack.HeaderField{Name: ":method", Value: "POST"})
 	headerFields = append(headerFields, hpack.HeaderField{Name: ":scheme", Value: t.scheme})
 	headerFields = append(headerFields, hpack.HeaderField{Name: ":path", Value: callHdr.Method})
-	headerFields = append(headerFields, hpack.HeaderField{Name: ":authority", Value: callHdr.Host})
+	headerFields = append(headerFields, hpack.HeaderField{Name: ":authority", Value: callHdr.Authority})
 	headerFields = append(headerFields, hpack.HeaderField{Name: "content-type", Value: grpcutil.ContentType(callHdr.ContentSubtype)})
 	headerFields = append(headerFields, hpack.HeaderField{Name: "user-agent", Value: t.userAgent})
 	headerFields = append(headerFields, hpack.HeaderField{Name: "te", Value: "trailers"})
diff --git a/internal/transport/transport.go b/internal/transport/transport.go
index 2e615ee20cc5..70e56a63c4e5 100644
--- a/internal/transport/transport.go
+++ b/internal/transport/transport.go
@@ -600,6 +600,9 @@ type CallHdr struct {
 	// Host specifies the peer's host.
 	Host string
 
+	// Authority specifies the :authority header.
+	Authority string
+
 	// Method specifies the operation to perform.
 	Method string
 
diff --git a/rpc_util.go b/rpc_util.go
index 934fc1aa015e..a43c6ee11330 100644
--- a/rpc_util.go
+++ b/rpc_util.go
@@ -160,6 +160,7 @@ type callInfo struct {
 	contentSubtype        string
 	codec                 baseCodec
 	maxRetryRPCBufferSize int
+	authorityModifier     func(ctx context.Context, authority string) string
 }
 
 func defaultCallInfo() *callInfo {
@@ -233,6 +234,29 @@ func (o TrailerCallOption) after(c *callInfo, attempt *csAttempt) {
 	*o.TrailerAddr = attempt.s.Trailer()
 }
 
+// Authority returns a CallOption that mofifies the :authority header for a RPC.
+func Authority(modifier func(ctx context.Context, authority string) string) CallOption {
+	return AuthorityCallOption{
+		modifier: modifier,
+	}
+}
+
+// AuthorityCallOption is a CallOption for modifying the :authority header.
+//
+// # Experimental
+//
+// Notice: This type is EXPERIMENTAL and may be changed or removed in a
+// later release.
+type AuthorityCallOption struct {
+	modifier func(ctx context.Context, authority string) string
+}
+
+func (o AuthorityCallOption) before(c *callInfo) error {
+	c.authorityModifier = o.modifier
+	return nil
+}
+func (o AuthorityCallOption) after(c *callInfo, attempt *csAttempt) {}
+
 // Peer returns a CallOption that retrieves peer information for a unary RPC.
 // The peer field will be populated *after* the RPC completes.
 func Peer(p *peer.Peer) CallOption {
diff --git a/stream.go b/stream.go
index 960c3e33dfd3..624f1e6a923f 100644
--- a/stream.go
+++ b/stream.go
@@ -262,8 +262,14 @@ func newClientStreamWithParams(ctx context.Context, desc *StreamDesc, cc *Client
 		return nil, err
 	}
 
+	authority := cc.authority
+	if c.authorityModifier != nil {
+		authority = c.authorityModifier(ctx, authority)
+	}
+
 	callHdr := &transport.CallHdr{
 		Host:           cc.authority,
+		Authority:      authority,
 		Method:         method,
 		ContentSubtype: c.contentSubtype,
 		DoneFunc:       doneFunc,
@@ -762,7 +768,6 @@ func (cs *clientStream) Header() (metadata.MD, error) {
 		}
 		return toRPCErr(err)
 	}, cs.commitAttemptLocked)
-
 	if err != nil {
 		cs.finish(err)
 		return nil, err
@@ -1181,8 +1186,14 @@ func newNonRetryClientStream(ctx context.Context, desc *StreamDesc, method strin
 		return nil, err
 	}
 
+	authority := ac.cc.authority
+	if c.authorityModifier != nil {
+		authority = c.authorityModifier(ctx, authority)
+	}
+
 	callHdr := &transport.CallHdr{
 		Host:           ac.cc.authority,
+		Authority:      authority,
 		Method:         method,
 		ContentSubtype: c.contentSubtype,
 	}