Agent-System-Backend/internal/vision/vision_test.go

package vision

import (
	"context"
	"encoding/base64"
	"net"
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"

	agentllm "agent-native/runtime/llm"
	core "agent-system-core"
)

func TestCheckIP(t *testing.T) {
	blocked := []string{"127.0.0.1", "10.1.2.3", "192.168.1.1", "172.16.0.1", "169.254.169.254", "100.64.0.1", "::1", "fc00::1", "0.0.0.0"}
	for _, s := range blocked {
		if err := checkIP(net.ParseIP(s)); err == nil {
			t.Errorf("%s should be blocked", s)
		}
	}
	if err := checkIP(net.ParseIP("8.8.8.8")); err != nil {
		t.Fatal(err)
	}
	if err := checkIP(net.ParseIP("1.1.1.1")); err != nil {
		t.Fatal(err)
	}
}

func TestValidateImageURLRejectsUnsafeTargets(t *testing.T) {
	ctx := context.Background()
	for _, raw := range []string{
		"http://127.0.0.1/a.png",
		"http://169.254.169.254/latest/a.png",
		"http://10.0.0.8/a.jpg",
		"http://user:pass@8.8.8.8/a.png",
		"file:///tmp/a.png",
		"http://localhost/a.png",
	} {
		if err := validateImageURL(ctx, raw); err == nil {
			t.Errorf("accepted %s", raw)
		}
	}
	if err := validateImageURL(ctx, "https://8.8.8.8/a.png"); err != nil {
		t.Fatal(err)
	}
}

func TestFetchImageRejectsLoopbackServer(t *testing.T) {
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Type", "image/png")
		_, _ = w.Write(onePixelPNG)
	}))
	defer srv.Close()
	_, err := FetchImage(context.Background(), srv.URL+"/pixel.png")
	if err == nil || !strings.Contains(err.Error(), "允许范围") {
		t.Fatalf("err = %v", err)
	}
}

func TestImageLinksAndAttachNote(t *testing.T) {
	text := "看这张 http://127.0.0.1/a.png。还有 https://cdn.example.com/b.JPG?x=1 和重复 http://127.0.0.1/a.png"
	links := ImageLinks(text)
	if len(links) != 2 || links[0] != "http://127.0.0.1/a.png" || !strings.Contains(links[1], "b.JPG") {
		t.Fatalf("%#v", links)
	}
	msgs := []agentllm.Message{{Role: agentllm.RoleUser, Content: "请看 http://127.0.0.1/secret.png"}}
	notes := AttachLinkedImages(context.Background(), msgs)
	if len(notes) == 0 || len(msgs[0].Images) != 0 {
		t.Fatalf("notes=%v images=%d", notes, len(msgs[0].Images))
	}
	if !strings.Contains(msgs[0].Content, "地址不在允许范围") {
		t.Fatalf("content missing reason: %s", msgs[0].Content)
	}
}

func TestPrepareImageAndBudget(t *testing.T) {
	block, warn, ok := PrepareImage("pixel.png", "image/png", onePixelPNG)
	if !ok || !strings.Contains(warn, "按图像发送") || block.Type != core.BlockImageRef || block.Ref == "" {
		t.Fatalf("ok=%v warn=%q block=%+v", ok, warn, block)
	}
	if _, warn, ok := PrepareImage("icon.svg", "image/svg+xml", []byte("<svg/>")); ok || !strings.Contains(warn, "不作为图像发送") {
		t.Fatalf("svg ok=%v warn=%q", ok, warn)
	}
	big := make([]byte, MaxImageBytes+1)
	if _, warn, ok := PrepareImage("big.png", "image/png", big); ok || !strings.Contains(warn, "超过") {
		t.Fatalf("big ok=%v warn=%q", ok, warn)
	}
	if _, warn, ok := PrepareImage("bad.png", "image/png", []byte("not a png")); ok || !strings.Contains(warn, "不符") {
		t.Fatalf("bad ok=%v warn=%q", ok, warn)
	}

	img := agentllm.Image{MIME: "image/png", Data: block.Ref}
	kept, notes := Admit(nil, []agentllm.Image{img})
	if len(kept) != 1 || len(notes) != 0 {
		t.Fatalf("kept=%d notes=%v", len(kept), notes)
	}
	many := make([]agentllm.Image, MaxImages+1)
	for i := range many {
		many[i] = img
	}
	kept, notes = Admit(nil, many)
	if len(kept) != MaxImages || len(notes) == 0 {
		t.Fatalf("kept=%d notes=%v", len(kept), notes)
	}
	raw, err := base64.StdEncoding.DecodeString(block.Ref)
	if err != nil || len(raw) != len(onePixelPNG) {
		t.Fatal(err)
	}
}

// 1x1 PNG。
var onePixelPNG = []byte{
	0x89, 0x50, 0x4e, 0x47, 0x0d, 0x0a, 0x1a, 0x0a, 0x00, 0x00, 0x00, 0x0d, 0x49, 0x48, 0x44, 0x52,
	0x00, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00, 0x01, 0x08, 0x02, 0x00, 0x00, 0x00, 0x90, 0x77, 0x53,
	0xde, 0x00, 0x00, 0x00, 0x0c, 0x49, 0x44, 0x41, 0x54, 0x08, 0xd7, 0x63, 0xf8, 0xcf, 0xc0, 0x50,
	0x0f, 0x00, 0x04, 0x85, 0x01, 0x80, 0x84, 0xa9, 0x8c, 0x21, 0x00, 0x00, 0x00, 0x00, 0x49, 0x45,
	0x4e, 0x44, 0xae, 0x42, 0x60, 0x82,
}