Agent-System/internal/server/idor_hotfixes_test.go

package server_test

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

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

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

	suffix := time.Now().Format("150405.000")
	aliceToken := createAndLoginUser(t, handler, adminToken, "pref_alice_"+suffix, "pass12345")
	bobToken := createAndLoginUser(t, handler, adminToken, "pref_bob_"+suffix, "pass12345")

	// resolve alice user id via /auth/me
	meReq := httptest.NewRequest(http.MethodGet, "/api/v1/auth/me", nil)
	withBearerAuth(meReq, aliceToken)
	meRec := httptest.NewRecorder()
	handler.ServeHTTP(meRec, meReq)
	if meRec.Code != http.StatusOK {
		t.Fatalf("auth/me: %d %s", meRec.Code, meRec.Body.String())
	}
	var me map[string]any
	_ = json.Unmarshal(meRec.Body.Bytes(), &me)
	userObj, _ := me["user"].(map[string]any)
	aliceID, _ := userObj["id"].(float64)
	if aliceID <= 0 {
		t.Fatalf("cannot resolve alice id from me: %s", meRec.Body.String())
	}

	body, _ := json.Marshal(map[string]any{
		"user_id":             int64(aliceID),
		"communication_style": "hijacked",
		"custom_instructions": "inject prompt as bob",
	})
	req := httptest.NewRequest(http.MethodPut, "/api/v1/context/user/preference", bytes.NewReader(body))
	req.Header.Set("Content-Type", "application/json")
	withBearerAuth(req, bobToken)
	rec := httptest.NewRecorder()
	handler.ServeHTTP(rec, req)
	if rec.Code != http.StatusForbidden {
		t.Fatalf("expected 403 cross-user preference write, got %d %s", rec.Code, rec.Body.String())
	}
}

func TestGetProject_OtherUserSessionDenied(t *testing.T) {
	cfg := &config.Config{DataDir: t.TempDir(), AgentBackend: "stub", AdminUsername: "admin", AdminPassword: "admin-test-pass"}
	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, "proj_victim_"+suffix, "pass12345")
	attackerToken := createAndLoginUser(t, handler, adminToken, "proj_attacker_"+suffix, "pass12345")
	sessionID := "sess-proj-idor-" + suffix

	chatBody, _ := json.Marshal(map[string]any{
		"session_id":   sessionID,
		"project_name": "IDOR会话项目",
		"message":      "私密立项上下文",
	})
	req := httptest.NewRequest(http.MethodPost, "/api/v1/chat", bytes.NewReader(chatBody))
	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())
	}

	getReq := httptest.NewRequest(http.MethodGet, "/api/v1/projects?session_id="+sessionID, nil)
	withBearerAuth(getReq, attackerToken)
	getRec := httptest.NewRecorder()
	handler.ServeHTTP(getRec, getReq)
	if getRec.Code != http.StatusForbidden {
		t.Fatalf("expected 403 other-user session project read, got %d %s", getRec.Code, getRec.Body.String())
	}
}

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

	suffix := time.Now().Format("150405.000")
	ownerToken := createAndLoginUser(t, handler, adminToken, "recat_owner_"+suffix, "pass12345")
	otherToken := createAndLoginUser(t, handler, adminToken, "recat_other_"+suffix, "pass12345")
	sessionID := "sess-recat-idor-" + suffix

	chatBody, _ := json.Marshal(map[string]any{"session_id": sessionID, "message": "需要重分类的消息"})
	req := httptest.NewRequest(http.MethodPost, "/api/v1/chat", bytes.NewReader(chatBody))
	req.Header.Set("Content-Type", "application/json")
	withBearerAuth(req, ownerToken)
	rec := httptest.NewRecorder()
	handler.ServeHTTP(rec, req)
	if rec.Code != http.StatusOK {
		t.Fatalf("owner chat: %d %s", rec.Code, rec.Body.String())
	}

	emptyBody := bytes.NewReader([]byte(`{}`))
	emptyReq := httptest.NewRequest(http.MethodPost, "/api/v1/chat/recategorize", emptyBody)
	emptyReq.Header.Set("Content-Type", "application/json")
	withBearerAuth(emptyReq, ownerToken)
	emptyRec := httptest.NewRecorder()
	handler.ServeHTTP(emptyRec, emptyReq)
	if emptyRec.Code != http.StatusBadRequest {
		t.Fatalf("expected 400 for empty session_id, got %d %s", emptyRec.Code, emptyRec.Body.String())
	}

	otherBody, _ := json.Marshal(map[string]string{"session_id": sessionID})
	otherReq := httptest.NewRequest(http.MethodPost, "/api/v1/chat/recategorize", bytes.NewReader(otherBody))
	otherReq.Header.Set("Content-Type", "application/json")
	withBearerAuth(otherReq, otherToken)
	otherRec := httptest.NewRecorder()
	handler.ServeHTTP(otherRec, otherReq)
	if otherRec.Code != http.StatusForbidden {
		t.Fatalf("expected 403 for other user's session, got %d %s", otherRec.Code, otherRec.Body.String())
	}

	ownBody, _ := json.Marshal(map[string]string{"session_id": sessionID})
	ownReq := httptest.NewRequest(http.MethodPost, "/api/v1/chat/recategorize", bytes.NewReader(ownBody))
	ownReq.Header.Set("Content-Type", "application/json")
	withBearerAuth(ownReq, ownerToken)
	ownRec := httptest.NewRecorder()
	handler.ServeHTTP(ownRec, ownReq)
	if ownRec.Code != http.StatusOK {
		t.Fatalf("expected 200 for owned session, got %d %s", ownRec.Code, ownRec.Body.String())
	}
}