go_blog/handlers/security_test.go

package handlers

import (
	"fmt"
	"io"
	"mime/multipart"
	"net/http"
	"net/http/httptest"
	"net/url"
	"os"
	"path/filepath"
	"regexp"
	"strings"
	"testing"

	"github.com/gin-contrib/sessions"
	"github.com/gin-contrib/sessions/cookie"
	"github.com/gin-gonic/gin"
	"github.com/glebarez/sqlite"
	"gorm.io/gorm"

	"go_blog/middleware"
	"go_blog/models"
)

// securityTestEnv wires a router that mirrors the production middleware chain
// (sessions -> CSRF -> user context) plus the routes under test.
type securityTestEnv struct {
	router     *gin.Engine
	db         *gorm.DB
	storageDir string
}

var csrfTokenRe = regexp.MustCompile(`name="_csrf" value="([^"]+)"`)

func newSecurityTestEnv(t *testing.T) *securityTestEnv {
	t.Helper()
	gin.SetMode(gin.TestMode)

	db, err := gorm.Open(sqlite.Open(filepath.Join(t.TempDir(), "test.db")), &gorm.Config{})
	if err != nil {
		t.Fatalf("open sqlite: %v", err)
	}
	if err := db.AutoMigrate(&models.User{}, &models.Article{}, &models.Attachment{}, &models.SiteSetting{},
		&models.UploadConfig{}, &models.UploadFileType{}, &models.CommentConfig{}, &models.NavLink{},
		&models.DownloadBaseURL{}); err != nil {
		t.Fatalf("migrate: %v", err)
	}

	storageDir := t.TempDir()

	// Seed the upload policy so ValidateUpload accepts .txt files.
	db.Create(&models.SiteSetting{ID: 1})
	db.Create(&models.CommentConfig{ID: 1, Enabled: true, AllowGuest: true})
	db.Create(&models.UploadConfig{ID: 1, Enabled: true, DefaultMaxSize: 1024 * 1024, StorageDir: "attachments"})
	db.Create(&models.UploadFileType{Extension: ".txt", MimeType: "text/plain", Category: models.CategoryDocument, Enabled: true})
	models.LoadConfigCache(db)

	// Seed users.
	mustUser(t, db, "admin", models.RoleAdmin)
	alice := mustUser(t, db, "alice", models.RoleAuthor)
	bob := mustUser(t, db, "bob", models.RoleAuthor)

	// Seed one article per author.
	aliceArt := models.Article{AuthorID: alice.ID, Title: "alice post", Slug: "alice-post", Content: "x", Status: models.ArticlePublished}
	bobArt := models.Article{AuthorID: bob.ID, Title: "bob post", Slug: "bob-post", Content: "x", Status: models.ArticlePublished}
	db.Create(&aliceArt)
	db.Create(&bobArt)

	r := gin.New()
	if err := r.SetTrustedProxies(nil); err != nil {
		t.Fatalf("set trusted proxies: %v", err)
	}
	r.LoadHTMLGlob("../templates/**/*.html")
	store := cookie.NewStore([]byte("test-secret"))
	r.Use(sessions.Sessions("blog_session", store))
	r.Use(middleware.CSRFProtect())
	r.Use(middleware.SetUserContext(db))

	r.GET("/login", LoginPage())
	r.POST("/login", Login(db))
	r.POST("/logout", Logout())

	protected := r.Group("/my", middleware.AuthRequired())
	{
		protected.POST("/articles/attachments", UploadAttachment(db, storageDir))
		protected.POST("/articles/attachments/:id/delete", DeleteAttachment(db, storageDir))
		protected.GET("/articles/:id/attachments", ListAttachments(db))
		protected.GET("/whoami", func(c *gin.Context) {
			uid, _ := sessionAuthorID(c)
			c.String(http.StatusOK, "uid=%d", uid)
		})
	}

	return &securityTestEnv{router: r, db: db, storageDir: storageDir}
}

