mirror of https://github.com/usememos/memos
chore: harden MCP access control and origin validation
boojack
2 months ago
parent
7c708ee27e
commit
d3f6e8ee31
@ -0,0 +1,113 @@
|
||||
package mcp
|
||||
|
||||
import (
|
||||
"context"
|
||||
"net/http"
|
||||
"net/url"
|
||||
"strconv"
|
||||
"strings"
|
||||
|
||||
"github.com/pkg/errors"
|
||||
|
||||
"github.com/usememos/memos/store"
|
||||
)
|
||||
|
||||
// checkMemoAccess returns an error if the caller cannot read the memo.
|
||||
// userID == 0 means anonymous.
|
||||
func checkMemoAccess(memo *store.Memo, userID int32) error {
|
||||
if memo.RowStatus == store.Archived && memo.CreatorID != userID {
|
||||
return errors.New("permission denied")
|
||||
}
|
||||
|
||||
switch memo.Visibility {
|
||||
case store.Protected:
|
||||
if userID == 0 {
|
||||
return errors.New("permission denied")
|
||||
}
|
||||
case store.Private:
|
||||
if memo.CreatorID != userID {
|
||||
return errors.New("permission denied")
|
||||
}
|
||||
default:
|
||||
// store.Public and any unknown visibility: allow.
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func checkMemoOwnership(memo *store.Memo, userID int32) error {
|
||||
if memo.CreatorID != userID {
|
||||
return errors.New("permission denied")
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// applyVisibilityFilter restricts find to memos the caller may see.
|
||||
func applyVisibilityFilter(find *store.FindMemo, userID int32, rowStatus *store.RowStatus) {
|
||||
if rowStatus != nil && *rowStatus == store.Archived {
|
||||
if userID == 0 {
|
||||
impossibleCreatorID := int32(-1)
|
||||
find.CreatorID = &impossibleCreatorID
|
||||
return
|
||||
}
|
||||
find.CreatorID = &userID
|
||||
return
|
||||
}
|
||||
if userID == 0 {
|
||||
find.VisibilityList = []store.Visibility{store.Public}
|
||||
return
|
||||
}
|
||||
find.Filters = append(find.Filters, "creator_id == "+itoa32(userID)+` || visibility in ["PUBLIC", "PROTECTED"]`)
|
||||
}
|
||||
|
||||
func (s *MCPService) checkAttachmentAccess(ctx context.Context, attachment *store.Attachment, userID int32) error {
|
||||
if attachment.CreatorID == userID {
|
||||
return nil
|
||||
}
|
||||
if attachment.MemoID == nil {
|
||||
return errors.New("permission denied")
|
||||
}
|
||||
|
||||
memo, err := s.store.GetMemo(ctx, &store.FindMemo{ID: attachment.MemoID})
|
||||
if err != nil {
|
||||
return errors.Wrap(err, "failed to get linked memo")
|
||||
}
|
||||
if memo == nil {
|
||||
return errors.New("linked memo not found")
|
||||
}
|
||||
return checkMemoAccess(memo, userID)
|
||||
}
|
||||
|
||||
func (s *MCPService) isAllowedOrigin(r *http.Request) bool {
|
||||
origin := r.Header.Get("Origin")
|
||||
if origin == "" {
|
||||
return true
|
||||
}
|
||||
|
||||
originURL, err := url.Parse(origin)
|
||||
if err != nil || originURL.Scheme == "" || originURL.Host == "" {
|
||||
return false
|
||||
}
|
||||
|
||||
if sameOriginHost(originURL.Host, r.Host) {
|
||||
return true
|
||||
}
|
||||
|
||||
if s.profile.InstanceURL == "" {
|
||||
return false
|
||||
}
|
||||
|
||||
instanceURL, err := url.Parse(s.profile.InstanceURL)
|
||||
if err != nil || instanceURL.Scheme == "" || instanceURL.Host == "" {
|
||||
return false
|
||||
}
|
||||
|
||||
return strings.EqualFold(originURL.Scheme, instanceURL.Scheme) && sameOriginHost(originURL.Host, instanceURL.Host)
|
||||
}
|
||||
|
||||
func sameOriginHost(a, b string) bool {
|
||||
return strings.EqualFold(a, b)
|
||||
}
|
||||
|
||||
func itoa32(v int32) string {
|
||||
return strconv.FormatInt(int64(v), 10)
|
||||
}
|
||||
@ -0,0 +1,275 @@
|
||||
package mcp
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"net/http/httptest"
|
||||
"testing"
|
||||
|
||||
"github.com/lithammer/shortuuid/v4"
|
||||
"github.com/mark3labs/mcp-go/mcp"
|
||||
"github.com/stretchr/testify/require"
|
||||
|
||||
"github.com/usememos/memos/internal/profile"
|
||||
storepb "github.com/usememos/memos/proto/gen/store"
|
||||
"github.com/usememos/memos/server/auth"
|
||||
"github.com/usememos/memos/store"
|
||||
teststore "github.com/usememos/memos/store/test"
|
||||
)
|
||||
|
||||
type testMCPService struct {
|
||||
service *MCPService
|
||||
store *store.Store
|
||||
}
|
||||
|
||||
func newTestMCPService(t *testing.T) *testMCPService {
|
||||
t.Helper()
|
||||
|
||||
ctx := context.Background()
|
||||
stores := teststore.NewTestingStore(ctx, t)
|
||||
t.Cleanup(func() {
|
||||
require.NoError(t, stores.Close())
|
||||
})
|
||||
|
||||
svc := NewMCPService(&profile.Profile{
|
||||
Driver: "sqlite",
|
||||
InstanceURL: "https://notes.example.com",
|
||||
}, stores, "test-secret")
|
||||
return &testMCPService{
|
||||
service: svc,
|
||||
store: stores,
|
||||
}
|
||||
}
|
||||
|
||||
func (s *testMCPService) createUser(t *testing.T, username string) *store.User {
|
||||
t.Helper()
|
||||
|
||||
user, err := s.store.CreateUser(context.Background(), &store.User{
|
||||
Username: username,
|
||||
Role: store.RoleUser,
|
||||
Email: username + "@example.com",
|
||||
})
|
||||
require.NoError(t, err)
|
||||
return user
|
||||
}
|
||||
|
||||
func (s *testMCPService) createMemo(t *testing.T, creatorID int32, visibility store.Visibility, content string) *store.Memo {
|
||||
t.Helper()
|
||||
|
||||
memo, err := s.store.CreateMemo(context.Background(), &store.Memo{
|
||||
UID: shortuuid.New(),
|
||||
CreatorID: creatorID,
|
||||
RowStatus: store.Normal,
|
||||
Visibility: visibility,
|
||||
Content: content,
|
||||
})
|
||||
require.NoError(t, err)
|
||||
return memo
|
||||
}
|
||||
|
||||
func (s *testMCPService) archiveMemo(t *testing.T, memoID int32) {
|
||||
t.Helper()
|
||||
|
||||
rowStatus := store.Archived
|
||||
require.NoError(t, s.store.UpdateMemo(context.Background(), &store.UpdateMemo{
|
||||
ID: memoID,
|
||||
RowStatus: &rowStatus,
|
||||
}))
|
||||
}
|
||||
|
||||
func (s *testMCPService) createAttachment(t *testing.T, creatorID int32, memoID *int32) *store.Attachment {
|
||||
t.Helper()
|
||||
|
||||
attachment, err := s.store.CreateAttachment(context.Background(), &store.Attachment{
|
||||
UID: shortuuid.New(),
|
||||
CreatorID: creatorID,
|
||||
Filename: "note.txt",
|
||||
Type: "text/plain",
|
||||
Size: 4,
|
||||
StorageType: storepb.AttachmentStorageType_ATTACHMENT_STORAGE_TYPE_UNSPECIFIED,
|
||||
Reference: "db://attachment/note.txt",
|
||||
MemoID: memoID,
|
||||
})
|
||||
require.NoError(t, err)
|
||||
return attachment
|
||||
}
|
||||
|
||||
func withUser(ctx context.Context, userID int32) context.Context {
|
||||
return context.WithValue(ctx, auth.UserIDContextKey, userID)
|
||||
}
|
||||
|
||||
func toolRequest(name string, arguments map[string]any) mcp.CallToolRequest {
|
||||
return mcp.CallToolRequest{
|
||||
Params: mcp.CallToolParams{
|
||||
Name: name,
|
||||
Arguments: arguments,
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
func firstText(t *testing.T, result *mcp.CallToolResult) string {
|
||||
t.Helper()
|
||||
require.NotEmpty(t, result.Content)
|
||||
text, ok := result.Content[0].(mcp.TextContent)
|
||||
require.True(t, ok)
|
||||
return text.Text
|
||||
}
|
||||
|
||||
func TestHandleGetMemoAndReadResourceDenyArchivedMemoToNonCreator(t *testing.T) {
|
||||
ts := newTestMCPService(t)
|
||||
owner := ts.createUser(t, "owner")
|
||||
other := ts.createUser(t, "other")
|
||||
|
||||
memo := ts.createMemo(t, owner.ID, store.Public, "archived")
|
||||
ts.archiveMemo(t, memo.ID)
|
||||
|
||||
ctx := withUser(context.Background(), other.ID)
|
||||
result, err := ts.service.handleGetMemo(ctx, toolRequest("get_memo", map[string]any{
|
||||
"name": "memos/" + memo.UID,
|
||||
}))
|
||||
require.NoError(t, err)
|
||||
require.True(t, result.IsError)
|
||||
require.Contains(t, firstText(t, result), "permission denied")
|
||||
|
||||
_, err = ts.service.handleReadMemoResource(ctx, mcp.ReadResourceRequest{
|
||||
Params: mcp.ReadResourceParams{
|
||||
URI: "memo://memos/" + memo.UID,
|
||||
},
|
||||
})
|
||||
require.ErrorContains(t, err, "permission denied")
|
||||
}
|
||||
|
||||
func TestHandleListMemosArchivedOnlyReturnsCreatorMemos(t *testing.T) {
|
||||
ts := newTestMCPService(t)
|
||||
owner := ts.createUser(t, "owner")
|
||||
other := ts.createUser(t, "other")
|
||||
|
||||
ownerMemo := ts.createMemo(t, owner.ID, store.Public, "owner archived")
|
||||
ts.archiveMemo(t, ownerMemo.ID)
|
||||
otherMemo := ts.createMemo(t, other.ID, store.Public, "other archived")
|
||||
ts.archiveMemo(t, otherMemo.ID)
|
||||
|
||||
result, err := ts.service.handleListMemos(withUser(context.Background(), owner.ID), toolRequest("list_memos", map[string]any{
|
||||
"state": "ARCHIVED",
|
||||
}))
|
||||
require.NoError(t, err)
|
||||
require.False(t, result.IsError)
|
||||
|
||||
var payload struct {
|
||||
Memos []memoJSON `json:"memos"`
|
||||
}
|
||||
require.NoError(t, json.Unmarshal([]byte(firstText(t, result)), &payload))
|
||||
require.Len(t, payload.Memos, 1)
|
||||
require.Equal(t, "memos/"+ownerMemo.UID, payload.Memos[0].Name)
|
||||
|
||||
anonResult, err := ts.service.handleListMemos(context.Background(), toolRequest("list_memos", map[string]any{
|
||||
"state": "ARCHIVED",
|
||||
}))
|
||||
require.NoError(t, err)
|
||||
require.NoError(t, json.Unmarshal([]byte(firstText(t, anonResult)), &payload))
|
||||
require.Empty(t, payload.Memos)
|
||||
}
|
||||
|
||||
func TestHandleListMemoRelationsFiltersUnreadableTargets(t *testing.T) {
|
||||
ts := newTestMCPService(t)
|
||||
owner := ts.createUser(t, "owner")
|
||||
privateUser := ts.createUser(t, "private-user")
|
||||
publicUser := ts.createUser(t, "public-user")
|
||||
|
||||
source := ts.createMemo(t, owner.ID, store.Public, "source")
|
||||
privateTarget := ts.createMemo(t, privateUser.ID, store.Private, "private")
|
||||
publicTarget := ts.createMemo(t, publicUser.ID, store.Public, "public")
|
||||
|
||||
_, err := ts.store.UpsertMemoRelation(context.Background(), &store.MemoRelation{
|
||||
MemoID: source.ID,
|
||||
RelatedMemoID: privateTarget.ID,
|
||||
Type: store.MemoRelationReference,
|
||||
})
|
||||
require.NoError(t, err)
|
||||
_, err = ts.store.UpsertMemoRelation(context.Background(), &store.MemoRelation{
|
||||
MemoID: source.ID,
|
||||
RelatedMemoID: publicTarget.ID,
|
||||
Type: store.MemoRelationReference,
|
||||
})
|
||||
require.NoError(t, err)
|
||||
|
||||
result, err := ts.service.handleListMemoRelations(context.Background(), toolRequest("list_memo_relations", map[string]any{
|
||||
"name": "memos/" + source.UID,
|
||||
}))
|
||||
require.NoError(t, err)
|
||||
require.False(t, result.IsError)
|
||||
|
||||
var relations []relationJSON
|
||||
require.NoError(t, json.Unmarshal([]byte(firstText(t, result)), &relations))
|
||||
require.Len(t, relations, 1)
|
||||
require.Equal(t, "memos/"+publicTarget.UID, relations[0].RelatedMemo)
|
||||
|
||||
denied, err := ts.service.handleListMemoRelations(context.Background(), toolRequest("list_memo_relations", map[string]any{
|
||||
"name": "memos/" + privateTarget.UID,
|
||||
}))
|
||||
require.NoError(t, err)
|
||||
require.True(t, denied.IsError)
|
||||
require.Contains(t, firstText(t, denied), "permission denied")
|
||||
}
|
||||
|
||||
func TestHandleLinkAttachmentToMemoRequiresMemoOwnership(t *testing.T) {
|
||||
ts := newTestMCPService(t)
|
||||
attachmentOwner := ts.createUser(t, "attachment-owner")
|
||||
memoOwner := ts.createUser(t, "memo-owner")
|
||||
|
||||
attachment := ts.createAttachment(t, attachmentOwner.ID, nil)
|
||||
memo := ts.createMemo(t, memoOwner.ID, store.Public, "target")
|
||||
|
||||
result, err := ts.service.handleLinkAttachmentToMemo(withUser(context.Background(), attachmentOwner.ID), toolRequest("link_attachment_to_memo", map[string]any{
|
||||
"name": "attachments/" + attachment.UID,
|
||||
"memo": "memos/" + memo.UID,
|
||||
}))
|
||||
require.NoError(t, err)
|
||||
require.True(t, result.IsError)
|
||||
require.Contains(t, firstText(t, result), "permission denied")
|
||||
}
|
||||
|
||||
func TestHandleGetAttachmentDeniesArchivedLinkedMemoToNonCreator(t *testing.T) {
|
||||
ts := newTestMCPService(t)
|
||||
owner := ts.createUser(t, "owner")
|
||||
other := ts.createUser(t, "other")
|
||||
|
||||
memo := ts.createMemo(t, owner.ID, store.Public, "memo")
|
||||
ts.archiveMemo(t, memo.ID)
|
||||
attachment := ts.createAttachment(t, owner.ID, &memo.ID)
|
||||
|
||||
result, err := ts.service.handleGetAttachment(withUser(context.Background(), other.ID), toolRequest("get_attachment", map[string]any{
|
||||
"name": "attachments/" + attachment.UID,
|
||||
}))
|
||||
require.NoError(t, err)
|
||||
require.True(t, result.IsError)
|
||||
require.Contains(t, firstText(t, result), "permission denied")
|
||||
}
|
||||
|
||||
func TestIsAllowedOrigin(t *testing.T) {
|
||||
ts := newTestMCPService(t)
|
||||
|
||||
t.Run("allow missing origin", func(t *testing.T) {
|
||||
req := httptest.NewRequest("POST", "http://localhost:5230/mcp", nil)
|
||||
require.True(t, ts.service.isAllowedOrigin(req))
|
||||
})
|
||||
|
||||
t.Run("allow same origin as request host", func(t *testing.T) {
|
||||
req := httptest.NewRequest("POST", "http://localhost:5230/mcp", nil)
|
||||
req.Header.Set("Origin", "http://localhost:5230")
|
||||
require.True(t, ts.service.isAllowedOrigin(req))
|
||||
})
|
||||
|
||||
t.Run("allow configured instance origin", func(t *testing.T) {
|
||||
req := httptest.NewRequest("POST", "http://127.0.0.1:5230/mcp", nil)
|
||||
req.Host = "127.0.0.1:5230"
|
||||
req.Header.Set("Origin", "https://notes.example.com")
|
||||
require.True(t, ts.service.isAllowedOrigin(req))
|
||||
})
|
||||
|
||||
t.Run("reject cross origin", func(t *testing.T) {
|
||||
req := httptest.NewRequest("POST", "http://localhost:5230/mcp", nil)
|
||||
req.Header.Set("Origin", "https://evil.example.com")
|
||||
require.False(t, ts.service.isAllowedOrigin(req))
|
||||
})
|
||||
}
|
||||
Loading…
Reference in New Issue