From f5a2e82437e4b5cb013a3669efb17dc7bf9ea110 Mon Sep 17 00:00:00 2001
From: Danil Petrov <danilpetrovvi@gmail.com>
Date: Fri, 22 Feb 2019 01:36:26 +1000
Subject: [PATCH] refactor websocket_wrapper.go with copyHeader func (#362)

This commit makes websocket_wrapper.go file consistent with the changes
in grpc_web_response.go file, PR #359.
---
 go/grpcweb/websocket_wrapper.go | 44 ++++++++++++---------------------
 1 file changed, 16 insertions(+), 28 deletions(-)

diff --git a/go/grpcweb/websocket_wrapper.go b/go/grpcweb/websocket_wrapper.go
index a1513b0c..f281f3d5 100644
--- a/go/grpcweb/websocket_wrapper.go
+++ b/go/grpcweb/websocket_wrapper.go
@@ -54,15 +54,11 @@ func (w *webSocketResponseWriter) writeHeaderFrame(headers http.Header) {
 }
 
 func (w *webSocketResponseWriter) copyFlushedHeaders() {
-	for k, vv := range w.headers {
-		// Skip the pre-annoucement of Trailer headers. Don't add them to the response headers.
-		if strings.ToLower(k) == "trailer" {
-			continue
-		}
-		for _, v := range vv {
-			w.flushedHeaders.Add(k, v)
-		}
-	}
+	copyHeader(
+		w.flushedHeaders, w.headers,
+		skipKeys("trailer"),
+		keyCase(http.CanonicalHeaderKey),
+	)
 }
 
 func (w *webSocketResponseWriter) WriteHeader(code int) {
@@ -73,25 +69,17 @@ func (w *webSocketResponseWriter) WriteHeader(code int) {
 }
 
 func (w *webSocketResponseWriter) extractTrailerHeaders() http.Header {
-	trailerHeaders := make(http.Header)
-	for k, vv := range w.headers {
-		// Skip the pre-annoucement of Trailer headers. Don't add them to the response headers.
-		if strings.ToLower(k) == "trailer" {
-			continue
-		}
-		// Skip existing headers that were already sent.
-		if _, exists := w.flushedHeaders[k]; exists {
-			continue
-		}
-		// Skip the Trailer prefix
-		if strings.HasPrefix(k, http2.TrailerPrefix) {
-			k = k[len(http2.TrailerPrefix):]
-		}
-		for _, v := range vv {
-			trailerHeaders.Add(k, v)
-		}
-	}
-	return trailerHeaders
+	th := make(http.Header)
+	copyHeader(
+		th, w.headers,
+		skipKeys(append([]string{"trailer"}, headerKeys(w.flushedHeaders)...)...),
+		replaceInKeys(http2.TrailerPrefix, ""),
+		// gRPC-Web spec says that must use lower-case header/trailer names.
+		// See "HTTP wire protocols" section in
+		// https://github.com/grpc/grpc/blob/master/doc/PROTOCOL-WEB.md#protocol-differences-vs-grpc-over-http2
+		keyCase(strings.ToLower),
+	)
+	return th
 }
 
 func (w *webSocketResponseWriter) FlushTrailers() {