feat: implement favicon download functionality with content type validation

Author: 0xJacky

internal/sitecheck/checker.go

@@ -616,24 +616,77 @@ func (sc *SiteChecker) downloadFavicon(ctx context.Context, faviconURL string) s
 return ""
 }
 
-// Get content type
-contentType := resp.Header.Get("Content-Type")
-if contentType == "" {
-// Try to determine from URL extension
-if strings.HasSuffix(faviconURL, ".png") {
-contentType = "image/png"
-} else if strings.HasSuffix(faviconURL, ".ico") {
-contentType = "image/x-icon"
-} else {
-contentType = "image/x-icon" // default
-}
+headerContentType := normalizeContentType(resp.Header.Get("Content-Type"))
+inferredContentType := inferContentTypeFromURL(faviconURL)
+sniffedContentType := normalizeContentType(http.DetectContentType(body))
+
+contentType := headerContentType
+if !isAllowedFaviconContentType(contentType) && isAllowedFaviconContentType(sniffedContentType) {
+contentType = sniffedContentType
+}
+if !isAllowedFaviconContentType(contentType) &&
+headerContentType == "" &&
+isUnknownContentType(sniffedContentType) &&
+isAllowedFaviconContentType(inferredContentType) {
+contentType = inferredContentType
+}
+if !isAllowedFaviconContentType(contentType) {
+return ""
 }
 
-// Encode as data URL
 encoded := base64.StdEncoding.EncodeToString(body)
 return fmt.Sprintf("data:%s;base64,%s", contentType, encoded)
 }
 
+func normalizeContentType(contentType string) string {
+if contentType == "" {
+return ""
+}
+if semi := strings.Index(contentType, ";"); semi != -1 {
+contentType = contentType[:semi]
+}
+return strings.TrimSpace(strings.ToLower(contentType))
+}
+
+func inferContentTypeFromURL(faviconURL string) string {
+lower := strings.ToLower(faviconURL)
+switch {
+case strings.HasSuffix(lower, ".png"):
+return "image/png"
+case strings.HasSuffix(lower, ".jpg"), strings.HasSuffix(lower, ".jpeg"):
+return "image/jpeg"
+case strings.HasSuffix(lower, ".webp"):
+return "image/webp"
+case strings.HasSuffix(lower, ".gif"):
+return "image/gif"
+case strings.HasSuffix(lower, ".svg"):
+return "image/svg+xml"
+case strings.HasSuffix(lower, ".ico"):
+return "image/x-icon"
+default:
+return ""
+}
+}
+
+func isAllowedFaviconContentType(contentType string) bool {
+switch contentType {
+case "image/png",
+"image/jpeg",
+"image/webp",
+"image/gif",
+"image/svg+xml",
+"image/x-icon",
+"image/vnd.microsoft.icon":
+return true
+default:
+return false
+}
+}
+
+func isUnknownContentType(contentType string) bool {
+return contentType == "" || contentType == "application/octet-stream"
+}
+
 // generateDisplayURL generates the URL to display in UI based on health check protocol
 func generateDisplayURL(originalURL, protocol string) string {
 parsed, err := url.Parse(originalURL)

internal/sitecheck/checker_test.go

@@ -2,7 +2,11 @@ package sitecheck
 
 import (
 "context"
+"encoding/base64"
+"io"
 "net/http"
+"net/http/httptest"
+"strings"
 "testing"
 
 "github.com/0xJacky/Nginx-UI/model"
@@ -46,3 +50,72 @@ func TestCheckSiteSkipsNetworkWhenDisabled(t *testing.T) {
 t.Fatalf("CheckSite returned error: %v", err)
 }
 }
+
+func TestDownloadFaviconAcceptsValidImage(t *testing.T) {
+pngData := []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, 0x0A, 0x49, 0x44, 0x41,
+0x54, 0x08, 0xD7, 0x63, 0xF8, 0x0F, 0x04, 0x00,
+0x09, 0xFB, 0x03, 0xFD, 0xA7, 0x89, 0x81, 0xB9,
+0x00, 0x00, 0x00, 0x00, 0x49, 0x45, 0x4E, 0x44,
+0xAE, 0x42, 0x60, 0x82,
+}
+
+server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+w.Header().Set("Content-Type", "image/png")
+w.Write(pngData)
+}))
+defer server.Close()
+
+checker := NewSiteChecker(DefaultCheckOptions())
+dataURL := checker.downloadFavicon(context.Background(), server.URL+"/favicon.png")
+if dataURL == "" {
+t.Fatal("expected data URL for valid favicon")
+}
+
+expectedPrefix := "data:image/png;base64,"
+if !strings.HasPrefix(dataURL, expectedPrefix) {
+t.Fatalf("unexpected data URL prefix: %s", dataURL)
+}
+
+expectedPayload := base64.StdEncoding.EncodeToString(pngData)
+if payload := strings.TrimPrefix(dataURL, expectedPrefix); payload != expectedPayload {
+t.Fatalf("unexpected base64 payload: got %s want %s", payload, expectedPayload)
+}
+}
+
+func TestDownloadFaviconRejectsHTMLContent(t *testing.T) {
+server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+w.Header().Set("Content-Type", "text/html; charset=utf-8")
+io.WriteString(w, "<html><body>not an image</body></html>")
+}))
+defer server.Close()
+
+checker := NewSiteChecker(DefaultCheckOptions())
+dataURL := checker.downloadFavicon(context.Background(), server.URL+"/favicon.ico")
+if dataURL != "" {
+t.Fatalf("expected empty data URL for non-image content, got %s", dataURL)
+}
+}
+
+func TestDownloadFaviconRejectsHTMLContentWithoutHeader(t *testing.T) {
+checker := NewSiteChecker(DefaultCheckOptions())
+checker.client = &http.Client{
+Transport: roundTripFunc(func(req *http.Request) (*http.Response, error) {
+return &http.Response{
+StatusCode: http.StatusOK,
+Header: make(http.Header),
+Body: io.NopCloser(strings.NewReader("<html><body>not an image</body></html>")),
+Request: req,
+}, nil
+}),
+}
+
+dataURL := checker.downloadFavicon(context.Background(), "http://example.com/favicon.ico")
+if dataURL != "" {
+t.Fatalf("expected empty data URL when header missing and content sniffing rejects, got %s", dataURL)
+}
+}