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("")); 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, }