Agent-System/internal/server/security_penetration_test.go

package server

import (
	"bytes"
	"io"
	"mime/multipart"
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"

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

// newStrictAuthServer 禁用 legacy 匿名 principal，用于安全负例测试。
func newStrictAuthServer(t *testing.T) *Server {
	t.Helper()
	cfg := &config.Config{
		Env:           "main",
		DataDir:       t.TempDir(),
		AgentBackend:  "stub",
		AllowInMemory: false,
	}
	svc, err := bizdev.NewService(cfg)
	if err != nil {
		t.Fatalf("NewService: %v", err)
	}
	return NewServer(cfg, svc)
}

func TestPenetration_AnonymousAdminTemplateUpload401(t *testing.T) {
	srv := newStrictAuthServer(t)

	var body bytes.Buffer
	w := multipart.NewWriter(&body)
	_ = w.WriteField("asset_key", "test-template")
	part, _ := w.CreateFormFile("file", "brief.md")
	_, _ = part.Write([]byte("# ok"))
	_ = w.Close()

	req := httptest.NewRequest(http.MethodPost, "/api/v1/admin/templates", &body)
	req.Header.Set("Content-Type", w.FormDataContentType())
	rec := httptest.NewRecorder()
	srv.Routes().ServeHTTP(rec, req)
	if rec.Code != http.StatusUnauthorized {
		t.Fatalf("expected 401, got %d body=%s", rec.Code, rec.Body.String())
	}
}

func TestPenetration_AnonymousTaskRunMutations401(t *testing.T) {
	srv := newStrictAuthServer(t)

	cases := []struct {
		method string
		path   string
		body   string
	}{
		{http.MethodPost, "/api/v1/tasks/runs", `{"task_type":"freeform"}`},
		{http.MethodPost, "/api/v1/tasks/runs/run_evil/cancel", ""},
		{http.MethodPost, "/api/v1/tasks/runs/run_evil/resume", `{"confirm":true}`},
		{http.MethodGet, "/api/v1/tasks/runs/run_evil/events", ""},
	}
	for _, tc := range cases {
		var reader *bytes.Reader
		if tc.body != "" {
			reader = bytes.NewReader([]byte(tc.body))
		} else {
			reader = bytes.NewReader(nil)
		}
		req := httptest.NewRequest(tc.method, tc.path, reader)
		if tc.body != "" {
			req.Header.Set("Content-Type", "application/json")
		}
		rec := httptest.NewRecorder()
		srv.Routes().ServeHTTP(rec, req)
		if rec.Code != http.StatusUnauthorized {
			t.Fatalf("%s %s expected 401 got %d", tc.method, tc.path, rec.Code)
		}
	}
}

func TestPenetration_OversizedJSONBody413(t *testing.T) {
	srv := &Server{}
	inner := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if _, err := io.ReadAll(r.Body); isRequestBodyTooLarge(err) {
			w.WriteHeader(http.StatusRequestEntityTooLarge)
			return
		}
		w.WriteHeader(http.StatusOK)
	})
	handler := srv.bodyLimitMiddleware(inner)

	body := bytes.Repeat([]byte("x"), int(defaultJSONBodyLimit)+2048)
	req := httptest.NewRequest(http.MethodPost, "/api/v1/chat", bytes.NewReader(body))
	rec := httptest.NewRecorder()
	handler.ServeHTTP(rec, req)
	if rec.Code != http.StatusRequestEntityTooLarge {
		t.Fatalf("expected 413, got %d", rec.Code)
	}
}

func TestPenetration_OversizedTemplateUpload413(t *testing.T) {
	srv := &Server{}
	inner := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if _, err := io.ReadAll(r.Body); isRequestBodyTooLarge(err) {
			w.WriteHeader(http.StatusRequestEntityTooLarge)
			return
		}
		w.WriteHeader(http.StatusOK)
	})
	handler := srv.bodyLimitMiddleware(inner)

	body := bytes.Repeat([]byte("a"), int(defaultUploadBodyLimit)+1024)
	req := httptest.NewRequest(http.MethodPost, "/api/v1/admin/templates", bytes.NewReader(body))
	rec := httptest.NewRecorder()
	handler.ServeHTTP(rec, req)
	if rec.Code != http.StatusRequestEntityTooLarge {
		t.Fatalf("expected 413 for oversized template upload, got %d", rec.Code)
	}
}

