ufile/app_test.go

package main

import (
	"bytes"
	"encoding/json"
	"io"
	"mime/multipart"
	"net/http"
	"net/http/httptest"
	"os"
	"path/filepath"
	"strings"
	"sync"
	"testing"
	"testing/fstest"
	"time"

	"golang.org/x/crypto/bcrypt"
)

func testApp(t *testing.T) (*App, string) {
	t.Helper()
	dir := t.TempDir()
	web := fstest.MapFS{"index.html": &fstest.MapFile{Data: []byte("test")}}
	app, password, err := newApp(filepath.Join(dir, "test.db"), dir, web)
	if err != nil {
		t.Fatal(err)
	}
	return app, password
}

func request(t *testing.T, router http.Handler, method, path string, body io.Reader, contentType string, cookie *http.Cookie) *httptest.ResponseRecorder {
	t.Helper()
	req := httptest.NewRequest(method, path, body)
	if contentType != "" {
		req.Header.Set("Content-Type", contentType)
	}
	if cookie != nil {
		req.AddCookie(cookie)
	}
	recorder := httptest.NewRecorder()
	router.ServeHTTP(recorder, req)
	return recorder
}

func loginAs(t *testing.T, router http.Handler, username, password string) *http.Cookie {
	t.Helper()
	body, _ := json.Marshal(map[string]string{"username": username, "password": password})
	response := request(t, router, http.MethodPost, "/api/auth/login", bytes.NewReader(body), "application/json", nil)
	if response.Code != http.StatusOK {
		t.Fatalf("login status = %d, body = %s", response.Code, response.Body.String())
	}
	for _, cookie := range response.Result().Cookies() {
		if cookie.Name == sessionCookie {
			return cookie
		}
	}
	t.Fatal("session cookie was not set")
	return nil
}

func uploadBody(t *testing.T, fields map[string]string, name string, content []byte) (*bytes.Buffer, string) {
	t.Helper()
	body := &bytes.Buffer{}
	writer := multipart.NewWriter(body)
	for key, value := range fields {
		if err := writer.WriteField(key, value); err != nil {
			t.Fatal(err)
		}
	}
	part, err := writer.CreateFormFile("file", name)
	if err != nil {
		t.Fatal(err)
	}
	if _, err := part.Write(content); err != nil {
		t.Fatal(err)
	}
	if err := writer.Close(); err != nil {
		t.Fatal(err)
	}
	return body, writer.FormDataContentType()
}

func TestBootstrapAdminAndForcePasswordChange(t *testing.T) {
	app, password := testApp(t)
	if password == "" {
		t.Fatal("initial password was not generated")
	}
	var admin User
	if err := app.db.First(&admin, 1).Error; err != nil {
		t.Fatal(err)
	}
	if admin.Username != "admin" || !admin.MustChangePassword || !admin.CanManageUsers || !admin.CanAdvancedUpload {
		t.Fatalf("unexpected admin: %+v", admin)
	}
	if bcrypt.CompareHashAndPassword([]byte(admin.PasswordHash), []byte(password)) != nil {
		t.Fatal("generated password does not match")
	}

	router := app.router()
	cookie := loginAs(t, router, "admin", password)
	blocked := request(t, router, http.MethodPost, "/api/files", nil, "", cookie)
	if blocked.Code != http.StatusForbidden || !strings.Contains(blocked.Body.String(), "PASSWORD_CHANGE_REQUIRED") {
		t.Fatalf("forced password status = %d, body = %s", blocked.Code, blocked.Body.String())
	}
	body, _ := json.Marshal(map[string]string{"currentPassword": password, "newPassword": "new-password-123"})
	changed := request(t, router, http.MethodPut, "/api/auth/password", bytes.NewReader(body), "application/json", cookie)
	if changed.Code != http.StatusOK {
		t.Fatalf("change password status = %d, body = %s", changed.Code, changed.Body.String())
	}
	if err := app.db.First(&admin, 1).Error; err != nil || admin.MustChangePassword {
		t.Fatal("password change flag was not cleared")
	}
}

