Agent-System/internal/server/chat_session_ownership_test.go

package server_test

import (
	"bytes"
	"encoding/json"
	"net/http"
	"net/http/httptest"
	"strconv"
	"strings"
	"testing"
	"time"

	"chat-agent-system/internal/bizdev"
	"chat-agent-system/internal/config"
	"chat-agent-system/internal/server"
)

func createAndLoginUser(t *testing.T, handler http.Handler, adminToken, username, password string) string {
	t.Helper()
	createBody, _ := json.Marshal(map[string]string{
		"username":         username,
		"name":             username,
		"initial_password": password,
	})
	reqCreate := httptest.NewRequest(http.MethodPost, "/api/v1/admin/users", bytes.NewReader(createBody))
	reqCreate.Header.Set("Content-Type", "application/json")
	withBearerAuth(reqCreate, adminToken)
	recCreate := httptest.NewRecorder()
	handler.ServeHTTP(recCreate, reqCreate)
	if recCreate.Code != http.StatusOK {
		t.Fatalf("admin create user %s: %d %s", username, recCreate.Code, recCreate.Body.String())
	}

	token := loginTestUser(t, handler, username, password)
	// 管理员创建的用户可能要求改密；改完后再登录拿稳定 token。
	changeBody, _ := json.Marshal(map[string]string{
		"old_password": password,
		"new_password": password + "_ok",
	})
	reqChange := httptest.NewRequest(http.MethodPost, "/api/v1/auth/change-password", bytes.NewReader(changeBody))
	reqChange.Header.Set("Content-Type", "application/json")
	withBearerAuth(reqChange, token)
	recChange := httptest.NewRecorder()
	handler.ServeHTTP(recChange, reqChange)
	if recChange.Code == http.StatusOK {
		return loginTestUser(t, handler, username, password+"_ok")
	}
	return token
}

func TestChatSessionOwnership_CrossUserHijackDenied_Memory(t *testing.T) {
	cfg := &config.Config{DataDir: t.TempDir(), AgentBackend: "stub"}
	svc := bizdev.MustNewService(t, cfg)
	srv := server.NewServer(cfg, svc)
	handler := srv.Routes()
	adminToken := loginTestAdmin(t, handler, cfg)

	suffix := time.Now().Format("150405.000")
	victimToken := createAndLoginUser(t, handler, adminToken, "hijack_victim_"+suffix, "pass12345")
	attackerToken := createAndLoginUser(t, handler, adminToken, "hijack_attacker_"+suffix, "pass12345")

	sessionID := "sess-hijack-" + suffix

	victimBody, _ := json.Marshal(map[string]any{"session_id": sessionID, "message": "机密立项讨论"})
	req := httptest.NewRequest(http.MethodPost, "/api/v1/chat", bytes.NewReader(victimBody))
	req.Header.Set("Content-Type", "application/json")
	withBearerAuth(req, victimToken)
	rec := httptest.NewRecorder()
	handler.ServeHTTP(rec, req)
	if rec.Code != http.StatusOK {
		t.Fatalf("victim chat: %d %s", rec.Code, rec.Body.String())
	}

	attackerBody, _ := json.Marshal(map[string]any{"session_id": sessionID, "message": "尝试接管"})
	req2 := httptest.NewRequest(http.MethodPost, "/api/v1/chat", bytes.NewReader(attackerBody))
	req2.Header.Set("Content-Type", "application/json")
	withBearerAuth(req2, attackerToken)
	rec2 := httptest.NewRecorder()
	handler.ServeHTTP(rec2, req2)
	if rec2.Code != http.StatusForbidden {
		t.Fatalf("attacker hijack expected 403, got %d: %s", rec2.Code, rec2.Body.String())
	}
	if !strings.Contains(rec2.Body.String(), "SESSION_FORBIDDEN") {
		t.Fatalf("expected SESSION_FORBIDDEN, got %s", rec2.Body.String())
	}

	req3 := httptest.NewRequest(http.MethodPost, "/api/v1/chat/stream", bytes.NewReader(attackerBody))
	req3.Header.Set("Content-Type", "application/json")
	withBearerAuth(req3, attackerToken)
	rec3 := httptest.NewRecorder()
	handler.ServeHTTP(rec3, req3)
	if rec3.Code != http.StatusForbidden {
		t.Fatalf("attacker stream hijack expected 403, got %d", rec3.Code)
	}
	if strings.Contains(rec3.Header().Get("Content-Type"), "text/event-stream") {
		t.Fatal("must reject before SSE headers are set")
	}

	upsertBody, _ := json.Marshal(map[string]any{"session_id": sessionID, "user_id": 999999, "title": "hijack via upsert"})
	req4 := httptest.NewRequest(http.MethodPost, "/api/v1/chat/sessions", bytes.NewReader(upsertBody))
	req4.Header.Set("Content-Type", "application/json")
	withBearerAuth(req4, attackerToken)
	rec4 := httptest.NewRecorder()
	handler.ServeHTTP(rec4, req4)
	if rec4.Code != http.StatusForbidden {
		t.Fatalf("attacker upsert hijack expected 403, got %d: %s", rec4.Code, rec4.Body.String())
	}

	req5 := httptest.NewRequest(http.MethodPost, "/api/v1/chat", bytes.NewReader(victimBody))
	req5.Header.Set("Content-Type", "application/json")
	withBearerAuth(req5, victimToken)
	rec5 := httptest.NewRecorder()
	handler.ServeHTTP(rec5, req5)
	if rec5.Code != http.StatusOK {
		t.Fatalf("victim continued access broken: %d %s", rec5.Code, rec5.Body.String())
	}
}

func TestListChatSessions_CrossUserIs403(t *testing.T) {
	cfg := &config.Config{DataDir: t.TempDir(), AgentBackend: "stub"}
	svc := bizdev.MustNewService(t, cfg)
	handler := server.NewServer(cfg, svc).Routes()
	adminToken := loginTestAdmin(t, handler, cfg)
	suffix := time.Now().Format("150405.000")
	victimToken := createAndLoginUser(t, handler, adminToken, "list_victim_"+suffix, "pass12345")
	attackerToken := createAndLoginUser(t, handler, adminToken, "list_attacker_"+suffix, "pass12345")
	victimID := fetchAuthUserID(t, handler, victimToken)

	req := httptest.NewRequest(http.MethodGet, "/api/v1/chat/sessions?user_id="+strconv.FormatInt(victimID, 10), nil)
	withBearerAuth(req, attackerToken)
	rec := httptest.NewRecorder()
	handler.ServeHTTP(rec, req)
	if rec.Code != http.StatusForbidden {
		t.Fatalf("cross-user session list expected 403, got %d: %s", rec.Code, rec.Body.String())
	}
}