package vision import ( "bytes" "encoding/base64" "errors" "image" "image/color" "image/png" "testing" ) // pngBytes 返回一张真实的最小 PNG(1x1),使「字节嗅探」在测试里 // 有真实依据,而不是靠手写常量。 func pngBytes(t *testing.T) []byte { t.Helper() img := image.NewRGBA(image.Rect(0, 0, 1, 1)) img.Set(0, 0, color.RGBA{R: 0x11, G: 0x22, B: 0x33, A: 0xff}) var buf bytes.Buffer if err := png.Encode(&buf, img); err != nil { t.Fatalf("encode png: %v", err) } return buf.Bytes() } // htmlBytes 模拟网关/PROXY 的错误页:非图像响应体的典型形态。 func htmlBytes() []byte { return []byte("
502 Bad Gateway") } // webpBytes 构造一个满足 RIFF/WEBPVP8 嗅探签名的 WebP 头。 func webpBytes() []byte { return append([]byte("RIFF\x00\x00\x00\x00WEBPVP8 "), make([]byte, 16)...) } // bmpBytes 构造一个满足 "BM" 嗅探签名的 BMP 头。 func bmpBytes() []byte { return append([]byte("BM"), make([]byte, 32)...) } // TestImageMIMEFor 覆盖「内容类型不能瞎猜」的裁决表:声明的类型只作参考, // 真实字节嗅探是否决权,判不出图像一律报错(绝不静默回退 image/png)。 func TestImageMIMEFor(t *testing.T) { t.Parallel() tests := []struct { name string declared string data []byte wantMIME string wantErr bool }{ {name: "png with correct type", declared: "image/png", data: pngBytes(t), wantMIME: "image/png"}, {name: "png with charset param", declared: "image/png; charset=binary", data: pngBytes(t), wantMIME: "image/png"}, {name: "png with missing type sniffs", declared: "", data: pngBytes(t), wantMIME: "image/png"}, {name: "bytes win over a wrong declaration", declared: "image/jpg", data: pngBytes(t), wantMIME: "image/png"}, {name: "webp sniffs", declared: "", data: webpBytes(), wantMIME: "image/webp"}, {name: "bmp sniffs", declared: "", data: bmpBytes(), wantMIME: "image/bmp"}, {name: "html declared as html", declared: "text/html; charset=utf-8", data: htmlBytes(), wantErr: true}, {name: "html with missing type", declared: "", data: htmlBytes(), wantErr: true}, {name: "html falsely declared png", declared: "image/png", data: htmlBytes(), wantErr: true}, {name: "svg rejected", declared: "image/svg+xml", data: []byte(``), wantErr: true}, {name: "empty payload", declared: "", data: nil, wantErr: true}, {name: "unknown binary without type", declared: "", data: []byte{0x00, 0x01, 0x02, 0x03}, wantErr: true}, {name: "declared image with unsniffable bytes trusts the declaration", declared: "image/avif", data: []byte{0x00, 0x00, 0x00, 0x20}, wantMIME: "image/avif"}, {name: "declared jpg alias normalises in the fallback", declared: "image/jpg", data: []byte{0x00, 0x01, 0x02, 0x03}, wantMIME: "image/jpeg"}, } for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { t.Parallel() got, err := imageMIMEFor(tt.declared, tt.data) if tt.wantErr { if err == nil { t.Fatalf("imageMIMEFor(%q) = %q, want error", tt.declared, got) } if !errors.Is(err, ErrNotImage) { t.Fatalf("error = %v, want ErrNotImage", err) } return } if err != nil { t.Fatalf("imageMIMEFor(%q) error: %v", tt.declared, err) } if got != tt.wantMIME { t.Fatalf("mime = %q, want %q", got, tt.wantMIME) } }) } } // TestMimeByExtLeavesUnknownEmpty 钉住未知扩展名不再被猜成 png —— 交给 // 字节嗅探裁决。 func TestMimeByExtLeavesUnknownEmpty(t *testing.T) { t.Parallel() if got := mimeByExt("/tmp/notes.txt"); got != "" { t.Fatalf("mimeByExt(.txt) = %q, want empty (let the bytes decide)", got) } if got := mimeByExt("/tmp/cat.JPEG"); got != "image/jpeg" { t.Fatalf("mimeByExt(.JPEG) = %q, want image/jpeg", got) } if got := mimeByExt("/tmp/shot.bmp"); got != "image/bmp" { t.Fatalf("mimeByExt(.bmp) = %q, want image/bmp", got) } } // TestReadFileRejectsRenamedHTML 钉住本地路径也不喂错类型:把 HTML 改名成 // .png 不能再被当成 PNG 读进模型。 func TestReadFileRejectsRenamedHTML(t *testing.T) { t.Parallel() path := t.TempDir() + "/liar.png" if err := writeFileBytes(path, htmlBytes()); err != nil { t.Fatalf("write: %v", err) } mime, data, err := readFile(path) if err == nil { t.Fatalf("readFile = (%q, %d bytes), want error", mime, len(data)) } if !errors.Is(err, ErrNotImage) { t.Fatalf("error = %v, want ErrNotImage", err) } } // TestReadFileKeepsKnownImage 钉住正常本地图片不退化(含新增的 bmp)。 func TestReadFileKeepsKnownImage(t *testing.T) { t.Parallel() path := t.TempDir() + "/cat.png" want := pngBytes(t) if err := writeFileBytes(path, want); err != nil { t.Fatalf("write: %v", err) } mime, data, err := readFile(path) if err != nil { t.Fatalf("readFile: %v", err) } if mime != "image/png" { t.Fatalf("mime = %q, want image/png", mime) } if !bytes.Equal(data, want) { t.Fatal("bytes changed on read") } } // TestDataURIStable 是 base64 路径逐字节不变的底线。 func TestDataURIStable(t *testing.T) { t.Parallel() data := pngBytes(t) want := "data:image/png;base64," + base64.StdEncoding.EncodeToString(data) if got := DataURI("image/png", data); got != want { t.Fatalf("DataURI = %q, want %q", got, want) } }