Vectrum
Sign In
219 lines · 7.0 KB · Go
package websearch

import (
	"context"
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"
)

// Every address the jail has to refuse. The model writes these URLs, so this
// list is the difference between a documentation reader and a way to read the
// cloud metadata endpoint out loud.
func TestCheckURLRefusesEverythingOffTheInternet(t *testing.T) {
	cases := []struct{ name, url string }{
		{"file scheme", "file:///etc/passwd"},
		{"ftp scheme", "ftp://example.com/x"},
		{"no scheme", "example.com/x"},
		{"bare path", "/etc/passwd"},
		{"loopback name", "http://localhost:8080/"},
		{"loopback v4", "http://127.0.0.1/"},
		{"loopback v6", "http://[::1]/"},
		{"private 10", "http://10.0.0.1/"},
		{"private 192", "http://192.168.1.1/"},
		{"private 172", "http://172.16.0.1/"},
		{"link local", "http://169.254.169.254/latest/meta-data/"},
		{"unique local v6", "http://[fc00::1]/"},
		{"unspecified", "http://0.0.0.0/"},
		{"no host", "http:///x"},
	}
	for _, tc := range cases {
		t.Run(tc.name, func(t *testing.T) {
			if _, err := CheckURL(context.Background(), tc.url); err == nil {
				t.Errorf("%s was allowed", tc.url)
			}
		})
	}
}

// The refusal has to name the address: the model reads this error, and one
// that does not say what was wrong invites it to try the same URL again.
func TestCheckURLSaysWhatItRefused(t *testing.T) {
	_, err := CheckURL(context.Background(), "http://169.254.169.254/")
	if err == nil {
		t.Fatal("the metadata endpoint was allowed")
	}
	if !strings.Contains(err.Error(), "169.254.169.254") {
		t.Errorf("the error does not name the address: %v", err)
	}
}

// A public URL has to survive the check intact, query string and all.
func TestCheckURLAllowsAPublicURL(t *testing.T) {
	u, err := CheckURL(context.Background(), "https://pkg.go.dev/search?q=json")
	if err != nil {
		t.Skipf("no DNS in this environment: %v", err)
	}
	if u.RawQuery != "q=json" {
		t.Errorf("the query string was lost: %q", u.RawQuery)
	}
}

// httptest servers listen on 127.0.0.1, which the jail exists to refuse — so
// every test here dials one through a client with the jail's checks removed,
// and the jail itself is tested above on addresses rather than on connections.
// The one thing that still has to be proven against a live server is the
// redirect hop, which is checked in the client rather than in CheckURL.
func TestRedirectsOffTheInternetAreRefused(t *testing.T) {
	private := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Write([]byte("secrets"))
	}))
	defer private.Close()

	var redirected bool
	public := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		redirected = true
		http.Redirect(w, r, private.URL, http.StatusFound)
	}))
	defer public.Close()

	client := SafeClient(FetchTimeout)
	// The transport's dial check would refuse the first hop too — this test is
	// about CheckRedirect, so the first hop is made with the default transport
	// and only the redirect policy is the jail's.
	client.Transport = http.DefaultTransport

	resp, err := client.Get(public.URL)
	if err == nil {
		resp.Body.Close()
		t.Fatal("a redirect to a loopback address was followed")
	}
	if !redirected {
		t.Fatal("the first hop never happened, so the redirect was not what was refused")
	}
	if !strings.Contains(err.Error(), "loopback") {
		t.Errorf("the error does not explain the refusal: %v", err)
	}
}

// The dial check is the backstop for a name that resolves differently between
// the check and the connection. It is the reason a URL that somehow got past
// CheckURL still cannot reach this machine.
func TestTheDialerRefusesPrivateAddresses(t *testing.T) {
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Write([]byte("secrets"))
	}))
	defer srv.Close()

	resp, err := SafeClient(FetchTimeout).Get(srv.URL)
	if err == nil {
		resp.Body.Close()
		t.Fatal("the dialer connected to a loopback address")
	}
	if !strings.Contains(err.Error(), "refusing to connect") {
		t.Errorf("unexpected error: %v", err)
	}
}

// fetchLocal runs Fetch's own logic against a test server, which lives on
// loopback. The jail is bypassed deliberately and only here: what is under
// test is the reading and reducing, and the refusing is tested above.
func fetchLocal(t *testing.T, srv *httptest.Server, maxBytes int) (*Page, error) {
	t.Helper()
	old := http.DefaultClient
	defer func() { http.DefaultClient = old }()
	return fetchWith(context.Background(), http.DefaultClient, srv.URL, maxBytes)
}

func TestFetchReadsWhatItCanAndRefusesWhatItCannot(t *testing.T) {
	cases := []struct {
		name        string
		ctype, body string
		status      int
		wantErr     string
		wantText    string
	}{{
		name:  "html is reduced",
		ctype: "text/html; charset=utf-8",
		// No furniture here: what a short page does with its <nav> is
		// text.go's business and is tested there. This case is about the
		// content type deciding whether the reducer runs at all.
		body:     "<html><head><style>p{color:red}</style></head><body><h1>Title</h1><p>Hello &amp; welcome</p><script>alert(1)</script></body></html>",
		wantText: "Title\nHello & welcome",
	}, {
		name:     "plain text passes through",
		ctype:    "text/plain",
		body:     "just text",
		wantText: "just text",
	}, {
		name:     "json passes through",
		ctype:    "application/json",
		body:     `{"a":1}`,
		wantText: `{"a":1}`,
	}, {
		name:    "binary is refused",
		ctype:   "application/pdf",
		body:    "%PDF-1.7",
		wantErr: "cannot be read as text",
	}, {
		name:    "a 404 is an error",
		ctype:   "text/html",
		body:    "gone",
		status:  http.StatusNotFound,
		wantErr: "404",
	}}

	for _, tc := range cases {
		t.Run(tc.name, func(t *testing.T) {
			srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
				w.Header().Set("Content-Type", tc.ctype)
				if tc.status != 0 {
					w.WriteHeader(tc.status)
				}
				w.Write([]byte(tc.body))
			}))
			defer srv.Close()

			page, err := fetchLocal(t, srv, 0)
			if tc.wantErr != "" {
				if err == nil {
					t.Fatalf("no error, got page %+v", page)
				}
				if !strings.Contains(err.Error(), tc.wantErr) {
					t.Errorf("error %v does not mention %q", err, tc.wantErr)
				}
				return
			}
			if err != nil {
				t.Fatal(err)
			}
			if page.Text != tc.wantText {
				t.Errorf("text = %q, want %q", page.Text, tc.wantText)
			}
		})
	}
}

// A page bigger than the budget must stop being read, and must say that it
// stopped — text that ends mid-sentence without a word about why is text the
// model will treat as the whole document.
func TestFetchStopsAtTheByteCeiling(t *testing.T) {
	big := strings.Repeat("word ", 10000)
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Type", "text/plain")
		w.Write([]byte(big))
	}))
	defer srv.Close()

	page, err := fetchLocal(t, srv, 100)
	if err != nil {
		t.Fatal(err)
	}
	if !page.Truncated {
		t.Error("an over-long page did not report being truncated")
	}
	if page.Bytes != 100 {
		t.Errorf("read %d bytes, want the 100 asked for", page.Bytes)
	}
}