feat: implement favicon download functionality with content type validation

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)
+	}
+}