Skip to content

Commit e1941e7

Browse files
committed
fix: attachment type checks
1 parent 874a4a7 commit e1941e7

File tree

2 files changed

+73
-1
lines changed

2 files changed

+73
-1
lines changed

server/router/api/v1/attachment_service.go

Lines changed: 15 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,8 @@ import (
66
"encoding/binary"
77
"fmt"
88
"io"
9+
"mime"
10+
"net/http"
911
"os"
1012
"path/filepath"
1113
"regexp"
@@ -63,7 +65,19 @@ func (s *APIV1Service) CreateAttachment(ctx context.Context, request *v1pb.Creat
6365
return nil, status.Errorf(codes.InvalidArgument, "filename contains invalid characters or format")
6466
}
6567
if request.Attachment.Type == "" {
66-
return nil, status.Errorf(codes.InvalidArgument, "type is required")
68+
ext := filepath.Ext(request.Attachment.Filename)
69+
mimeType := mime.TypeByExtension(ext)
70+
if mimeType == "" {
71+
mimeType = http.DetectContentType(request.Attachment.Content)
72+
}
73+
// ParseMediaType to strip parameters
74+
mediaType, _, err := mime.ParseMediaType(mimeType)
75+
if err == nil {
76+
request.Attachment.Type = mediaType
77+
}
78+
}
79+
if request.Attachment.Type == "" {
80+
request.Attachment.Type = "application/octet-stream"
6781
}
6882
if !isValidMimeType(request.Attachment.Type) {
6983
return nil, status.Errorf(codes.InvalidArgument, "invalid MIME type format")
Lines changed: 58 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,58 @@
1+
package test
2+
3+
import (
4+
"context"
5+
"testing"
6+
7+
"github.com/stretchr/testify/require"
8+
v1pb "github.com/usememos/memos/proto/gen/api/v1"
9+
)
10+
11+
func TestCreateAttachment(t *testing.T) {
12+
ts := NewTestService(t)
13+
defer ts.Cleanup()
14+
ctx := context.Background()
15+
16+
user, err := ts.CreateRegularUser(ctx, "test_user")
17+
require.NoError(t, err)
18+
userCtx := ts.CreateUserContext(ctx, user.ID)
19+
20+
// Test case 1: Create attachment with empty type but known extension
21+
t.Run("EmptyType_KnownExtension", func(t *testing.T) {
22+
attachment, err := ts.Service.CreateAttachment(userCtx, &v1pb.CreateAttachmentRequest{
23+
Attachment: &v1pb.Attachment{
24+
Filename: "test.png",
25+
Content: []byte("fake png content"),
26+
},
27+
})
28+
require.NoError(t, err)
29+
require.Equal(t, "image/png", attachment.Type)
30+
})
31+
32+
// Test case 2: Create attachment with empty type and unknown extension, but detectable content
33+
t.Run("EmptyType_UnknownExtension_ContentSniffing", func(t *testing.T) {
34+
// PNG magic header: 89 50 4E 47 0D 0A 1A 0A
35+
pngContent := []byte{0x89, 0x50, 0x4E, 0x47, 0x0D, 0x0A, 0x1A, 0x0A}
36+
attachment, err := ts.Service.CreateAttachment(userCtx, &v1pb.CreateAttachmentRequest{
37+
Attachment: &v1pb.Attachment{
38+
Filename: "test.unknown",
39+
Content: pngContent,
40+
},
41+
})
42+
require.NoError(t, err)
43+
require.Equal(t, "image/png", attachment.Type)
44+
})
45+
46+
// Test case 3: Empty type, unknown extension, random content -> fallback to application/octet-stream
47+
t.Run("EmptyType_Fallback", func(t *testing.T) {
48+
randomContent := []byte{0x00, 0x01, 0x02, 0x03}
49+
attachment, err := ts.Service.CreateAttachment(userCtx, &v1pb.CreateAttachmentRequest{
50+
Attachment: &v1pb.Attachment{
51+
Filename: "test.data",
52+
Content: randomContent,
53+
},
54+
})
55+
require.NoError(t, err)
56+
require.Equal(t, "application/octet-stream", attachment.Type)
57+
})
58+
}

0 commit comments

Comments
 (0)