func TestPenetration_ProjectAssetRejectsPathTraversalName(t *testing.T) {
	cfg := &config.Config{Env: "test", DataDir: t.TempDir()}
	srv := NewServer(cfg, bizdev.MustNewService(t, cfg))
	ts := httptest.NewServer(srv.Routes())
	defer ts.Close()

	var body bytes.Buffer
	w := multipart.NewWriter(&body)
	part, _ := w.CreateFormFile("file", "../../etc/passwd.txt")
	_, _ = part.Write([]byte("not html"))
	_ = w.Close()

	req, _ := http.NewRequest(http.MethodPost, ts.URL+"/api/v1/projects/7/assets", &body)
	req.Header.Set("Content-Type", w.FormDataContentType())
	withBearerAuth(req, loginTestAdmin(t, srv))
	resp, err := http.DefaultClient.Do(req)
	if err != nil {
		t.Fatal(err)
	}
	defer resp.Body.Close()
	if resp.StatusCode != http.StatusOK {
		t.Fatalf("expected 200 with sanitized basename, got %d", resp.StatusCode)
	}
}

func TestPenetration_ProjectAssetRejectsScriptContent(t *testing.T) {
	cfg := &config.Config{Env: "test", DataDir: t.TempDir()}
	srv := NewServer(cfg, bizdev.MustNewService(t, cfg))
	ts := httptest.NewServer(srv.Routes())
	defer ts.Close()

	var body bytes.Buffer
	w := multipart.NewWriter(&body)
	part, _ := w.CreateFormFile("file", "payload.txt")
	_, _ = part.Write([]byte("#!/bin/bash\nrm -rf /"))
	_ = w.Close()

	req, _ := http.NewRequest(http.MethodPost, ts.URL+"/api/v1/projects/7/assets", &body)
	req.Header.Set("Content-Type", w.FormDataContentType())
	withBearerAuth(req, loginTestAdmin(t, srv))
	resp, err := http.DefaultClient.Do(req)
	if err != nil {
		t.Fatal(err)
	}
	defer resp.Body.Close()
	if resp.StatusCode != http.StatusUnsupportedMediaType {
		t.Fatalf("expected 415 for script content, got %d", resp.StatusCode)
	}
}

func TestPenetration_ReadyStrictWhenDBConfiguredButDown(t *testing.T) {
	cfg := &config.Config{
		DataDir:           t.TempDir(),
		AgentBackend:      "stub",
		ContextDbHost:     "127.0.0.1",
		ContextDbPort:     59999,
		AllowInMemory:     false,
	}
	srv := NewServer(cfg, bizdev.MustNewService(t, &config.Config{DataDir: t.TempDir(), AgentBackend: "stub"}))
	srv.cfg = cfg

	req := httptest.NewRequest(http.MethodGet, "/api/v1/ready", nil)
	rec := httptest.NewRecorder()
	srv.handleReadiness(rec, req)
	if rec.Code != http.StatusServiceUnavailable {
		t.Fatalf("expected 503 ready when DB down, got %d %s", rec.Code, rec.Body.String())
	}
}

func TestPenetration_DisallowedOriginNoWildcardCORS(t *testing.T) {
	cfg := &config.Config{
		DataDir:            t.TempDir(),
		AgentBackend:       "stub",
		CORSAllowedOrigins: []string{"https://trusted.example.com"},
	}
	srv := NewServer(cfg, bizdev.MustNewService(t, cfg))

	req := httptest.NewRequest(http.MethodOptions, "/api/v1/projects", nil)
	req.Header.Set("Origin", "https://attacker.example.com")
	req.Header.Set("Access-Control-Request-Method", "GET")
	rec := httptest.NewRecorder()
	srv.Routes().ServeHTTP(rec, req)
	allowOrigin := rec.Header().Get("Access-Control-Allow-Origin")
	if allowOrigin == "*" || allowOrigin == "https://attacker.example.com" {
		t.Fatalf("CORS must not reflect attacker origin, got %q", allowOrigin)
	}
}

func TestPenetration_LongFilenameRejected415(t *testing.T) {
	cfg := &config.Config{Env: "test", DataDir: t.TempDir()}
	srv := NewServer(cfg, bizdev.MustNewService(t, cfg))
	ts := httptest.NewServer(srv.Routes())
	defer ts.Close()

	longName := strings.Repeat("z", 201) + ".md"
	var body bytes.Buffer
	w := multipart.NewWriter(&body)
	part, _ := w.CreateFormFile("file", longName)
	_, _ = part.Write([]byte("# x"))
	_ = w.Close()

	req, _ := http.NewRequest(http.MethodPost, ts.URL+"/api/v1/projects/1/assets", &body)
	req.Header.Set("Content-Type", w.FormDataContentType())
	withBearerAuth(req, loginTestAdmin(t, srv))
	resp, err := http.DefaultClient.Do(req)
	if err != nil {
		t.Fatal(err)
	}
	defer resp.Body.Close()
	if resp.StatusCode != http.StatusUnsupportedMediaType {
		t.Fatalf("expected 415, got %d", resp.StatusCode)
	}
}