Combine header functions into one file

Summary: They all deal with context headers.

Reviewed By: leoleovich

Differential Revision: D70210416

fbshipit-source-id: 00ade391a16efc1bcfbe3b326969f60804124fbd

thrift/lib/go/thrift/context_headers.go

@@ -22,6 +22,15 @@ import (
 	"github.com/facebook/fbthrift/thrift/lib/go/thrift/types"
 )
 
+// The headersKeyType type is unexported to prevent collisions with context keys.
+type headersKeyType int
+
+// RequestHeadersKey is a context key.
+const RequestHeadersKey headersKeyType = 0
+
+// ResponseHeadersKey is a context key.
+const ResponseHeadersKey headersKeyType = 1
+
 // AddHeader adds a header to the context, which will be sent as part of the request.
 // AddHeader can be called multiple times to add multiple headers.
 // These headers are not persistent and will only be sent with the current request.
@@ -106,3 +115,42 @@ func RequestHeadersFromContext(ctx context.Context) map[string]string {
 	}
 	return requestHeadersMap
 }
+
+// SetRequestHeaders sets the Headers in the protocol to send with the request.
+// These headers will be written via the Write method, inside the Call method for each generated request.
+// These Headers will be cleared with Flush, as they are not persistent.
+func SetRequestHeaders(ctx context.Context, protocol Protocol) error {
+	if ctx == nil {
+		return nil
+	}
+	headers := ctx.Value(RequestHeadersKey)
+	if headers == nil {
+		return nil
+	}
+	headersMap, ok := headers.(map[string]string)
+	if !ok {
+		return NewTransportException(INVALID_HEADERS_TYPE, "Headers key in context value is not map[string]string")
+	}
+	for k, v := range headersMap {
+		protocol.setRequestHeader(k, v)
+	}
+	return nil
+}
+
+func setResponseHeaders(ctx context.Context, responseHeaders map[string]string) {
+	if ctx == nil {
+		return
+	}
+	responseHeadersMapIf := ctx.Value(ResponseHeadersKey)
+	if responseHeadersMapIf == nil {
+		return
+	}
+	responseHeadersMap, ok := responseHeadersMapIf.(map[string]string)
+	if !ok {
+		return
+	}
+
+	for k, v := range responseHeaders {
+		responseHeadersMap[k] = v
+	}
+}