rpc: support inject http headers using outgoing context

rpc/http.go

@@ -214,6 +214,8 @@ func (hc *httpConn) doRequest(ctx context.Context, msg interface{}) (io.ReadClos
 	hc.mu.Lock()
 	req.Header = hc.headers.Clone()
 	hc.mu.Unlock()
+	mergeHeadersFromOutgoingContext(ctx, req.Header)
+
 	if hc.auth != nil {
 		if err := hc.auth(req.Header); err != nil {
 			return nil, err

rpc/http_test.go

@@ -17,6 +17,7 @@
 package rpc
 
 import (
+	"context"
 	"net/http"
 	"net/http/httptest"
 	"strings"
@@ -198,3 +199,29 @@ func TestHTTPPeerInfo(t *testing.T) {
 		t.Errorf("wrong HTTP.Origin %q", info.HTTP.UserAgent)
 	}
 }
+
+func TestNewOutgoingContext(t *testing.T) {
+	const (
+		testHeaderKey   = "test-header-key"
+		testHeaderValue = "test-header-value"
+	)
+	server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) {
+		if got := request.Header.Get(testHeaderKey); got != testHeaderValue {
+			t.Errorf("wrong request headers for %s, expected: %s, actual: %s", testHeaderKey, testHeaderValue, got)
+		}
+		writer.WriteHeader(http.StatusOK)
+		_, _ = writer.Write([]byte(`{}`))
+	}))
+	defer server.Close()
+	client, err := Dial(server.URL)
+	if err != nil {
+		panic(err)
+	}
+	defer client.Close()
+	header := http.Header{}
+	header.Set(testHeaderKey, testHeaderValue)
+	ctx := NewOutgoingContext(context.TODO(), header)
+	if err := client.CallContext(ctx, &struct{}{}, "test"); err != ErrNoResult {
+		t.Errorf("failed to call context: %s", err)
+	}
+}

rpc/metadata.go

@@ -0,0 +1,42 @@
+package rpc
+
+import (
+	"context"
+	"net/http"
+)
+
+type rawMd struct {
+	headers http.Header
+}
+
+type mdOutgoingKey struct{}
+
+// NewOutgoingContext is used to attach http headers into the context
+func NewOutgoingContext(ctx context.Context, header http.Header) context.Context {
+	return context.WithValue(ctx, mdOutgoingKey{}, rawMd{headers: header})
+}
+
+// HeadersFromOutgoingContext is used to extract http headers from the context
+func HeadersFromOutgoingContext(ctx context.Context) (http.Header, bool) {
+	value := ctx.Value(mdOutgoingKey{})
+	if value == nil {
+		return nil, false
+	}
+	headers := value.(rawMd).headers
+	if headers == nil {
+		return nil, false
+	}
+	return headers, true
+}
+
+// mergeHeadersFromOutgoingContext is used to extract http headers from the context and inject it into the provided headers
+func mergeHeadersFromOutgoingContext(ctx context.Context, headers http.Header) {
+	if kvs, ok := HeadersFromOutgoingContext(ctx); ok {
+		for key, values := range kvs {
+			headers.Del(key)
+			for _, val := range values {
+				headers.Add(key, val)
+			}
+		}
+	}
+}