func TestPickupCodePreviewConsumesFinalAccessAndDeletesFile(t *testing.T) {
	app, _ := testApp(t)
	hash, _ := bcrypt.GenerateFromPassword([]byte("password-123"), bcrypt.DefaultCost)
	user := User{Username: "sender", PasswordHash: string(hash), CanAdvancedUpload: true}
	if err := app.db.Create(&user).Error; err != nil {
		t.Fatal(err)
	}
	router := app.router()
	cookie := loginAs(t, router, user.Username, "password-123")
	body, contentType := uploadBody(t, map[string]string{"retention": "1h", "downloads": "1", "pickupCode": "123456"}, "hello.txt", []byte("hello"))
	uploaded := request(t, router, http.MethodPost, "/api/files", body, contentType, cookie)
	if uploaded.Code != http.StatusCreated {
		t.Fatalf("upload status = %d, body = %s", uploaded.Code, uploaded.Body.String())
	}
	var file fileResponse
	if err := json.Unmarshal(uploaded.Body.Bytes(), &file); err != nil {
		t.Fatal(err)
	}

	missingCode := request(t, router, http.MethodGet, "/api/pickup/"+file.Hash, nil, "", nil)
	if missingCode.Code != http.StatusForbidden || !strings.Contains(missingCode.Body.String(), "请输入取件码") {
		t.Fatalf("missing code status = %d, body = %s", missingCode.Code, missingCode.Body.String())
	}
	wrongCode := request(t, router, http.MethodGet, "/api/pickup/"+file.Hash+"?code=654321", nil, "", nil)
	if wrongCode.Code != http.StatusNotFound || !strings.Contains(wrongCode.Body.String(), publicFileOrCodeError) {
		t.Fatalf("wrong code status = %d, body = %s", wrongCode.Code, wrongCode.Body.String())
	}
	missingFile := request(t, router, http.MethodGet, "/api/pickup/notfound?code=654321", nil, "", nil)
	if missingFile.Code != wrongCode.Code || missingFile.Body.String() != wrongCode.Body.String() {
		t.Fatalf("missing file and wrong code responses differ: missing = %d %s, wrong = %d %s", missingFile.Code, missingFile.Body.String(), wrongCode.Code, wrongCode.Body.String())
	}
	verified := request(t, router, http.MethodGet, "/api/pickup/"+file.Hash+"?code=123456", nil, "", nil)
	if verified.Code != http.StatusOK {
		t.Fatalf("verify status = %d, body = %s", verified.Code, verified.Body.String())
	}
	var stored SharedFile
	if err := app.db.Where("hash = ?", file.Hash).First(&stored).Error; err != nil || stored.DownloadCount != 0 || stored.FailedCodeAttempts != 0 {
		t.Fatalf("verification state: downloads = %d, code failures = %d, err = %v", stored.DownloadCount, stored.FailedCodeAttempts, err)
	}
	preview := request(t, router, http.MethodGet, "/uf/"+file.Hash+"?preview=true&code=123456", nil, "", nil)
	if preview.Code != http.StatusOK || preview.Body.String() != "hello" {
		t.Fatalf("preview status = %d, body = %s", preview.Code, preview.Body.String())
	}
	if disposition := preview.Header().Get("Content-Disposition"); disposition != "inline" {
		t.Fatalf("disposition = %q", disposition)
	}
	var count int64
	app.db.Model(&SharedFile{}).Where("hash = ?", file.Hash).Count(&count)
	if count != 0 {
		t.Fatal("file record remains after final access")
	}
	if _, err := os.Stat(filepath.Join(app.filesDir, file.Hash)); !os.IsNotExist(err) {
		t.Fatal("stored file remains after final access")
	}
	again := request(t, router, http.MethodGet, "/uf/"+file.Hash+"?code=123456", nil, "", nil)
	if again.Code != http.StatusNotFound {
		t.Fatalf("second access status = %d", again.Code)
	}
}

