From ba1fc71f165c06bddbcbc88a8cb555015aae20a6 Mon Sep 17 00:00:00 2001 From: weichen Date: Fri, 18 Sep 2026 03:39:08 +0000 Subject: [PATCH] =?UTF-8?q?test(server):=20=E8=B7=A8=E7=94=A8=E6=88=B7?= =?UTF-8?q?=E4=BC=9A=E8=AF=9D=E6=8C=82=E9=9D=A0=E8=B4=9F=E4=BE=8B=E6=B5=8B?= =?UTF-8?q?=E8=AF=95=EF=BC=88PG=20+=20=E5=86=85=E5=AD=98=E5=8F=8C=E5=AE=9E?= =?UTF-8?q?=E7=8E=B0=EF=BC=89?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 覆盖 chat / chat/stream / chat/sessions 三条路径；受害者续用不受影响。 Co-authored-by: Cursor --- .../server/chat_session_ownership_pg_test.go | 79 ++++++++++++ .../server/chat_session_ownership_test.go | 116 ++++++++++++++++++ 2 files changed, 195 insertions(+) create mode 100644 internal/server/chat_session_ownership_pg_test.go create mode 100644 internal/server/chat_session_ownership_test.go diff --git a/internal/server/chat_session_ownership_pg_test.go b/internal/server/chat_session_ownership_pg_test.go new file mode 100644 index 00000000..011aca25 --- /dev/null +++ b/internal/server/chat_session_ownership_pg_test.go @@ -0,0 +1,79 @@ +package server_test + +import ( + "bytes" + "encoding/json" + "net/http" + "net/http/httptest" + "strings" + "testing" + "time" +) + +func TestChatSessionOwnership_CrossUserHijackDenied_PG(t *testing.T) { + pool := openPGPoolWithMigrations(t) + if pool == nil { + return + } + defer pool.Close() + + handler, _ := setupPGGovernanceServer(t, pool) + victimToken := loginTestUser(t, handler, "zhanggong", "123456") + attackerToken := loginTestUser(t, handler, "lizhuguan", "123456") + + sessionID := "sess-hijack-pg-" + time.Now().Format("150405.000") + + 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()) + } +} diff --git a/internal/server/chat_session_ownership_test.go b/internal/server/chat_session_ownership_test.go new file mode 100644 index 00000000..7b5cc149 --- /dev/null +++ b/internal/server/chat_session_ownership_test.go @@ -0,0 +1,116 @@ +package server_test + +import ( + "bytes" + "encoding/json" + "net/http" + "net/http/httptest" + "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()) + } +}