forked from hswaw/hscloud
84 lines
1.9 KiB
Go
84 lines
1.9 KiB
Go
package main
|
|
|
|
import (
|
|
"fmt"
|
|
"io"
|
|
"net/http"
|
|
"net/http/httptest"
|
|
"net/url"
|
|
"testing"
|
|
)
|
|
|
|
func TestForward(t *testing.T) {
|
|
// Test backend which proudly proclaims the value of the X-Forwarded-For header it received.
|
|
backendServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
fmt.Fprintf(w, "hello %s %s\n", r.Host, r.Header.Get("X-Forwarded-For"))
|
|
}))
|
|
defer backendServer.Close()
|
|
rpURL, err := url.Parse(backendServer.URL)
|
|
if err != nil {
|
|
t.Fatalf("parsing test backend URL failed: %v", err)
|
|
}
|
|
|
|
// Configure and run proxy.
|
|
flagUpstream = rpURL.Host
|
|
flagUpstreamHost = "example.com"
|
|
flagDownstreamHost = "matrix.example.com"
|
|
proxy := httptest.NewServer(newProxy())
|
|
defer proxy.Close()
|
|
|
|
// Run through a few tests.
|
|
for i, te := range []struct {
|
|
headers map[string]string
|
|
host string
|
|
want string
|
|
}{
|
|
{
|
|
// 0: expected to succeed
|
|
headers: map[string]string{
|
|
"Hscloud-Nic-Source-IP": "1.2.3.4",
|
|
"Hscloud-Nic-Source-Port": "1337",
|
|
},
|
|
host: "matrix.example.com",
|
|
want: "hello example.com 1.2.3.4:1337, 127.0.0.1\n",
|
|
},
|
|
{
|
|
// 1: expected to succeed
|
|
host: "matrix.example.com",
|
|
want: "hello example.com 127.0.0.1\n",
|
|
},
|
|
{
|
|
// 2: expected to succeed
|
|
host: "matrix.example.com:443",
|
|
want: "hello example.com 127.0.0.1\n",
|
|
},
|
|
{
|
|
// 3: expected to fail
|
|
host: "example.com",
|
|
want: "invalid host\n",
|
|
},
|
|
} {
|
|
req, _ := http.NewRequest("GET", proxy.URL, nil)
|
|
req.Host = te.host
|
|
for k, v := range te.headers {
|
|
req.Header.Set(k, v)
|
|
}
|
|
|
|
resp, err := http.DefaultClient.Do(req)
|
|
if err != nil {
|
|
t.Fatalf("Get failed: %v", err)
|
|
}
|
|
|
|
b, err := io.ReadAll(resp.Body)
|
|
if err != nil {
|
|
t.Fatalf("Read failed: %v", err)
|
|
}
|
|
resp.Body.Close()
|
|
|
|
if want, got := te.want, string(b); want != got {
|
|
t.Errorf("%d: wrong response from upstream, wanted %q, got %q", i, want, got)
|
|
}
|
|
}
|
|
|
|
}
|