[tailscale] net/http: add enforcement hook to Transport.RoundTrip, li…

…ke our net ones

Updates #55
Updates tailscale/corp#12702

Signed-off-by: Brad Fitzpatrick <bradfitz@tailscale.com>
(cherry picked from commit 8df9488)

api/go1.99999.txt

@@ -8,3 +8,4 @@ pkg net, type SockTrace struct, DidRead func(int)
 pkg net, type SockTrace struct, DidWrite func(int)
 pkg net, type SockTrace struct, WillCloseTCPConn func(syscall.RawConn)
 pkg net, type SockTrace struct, WillOverwrite func(*SockTrace)
+pkg net/http, func SetRoundTripEnforcer(func(*Request) error)

src/net/http/tailscale.go

@@ -0,0 +1,21 @@
+package http
+
+var roundTripEnforcer func(*Request) error
+
+// SetRoundTripEnforcer set a program-global resolver enforcer that can cause
+// RoundTrip calls to fail based on the request and its context.
+//
+// f must be non-nil.
+//
+// SetRoundTripEnforcer can only be called once, and must not be called
+// concurrent with any RoundTrip call; it's expected to be registered during
+// init.
+func SetRoundTripEnforcer(f func(*Request) error) {
+	if f == nil {
+		panic("nil func")
+	}
+	if roundTripEnforcer != nil {
+		panic("already called")
+	}
+	roundTripEnforcer = f
+}

src/net/http/transport.go

@@ -528,6 +528,12 @@ func validateHeaders(hdrs Header) string {
 
 // roundTrip implements a RoundTripper over HTTP.
 func (t *Transport) roundTrip(req *Request) (_ *Response, err error) {
+	if roundTripEnforcer != nil {
+		if err := roundTripEnforcer(req); err != nil {
+			return nil, err
+		}
+	}
+
 	t.nextProtoOnce.Do(t.onceSetNextProtoDefaults)
 	ctx := req.Context()
 	trace := httptrace.ContextClientTrace(ctx)