diff --git a/runtime/context.go b/runtime/context.go index 571e353dd5c..6e0eb27e285 100644 --- a/runtime/context.go +++ b/runtime/context.go @@ -61,10 +61,8 @@ func AnnotateContext(ctx context.Context, mux *ServeMux, req *http.Request) (con if strings.ToLower(key) == "authorization" { pairs = append(pairs, "authorization", val) } - if mux.incomingHeaderMatcher != nil { - if h, ok := mux.incomingHeaderMatcher(key); ok { - pairs = append(pairs, h, val) - } + if h, ok := mux.incomingHeaderMatcher(key); ok { + pairs = append(pairs, h, val) } } } diff --git a/runtime/handler.go b/runtime/handler.go index 22eb1b3dbad..0f4750f4c80 100644 --- a/runtime/handler.go +++ b/runtime/handler.go @@ -68,15 +68,11 @@ func ForwardResponseStream(ctx context.Context, mux *ServeMux, marshaler Marshal func handleForwardResponseServerMetadata(w http.ResponseWriter, mux *ServeMux, md ServerMetadata) { for k, vs := range md.HeaderMD { - hKey := fmt.Sprintf("%s%s", MetadataHeaderPrefix, k) - if mux.outgoingHeaderMatcher != nil { - if h, ok := mux.outgoingHeaderMatcher(k); ok { - hKey = h + if h, ok := mux.outgoingHeaderMatcher(k); ok { + for _, v := range vs { + w.Header().Add(h, v) } } - for i := range vs { - w.Header().Add(hKey, vs[i]) - } } } @@ -90,8 +86,8 @@ func handleForwardResponseTrailerHeader(w http.ResponseWriter, md ServerMetadata func handleForwardResponseTrailer(w http.ResponseWriter, md ServerMetadata) { for k, vs := range md.TrailerMD { tKey := fmt.Sprintf("%s%s", MetadataTrailerPrefix, k) - for i := range vs { - w.Header().Add(tKey, vs[i]) + for _, v := range vs { + w.Header().Add(tKey, v) } } } diff --git a/runtime/mux.go b/runtime/mux.go index dda3db205fd..f7a228031dc 100644 --- a/runtime/mux.go +++ b/runtime/mux.go @@ -1,6 +1,7 @@ package runtime import ( + "fmt" "net/http" "net/textproto" "strings" @@ -107,7 +108,6 @@ func NewServeMux(opts ...ServeMuxOption) *ServeMux { handlers: make(map[string][]handler), forwardResponseOptions: make([]func(context.Context, http.ResponseWriter, proto.Message) error, 0), marshalers: makeMarshalerMIMERegistry(), - incomingHeaderMatcher: DefaultHeaderMatcher, } for _, opt := range opts { @@ -126,6 +126,16 @@ func NewServeMux(opts ...ServeMuxOption) *ServeMux { } } + if serveMux.incomingHeaderMatcher == nil { + serveMux.incomingHeaderMatcher = DefaultHeaderMatcher + } + + if serveMux.outgoingHeaderMatcher == nil { + serveMux.outgoingHeaderMatcher = func(key string) (string, bool) { + return fmt.Sprintf("%s%s", MetadataHeaderPrefix, key), true + } + } + return serveMux }