func TestFiveConsecutivePickupCodeFailuresDeleteFile(t *testing.T) {
	app, _ := testApp(t)
	codeHash, _ := bcrypt.GenerateFromPassword([]byte("123456"), bcrypt.DefaultCost)
	max := 10
	file := SharedFile{Hash: "badcode5", OriginalName: "secret.txt", MimeType: "text/plain", OwnerID: 1, MaxDownloads: &max, PickupCodeHash: string(codeHash), Ready: true, Size: 6}
	if err := app.db.Create(&file).Error; err != nil {
		t.Fatal(err)
	}
	path := filepath.Join(app.filesDir, file.Hash)
	if err := os.WriteFile(path, []byte("secret"), 0o640); err != nil {
		t.Fatal(err)
	}
	router := app.router()
	requests := []string{
		"/api/pickup/" + file.Hash,
		"/uf/" + file.Hash + "?code=000001",
		"/api/pickup/" + file.Hash + "?code=000002",
		"/uf/" + file.Hash + "?code=000003",
		"/api/pickup/" + file.Hash + "?code=000004",
	}
	for index, path := range requests {
		response := request(t, router, http.MethodGet, path, nil, "", nil)
		want := http.StatusNotFound
		if index == 0 {
			want = http.StatusForbidden
		}
		if response.Code != want {
			t.Fatalf("failure %d status = %d, want %d, body = %s", index+1, response.Code, want, response.Body.String())
		}
	}
	var count int64
	app.db.Model(&SharedFile{}).Where("id = ?", file.ID).Count(&count)
	if count != 0 {
		t.Fatal("file record remains after five code failures")
	}
	if _, err := os.Stat(path); !os.IsNotExist(err) {
		t.Fatal("file remains on disk after five code failures")
	}
}

func TestFiniteDownloadCountDeletesAfterMoreThanOneAccess(t *testing.T) {
	app, _ := testApp(t)
	max := 3
	file := SharedFile{Hash: "three123", OriginalName: "three.txt", MimeType: "text/plain", OwnerID: 1, MaxDownloads: &max, Ready: true, Size: 1}
	if err := app.db.Create(&file).Error; err != nil {
		t.Fatal(err)
	}
	if err := os.WriteFile(filepath.Join(app.filesDir, file.Hash), []byte("x"), 0o640); err != nil {
		t.Fatal(err)
	}
	router := app.router()
	for i := 0; i < max; i++ {
		response := request(t, router, http.MethodGet, "/uf/"+file.Hash, nil, "", nil)
		if response.Code != http.StatusOK {
			t.Fatalf("access %d status = %d", i+1, response.Code)
		}
	}
	response := request(t, router, http.MethodGet, "/uf/"+file.Hash, nil, "", nil)
	if response.Code != http.StatusNotFound {
		t.Fatalf("access after limit status = %d", response.Code)
	}
	if _, err := os.Stat(filepath.Join(app.filesDir, file.Hash)); !os.IsNotExist(err) {
		t.Fatal("file remains after third access")
	}
}

func TestConcurrentAccessCannotExceedDownloadLimit(t *testing.T) {
	app, _ := testApp(t)
	max := 5
	file := SharedFile{Hash: "parallel", OriginalName: "parallel.txt", MimeType: "text/plain", OwnerID: 1, MaxDownloads: &max, Ready: true, Size: 1}
	if err := app.db.Create(&file).Error; err != nil {
		t.Fatal(err)
	}
	if err := os.WriteFile(filepath.Join(app.filesDir, file.Hash), []byte("x"), 0o640); err != nil {
		t.Fatal(err)
	}
	router := app.router()
	statuses := make(chan int, 20)
	var group sync.WaitGroup
	for i := 0; i < 20; i++ {
		group.Add(1)
		go func() {
			defer group.Done()
			statuses <- request(t, router, http.MethodGet, "/uf/"+file.Hash, nil, "", nil).Code
		}()
	}
	group.Wait()
	close(statuses)
	successes := 0
	for status := range statuses {
		if status == http.StatusOK {
			successes++
		}
	}
	if successes != max {
		t.Fatalf("successful accesses = %d, want %d", successes, max)
	}
}

