eleven / ocho Public

Download this file

1	package bot
2
3	import (
4	"net/http"
5	"net/http/httptest"
6	"strings"
7	"testing"
8	)
9
10	func TestCheckFetchURL(t *testing.T) {
11	allow := []string{"https://www.rte.ie/feeds/", "https://rss.nytimes.com/"}
12	if err := checkFetchURL("https://www.rte.ie/feeds/rss/news.xml", allow); err != nil {
13	t.Errorf("allowlisted url rejected: %v", err)
14	}
15	if err := checkFetchURL("https://evil.example.com/feeds/", allow); err == nil {
16	t.Error("off-list url accepted")
17	}
18	if err := checkFetchURL("http://www.rte.ie/feeds/rss/news.xml", allow); err == nil {
19	t.Error("plain http accepted outside test mode")
20	}
21	// Prefix means prefix: a lookalike host that merely CONTAINS the prefix
22	// elsewhere must fail.
23	if err := checkFetchURL("https://phish.example/https://www.rte.ie/feeds/", allow); err == nil {
24	t.Error("embedded-prefix url accepted")
25	}
26	}
27
28	// TestCheckFetchURLWWW pins the apex↔www tolerance: a homepage-fallback allow
29	// entry survives the site's own redirect to (or from) www — the live failure
30	// that broke "switch to the Irish Times" (irishtimes.com 301s to www).
31	func TestCheckFetchURLWWW(t *testing.T) {
32	cases := []struct {
33	raw string
34	allow string
35	ok bool
36	}{
37	{"https://www.irishtimes.com/", "https://irishtimes.com", true},
38	{"https://irishtimes.com/news", "https://www.irishtimes.com", true},
39	{"https://www.irishtimes.com/news?page=2", "https://irishtimes.com", true},
40	// www tolerance is not subdomain tolerance.
41	{"https://api.irishtimes.com/", "https://irishtimes.com", false},
42	{"https://wwwx.irishtimes.com/", "https://irishtimes.com", false},
43	}
44	for _, c := range cases {
45	err := checkFetchURL(c.raw, []string{c.allow})
46	if (err == nil) != c.ok {
47	t.Errorf("checkFetchURL(%q, allow %q) = %v, want ok=%v", c.raw, c.allow, err, c.ok)
48	}
49	}
50	}
51
52	// TestAllowPrefixMatchBoundary pins the boundary rule: a bare-host entry must
53	// not match a LONGER host that merely starts with the same bytes — the model
54	// picks fetch URLs, so this edge is load-bearing.
55	func TestAllowPrefixMatchBoundary(t *testing.T) {
56	cases := []struct {
57	raw string
58	entry string
59	ok bool
60	}{
61	{"https://irishtimes.com", "https://irishtimes.com", true},
62	{"https://irishtimes.com/", "https://irishtimes.com", true},
63	{"https://irishtimes.com/news", "https://irishtimes.com", true},
64	{"https://irishtimes.com?utm=x", "https://irishtimes.com", true},
65	{"https://irishtimes.com.evil.com/", "https://irishtimes.com", false},
66	{"https://irishtimes.community/", "https://irishtimes.com", false},
67	{"https://irishtimes.com:8443/", "https://irishtimes.com", false},
68	// Entries ending in a slash or already inside a query keep plain
69	// prefix semantics.
70	{"https://www.rte.ie/feeds/rss/news.xml", "https://www.rte.ie/feeds/", true},
71	{"https://a.com/feed?index=/news/&limit=100", "https://a.com/feed?index=/news/", true},
72	{"https://a.com/feedsecret", "https://a.com/feeds", false},
73	}
74	for _, c := range cases {
75	if got := allowPrefixMatch(c.raw, c.entry); got != c.ok {
76	t.Errorf("allowPrefixMatch(%q, %q) = %v, want %v", c.raw, c.entry, got, c.ok)
77	}
78	}
79	}
80
81	// TestFetchBlocksPrivateAddresses pins the SSRF guard: without the test-only
82	// escape hatch, a loopback URL — even one on the allowlist — must not connect.
83	func TestFetchBlocksPrivateAddresses(t *testing.T) {
84	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
85	w.Write([]byte("secret"))
86	}))
87	defer srv.Close()
88
89	// The dialer's Control hook must reject the loopback address (scheme check
90	// is bypassed by allowlisting the URL itself — the layers are independent).
91	if _, err := fetchURL(srv.URL+"/feed", []string{srv.URL}); err == nil {
92	t.Fatal("loopback fetch succeeded without LCHAT_BOT_ALLOW_PRIVATE_FETCH")
93	}
94
95	t.Setenv("LCHAT_BOT_ALLOW_PRIVATE_FETCH", "1")
96	body, err := fetchURL(srv.URL+"/feed", []string{srv.URL})
97	if err != nil {
98	t.Fatalf("test-mode loopback fetch: %v", err)
99	}
100	if body != "secret" {
101	t.Fatalf("got %q", body)
102	}
103	}
104
105	func TestValidateTaskSpec(t *testing.T) {
106	ok := &TaskSpec{Prompt: "summarise", Tools: []string{"fetch"},
107	Fetch: &FetchSpec{Allow: []string{"https://rss.nytimes.com/"}}}
108	if err := validateTaskSpec(0, ok); err != nil {
109	t.Errorf("valid spec rejected: %v", err)
110	}
111
112	bad := []struct {
113	name string
114	spec *TaskSpec
115	want string
116	}{
117	{"nil", nil, "needs a prompt"},
118	{"no prompt", &TaskSpec{}, "needs a prompt"},
119	{"fetch without allowlist", &TaskSpec{Prompt: "x", Tools: []string{"fetch"}}, "non-empty allowlist"},
120	{"unknown tool", &TaskSpec{Prompt: "x", Tools: []string{"shell"}}, "unknown task tool"},
121	{"http allow entry", &TaskSpec{Prompt: "x", Tools: []string{"fetch"},
122	Fetch: &FetchSpec{Allow: []string{"http://a.example/"}}}, "https:// prefixes"},
123	{"fetch config, no tool", &TaskSpec{Prompt: "x", Fetch: &FetchSpec{Allow: []string{"https://a.example/"}}}, "without the fetch tool"},
124	{"notice too long", &TaskSpec{Prompt: "x", Post: &PostSpec{Notice: &NoticeSpec{Title: strings.Repeat("t", 200)}}}, "notice too long"},
125	{"budget out of range", &TaskSpec{Prompt: "x", MaxToolCalls: 99}, "out of range"},
126	}
127	for _, tc := range bad {
128	err := validateTaskSpec(0, tc.spec)
129	if err == nil || !strings.Contains(err.Error(), tc.want) {
130	t.Errorf("%s: want error containing %q, got %v", tc.name, tc.want, err)
131	}
132	}
133	}
134