app/matrix/media-repo-proxy/main_test.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)
		}
	}

}