feat: HTTP API, WebSocket streaming with token auth and origin check, server entrypoint

internal/server/server_test.go

@@ -0,0 +1,156 @@
+package server
+
+import (
+	"context"
+	"encoding/json"
+	"io"
+	"log/slog"
+	"net/http"
+	"net/http/httptest"
+	"path/filepath"
+	"strings"
+	"testing"
+	"time"
+
+	"github.com/coder/websocket"
+	"github.com/coder/websocket/wsjson"
+
+	"github.com/crypto/openclaw/internal/agent"
+	"github.com/crypto/openclaw/internal/llm"
+	"github.com/crypto/openclaw/internal/store"
+	"github.com/crypto/openclaw/internal/tools"
+)
+
+const tok = "test-token-0123456789"
+
+type echo struct{}
+
+func (echo) Chat(_ context.Context, req llm.Request, d llm.DeltaFunc) (llm.Response, error) {
+	last := req.Messages[len(req.Messages)-1].Content
+	for _, w := range strings.Fields("echo: " + last) {
+		d(w + " ")
+	}
+	return llm.Response{Content: "echo: " + last}, nil
+}
+
+func setup(t *testing.T, origins ...string) *httptest.Server {
+	st, err := store.Open(filepath.Join(t.TempDir(), "s.db"))
+	if err != nil {
+		t.Fatal(err)
+	}
+	t.Cleanup(func() { st.Close() })
+	s := &Server{Store: st, Token: tok, AllowedOrigins: origins, Log: slog.New(slog.NewTextHandler(io.Discard, nil)), RunTimeout: 5 * time.Second,
+		Agent: &agent.Agent{LLM: echo{}, Store: st, Tools: tools.NewRegistry(), MaxIterations: 3}}
+	ts := httptest.NewServer(s.Handler())
+	t.Cleanup(ts.Close)
+	return ts
+}
+
+func do(t *testing.T, method, url, body string, hdr map[string]string) *http.Response {
+	req, _ := http.NewRequest(method, url, strings.NewReader(body))
+	for k, v := range hdr {
+		req.Header.Set(k, v)
+	}
+	resp, err := http.DefaultClient.Do(req)
+	if err != nil {
+		t.Fatal(err)
+	}
+	t.Cleanup(func() { resp.Body.Close() })
+	return resp
+}
+
+var authH = map[string]string{"Authorization": "Bearer " + tok}
+
+func TestAuthAndOrigin(t *testing.T) {
+	ts := setup(t, "https://app.example")
+	if r := do(t, "GET", ts.URL+"/healthz", "", nil); r.StatusCode != 200 {
+		t.Fatal("healthz", r.StatusCode)
+	}
+	if r := do(t, "GET", ts.URL+"/api/sessions", "", nil); r.StatusCode != 401 {
+		t.Fatal("no token", r.StatusCode)
+	}
+	if r := do(t, "GET", ts.URL+"/api/sessions", "", map[string]string{"Authorization": "Bearer wrong"}); r.StatusCode != 401 {
+		t.Fatal("bad token", r.StatusCode)
+	}
+	evil := map[string]string{"Authorization": "Bearer " + tok, "Origin": "https://evil.example"}
+	if r := do(t, "GET", ts.URL+"/api/sessions", "", evil); r.StatusCode != 403 {
+		t.Fatal("evil origin", r.StatusCode)
+	}
+	ok := map[string]string{"Authorization": "Bearer " + tok, "Origin": "https://app.example"}
+	if r := do(t, "GET", ts.URL+"/api/sessions", "", ok); r.StatusCode != 200 {
+		t.Fatal("allowed origin", r.StatusCode)
+	}
+	same := map[string]string{"Authorization": "Bearer " + tok, "Origin": ts.URL}
+	if r := do(t, "GET", ts.URL+"/api/sessions", "", same); r.StatusCode != 200 {
+		t.Fatal("same origin", r.StatusCode)
+	}
+}
+
+func TestChatAndSessions(t *testing.T) {
+	ts := setup(t)
+	r := do(t, "POST", ts.URL+"/api/chat", `{"text":"hello world"}`, authH)
+	var out map[string]string
+	json.NewDecoder(r.Body).Decode(&out)
+	if r.StatusCode != 200 || out["reply"] != "echo: hello world" || out["session_id"] == "" {
+		t.Fatalf("%d %v", r.StatusCode, out)
+	}
+	r = do(t, "GET", ts.URL+"/api/sessions", "", authH)
+	var list []store.Session
+	json.NewDecoder(r.Body).Decode(&list)
+	if len(list) != 1 || list[0].Title != "hello world" {
+		t.Fatalf("%+v", list)
+	}
+	r = do(t, "GET", ts.URL+"/api/sessions/"+out["session_id"]+"/messages", "", authH)
+	var msgs []store.StoredMessage
+	json.NewDecoder(r.Body).Decode(&msgs)
+	if len(msgs) != 2 {
+		t.Fatalf("msgs %+v", msgs)
+	}
+	if r := do(t, "DELETE", ts.URL+"/api/sessions/"+out["session_id"], "", authH); r.StatusCode != 204 {
+		t.Fatal(r.StatusCode)
+	}
+	if r := do(t, "GET", ts.URL+"/api/sessions/nope/messages", "", authH); r.StatusCode != 404 {
+		t.Fatal(r.StatusCode)
+	}
+}
+
+func TestWebSocketStreaming(t *testing.T) {
+	ts := setup(t)
+	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+	defer cancel()
+	wsURL := "ws" + strings.TrimPrefix(ts.URL, "http") + "/api/ws"
+	if _, _, err := websocket.Dial(ctx, wsURL, nil); err == nil {
+		t.Fatal("ws without token must fail")
+	}
+	if _, _, err := websocket.Dial(ctx, wsURL, &websocket.DialOptions{HTTPHeader: http.Header{"Authorization": {"Bearer " + tok}, "Origin": {"https://evil.example"}}}); err == nil {
+		t.Fatal("ws from foreign origin must fail")
+	}
+	c, _, err := websocket.Dial(ctx, wsURL, &websocket.DialOptions{HTTPHeader: http.Header{"Authorization": {"Bearer " + tok}}})
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer c.CloseNow()
+	wsjson.Write(ctx, c, map[string]string{"type": "message", "text": "hi there"})
+	var deltas strings.Builder
+	var sid string
+	for {
+		var e agent.Event
+		if err := wsjson.Read(ctx, c, &e); err != nil {
+			t.Fatal(err)
+		}
+		switch e.Type {
+		case "session":
+			sid = e.Session
+		case "delta":
+			deltas.WriteString(e.Text)
+		case "error":
+			t.Fatal(e.Text)
+		}
+		if e.Type == "done" {
+			if e.Text != "echo: hi there" || sid == "" || e.Session != sid || strings.TrimSpace(deltas.String()) != "echo: hi there" {
+				t.Fatalf("done %+v deltas %q", e, deltas.String())
+			}
+			return
+		}
+	}
+}