-
Notifications
You must be signed in to change notification settings - Fork 0
/
main.go
105 lines (96 loc) · 2.36 KB
/
main.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
package main
import (
"fmt"
"github.com/labstack/echo/v4"
"net/http"
"net/http/httptest"
"net/http/httputil"
"os"
"path/filepath"
"strconv"
"time"
"mime"
)
func main() {
// Config
host := os.Getenv("HOST")
if host == "" {
fmt.Printf("missing host")
}
timeout := os.Getenv("TIMEOUT")
scheme := os.Getenv("SCHEME")
if scheme == "" {
scheme = "http"
}
proxyTimeout, err := strconv.Atoi(timeout)
if err != nil {
fmt.Printf("invalid timeout: %v\n", err)
return
}
// Server
e := echo.New()
e.Any("/*", timeoutFixerFor(host, scheme, proxyTimeout))
e.Start(":8080")
}
func timeoutFixerFor(host, scheme string, timeout int) echo.HandlerFunc {
return func(c echo.Context) error {
director := func(req *http.Request) {
r := c.Request()
req = r
req.URL.Scheme = scheme
req.URL.Host = r.URL.Host
}
req := c.Request()
req.URL.Host = host
req.Host = host
req.URL.Scheme = scheme
proxy := &httputil.ReverseProxy{Director: director}
recorder := httptest.NewRecorder()
done := make(chan bool)
// Replace accept-encoding header
req.Header.Set("Accept-Encoding", "identity")
// This might take up to an hour, so spawn a go routine
go func() {
proxy.ServeHTTP(recorder, req)
done <- true
}()
// Set correct mime-type
ext := filepath.Ext(req.URL.Path)
mimeType := mime.TypeByExtension(ext)
if mimeType == "" {
mimeType = "text/html" // Default
}
c.Response().Header().Set("Content-Type", mimeType)
c.Response().Header().Set("X-Proxy-Pass", "timeout-tricker")
// Write headers
// We may have to use some heuristics based on the request
// to send the correct headers
headersSent := false
for {
select {
case <-done:
// Upstream request is done!
// Write out the original body
fmt.Printf("remote request finished\n")
// and header if we did not sent anything yet!
if !headersSent {
writer := c.Response().Writer
for k, v := range recorder.Header() {
writer.Header()[k] = v
}
c.Response().WriteHeader(recorder.Result().StatusCode)
}
_, err := c.Response().Writer.Write(recorder.Body.Bytes())
return err
case <-time.After(time.Duration(timeout) * time.Second):
writer := c.Response().Writer
_, _ = writer.Write([]byte(" "))
if f, ok := writer.(http.Flusher); ok {
fmt.Printf("flushing\n")
f.Flush()
}
headersSent = true
}
}
}
}