| 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) |
| } |
| } |
| |
| } |