add allowHeaders

Author: molon

server/cors.go

@@ -5,7 +5,11 @@ import (
 	"strings"
 )
 
-func allowCORS(h http.Handler, allowOrigins []string) http.Handler {
+var (
+	defaultAllowHeaders = []string{"Content-Type", "Accept", "Authorization", "Origin"}
+)
+
+func allowCORS(h http.Handler, allowOrigins []string, extraAllowHeaders []string) http.Handler {
 	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 		if origin := r.Header.Get("Origin"); origin != "" {
 			if len(allowOrigins) > 0 {
@@ -20,16 +24,34 @@ func allowCORS(h http.Handler, allowOrigins []string) http.Handler {
 			}
 
 			if r.Method == "OPTIONS" && r.Header.Get("Access-Control-Request-Method") != "" {
-				preflightHandler(w, r)
+				preflightHandler(w, r, extraAllowHeaders)
 				return
 			}
 		}
 		h.ServeHTTP(w, r)
 	})
 }
 
-func preflightHandler(w http.ResponseWriter, r *http.Request) {
-	headers := []string{"Content-Type", "Accept", "Authorization", "Origin"}
+func evaluateExtraAllowHeaders(allowHeaders []string) []string {
+	m := map[string]bool{}
+	for _, h := range defaultAllowHeaders {
+		m[h] = true
+	}
+
+	extraAllowHeaders := []string{}
+	for _, h := range allowHeaders {
+		if m[h] == false {
+			extraAllowHeaders = append(extraAllowHeaders, h)
+		}
+	}
+	return extraAllowHeaders
+}
+
+func preflightHandler(w http.ResponseWriter, r *http.Request, extraAllowHeaders []string) {
+	headers := defaultAllowHeaders
+	if len(extraAllowHeaders) > 0 {
+		headers = append(headers, extraAllowHeaders...)
+	}
 	w.Header().Set("Access-Control-Allow-Headers", strings.Join(headers, ","))
 	methods := []string{"GET", "HEAD", "POST", "PUT", "DELETE"}
 	w.Header().Set("Access-Control-Allow-Methods", strings.Join(methods, ","))

server/server.go

@@ -32,6 +32,8 @@ type Server struct {
 	initializeTimeout time.Duration
 	registrars        []func(mux *http.ServeMux) error
 	allowOrigins      []string
+	allowHeaders      []string
+	extraAllowHeaders []string
 
 	// GRPCServer will be started whenever this is served
 	GRPCServer *grpc.Server
@@ -66,7 +68,7 @@ func NewServer(opts ...Option) (*Server, error) {
 			return nil, err
 		}
 	}
-	s.HTTPServer.Handler = allowCORS(mux, s.allowOrigins)
+	s.HTTPServer.Handler = allowCORS(mux, s.allowOrigins, s.extraAllowHeaders)
 
 	return s, nil
 }
@@ -137,6 +139,14 @@ func WithAllowOrigins(allowOrigins []string) Option {
 	}
 }
 
+func WithAllowHeaders(allowHeaders []string) Option {
+	return func(s *Server) error {
+		s.allowHeaders = allowHeaders
+		s.extraAllowHeaders = evaluateExtraAllowHeaders(allowHeaders)
+		return nil
+	}
+}
+
 // Serve invokes all initializers then serves on the given listeners.
 //
 // If a listener is left blank, then that particular part will not be served.