Cyrene_v2/internal/runtime/multimodal/vision/limits_test.go

package vision

import (
	"bytes"
	"context"
	"encoding/base64"
	"errors"
	"strings"
	"testing"
)

// pngSignature 是 PNG 的 8 字节魔数（http.DetectContentType 的判据）。
// 体积口径类用例要在载荷头部带上它，才不会被内容嗅探当成文本否决。
var pngSignature = []byte("\x89PNG\r\n\x1a\n")

// TestInlineImageLengthBoundary 钉住长度预检的边界与错误文案：恰好等于
// 上限通过，多一个字符就报 ErrImageTooLarge。比较运算符一旦被挪位
// （>= / <），或口径被换成另一个数，这里立刻红。
func TestInlineImageLengthBoundary(t *testing.T) {
	t.Parallel()
	if err := checkInlineImageLength(strings.Repeat("A", MaxInlineImageChars)); err != nil {
		t.Fatalf("恰好等于上限被拒：%v", err)
	}
	err := checkInlineImageLength(strings.Repeat("A", MaxInlineImageChars+1))
	if !errors.Is(err, ErrImageTooLarge) {
		t.Fatalf("error = %v, want ErrImageTooLarge", err)
	}
	msg := err.Error()
	if !strings.Contains(msg, "上限") || !strings.Contains(msg, "699") {
		t.Fatalf("文案没有带上限：%v", err)
	}
	if strings.Contains(msg, "AAAA") {
		t.Fatalf("文案泄漏了载荷内容：%v", err)
	}
}

// TestMaxInlineImageCharsMatchesByteBudget 钉住换算口径：字符上限换算回
// 解码字节必须装得下一张满额 5 MiB 图，余量只能是 data: 前缀那点字符，
// 不能顺带塞进第二张图 —— 与网关入口 checkChatImageSize 同一口径。
func TestMaxInlineImageCharsMatchesByteBudget(t *testing.T) {
	t.Parallel()
	if got := MaxInlineImageChars / 4 * 3; got < MaxInlineImageBytes {
		t.Fatalf("字符上限只装得下 %d 字节，少于预算 %d", got, MaxInlineImageBytes)
	}
	if want := ((MaxInlineImageBytes+2)/3)*4 + InlineImagePrefixAllowance; MaxInlineImageChars != want {
		t.Fatalf("MaxInlineImageChars = %d, want %d（base64 膨胀 4/3 + 前缀余量）", MaxInlineImageChars, want)
	}
	if MaxInlineImageChars > ((MaxInlineImageBytes+2)/3)*4+4096 {
		t.Fatalf("余量过大（应只是 data: 前缀余量）：%d", MaxInlineImageChars)
	}
	if MaxInlineImageBytes != 5<<20 {
		t.Fatalf("MaxInlineImageBytes = %d, want 5 MiB（与上传/网关口徑同值）", MaxInlineImageBytes)
	}
}

// TestDecodeBase64_RejectsOversizeBeforeDecoding 钉住「预检先于解码」：
// 用一段既超限又非法 base64 的载荷，错误必须是 ErrImageTooLarge 而不是
// ErrBadImage —— 后者只可能来自 base64.StdEncoding.DecodeString，
// 也就是「先按约 4/3 分配再报错」，正是要避免的那一步。
func TestDecodeBase64_RejectsOversizeBeforeDecoding(t *testing.T) {
	t.Parallel()
	oversize := strings.Repeat("!", MaxInlineImageChars+1)
	if _, _, err := decodeBase64(oversize); !errors.Is(err, ErrImageTooLarge) {
		t.Fatalf("裸载荷 error = %v, want ErrImageTooLarge（预检没有先于解码）", err)
	}
	if _, _, err := decodeBase64("data:image/png;base64," + oversize); !errors.Is(err, ErrImageTooLarge) {
		t.Fatalf("data: URI error = %v, want ErrImageTooLarge", err)
	}
}

