Add headerbp integration

httpbp/client_middlewares.go

@@ -15,6 +15,7 @@ import (
 	"github.com/prometheus/client_golang/prometheus"
 
 	"github.com/reddit/baseplate.go/breakerbp"
+	"github.com/reddit/baseplate.go/internal/headerbp"
 	//lint:ignore SA1019 This library is internal only, not actually deprecated
 	"github.com/reddit/baseplate.go/internalv2compat"
 	"github.com/reddit/baseplate.go/retrybp"
@@ -88,6 +89,11 @@ func NewClient(config ClientConfig, middleware ...ClientMiddleware) (*http.Clien
 	if config.CircuitBreaker != nil {
 		defaults = append([]ClientMiddleware{CircuitBreaker(*config.CircuitBreaker)}, defaults...)
 	}
+
+	// only add the middleware to forward baseplate headers if the client is for internal calls
+	if config.InternalOnly {
+		defaults = append(defaults, ClientBaseplateHeadersMiddleware(config.Slug))
+	}
 	middleware = append(middleware, defaults...)
 
 	return &http.Client{
@@ -349,3 +355,26 @@ func PrometheusClientMetrics(serverSlug string) ClientMiddleware {
 		})
 	}
 }
+
+// ClientBaseplateHeadersMiddleware is a middleware that forwards baseplate headers from the context to the outgoing request.
+//
+// If it detects any new baseplate headers set on the request, it will reject the request and return an error.
+func ClientBaseplateHeadersMiddleware(client string) ClientMiddleware {
+	return func(next http.RoundTripper) http.RoundTripper {
+		return roundTripperFunc(func(req *http.Request) (*http.Response, error) {
+			for k := range req.Header {
+				if err := headerbp.CheckClientHeader(k,
+					headerbp.WithHTTPClient("", client, ""),
+				); err != nil {
+					return nil, err
+				}
+			}
+			headerbp.SetOutgoingHeaders(
+				req.Context(),
+				headerbp.WithHTTPClient("", client, ""),
+				headerbp.WithHeaderSetter(req.Header.Set),
+			)
+			return next.RoundTrip(req)
+		})
+	}
+}

httpbp/config.go

@@ -16,6 +16,7 @@ type ClientConfig struct {
 	MaxConnections    int               `yaml:"maxConnections"`
 	CircuitBreaker    *breakerbp.Config `yaml:"circuitBreaker"`
 	RetryOptions      []retry.Option
+	InternalOnly      bool
 }
 
 // Validate checks ClientConfig for any missing or erroneous values.

httpbp/middlewares.go

@@ -15,6 +15,7 @@ import (
 
 	"github.com/reddit/baseplate.go/ecinterface"
 	"github.com/reddit/baseplate.go/errorsbp"
+	"github.com/reddit/baseplate.go/internal/headerbp"
 	//lint:ignore SA1019 This library is internal only, not actually deprecated
 	"github.com/reddit/baseplate.go/internalv2compat"
 	"github.com/reddit/baseplate.go/log"
@@ -517,3 +518,48 @@ func (rr *responseRecorder) WriteHeader(code int) {
 	rr.ResponseWriter.WriteHeader(code)
 	rr.responseCode = code
 }
+
+type untrustedHeadersKey struct{}
+
+func setUntrustedHeaders(ctx context.Context, h map[string]string) context.Context {
+	return context.WithValue(ctx, untrustedHeadersKey{}, h)
+}
+
+func GetUntrustedBaseplateHeaders(ctx context.Context) (map[string]string, bool) {
+	h, ok := ctx.Value(untrustedHeadersKey{}).(map[string]string)
+	return h, ok
+}
+
+// ServerBaseplateHeadersMiddleware is a middleware that extracts baseplate headers from the incoming request and adds them to the context.
+//
+// If the request is flagged as untrusted, it will remove the baseplate headers from the request and add them to the
+// context. These can be retrieved using GetUntrustedBaseplateHeaders.
+func ServerBaseplateHeadersMiddleware(service string) Middleware {
+	return func(name string, next HandlerFunc) HandlerFunc {
+		return func(ctx context.Context, w http.ResponseWriter, r *http.Request) (err error) {
+			if r.Header.Get(headerbp.IsUntrustedRequestHeaderCanonicalHTTP) != "" {
+				untrusted := make(map[string]string)
+				for k, v := range r.Header {
+					if headerbp.IsBaseplateHeader(k) {
+						if len(v) > 0 {
+							untrusted[strings.ToLower(k)] = v[0]
+						}
+						r.Header.Del(k)
+					}
+				}
+				ctx = setUntrustedHeaders(ctx, untrusted)
+				return next(ctx, w, r)
+			}
+			headers := headerbp.NewIncomingHeaders(
+				headerbp.WithHTTPService(service, name),
+			)
+			for k, v := range r.Header {
+				if len(v) > 0 {
+					headers.RecordHeader(k, v[0])
+				}
+			}
+			ctx = headers.SetOnContext(ctx)
+			return next(ctx, w, r.WithContext(ctx))
+		}
+	}
+}