MailUI4Agents/gateway/internal/blob/store_test.go

package blob

import (
	"bytes"
	"crypto/sha256"
	"encoding/hex"
	"errors"
	"io"
	"os"
	"path/filepath"
	"strings"
	"testing"
)

func newStore(t *testing.T) *Store {
	t.Helper()
	s, err := New(t.TempDir())
	if err != nil {
		t.Fatal(err)
	}
	return s
}

func TestPutAndOpenRoundTrip(t *testing.T) {
	s := newStore(t)
	data := []byte("附件内容 with bytes \x00\x01")

	sum, n, err := s.Put(bytes.NewReader(data), 0)
	if err != nil {
		t.Fatal(err)
	}
	if n != int64(len(data)) {
		t.Errorf("写入 %d 字节，报告 %d", len(data), n)
	}

	h := sha256.Sum256(data)
	if sum != hex.EncodeToString(h[:]) {
		t.Errorf("sha256 = %s，与内容不符", sum)
	}

	f, err := s.Open(sum)
	if err != nil {
		t.Fatal(err)
	}
	defer f.Close()
	got, _ := io.ReadAll(f)
	if !bytes.Equal(got, data) {
		t.Error("读回的内容与写入不一致")
	}
}

// 相同内容重复上传必须复用同一个文件，不占额外空间。
func TestPutDeduplicates(t *testing.T) {
	s := newStore(t)
	data := []byte("same content")

	sum1, _, err := s.Put(bytes.NewReader(data), 0)
	if err != nil {
		t.Fatal(err)
	}
	sum2, _, err := s.Put(bytes.NewReader(data), 0)
	if err != nil {
		t.Fatal(err)
	}
	if sum1 != sum2 {
		t.Fatalf("同内容得到不同哈希: %s vs %s", sum1, sum2)
	}

	// 目录里应当只有一个内容文件（外加两级目录）
	var files int
	filepath.Walk(s.Root(), func(_ string, info os.FileInfo, _ error) error {
		if info != nil && !info.IsDir() {
			files++
		}
		return nil
	})
	if files != 1 {
		t.Errorf("去重后应只剩 1 个文件，实际 %d", files)
	}
}

// 超限必须拒绝，且不能留下半个临时文件。
func TestPutTooLargeLeavesNoGarbage(t *testing.T) {
	s := newStore(t)
	data := bytes.Repeat([]byte("x"), 1024)

	_, _, err := s.Put(bytes.NewReader(data), 512)
	if !errors.Is(err, ErrTooLarge) {
		t.Fatalf("期望 ErrTooLarge，得到 %v", err)
	}

	entries, _ := os.ReadDir(s.Root())
	for _, e := range entries {
		if strings.HasPrefix(e.Name(), ".upload-") {
			t.Errorf("超限后残留临时文件 %s", e.Name())
		}
	}
}

// 恰好等于上限应当通过 —— 边界不能误杀。
func TestPutExactlyAtLimit(t *testing.T) {
	s := newStore(t)
	data := bytes.Repeat([]byte("y"), 512)

	if _, n, err := s.Put(bytes.NewReader(data), 512); err != nil {
		t.Fatalf("恰好等于上限被拒: %v", err)
	} else if n != 512 {
		t.Errorf("字节数 = %d，want 512", n)
	}
}

// 路径完全由 sha256 派生，任何非法 sum 都不能落到库外。
func TestPathTraversalRejected(t *testing.T) {
	s := newStore(t)

	for _, bad := range []string{
		"../../etc/passwd",
		"..",
		"/etc/passwd",
		"ABCDEF", // 大写非法
		strings.Repeat("g", 64), // 非十六进制
		strings.Repeat("a", 63), // 长度不足
		"",
	} {
		if _, err := s.pathFor(bad); err == nil {
			t.Errorf("pathFor(%q) 应报错", bad)
		}
		if _, err := s.Open(bad); err == nil {
			t.Errorf("Open(%q) 应报错", bad)
		}
		if s.Exists(bad) {
			t.Errorf("Exists(%q) 应为 false", bad)
		}
	}
}

// 生成的路径必须落在库根目录之内。
func TestPathStaysInsideRoot(t *testing.T) {
	s := newStore(t)
	sum := strings.Repeat("ab", 32)

	p, err := s.pathFor(sum)
	if err != nil {
		t.Fatal(err)
	}
	rel, err := filepath.Rel(s.Root(), p)
	if err != nil || strings.HasPrefix(rel, "..") {
		t.Errorf("路径逃出库根: %s", p)
	}
}

func TestRemove(t *testing.T) {
	s := newStore(t)
	sum, _, err := s.Put(bytes.NewReader([]byte("z")), 0)
	if err != nil {
		t.Fatal(err)
	}
	if !s.Exists(sum) {
		t.Fatal("写入后应存在")
	}
	if err := s.Remove(sum); err != nil {
		t.Fatal(err)
	}
	if s.Exists(sum) {
		t.Error("删除后仍存在")
	}
	// 重复删除应当幂等，不报错
	if err := s.Remove(sum); err != nil {
		t.Errorf("重复删除报错: %v", err)
	}
}