all: make use of ctxkey everywhere (#10846)

Also perform minor cleanups on the ctxkey package itself.
Provide guidance on when to use ctxkey.Key[T] over ctxkey.New.
Also, allow for interface kinds because the value wrapping trick
also happens to fix edge cases with interfaces in Go.

Updates #cleanup

Signed-off-by: Joe Tsai <joetsai@digital-static.net>
This commit is contained in:
Joe Tsai
2024-01-16 13:56:23 -08:00
committed by GitHub
parent 7732377cd7
commit c25968e1c5
13 changed files with 97 additions and 85 deletions

View File

@@ -34,6 +34,7 @@ import (
"tailscale.com/tailcfg"
"tailscale.com/types/lazy"
"tailscale.com/types/logger"
"tailscale.com/util/ctxkey"
"tailscale.com/util/mak"
"tailscale.com/version"
)
@@ -48,8 +49,7 @@ const (
// current etag of a resource.
var ErrETagMismatch = errors.New("etag mismatch")
// serveHTTPContextKey is the context.Value key for a *serveHTTPContext.
type serveHTTPContextKey struct{}
var serveHTTPContextKey ctxkey.Key[*serveHTTPContext]
type serveHTTPContext struct {
SrcAddr netip.AddrPort
@@ -433,7 +433,7 @@ func (b *LocalBackend) tcpHandlerForServe(dport uint16, srcAddr netip.AddrPort)
hs := &http.Server{
Handler: http.HandlerFunc(b.serveWebHandler),
BaseContext: func(_ net.Listener) context.Context {
return context.WithValue(context.Background(), serveHTTPContextKey{}, &serveHTTPContext{
return serveHTTPContextKey.WithValue(context.Background(), &serveHTTPContext{
SrcAddr: srcAddr,
DestPort: dport,
})
@@ -500,11 +500,6 @@ func (b *LocalBackend) tcpHandlerForServe(dport uint16, srcAddr netip.AddrPort)
return nil
}
func getServeHTTPContext(r *http.Request) (c *serveHTTPContext, ok bool) {
c, ok = r.Context().Value(serveHTTPContextKey{}).(*serveHTTPContext)
return c, ok
}
func (b *LocalBackend) getServeHandler(r *http.Request) (_ ipn.HTTPHandlerView, at string, ok bool) {
var z ipn.HTTPHandlerView // zero value
@@ -521,7 +516,7 @@ func (b *LocalBackend) getServeHandler(r *http.Request) (_ ipn.HTTPHandlerView,
hostname = r.TLS.ServerName
}
sctx, ok := getServeHTTPContext(r)
sctx, ok := serveHTTPContextKey.ValueOk(r.Context())
if !ok {
b.logf("[unexpected] localbackend: no serveHTTPContext in request")
return z, "", false
@@ -684,7 +679,7 @@ func addProxyForwardedHeaders(r *httputil.ProxyRequest) {
if r.In.TLS != nil {
r.Out.Header.Set("X-Forwarded-Proto", "https")
}
if c, ok := getServeHTTPContext(r.Out); ok {
if c, ok := serveHTTPContextKey.ValueOk(r.Out.Context()); ok {
r.Out.Header.Set("X-Forwarded-For", c.SrcAddr.Addr().String())
}
}
@@ -696,7 +691,7 @@ func (b *LocalBackend) addTailscaleIdentityHeaders(r *httputil.ProxyRequest) {
r.Out.Header.Del("Tailscale-User-Profile-Pic")
r.Out.Header.Del("Tailscale-Headers-Info")
c, ok := getServeHTTPContext(r.Out)
c, ok := serveHTTPContextKey.ValueOk(r.Out.Context())
if !ok {
return
}

View File

@@ -158,7 +158,7 @@ func TestGetServeHandler(t *testing.T) {
TLS: &tls.ConnectionState{ServerName: serverName},
}
port := cmpx.Or(tt.port, 443)
req = req.WithContext(context.WithValue(req.Context(), serveHTTPContextKey{}, &serveHTTPContext{
req = req.WithContext(serveHTTPContextKey.WithValue(req.Context(), &serveHTTPContext{
DestPort: port,
}))
@@ -428,7 +428,7 @@ func TestServeHTTPProxy(t *testing.T) {
URL: &url.URL{Path: "/"},
TLS: &tls.ConnectionState{ServerName: "example.ts.net"},
}
req = req.WithContext(context.WithValue(req.Context(), serveHTTPContextKey{}, &serveHTTPContext{
req = req.WithContext(serveHTTPContextKey.WithValue(req.Context(), &serveHTTPContext{
DestPort: 443,
SrcAddr: netip.MustParseAddrPort(tt.srcIP + ":1234"), // random src port for tests
}))