/
proxy.go
120 lines (104 loc) · 2.66 KB
/
proxy.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
/*
Copyright (c) 2022 - Present. Blend Labs, Inc. All rights reserved
Use of this source code is governed by a MIT license that can be found in the LICENSE file.
*/
package reverseproxy
import (
"fmt"
"net/http"
"os"
"time"
"github.com/blend/go-sdk/ex"
"github.com/blend/go-sdk/logger"
"github.com/blend/go-sdk/webutil"
)
const (
// FlagProxyRequest is a logger flag.
FlagProxyRequest = "proxy.request"
)
// NewProxy returns a new proxy.
func NewProxy(opts ...ProxyOption) (*Proxy, error) {
var err error
p := Proxy{
Headers: http.Header{},
}
for _, opt := range opts {
if err = opt(&p); err != nil {
return nil, err
}
}
return &p, nil
}
// Proxy is a factory for a simple reverse proxy.
type Proxy struct {
Headers http.Header
Log logger.Log
Upstreams []*Upstream
Resolver Resolver
Tracer webutil.HTTPTracer
TransformRequest TransformRequest
Timeout time.Duration
}
// ServeHTTP is the http entrypoint.
func (p *Proxy) ServeHTTP(rw http.ResponseWriter, req *http.Request) {
var err error
var tf webutil.HTTPTraceFinisher
srw := webutil.NewStatusResponseWriter(rw)
defer func() {
// NOTE: This uses the outer scope's `err` by design. This way updates
// to `err` will be reflected on (deferred) exit.
r := recover()
// see: https://golang.org/pkg/net/http/#ErrAbortHandler
if r != nil && r != http.ErrAbortHandler {
// Wrap the error with the reason for the panic.
err = ex.Nest(err, ex.New(r))
}
if tf != nil {
tf.Finish(srw.StatusCode(), err)
}
if err != nil {
if p.Log != nil {
p.Log.Fatalf("%v", r)
} else {
fmt.Fprintf(os.Stderr, "%v\n", r)
}
}
}()
if p.Tracer != nil {
tf, req = p.Tracer.Start(req)
}
// set the default resolver if unset.
if p.Resolver == nil {
p.Resolver = RoundRobinResolver(p.Upstreams)
}
upstream, err := p.Resolver(req, p.Upstreams)
if err != nil {
logger.MaybeError(p.Log, err)
srw.WriteHeader(http.StatusBadGateway)
return
}
if upstream == nil {
srw.WriteHeader(http.StatusBadGateway)
return
}
// Add extra forwarded headers.
// these are required for a majority of services to function correctly behind
// a reverse proxy.
// They are "Add" vs. "Set" in case there are existing values.
if port := webutil.GetPort(req); port != "" {
req.Header.Add("X-Forwarded-Port", port)
}
if proto := webutil.GetProto(req); proto != "" {
req.Header.Add("X-Forwarded-Proto", proto)
}
// add upstream headers.
for key, values := range p.Headers {
for _, value := range values {
req.Header.Add(key, value)
}
}
if p.TransformRequest != nil {
p.TransformRequest(req)
}
upstream.ServeHTTP(srw, req)
}