func TestCleanupRemovesExpiredFileFromDatabaseAndDisk(t *testing.T) {
	app, _ := testApp(t)
	expired := time.Now().Add(-time.Minute)
	max := 10
	file := SharedFile{Hash: "expired1", OriginalName: "expired.txt", MimeType: "text/plain", OwnerID: 1, ExpiresAt: &expired, MaxDownloads: &max, Ready: true, Size: 1}
	if err := app.db.Create(&file).Error; err != nil {
		t.Fatal(err)
	}
	path := filepath.Join(app.filesDir, file.Hash)
	if err := os.WriteFile(path, []byte("x"), 0o640); err != nil {
		t.Fatal(err)
	}
	app.cleanup()
	var count int64
	app.db.Model(&SharedFile{}).Where("id = ?", file.ID).Count(&count)
	if count != 0 {
		t.Fatal("expired database record remains")
	}
	if _, err := os.Stat(path); !os.IsNotExist(err) {
		t.Fatal("expired file remains on disk")
	}
}

func TestBasicUploadLimitAndAdvancedInheritance(t *testing.T) {
	app, _ := testApp(t)
	passwordHash, _ := bcrypt.GenerateFromPassword([]byte("password-123"), bcrypt.DefaultCost)
	basic := User{Username: "basic", PasswordHash: string(passwordHash), CanUpload: true}
	if err := app.db.Create(&basic).Error; err != nil {
		t.Fatal(err)
	}
	expires := time.Now().Add(time.Hour)
	max := 1
	for i := 0; i < 10; i++ {
		file := SharedFile{Hash: "A" + strings.Repeat("0", 6) + string(rune('0'+i)), OriginalName: "x", MimeType: "text/plain", OwnerID: basic.ID, ExpiresAt: &expires, MaxDownloads: &max, Ready: true}
		if err := app.db.Create(&file).Error; err != nil {
			t.Fatal(err)
		}
	}
	router := app.router()
	cookie := loginAs(t, router, basic.Username, "password-123")
	body, contentType := uploadBody(t, map[string]string{"retention": "1h", "downloads": "1"}, "blocked.txt", []byte("x"))
	response := request(t, router, http.MethodPost, "/api/files", body, contentType, cookie)
	if response.Code != http.StatusForbidden || !strings.Contains(response.Body.String(), "10") {
		t.Fatalf("limit status = %d, body = %s", response.Code, response.Body.String())
	}

	advanced := User{Username: "advanced", PasswordHash: string(passwordHash), CanAdvancedUpload: true}
	if err := app.db.Create(&advanced).Error; err != nil {
		t.Fatal(err)
	}
	cookie = loginAs(t, router, advanced.Username, "password-123")
	body, contentType = uploadBody(t, map[string]string{"retention": "forever", "downloads": "unlimited"}, "advanced.txt", []byte("x"))
	response = request(t, router, http.MethodPost, "/api/files", body, contentType, cookie)
	if response.Code != http.StatusCreated {
		t.Fatalf("advanced upload status = %d, body = %s", response.Code, response.Body.String())
	}
}

