Add header injector middlewares

pkg/middleware/headers.go

@@ -0,0 +1,102 @@
+package middleware
+
+import (
+	"fmt"
+	"net/http"
+
+	"github.com/justinas/alice"
+	"github.com/oauth2-proxy/oauth2-proxy/v7/pkg/apis/options"
+	"github.com/oauth2-proxy/oauth2-proxy/v7/pkg/header"
+)
+
+func NewRequestHeaderInjector(headers []options.Header) (alice.Constructor, error) {
+	headerInjector, err := newRequestHeaderInjector(headers)
+	if err != nil {
+		return nil, fmt.Errorf("error building request header injector: %v", err)
+	}
+
+	strip := newStripHeaders(headers)
+	if strip != nil {
+		return alice.New(strip, headerInjector).Then, nil
+	}
+	return headerInjector, nil
+}
+
+func newStripHeaders(headers []options.Header) alice.Constructor {
+	headersToStrip := []string{}
+	for _, header := range headers {
+		if !header.PreserveRequestValue {
+			headersToStrip = append(headersToStrip, header.Name)
+		}
+	}
+
+	if len(headersToStrip) == 0 {
+		return nil
+	}
+
+	return func(next http.Handler) http.Handler {
+		return stripHeaders(headersToStrip, next)
+	}
+}
+
+func stripHeaders(headers []string, next http.Handler) http.Handler {
+	return http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {
+		for _, header := range headers {
+			req.Header.Del(header)
+		}
+		next.ServeHTTP(rw, req)
+	})
+}
+
+func newRequestHeaderInjector(headers []options.Header) (alice.Constructor, error) {
+	injector, err := header.NewInjector(headers)
+	if err != nil {
+		return nil, fmt.Errorf("error building request injector: %v", err)
+	}
+
+	return func(next http.Handler) http.Handler {
+		return injectRequestHeaders(injector, next)
+	}, nil
+}
+
+func injectRequestHeaders(injector header.Injector, next http.Handler) http.Handler {
+	return http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {
+		scope := GetRequestScope(req)
+
+		// If scope is nil, this will panic.
+		// A scope should always be injected before this handler is called.
+		injector.Inject(req.Header, scope.Session)
+		next.ServeHTTP(rw, req)
+	})
+}
+
+func NewResponseHeaderInjector(headers []options.Header) (alice.Constructor, error) {
+	headerInjector, err := newResponseHeaderInjector(headers)
+	if err != nil {
+		return nil, fmt.Errorf("error building response header injector: %v", err)
+	}
+
+	return headerInjector, nil
+}
+
+func newResponseHeaderInjector(headers []options.Header) (alice.Constructor, error) {
+	injector, err := header.NewInjector(headers)
+	if err != nil {
+		return nil, fmt.Errorf("error building response injector: %v", err)
+	}
+
+	return func(next http.Handler) http.Handler {
+		return injectResponseHeaders(injector, next)
+	}, nil
+}
+
+func injectResponseHeaders(injector header.Injector, next http.Handler) http.Handler {
+	return http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {
+		scope := GetRequestScope(req)
+
+		// If scope is nil, this will panic.
+		// A scope should always be injected before this handler is called.
+		injector.Inject(rw.Header(), scope.Session)
+		next.ServeHTTP(rw, req)
+	})
+}