From e1941e78435eab82871989446235f9dd5c6fe143 Mon Sep 17 00:00:00 2001
From: Steven <stevenlgtm@gmail.com>
Date: Tue, 6 Jan 2026 20:37:04 +0800
Subject: [PATCH] fix: attachment type checks

---
 server/router/api/v1/attachment_service.go    | 16 ++++-
 .../api/v1/test/attachment_service_test.go    | 58 +++++++++++++++++++
 2 files changed, 73 insertions(+), 1 deletion(-)
 create mode 100644 server/router/api/v1/test/attachment_service_test.go

diff --git a/server/router/api/v1/attachment_service.go b/server/router/api/v1/attachment_service.go
index 3226b9f42..a87140f75 100644
--- a/server/router/api/v1/attachment_service.go
+++ b/server/router/api/v1/attachment_service.go
@@ -6,6 +6,8 @@ import (
 	"encoding/binary"
 	"fmt"
 	"io"
+	"mime"
+	"net/http"
 	"os"
 	"path/filepath"
 	"regexp"
@@ -63,7 +65,19 @@ func (s *APIV1Service) CreateAttachment(ctx context.Context, request *v1pb.Creat
 		return nil, status.Errorf(codes.InvalidArgument, "filename contains invalid characters or format")
 	}
 	if request.Attachment.Type == "" {
-		return nil, status.Errorf(codes.InvalidArgument, "type is required")
+		ext := filepath.Ext(request.Attachment.Filename)
+		mimeType := mime.TypeByExtension(ext)
+		if mimeType == "" {
+			mimeType = http.DetectContentType(request.Attachment.Content)
+		}
+		// ParseMediaType to strip parameters
+		mediaType, _, err := mime.ParseMediaType(mimeType)
+		if err == nil {
+			request.Attachment.Type = mediaType
+		}
+	}
+	if request.Attachment.Type == "" {
+		request.Attachment.Type = "application/octet-stream"
 	}
 	if !isValidMimeType(request.Attachment.Type) {
 		return nil, status.Errorf(codes.InvalidArgument, "invalid MIME type format")
diff --git a/server/router/api/v1/test/attachment_service_test.go b/server/router/api/v1/test/attachment_service_test.go
new file mode 100644
index 000000000..067df14b4
--- /dev/null
+++ b/server/router/api/v1/test/attachment_service_test.go
@@ -0,0 +1,58 @@
+package test
+
+import (
+	"context"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+	v1pb "github.com/usememos/memos/proto/gen/api/v1"
+)
+
+func TestCreateAttachment(t *testing.T) {
+	ts := NewTestService(t)
+	defer ts.Cleanup()
+	ctx := context.Background()
+
+	user, err := ts.CreateRegularUser(ctx, "test_user")
+	require.NoError(t, err)
+	userCtx := ts.CreateUserContext(ctx, user.ID)
+
+	// Test case 1: Create attachment with empty type but known extension
+	t.Run("EmptyType_KnownExtension", func(t *testing.T) {
+		attachment, err := ts.Service.CreateAttachment(userCtx, &v1pb.CreateAttachmentRequest{
+			Attachment: &v1pb.Attachment{
+				Filename: "test.png",
+				Content:  []byte("fake png content"),
+			},
+		})
+		require.NoError(t, err)
+		require.Equal(t, "image/png", attachment.Type)
+	})
+
+	// Test case 2: Create attachment with empty type and unknown extension, but detectable content
+	t.Run("EmptyType_UnknownExtension_ContentSniffing", func(t *testing.T) {
+		// PNG magic header: 89 50 4E 47 0D 0A 1A 0A
+		pngContent := []byte{0x89, 0x50, 0x4E, 0x47, 0x0D, 0x0A, 0x1A, 0x0A}
+		attachment, err := ts.Service.CreateAttachment(userCtx, &v1pb.CreateAttachmentRequest{
+			Attachment: &v1pb.Attachment{
+				Filename: "test.unknown",
+				Content:  pngContent,
+			},
+		})
+		require.NoError(t, err)
+		require.Equal(t, "image/png", attachment.Type)
+	})
+
+	// Test case 3: Empty type, unknown extension, random content -> fallback to application/octet-stream
+	t.Run("EmptyType_Fallback", func(t *testing.T) {
+		randomContent := []byte{0x00, 0x01, 0x02, 0x03}
+		attachment, err := ts.Service.CreateAttachment(userCtx, &v1pb.CreateAttachmentRequest{
+			Attachment: &v1pb.Attachment{
+				Filename: "test.data",
+				Content:  randomContent,
+			},
+		})
+		require.NoError(t, err)
+		require.Equal(t, "application/octet-stream", attachment.Type)
+	})
+}