-
Notifications
You must be signed in to change notification settings - Fork 22
/
handler_go119.go
67 lines (56 loc) · 1.78 KB
/
handler_go119.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
//go:build !go1.20
package listenerutil
import (
"bufio"
"net"
"net/http"
)
type ResponseWriter struct {
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
}
// Implement http.Flusher if available.
func (w *ResponseWriter) Flush() {
f, ok := w.wrapped.(http.Flusher)
if ok {
f.Flush()
}
}
// Implement http.Hijacker if available.
func (w *ResponseWriter) Hijack() (net.Conn, *bufio.ReadWriter, error) {
h, ok := w.wrapped.(http.Hijacker)
if ok {
return h.Hijack()
}
return nil, nil, http.ErrNotSupported
}
// 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{
wrapped: w,
headers: headers,
}
h.ServeHTTP(wrappedWriter, req)
if !wrappedWriter.headerWritten {
wrappedWriter.WriteHeader(http.StatusOK)
}
})
}