eleven / ocho Public

Download this file

1	package bot
2
3	import (
4	"crypto/rand"
5	"encoding/json"
6	"net/http"
7	"net/http/httptest"
8	"os"
9	"path/filepath"
10	"strings"
11	"sync/atomic"
12	"testing"
13
14	"lchat/internal/e2ee"
15	)
16
17	// fakeInviteServer implements the three endpoints Provision walks, the same
18	// wire shape agent/provision.mjs pins. It hands out one invite (the given
19	// fragment) and records how many claims it saw.
20	func fakeInviteServer(t *testing.T, fragment string, claims *atomic.Int32) *httptest.Server {
21	t.Helper()
22	inv, err := e2ee.DeriveInvite(fragment)
23	if err != nil {
24	t.Fatalf("derive invite: %v", err)
25	}
26	mux := http.NewServeMux()
27	mux.HandleFunc("GET /api/key-invite/"+inv.ID, func(w http.ResponseWriter, r *http.Request) {
28	json.NewEncoder(w).Encode(map[string]any{"ok": true, "wrapped": ""})
29	})
30	mux.HandleFunc("POST /api/key-invite/"+inv.ID+"/claim", func(w http.ResponseWriter, r *http.Request) {
31	var body struct {
32	ClaimSecret string `json:"claim_secret"`
33	}
34	json.NewDecoder(r.Body).Decode(&body)
35	if body.ClaimSecret != inv.ClaimSecret {
36	json.NewEncoder(w).Encode(map[string]any{"ok": false})
37	return
38	}
39	claims.Add(1)
40	json.NewEncoder(w).Encode(map[string]any{"ok": true, "token": "tok-ensure-test"})
41	})
42	mux.HandleFunc("POST /api/pubkey", func(w http.ResponseWriter, r *http.Request) {
43	var body struct {
44	Token string `json:"token"`
45	PublicKey string `json:"public_key"`
46	}
47	json.NewDecoder(r.Body).Decode(&body)
48	if body.Token == "" || body.PublicKey == "" {
49	http.Error(w, "missing", http.StatusBadRequest)
50	return
51	}
52	json.NewEncoder(w).Encode(map[string]any{"ok": true})
53	})
54	return httptest.NewServer(mux)
55	}
56
57	func testFragment(t *testing.T) string {
58	t.Helper()
59	raw := make([]byte, 32)
60	if _, err := rand.Read(raw); err != nil {
61	t.Fatal(err)
62	}
63	return e2ee.B64uEncode(raw)
64	}
65
66	// The #806 shape: session file absent + invite present → redeem, write the
67	// session, and hand it back — so a fresh add-bot'd sidecar boots zero-touch.
68	func TestEnsureSessionProvisionsFromInvite(t *testing.T) {
69	fragment := testFragment(t)
70	var claims atomic.Int32
71	srv := fakeInviteServer(t, fragment, &claims)
72	defer srv.Close()
73
74	path := filepath.Join(t.TempDir(), "eleven-session.json")
75	sess, err := EnsureSession(srv.URL, path, fragment, "Eleven", t.Logf)
76	if err != nil {
77	t.Fatalf("EnsureSession: %v", err)
78	}
79	if sess == nil {
80	t.Fatal("EnsureSession returned nil session despite an invite")
81	}
82	if sess.Name != "Eleven" || sess.Token != "tok-ensure-test" || sess.UserPriv == "" {
83	t.Fatalf("unexpected session: name=%q token=%q priv-empty=%v", sess.Name, sess.Token, sess.UserPriv == "")
84	}
85	if claims.Load() != 1 {
86	t.Fatalf("expected exactly one claim, got %d", claims.Load())
87	}
88	if _, err := os.Stat(path); err != nil {
89	t.Fatalf("session file not persisted: %v", err)
90	}
91
92	// Second boot: the persisted session wins; the (now spent) invite is never
93	// re-redeemed. The server is closed to prove no network is touched.
94	srv.Close()
95	again, err := EnsureSession(srv.URL, path, fragment, "Eleven", nil)
96	if err != nil {
97	t.Fatalf("EnsureSession reload: %v", err)
98	}
99	if again == nil || again.Token != "tok-ensure-test" {
100	t.Fatalf("reload did not return the persisted session: %+v", again)
101	}
102	}
103
104	// No session and no invite stays the normal "not provisioned" state — the
105	// caller keeps its clear operator-facing error.
106	func TestEnsureSessionMissingWithoutInvite(t *testing.T) {
107	path := filepath.Join(t.TempDir(), "session.json")
108	sess, err := EnsureSession("http://unused.invalid", path, " ", "Claude", nil)
109	if err != nil {
110	t.Fatalf("EnsureSession: %v", err)
111	}
112	if sess != nil {
113	t.Fatalf("expected nil session, got %+v", sess)
114	}
115	}
116
117	// An existing session short-circuits: the invite is ignored and nothing dials
118	// out (the base URL is unresolvable on purpose).
119	func TestEnsureSessionPrefersExistingFile(t *testing.T) {
120	path := filepath.Join(t.TempDir(), "session.json")
121	existing := &Session{Token: "tok-existing", Name: "Eleven", UserPub: "pub", UserPriv: "priv"}
122	if err := existing.Save(path); err != nil {
123	t.Fatal(err)
124	}
125	sess, err := EnsureSession("http://unused.invalid", path, testFragment(t), "Eleven", nil)
126	if err != nil {
127	t.Fatalf("EnsureSession: %v", err)
128	}
129	if sess == nil || sess.Token != "tok-existing" {
130	t.Fatalf("expected the existing session, got %+v", sess)
131	}
132	}
133
134	// A spent invite (the crash-loop-after-burn case) fails with a provision
135	// error rather than looping silently.
136	func TestEnsureSessionSpentInvite(t *testing.T) {
137	fragment := testFragment(t)
138	inv, err := e2ee.DeriveInvite(fragment)
139	if err != nil {
140	t.Fatal(err)
141	}
142	mux := http.NewServeMux()
143	mux.HandleFunc("GET /api/key-invite/"+inv.ID, func(w http.ResponseWriter, r *http.Request) {
144	json.NewEncoder(w).Encode(map[string]any{"ok": false})
145	})
146	srv := httptest.NewServer(mux)
147	defer srv.Close()
148
149	path := filepath.Join(t.TempDir(), "session.json")
150	_, err = EnsureSession(srv.URL, path, fragment, "Claude", nil)
151	if err == nil || !strings.Contains(err.Error(), "invite is used or invalid") {
152	t.Fatalf("expected spent-invite error, got %v", err)
153	}
154	if _, statErr := os.Stat(path); !os.IsNotExist(statErr) {
155	t.Fatalf("no session file should be written on failure; stat: %v", statErr)
156	}
157	}
158