func mustUser(t *testing.T, db *gorm.DB, username, role string) models.User {
	t.Helper()
	u := models.User{Username: username, DisplayName: username, Role: role, Status: models.StatusNormal}
	if err := u.SetPassword("pw-" + username); err != nil {
		t.Fatalf("set password: %v", err)
	}
	if err := db.Create(&u).Error; err != nil {
		t.Fatalf("create user %s: %v", username, err)
	}
	return u
}

// login performs the full login flow (GET the form for a CSRF token, then POST
// credentials) and returns the authenticated session cookie.
func (e *securityTestEnv) login(t *testing.T, username string) string {
	t.Helper()

	// Anonymous GET to obtain CSRF token + session cookie.
	req := httptest.NewRequest(http.MethodGet, "/login", nil)
	w := httptest.NewRecorder()
	e.router.ServeHTTP(w, req)
	if w.Code != http.StatusOK {
		t.Fatalf("GET /login: status = %d", w.Code)
	}
	m := csrfTokenRe.FindStringSubmatch(w.Body.String())
	if m == nil {
		t.Fatal("login page did not render a CSRF token")
	}
	cookie := e.sessionCookie(w)

	// POST credentials with the token.
	form := url.Values{}
	form.Set("username", username)
	form.Set("password", "pw-"+username)
	form.Set("_csrf", m[1])
	req = httptest.NewRequest(http.MethodPost, "/login", strings.NewReader(form.Encode()))
	req.Header.Set("Content-Type", "application/x-www-form-urlencoded")
	if cookie != "" {
		req.Header.Set("Cookie", cookie)
	}
	w = httptest.NewRecorder()
	e.router.ServeHTTP(w, req)
	if w.Code != http.StatusFound {
		t.Fatalf("POST /login (%s): status = %d, body %s", username, w.Code, w.Body.String())
	}
	authCookie := e.sessionCookie(w)
	if authCookie == "" {
		t.Fatal("login did not set a session cookie")
	}
	return authCookie
}

// sessionCookie extracts the blog_session cookie from a recorder. When
// several Set-Cookie headers are present (e.g. middleware and handler both
// save the session), the LAST one is the effective value - browsers apply
// them in order.
func (e *securityTestEnv) sessionCookie(w *httptest.ResponseRecorder) string {
	cookie := ""
	for _, c := range w.Result().Cookies() {
		if c.Name == "blog_session" {
			cookie = c.Name + "=" + c.Value
		}
	}
	return cookie
}

func (e *securityTestEnv) do(method, path, cookie string, body io.Reader, contentType string) *httptest.ResponseRecorder {
	req := httptest.NewRequest(method, path, body)
	if contentType != "" {
		req.Header.Set("Content-Type", contentType)
	}
	if cookie != "" {
		req.Header.Set("Cookie", cookie)
	}
	w := httptest.NewRecorder()
	e.router.ServeHTTP(w, req)
	return w
}

func (e *securityTestEnv) upload(t *testing.T, cookie, csrfToken, articleID string) *httptest.ResponseRecorder {
	t.Helper()
	var buf strings.Builder
	mw := multipart.NewWriter(&buf)
	if articleID != "" {
		mw.WriteField("article_id", articleID)
	} else {
		mw.WriteField("session_token", "test-pending-token")
	}
	mw.WriteField("_csrf", csrfToken)
	fw, _ := mw.CreateFormFile("file", "hello.txt")
	fw.Write([]byte("hello world"))
	mw.Close()
	return e.do(http.MethodPost, "/my/articles/attachments", cookie, strings.NewReader(buf.String()), mw.FormDataContentType())
}

