improve {incoming,outgoing}HeaderMatcher logic (#408)

* default implementation over nil-check
* really exclude header when outgoingHeaderMatcher returns false

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)
 			}
 		}
 	}

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)
 		}
 	}
 }

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
 }