func TestDeleteUserRemovesOwnedFilesAndProtectsIDOne(t *testing.T) {
	app, password := testApp(t)
	var admin User
	app.db.First(&admin, 1)
	app.db.Model(&admin).Update("must_change_password", false)
	user := User{Username: "doomed", PasswordHash: admin.PasswordHash}
	if err := app.db.Create(&user).Error; err != nil {
		t.Fatal(err)
	}
	file := SharedFile{Hash: "deadfile", OriginalName: "x", MimeType: "text/plain", OwnerID: user.ID, Ready: true}
	if err := app.db.Create(&file).Error; err != nil {
		t.Fatal(err)
	}
	if err := os.WriteFile(filepath.Join(app.filesDir, file.Hash), []byte("x"), 0o640); err != nil {
		t.Fatal(err)
	}
	router := app.router()
	cookie := loginAs(t, router, "admin", password)
	protected := request(t, router, http.MethodDelete, "/api/users/1", nil, "", cookie)
	if protected.Code != http.StatusForbidden {
		t.Fatalf("delete admin status = %d", protected.Code)
	}
	deleted := request(t, router, http.MethodDelete, "/api/users/"+jsonNumber(user.ID), nil, "", cookie)
	if deleted.Code != http.StatusNoContent {
		t.Fatalf("delete user status = %d, body = %s", deleted.Code, deleted.Body.String())
	}
	if _, err := os.Stat(filepath.Join(app.filesDir, file.Hash)); !os.IsNotExist(err) {
		t.Fatal("owned file remains")
	}
}

func TestUserManagerCannotGrantOrTakeOverPermissionsTheyDoNotHave(t *testing.T) {
	app, _ := testApp(t)
	passwordHash, _ := bcrypt.GenerateFromPassword([]byte("password-123"), bcrypt.DefaultCost)
	manager := User{Username: "manager", PasswordHash: string(passwordHash), CanManageUsers: true, CanUpload: true}
	higherUser := User{Username: "higher", PasswordHash: string(passwordHash), CanAdvancedUpload: true}
	basicUser := User{Username: "editable", PasswordHash: string(passwordHash), CanUpload: true}
	for _, user := range []*User{&manager, &higherUser, &basicUser} {
		if err := app.db.Create(user).Error; err != nil {
			t.Fatal(err)
		}
	}
	router := app.router()
	cookie := loginAs(t, router, manager.Username, "password-123")

	createBody, _ := json.Marshal(map[string]any{"username": "escalated", "canAdvancedUpload": true})
	created := request(t, router, http.MethodPost, "/api/users", bytes.NewReader(createBody), "application/json", cookie)
	if created.Code != http.StatusForbidden {
		t.Fatalf("create advanced user status = %d, body = %s", created.Code, created.Body.String())
	}

	updateBody, _ := json.Marshal(map[string]any{"username": basicUser.Username, "canUpload": true, "canAdvancedUpload": true})
	updated := request(t, router, http.MethodPut, "/api/users/"+jsonNumber(basicUser.ID), bytes.NewReader(updateBody), "application/json", cookie)
	if updated.Code != http.StatusForbidden {
		t.Fatalf("grant advanced permission status = %d, body = %s", updated.Code, updated.Body.String())
	}

	reset := request(t, router, http.MethodPost, "/api/users/"+jsonNumber(higherUser.ID)+"/reset-password", nil, "", cookie)
	if reset.Code != http.StatusForbidden {
		t.Fatalf("reset higher user status = %d, body = %s", reset.Code, reset.Body.String())
	}

	downgradeBody, _ := json.Marshal(map[string]any{"username": higherUser.Username, "canUpload": true})
	downgraded := request(t, router, http.MethodPut, "/api/users/"+jsonNumber(higherUser.ID), bytes.NewReader(downgradeBody), "application/json", cookie)
	if downgraded.Code != http.StatusForbidden {
		t.Fatalf("edit higher user status = %d, body = %s", downgraded.Code, downgraded.Body.String())
	}

	deleted := request(t, router, http.MethodDelete, "/api/users/"+jsonNumber(higherUser.ID), nil, "", cookie)
	if deleted.Code != http.StatusForbidden {
		t.Fatalf("delete higher user status = %d, body = %s", deleted.Code, deleted.Body.String())
	}

	createBody, _ = json.Marshal(map[string]any{"username": "allowed", "canUpload": true})
	created = request(t, router, http.MethodPost, "/api/users", bytes.NewReader(createBody), "application/json", cookie)
	if created.Code != http.StatusCreated {
		t.Fatalf("create user with owned permission status = %d, body = %s", created.Code, created.Body.String())
	}
}

func jsonNumber(id uint) string {
	data, _ := json.Marshal(id)
	return string(data)
}