// csrfTokenFor fetches a fresh CSRF token for an authenticated session.
func (e *securityTestEnv) csrfTokenFor(t *testing.T, cookie string) string {
	t.Helper()
	w := e.do(http.MethodGet, "/login", cookie, nil, "")
	if w.Code != http.StatusOK {
		t.Fatalf("GET /login: status = %d", w.Code)
	}
	m := csrfTokenRe.FindStringSubmatch(w.Body.String())
	if m == nil {
		t.Fatal("login page did not render a CSRF token")
	}
	return m[1]
}

func TestLoginRotatesSession(t *testing.T) {
	e := newSecurityTestEnv(t)

	// Obtain an anonymous session (pre-login cookie).
	req := httptest.NewRequest(http.MethodGet, "/login", nil)
	w := httptest.NewRecorder()
	e.router.ServeHTTP(w, req)
	preLoginCookie := e.sessionCookie(w)
	if preLoginCookie == "" {
		t.Fatal("expected anonymous session cookie")
	}

	authCookie := e.login(t, "alice")
	if authCookie == preLoginCookie {
		t.Fatal("session cookie was not rotated on login (fixation risk)")
	}

	// The authenticated session works.
	w = e.do(http.MethodGet, "/my/whoami", authCookie, nil, "")
	if w.Code != http.StatusOK || !strings.Contains(w.Body.String(), "uid=") {
		t.Fatalf("authenticated request failed: status=%d body=%s", w.Code, w.Body.String())
	}

	// The old (fixated) session must NOT carry the login.
	w = e.do(http.MethodGet, "/my/whoami", preLoginCookie, nil, "")
	if w.Code != http.StatusFound {
		t.Fatalf("pre-login session still authenticated after login: status=%d", w.Code)
	}
}

func TestAttachmentListRequiresOwnership(t *testing.T) {
	e := newSecurityTestEnv(t)
	var aliceArt, bobArt models.Article
	e.db.Where("slug = ?", "alice-post").First(&aliceArt)
	e.db.Where("slug = ?", "bob-post").First(&bobArt)

	alice := e.login(t, "alice")

	// Own article: allowed.
	w := e.do(http.MethodGet, fmt.Sprintf("/my/articles/%d/attachments", aliceArt.ID), alice, nil, "")
	if w.Code != http.StatusOK {
		t.Fatalf("list own attachments: status = %d, want 200", w.Code)
	}

	// Someone else's article: forbidden.
	w = e.do(http.MethodGet, fmt.Sprintf("/my/articles/%d/attachments", bobArt.ID), alice, nil, "")
	if w.Code != http.StatusForbidden {
		t.Fatalf("list other user's attachments: status = %d, want 403", w.Code)
	}
}

func TestAttachmentUploadRequiresOwnership(t *testing.T) {
	e := newSecurityTestEnv(t)
	var bobArt models.Article
	e.db.Where("slug = ?", "bob-post").First(&bobArt)

	alice := e.login(t, "alice")
	token := e.csrfTokenFor(t, alice)

	// Upload pending (article_id=0 + session token): allowed.
	w := e.upload(t, alice, token, "")
	if w.Code != http.StatusOK {
		t.Fatalf("pending upload: status = %d, body %s", w.Code, w.Body.String())
	}

	// Upload to someone else's article: forbidden.
	w = e.upload(t, alice, token, fmt.Sprint(bobArt.ID))
	if w.Code != http.StatusForbidden {
		t.Fatalf("upload to other user's article: status = %d, want 403", w.Code)
	}
}

