-
Notifications
You must be signed in to change notification settings - Fork 22
/
handler_go120.go
61 lines (50 loc) · 1.88 KB
/
handler_go120.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
// Copyright (c) HashiCorp, Inc.
// SPDX-License-Identifier: MPL-2.0
//go:build go1.20
package listenerutil
import (
"net/http"
)
type ResponseWriter struct {
// Embed ResponseController so we automatically implement
// http.Hijacker, SetReadDeadline and SetWriteDeadline.
*http.ResponseController
wrapped http.ResponseWriter
// headers contain a map of response code to header map such that
// headers[status][header name] = header value
// this map also contains values for hundred-level values in the format 1: "1xx", 2: "2xx", etc
// defaults are set to 0
headers map[int]http.Header
headerWritten bool
}
// We need to wrap the ResponseController Flush method to implement http.Flusher,
// since it doesn't normally return an error.
func (w *ResponseWriter) Flush() {
_ = w.ResponseController.Flush()
}
// WrapCustomHeadersHandler wraps the handler to pass a custom ResponseWriter struct to all
// later wrappers and handlers to assign custom headers by status code. This wrapper must
// be the outermost wrapper to function correctly.
func WrapCustomHeadersHandler(h http.Handler, config *ListenerConfig, isUiRequest uiRequestFunc) http.Handler {
uiHeaders := config.CustomUiResponseHeaders
apiHeaders := config.CustomApiResponseHeaders
return http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
// this function is extremely generic as all we want to do is wrap the http.ResponseWriter
// in our own ResponseWriter above, which will then perform all the logic we actually want
var headers map[int]http.Header
if isUiRequest(req) {
headers = uiHeaders
} else {
headers = apiHeaders
}
wrappedWriter := &ResponseWriter{
ResponseController: http.NewResponseController(w),
wrapped: w,
headers: headers,
}
h.ServeHTTP(wrappedWriter, req)
if !wrappedWriter.headerWritten {
wrappedWriter.WriteHeader(http.StatusOK)
}
})
}