vproxy/simpleproxy.go at master · jittering/vproxy · GitHub

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
package vproxy

import (
	"bytes"
	"fmt"
	"io/ioutil"
	"log"
	"net/http"
	"net/http/httputil"
	"net/url"
	"time"

	"github.com/cenkalti/backoff/v4"
)

var badGatewayMessage = `<html>
<body>
<h1>503 Service Unavailable</h1>
<p>Can't connect to upstream server (%s &mdash;&gt; %s), please try again later.</p>
</body>
</html>`

// proxyTransport is a simple http.RoundTripper implementation which returns a
// 503 on any error making a request to the upstream (backend) service
type proxyTransport struct {
	transport *http.Transport
	errMsg    string
}

func (t *proxyTransport) RoundTrip(request *http.Request) (*http.Response, error) {

	var (
		response *http.Response
		err      error
	)

	operation := func() error {
		response, err = t.transport.RoundTrip(request)
		// Handle Retry-After here, if you wish...
		// If err is nil, no retry will occur
		return err
	}

	be := backoff.NewExponentialBackOff()
	be.MaxElapsedTime = time.Second * 30
	err = backoff.Retry(operation, be)
	if err != nil {
		resp := &http.Response{
			StatusCode: http.StatusServiceUnavailable,
			Body:       ioutil.NopCloser(bytes.NewBufferString(t.errMsg)),
		}

		resp.StatusCode = 503
		resp.Status = "Can't connect to upstream server"
		log.Println("proxy: error fetching from upstream:", err)
		return resp, nil
	}

	return response, err
}

func createProxyTransport(targetURL url.URL, vhost string) *proxyTransport {
	t := &proxyTransport{errMsg: fmt.Sprintf(badGatewayMessage, targetURL.String(), vhost)}
	t.transport = http.DefaultTransport.(*http.Transport).Clone()
	t.transport.MaxConnsPerHost = 0 // unlim
	t.transport.MaxIdleConns = 800
	t.transport.MaxIdleConnsPerHost = 8
	return t
}

// CreateProxy with custom http.RoundTripper impl. Sets proper host headers
// using given vhost name.
func CreateProxy(targetURL url.URL, vhost string) *httputil.ReverseProxy {
	return &httputil.ReverseProxy{
		Director: func(r *http.Request) {
			p, q := r.URL.Path, r.URL.RawQuery
			*r.URL = targetURL
			r.URL.Path, r.URL.RawQuery = p, q
			if vhost != "" {
				r.Host = vhost
				r.Header.Add("X-Forwarded-Host", vhost)
			} else {
				r.Host = targetURL.Host
			}
			scheme := "http"
			if r.TLS != nil {
				scheme = "https"
			}
			r.Header.Add("X-Forwarded-Proto", scheme)
		},
		Transport: createProxyTransport(targetURL, vhost),
	}
}