func TestAttachmentDeleteRequiresOwnership(t *testing.T) {
	e := newSecurityTestEnv(t)
	var aliceArt, bobArt models.Article
	e.db.Where("slug = ?", "alice-post").First(&aliceArt)
	e.db.Where("slug = ?", "bob-post").First(&bobArt)

	alice := e.login(t, "alice")
	token := e.csrfTokenFor(t, alice)

	// Alice uploads an attachment to her own article.
	w := e.upload(t, alice, token, fmt.Sprint(aliceArt.ID))
	if w.Code != http.StatusOK {
		t.Fatalf("upload: status = %d, body %s", w.Code, w.Body.String())
	}

	// Bob uploads an attachment to his own article.
	bob := e.login(t, "bob")
	bobToken := e.csrfTokenFor(t, bob)
	w = e.upload(t, bob, bobToken, fmt.Sprint(bobArt.ID))
	if w.Code != http.StatusOK {
		t.Fatalf("upload (bob): status = %d, body %s", w.Code, w.Body.String())
	}

	var bobAtt models.Attachment
	if err := e.db.Where("uploader_id = ?", userIDByUsername(t, e.db, "bob")).First(&bobAtt).Error; err != nil {
		t.Fatalf("bob attachment not found: %v", err)
	}

	// Alice cannot delete Bob's attachment.
	form := url.Values{}
	form.Set("_csrf", token)
	w = e.do(http.MethodPost, fmt.Sprintf("/my/articles/attachments/%d/delete", bobAtt.ID), alice,
		strings.NewReader(form.Encode()), "application/x-www-form-urlencoded")
	if w.Code != http.StatusForbidden {
		t.Fatalf("delete other user's attachment: status = %d, want 403", w.Code)
	}

	// Bob can delete his own.
	form.Set("_csrf", bobToken)
	w = e.do(http.MethodPost, fmt.Sprintf("/my/articles/attachments/%d/delete", bobAtt.ID), bob,
		strings.NewReader(form.Encode()), "application/x-www-form-urlencoded")
	if w.Code != http.StatusOK {
		t.Fatalf("delete own attachment: status = %d, want 200 (body %s)", w.Code, w.Body.String())
	}

	// Bob's attachment record should be gone.
	var count int64
	e.db.Model(&models.Attachment{}).Where("id = ?", bobAtt.ID).Count(&count)
	if count != 0 {
		t.Fatal("attachment was not deleted")
	}
}

func TestAttachmentCSRFEnforced(t *testing.T) {
	e := newSecurityTestEnv(t)
	alice := e.login(t, "alice")

	// POST without a CSRF token must be rejected before reaching the handler.
	var buf strings.Builder
	mw := multipart.NewWriter(&buf)
	fw, _ := mw.CreateFormFile("file", "hello.txt")
	fw.Write([]byte("hello"))
	mw.Close()
	w := e.do(http.MethodPost, "/my/articles/attachments", alice, strings.NewReader(buf.String()), mw.FormDataContentType())
	if w.Code != http.StatusForbidden {
		t.Fatalf("upload without CSRF token: status = %d, want 403", w.Code)
	}
}

func TestAttachmentAdminOverride(t *testing.T) {
	e := newSecurityTestEnv(t)
	var bobArt models.Article
	e.db.Where("slug = ?", "bob-post").First(&bobArt)

	admin := e.login(t, "admin")
	token := e.csrfTokenFor(t, admin)

	// Admin may list and upload to any article.
	w := e.do(http.MethodGet, fmt.Sprintf("/my/articles/%d/attachments", bobArt.ID), admin, nil, "")
	if w.Code != http.StatusOK {
		t.Fatalf("admin list: status = %d, want 200", w.Code)
	}
	w = e.upload(t, admin, token, fmt.Sprint(bobArt.ID))
	if w.Code != http.StatusOK {
		t.Fatalf("admin upload: status = %d, want 200 (body %s)", w.Code, w.Body.String())
	}

	// Clean up files created during the test (best effort).
	entries, _ := os.ReadDir(filepath.Join(e.storageDir, "attachments"))
	for _, ent := range entries {
		os.Remove(filepath.Join(e.storageDir, "attachments", ent.Name()))
	}
}

func userIDByUsername(t *testing.T, db *gorm.DB, username string) uint {
	t.Helper()
	var u models.User
	if err := db.Where("username = ?", username).First(&u).Error; err != nil {
		t.Fatalf("user %s not found: %v", username, err)
	}
	return u.ID
}