91 lines
2.0 KiB
Go
91 lines
2.0 KiB
Go
package main
|
|
|
|
import (
|
|
"fmt"
|
|
"net/http"
|
|
"net/http/httptest"
|
|
"testing"
|
|
)
|
|
|
|
func init() {
|
|
allowedHosts = []string{"teapot-dummy-target.example.com"}
|
|
}
|
|
|
|
func TestHttpReverseProxy(t *testing.T) {
|
|
var header map[string][]string
|
|
|
|
teapot := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
header = r.Header.Clone()
|
|
w.WriteHeader(http.StatusTeapot)
|
|
}))
|
|
defer teapot.Close()
|
|
|
|
targetHost = teapot.Listener.Addr().String()
|
|
targetScheme = "http"
|
|
reverseProxy := newHttpReverseProxy(":8081", false)
|
|
|
|
fmt.Println("Proxy server listening on " + reverseProxy.Addr)
|
|
fmt.Println("Dummy target server listening on " + targetHost)
|
|
|
|
var tests = []struct {
|
|
host string
|
|
wantXForwardedHostHeader string
|
|
wantResponseCode int
|
|
}{
|
|
{
|
|
"teapot-dummy-target.example.com",
|
|
"teapot-dummy-target.example.com",
|
|
http.StatusTeapot,
|
|
},
|
|
{
|
|
myPtrAddr,
|
|
"",
|
|
http.StatusOK,
|
|
},
|
|
{
|
|
reverseProxy.Addr,
|
|
"",
|
|
http.StatusGatewayTimeout,
|
|
},
|
|
}
|
|
|
|
for _, test := range tests {
|
|
t.Run(fmt.Sprintf("requesting host %q returns status %d and header X-Forwarded-Host set to %q",
|
|
test.host,
|
|
test.wantResponseCode,
|
|
test.wantXForwardedHostHeader),
|
|
func(t *testing.T) {
|
|
request, _ := http.NewRequest("GET", "/", nil)
|
|
request.Host = test.host
|
|
response := httptest.NewRecorder()
|
|
reverseProxy.Handler.ServeHTTP(response, request)
|
|
|
|
gotHeader := header
|
|
header = nil
|
|
|
|
assertStatus(t, response.Code, test.wantResponseCode)
|
|
if len(test.wantXForwardedHostHeader) > 0 {
|
|
assertHeader(t, gotHeader, "X-Forwarded-Host", test.wantXForwardedHostHeader)
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
func assertHeader(t *testing.T, gotHeader map[string][]string, headerName, want string) {
|
|
t.Helper()
|
|
var got string
|
|
if lookup, ok := gotHeader[headerName]; ok {
|
|
got = lookup[0]
|
|
}
|
|
if got != want {
|
|
t.Errorf("got %s, want %s", got, want)
|
|
}
|
|
}
|
|
|
|
func assertStatus(t *testing.T, got, want int) {
|
|
t.Helper()
|
|
if got != want {
|
|
t.Errorf("got %d, want %d", got, want)
|
|
}
|
|
}
|