// TestInlineImageSrc_RejectsOversizeDataURI 覆盖短路路径：data: URI 在
// 原样返回**之前**也要过预检（这条路径不解码，此前完全不看体积），
// 而恰好等于上限的合法串必须逐字节原样返回、不受影响。
func TestInlineImageSrc_RejectsOversizeDataURI(t *testing.T) {
	t.Parallel()
	prefix := "data:image/png;base64,"
	oversize := prefix + strings.Repeat("A", MaxInlineImageChars)
	if _, err := inlineImageSrc(context.Background(), oversize); !errors.Is(err, ErrImageTooLarge) {
		t.Fatalf("error = %v, want ErrImageTooLarge", err)
	}

	exact := prefix + strings.Repeat("A", MaxInlineImageChars-len(prefix))
	got, err := inlineImageSrc(context.Background(), exact)
	if err != nil {
		t.Fatalf("恰好等于上限的 data: URI 被拒：%v", err)
	}
	if got != exact {
		t.Fatal("短路路径改动了原串（客户端 base64 路径必须逐字节不变）")
	}
}

// TestDecodeBase64_ExactBudgetImagePasses 用一张**恰好 5 MiB** 的载荷把口径
// 从字符数真的换算回解码字节数：满额图片编码后的字符串必须仍在字符上限内
// （否则网关入口与解码层对「一张图能有多大」给出了两个答案），解码结果
// 逐字节一致，绝不能误拒 —— 这正是「恰好 = 上限通过」的字节侧半边。
//
// 载荷头部带上 PNG 签名：本用例量的是体积口径，但 data: URI 解码后要过内容嗅探
// （image.go 的 checkInlinePayload），纯填充字节会被嗅成 text/plain 而否决 ——
// 那测的就是另一件事了。字节总长不变，故字符侧的口径换算不受影响。
//
// 判定口径是**字符数**（与网关入口一致，含 64 字符的 data: 前缀余量），
// 因此多出几十字节的载荷仍可能落在余量内通过；字符侧的 +1 越界由
// TestInlineImageLengthBoundary 钉住，余量本身不超过 64 字符
// （TestMaxInlineImageCharsMatchesByteBudget）。
func TestDecodeBase64_ExactBudgetImagePasses(t *testing.T) {
	t.Parallel()
	filler := bytes.Repeat([]byte{0x5a}, MaxInlineImageBytes-len(pngSignature))
	exact := append(append([]byte{}, pngSignature...), filler...)
	dataURI := DataURI("image/png", exact)
	if len(dataURI) > MaxInlineImageChars {
		t.Fatalf("满额图片编码后 %d 字符，超过上限 %d（口径不一致）", len(dataURI), MaxInlineImageChars)
	}
	mime, data, err := decodeBase64(dataURI)
	if err != nil {
		t.Fatalf("满额图片被拒：%v", err)
	}
	if mime != "image/png" || !bytes.Equal(data, exact) {
		t.Fatalf("满额图片解码结果不一致：mime=%q bytes=%d", mime, len(data))
	}
}

// TestDecodeBase64_SmallPayloadUnchanged 钉住预检不伤正常载荷：一张小图的
// data: URI 解码后 MIME 与字节都不变。
func TestDecodeBase64_SmallPayloadUnchanged(t *testing.T) {
	t.Parallel()
	png := pngBytes(t)
	mime, data, err := decodeBase64(DataURI("image/png", png))
	if err != nil {
		t.Fatalf("decodeBase64: %v", err)
	}
	if mime != "image/png" || !bytes.Equal(data, png) {
		t.Fatalf("decodeBase64 = (%q, %d bytes), want (image/png, %d bytes)", mime, len(data), len(png))
	}
	if _, _, err := decodeBase64(base64.StdEncoding.EncodeToString(png)); err != nil {
		t.Fatalf("裸 base64 被拒：%v", err)
	}
}