forked from stackrox/go-grpc-http1
-
Notifications
You must be signed in to change notification settings - Fork 0
/
response_writer.go
145 lines (123 loc) · 4.66 KB
/
response_writer.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
// Copyright (c) 2020 StackRox Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License
package grpcweb
import (
"bytes"
"encoding/binary"
"net/http"
"strings"
"github.com/mbanikazemi/go-grpc-http1/internal/sliceutils"
"github.com/mbanikazemi/go-grpc-http1/internal/stringutils"
)
type responseWriter struct {
w http.ResponseWriter
// List of trailers that were announced via the `Trailer` header at the time headers were written. Also used to keep
// track of whether headers were already written (in which case this is non-nil, even if it is the empty slice).
announcedTrailers []string
}
// NewResponseWriter returns a response writer that transparently transcodes an gRPC HTTP/2 response to a gRPC-Web
// response. It can be used as the response writer in the `ServeHTTP` method of a `grpc.Server`.
// The second return value is a finalization function that takes care of sending the data frame with trailers. It
// *needs* to be called before the response handler exits successfully (the returned error is simply any error of the
// underlying response writer passed through).
func NewResponseWriter(w http.ResponseWriter) (http.ResponseWriter, func() error) {
rw := &responseWriter{
w: w,
}
return rw, rw.Finalize
}
// Header returns the HTTP Header of the underlying response writer.
func (w *responseWriter) Header() http.Header {
return w.w.Header()
}
// Flush flushes any data not yet written. In contrast to most `http.ResponseWriter` implementations, it does not send
// headers if no data has been written yet.
func (w *responseWriter) Flush() {
if w.announcedTrailers == nil {
return // don't send headers needlessly, otherwise trailer-only responses won't work.
}
w.prepareHeadersIfNecessary()
if flusher, _ := w.w.(http.Flusher); flusher != nil {
flusher.Flush()
}
}
// prepareHeadersIfNecessary is called internally on any action that might cause headers to be sent.
func (w *responseWriter) prepareHeadersIfNecessary() {
if w.announcedTrailers != nil {
return
}
hdr := w.w.Header()
w.announcedTrailers = sliceutils.StringClone(hdr["Trailer"])
// Trailers are sent in a data frame, so don't announce trailers as otherwise downstream proxies might get confused.
hdr.Del("Trailer")
// "Downgrade" response content type to grpc-web.
contentType, contentSubtype := stringutils.Split2(hdr.Get("Content-Type"), "+")
respContentType := "application/grpc-web"
if contentType == "application/grpc" && contentSubtype != "" {
respContentType += "+" + contentSubtype
}
hdr.Set("Content-Type", respContentType)
// Any content length that might be set is no longer accurate because of trailers.
hdr.Del("Content-Length")
}
// WriteHeader sends HTTP headers to the client, along with the given status code.
func (w *responseWriter) WriteHeader(statusCode int) {
w.prepareHeadersIfNecessary()
w.w.WriteHeader(statusCode)
}
// Write writes a chunk of data.
func (w *responseWriter) Write(buf []byte) (int, error) {
w.prepareHeadersIfNecessary()
return w.w.Write(buf)
}
// Finalize sends trailer data in a data frame. It *needs* to be called
func (w *responseWriter) Finalize() error {
hdr := w.w.Header()
var trailers http.Header
if w.announcedTrailers == nil {
// Trailer-only response! Send trailers as headers...
trailers = hdr
delete(trailers, "Trailer")
} else {
trailers = make(http.Header)
for _, at := range w.announcedTrailers {
at = http.CanonicalHeaderKey(at)
trailers[at] = hdr[at]
}
}
for k, vs := range hdr {
if !strings.HasPrefix(k, http.TrailerPrefix) {
continue
}
trailerName := http.CanonicalHeaderKey(k[len(http.TrailerPrefix):])
trailers[trailerName] = append(trailers[trailerName], vs...)
delete(hdr, k)
}
if w.announcedTrailers == nil {
return nil // trailer-only response, don't send data frame.
}
var buf bytes.Buffer
if err := trailers.Write(&buf); err != nil {
return err // should not happen, only errors if (*bytes.Buffer).Write errors.
}
trailerFrameHeader := []byte{trailerMessageFlag, 0, 0, 0, 0}
binary.BigEndian.PutUint32(trailerFrameHeader[1:], uint32(buf.Len()))
if _, err := w.w.Write(trailerFrameHeader); err != nil {
return err
}
if _, err := w.w.Write(buf.Bytes()); err != nil {
return err
}
return nil
}