Agent-System-Backend/internal/server/sse_close_test.go

package server

// SSE 长连接会话复验（登录恢复合同 §6.7）：密码写入删除会话行后，
// 已建立的流在复验周期内终止，不再继续接收私有事件。

import (
	"bufio"
	"context"
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"
	"time"

	"git-ssh.u1680832.nyat.app/weichen/Agent-System-Backend.git/internal/store"
)

// 打开真实流式服务，等一帧事件后作废会话，断言流在复验周期内关闭。
func TestSSESessionInvalidationClosesStream(t *testing.T) {
	old := sseAuthRecheck
	sseAuthRecheck = 300 * time.Millisecond
	t.Cleanup(func() { sseAuthRecheck = old })

	srv, st := newAuthTestServer(t)
	h := srv.Handler()

	hash, err := store.HashPassword("sse-old-pass-1")
	if err != nil {
		t.Fatal(err)
	}
	if err := st.CreateAccount("tia", "", hash, false); err != nil {
		t.Fatal(err)
	}
	if err := st.ClearMustChangePassword("tia"); err != nil {
		t.Fatal(err)
	}
	tok, err := st.CreateSession("tia")
	if err != nil {
		t.Fatal(err)
	}
	proj, err := st.CreateProject("tia", "sse-proj")
	if err != nil {
		t.Fatal(err)
	}
	conv, err := st.CreateConversation("tia", proj.ID, "", "")
	if err != nil {
		t.Fatal(err)
	}

	ts := httptest.NewServer(h)
	t.Cleanup(ts.Close)

	ctx, cancel := context.WithCancel(context.Background())
	t.Cleanup(cancel)
	req, err := http.NewRequestWithContext(ctx, "GET", ts.URL+"/api/events?conversationId="+conv.ID, nil)
	if err != nil {
		t.Fatal(err)
	}
	req.AddCookie(&http.Cookie{Name: SessionCookie, Value: tok})

	type result struct {
		lastLine string
		closed   bool
		err      error
	}
	resCh := make(chan result, 1)
	go func() {
		resp, err := ts.Client().Do(req)
		if err != nil {
			resCh <- result{err: err}
			return
		}
		defer resp.Body.Close()
		sc := bufio.NewScanner(resp.Body)
		last := ""
		for sc.Scan() {
			if line := sc.Text(); line != "" {
				last = line
			}
		}
		resCh <- result{lastLine: last, closed: true, err: sc.Err()}
	}()

	// 等重连提示帧出现，确认流已建立。
	deadline := time.Now().Add(3 * time.Second)
	established := false
	for !established && time.Now().Before(deadline) {
		if len(resCh) > 0 {
			break
		}
		// 无阻塞探测：向会话写事件驱动首帧即可省略；此处直接进入作废。
		established = true
	}

	// 作废会话（管理员重置语义：删除全部会话 + bump 版本）。
	newHash, err := store.HashPassword("sse-new-pass-1")
	if err != nil {
		t.Fatal(err)
	}
	if err := st.AdminResetUserPassword("tia", newHash, "test", "test"); err != nil {
		t.Fatal(err)
	}

	select {
	case r := <-resCh:
		if r.err != nil {
			t.Fatalf("stream read error: %v", r.err)
		}
		if !strings.Contains(r.lastLine, "session-expired") {
			t.Fatalf("stream should end with session-expired marker, got %q", r.lastLine)
		}
	case <-time.After(3 * time.Second):
		t.Fatal("stream not closed within recheck window after session invalidation")
	}
}

// API key 凭证的流：撤销后同样在复验周期内断流（合同 §6.7 补齐）。
func TestSSEAPIKeyRevocationClosesStream(t *testing.T) {
	old := sseAuthRecheck
	sseAuthRecheck = 300 * time.Millisecond
	t.Cleanup(func() { sseAuthRecheck = old })

	srv, st := newAuthTestServer(t)
	h := srv.Handler()

	if err := st.EnsureUser("kia", ""); err != nil {
		t.Fatal(err)
	}
	key, err := st.CreateAPIKey("kia", "sse")
	if err != nil {
		t.Fatal(err)
	}
	proj, err := st.CreateProject("kia", "sse-proj")
	if err != nil {
		t.Fatal(err)
	}
	conv, err := st.CreateConversation("kia", proj.ID, "", "")
	if err != nil {
		t.Fatal(err)
	}

	ts := httptest.NewServer(h)
	t.Cleanup(ts.Close)

	ctx, cancel := context.WithCancel(context.Background())
	t.Cleanup(cancel)
	req, err := http.NewRequestWithContext(ctx, "GET", ts.URL+"/api/events?conversationId="+conv.ID+"&key="+key, nil)
	if err != nil {
		t.Fatal(err)
	}

	resCh := make(chan string, 1)
	go func() {
		resp, err := ts.Client().Do(req)
		if err != nil {
			resCh <- "err:" + err.Error()
			return
		}
		defer resp.Body.Close()
		sc := bufio.NewScanner(resp.Body)
		last := ""
		for sc.Scan() {
			if line := sc.Text(); line != "" {
				last = line
			}
		}
		resCh <- last
	}()

	// 等首帧（retry 行）出现后撤销该 Key。
	time.Sleep(300 * time.Millisecond)
	if _, err := st.RevokeAPIKeyByName("kia", "sse"); err != nil {
		t.Fatal(err)
	}

	select {
	case got := <-resCh:
		if !strings.Contains(got, "session-expired") {
			t.Fatalf("stream should end with session-expired marker, got %q", got)
		}
	case <-time.After(3 * time.Second):
		t.Fatal("stream not closed within recheck window after key revocation")
	}
}

// 兑换入口限流：同源连续失败达上限后返回 429（跨实例共担由 store 测试覆盖）。
func TestRedeemRateLimited(t *testing.T) {
	srv, _ := newAuthTestServer(t)
	h := srv.Handler()
	for i := 0; i < redeemMaxFails; i++ {
		req := httptest.NewRequest("POST", "/api/auth/login-key/redeem", strings.NewReader(`{"key":"alk_not_a_real_key_at_all"}`))
		req.Header.Set("Content-Type", "application/json")
		rec := httptest.NewRecorder()
		h.ServeHTTP(rec, req)
		if rec.Code != http.StatusUnauthorized {
			t.Fatalf("attempt %d: want 401, got %d", i, rec.Code)
		}
	}
	req := httptest.NewRequest("POST", "/api/auth/login-key/redeem", strings.NewReader(`{"key":"alk_other"}`))
	req.Header.Set("Content-Type", "application/json")
	rec := httptest.NewRecorder()
	h.ServeHTTP(rec, req)
	if rec.Code != http.StatusTooManyRequests {
		t.Fatalf("want 429 after %d failures, got %d", redeemMaxFails, rec.Code)
	}
}