From ea14280cb3cb138606dc71cbde94c66bbc598545 Mon Sep 17 00:00:00 2001
From: Johnny <yourselfhosted@gmail.com>
Date: Wed, 17 Dec 2025 08:58:43 +0800
Subject: [PATCH] feat: enhance attachment handling with MIME type validation

---
 server/router/api/v1/attachment_service.go | 15 +++++++
 server/router/fileserver/fileserver.go     | 48 ++++++++++++++++++++--
 2 files changed, 59 insertions(+), 4 deletions(-)

diff --git a/server/router/api/v1/attachment_service.go b/server/router/api/v1/attachment_service.go
index ab5149237..8dbc8d705 100644
--- a/server/router/api/v1/attachment_service.go
+++ b/server/router/api/v1/attachment_service.go
@@ -64,6 +64,9 @@ func (s *APIV1Service) CreateAttachment(ctx context.Context, request *v1pb.Creat
 	if request.Attachment.Type == "" {
 		return nil, status.Errorf(codes.InvalidArgument, "type is required")
 	}
+	if !isValidMimeType(request.Attachment.Type) {
+		return nil, status.Errorf(codes.InvalidArgument, "invalid MIME type format")
+	}
 
 	// Use provided attachment_id or generate a new one
 	attachmentUID := request.AttachmentId
@@ -457,3 +460,15 @@ func validateFilename(filename string) bool {
 
 	return true
 }
+
+func isValidMimeType(mimeType string) bool {
+	// Reject empty or excessively long MIME types
+	if mimeType == "" || len(mimeType) > 255 {
+		return false
+	}
+
+	// MIME type must match the pattern: type/subtype
+	// Allow common characters in MIME types per RFC 2045
+	matched, _ := regexp.MatchString(`^[a-zA-Z0-9][a-zA-Z0-9!#$&^_.+-]{0,126}/[a-zA-Z0-9][a-zA-Z0-9!#$&^_.+-]{0,126}$`, mimeType)
+	return matched
+}
diff --git a/server/router/fileserver/fileserver.go b/server/router/fileserver/fileserver.go
index 177754783..5407adbba 100644
--- a/server/router/fileserver/fileserver.go
+++ b/server/router/fileserver/fileserver.go
@@ -118,15 +118,39 @@ func (s *FileServerService) serveAttachmentFile(c echo.Context) error {
 		contentType += "; charset=utf-8"
 	}
 	// Prevent XSS attacks by serving potentially unsafe files as octet-stream
-	if strings.EqualFold(contentType, "image/svg+xml") ||
-		strings.EqualFold(contentType, "text/html") ||
-		strings.EqualFold(contentType, "application/xhtml+xml") {
-		contentType = "application/octet-stream"
+	unsafeTypes := []string{
+		"text/html",
+		"text/javascript",
+		"application/javascript",
+		"application/x-javascript",
+		"text/xml",
+		"application/xml",
+		"application/xhtml+xml",
+		"image/svg+xml",
+	}
+	for _, unsafeType := range unsafeTypes {
+		if strings.EqualFold(contentType, unsafeType) {
+			contentType = "application/octet-stream"
+			break
+		}
 	}
 
 	// Set common headers
 	c.Response().Header().Set("Content-Type", contentType)
 	c.Response().Header().Set("Cache-Control", "public, max-age=3600")
+	// Prevent MIME-type sniffing which could lead to XSS
+	c.Response().Header().Set("X-Content-Type-Options", "nosniff")
+	// Defense-in-depth: prevent embedding in frames and restrict content loading
+	c.Response().Header().Set("X-Frame-Options", "DENY")
+	c.Response().Header().Set("Content-Security-Policy", "default-src 'none'; style-src 'unsafe-inline';")
+
+	// Force download for non-media files to prevent XSS execution
+	if !strings.HasPrefix(contentType, "image/") &&
+		!strings.HasPrefix(contentType, "video/") &&
+		!strings.HasPrefix(contentType, "audio/") &&
+		contentType != "application/pdf" {
+		c.Response().Header().Set("Content-Disposition", fmt.Sprintf("attachment; filename=%q", attachment.Filename))
+	}
 
 	// For video/audio: Use http.ServeContent for automatic range request support
 	// This is critical for Safari which REQUIRES range request support
@@ -169,6 +193,18 @@ func (s *FileServerService) serveUserAvatar(c echo.Context) error {
 		return echo.NewHTTPError(http.StatusInternalServerError, "failed to extract image info").SetInternal(err)
 	}
 
+	// Validate avatar MIME type to prevent XSS
+	allowedAvatarTypes := map[string]bool{
+		"image/png":  true,
+		"image/jpeg": true,
+		"image/jpg":  true,
+		"image/gif":  true,
+		"image/webp": true,
+	}
+	if !allowedAvatarTypes[imageType] {
+		return echo.NewHTTPError(http.StatusBadRequest, "invalid avatar image type")
+	}
+
 	// Decode base64 data
 	imageData, err := base64.StdEncoding.DecodeString(base64Data)
 	if err != nil {
@@ -178,6 +214,10 @@ func (s *FileServerService) serveUserAvatar(c echo.Context) error {
 	// Set cache headers for avatars
 	c.Response().Header().Set("Content-Type", imageType)
 	c.Response().Header().Set("Cache-Control", "public, max-age=3600")
+	c.Response().Header().Set("X-Content-Type-Options", "nosniff")
+	// Defense-in-depth: prevent embedding in frames
+	c.Response().Header().Set("X-Frame-Options", "DENY")
+	c.Response().Header().Set("Content-Security-Policy", "default-src 'none'; style-src 'unsafe-inline';")
 
 	return c.Blob(http.StatusOK, imageType, imageData)
 }