rpc: support injecting HTTP headers through context (#26023)

This adds a way to specify HTTP headers per request.

Co-authored-by: Martin Holst Swende <martin@swende.se>
Co-authored-by: Felix Lange <fjl@twurst.com>
This commit is contained in:
storyicon 2022-11-16 22:22:12 +08:00 committed by GitHub
parent b4ea2bf7dd
commit add337e0f7
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
3 changed files with 100 additions and 0 deletions

56
rpc/context_headers.go Normal file
View File

@ -0,0 +1,56 @@
// Copyright 2022 The go-ethereum Authors
// This file is part of the go-ethereum library.
//
// The go-ethereum library is free software: you can redistribute it and/or modify
// it under the terms of the GNU Lesser General Public License as published by
// the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// The go-ethereum library is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
// GNU Lesser General Public License for more details.
//
// You should have received a copy of the GNU Lesser General Public License
// along with the go-ethereum library. If not, see <http://www.gnu.org/licenses/>.
package rpc
import (
"context"
"net/http"
)
type mdHeaderKey struct{}
// NewContextWithHeaders wraps the given context, adding HTTP headers. These headers will
// be applied by Client when making a request using the returned context.
func NewContextWithHeaders(ctx context.Context, h http.Header) context.Context {
if len(h) == 0 {
// This check ensures the header map set in context will never be nil.
return ctx
}
var ctxh http.Header
prev, ok := ctx.Value(mdHeaderKey{}).(http.Header)
if ok {
ctxh = setHeaders(prev.Clone(), h)
} else {
ctxh = h.Clone()
}
return context.WithValue(ctx, mdHeaderKey{}, ctxh)
}
// headersFromContext is used to extract http.Header from context.
func headersFromContext(ctx context.Context) http.Header {
source, _ := ctx.Value(mdHeaderKey{}).(http.Header)
return source
}
// setHeaders sets all headers from src in dst.
func setHeaders(dst http.Header, src http.Header) http.Header {
for key, values := range src {
dst[http.CanonicalHeaderKey(key)] = values
}
return dst
}

View File

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

View File

@ -17,6 +17,8 @@
package rpc package rpc
import ( import (
"context"
"fmt"
"net/http" "net/http"
"net/http/httptest" "net/http/httptest"
"strings" "strings"
@ -198,3 +200,43 @@ func TestHTTPPeerInfo(t *testing.T) {
t.Errorf("wrong HTTP.Origin %q", info.HTTP.UserAgent) t.Errorf("wrong HTTP.Origin %q", info.HTTP.UserAgent)
} }
} }
func TestNewContextWithHeaders(t *testing.T) {
expectedHeaders := 0
server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) {
for i := 0; i < expectedHeaders; i++ {
key, want := fmt.Sprintf("key-%d", i), fmt.Sprintf("val-%d", i)
if have := request.Header.Get(key); have != want {
t.Errorf("wrong request headers for %s, want: %s, have: %s", key, want, have)
}
}
writer.WriteHeader(http.StatusOK)
_, _ = writer.Write([]byte(`{}`))
}))
defer server.Close()
client, err := Dial(server.URL)
if err != nil {
t.Fatalf("failed to dial: %s", err)
}
defer client.Close()
newHdr := func(k, v string) http.Header {
header := http.Header{}
header.Set(k, v)
return header
}
ctx1 := NewContextWithHeaders(context.Background(), newHdr("key-0", "val-0"))
ctx2 := NewContextWithHeaders(ctx1, newHdr("key-1", "val-1"))
ctx3 := NewContextWithHeaders(ctx2, newHdr("key-2", "val-2"))
expectedHeaders = 3
if err := client.CallContext(ctx3, nil, "test"); err != ErrNoResult {
t.Error("call failed", err)
}
expectedHeaders = 2
if err := client.CallContext(ctx2, nil, "test"); err != ErrNoResult {
t.Error("call failed:", err)
}
}