mirror of
https://github.com/coder/coder.git
synced 2026-09-24 15:04:27 +08:00
fix: allow agents to attach any file type (#26560)
Agents can now attach any file type as a downloadable chat artifact, where previously the stored-file allowlist rejected types like `.zip`. The reason arbitrary types were blocked is that a single media-type list (`codersdk.AllChatAttachmentMediaTypes`) was doing three different jobs at once: gating what users may upload as prompt input, deciding what is safe to render inline in the browser, and admitting what the agent's `attach_file` could store. Because the agent storage path reused that same list as an admission gate, any artifact outside it was rejected even though agent artifacts are only ever downloaded by the user and are never forwarded to the model, so the prompt-input and inline-render constraints did not actually apply to them. This splits those concerns. `PrepareStoredFile` now only normalizes the name and classifies the bytes, and the prompt-input allowlist is enforced inline at `postChatFile` instead, which is the correct layer for user-provided input. User uploads are unchanged and still limited to the allowed prompt-input media types, and unsafe or unknown types remain download-only because `IsInlineRenderableStoredMediaType` still refuses to render them inline. Model replay is also unchanged: assistant and tool attachments are never forwarded to the LLM. Closes CODAGT-654
This commit is contained in:
+16
-8
@@ -6165,10 +6165,10 @@ func (api *API) postChatFile(rw http.ResponseWriter, r *http.Request) {
|
||||
}
|
||||
// application/octet-stream means the client could not classify the file
|
||||
// ahead of time, so we defer to byte classification below.
|
||||
if contentType != "application/octet-stream" && !chatfiles.IsAllowedStoredMediaType(contentType) {
|
||||
if contentType != "application/octet-stream" && !chatfiles.IsAllowedPromptInputMediaType(contentType) {
|
||||
httpapi.Write(ctx, rw, http.StatusBadRequest, codersdk.Response{
|
||||
Message: "Unsupported file type.",
|
||||
Detail: fmt.Sprintf("Allowed types: %s.", chatfiles.AllowedStoredMediaTypesString()),
|
||||
Detail: fmt.Sprintf("Allowed types: %s.", chatfiles.AllowedPromptInputMediaTypesString()),
|
||||
})
|
||||
return
|
||||
}
|
||||
@@ -6199,7 +6199,7 @@ func (api *API) postChatFile(rw http.ResponseWriter, r *http.Request) {
|
||||
return
|
||||
}
|
||||
|
||||
// Verify the actual content matches an allowed file type so that
|
||||
// Classify the actual content before applying the upload policy so
|
||||
// a client cannot spoof Content-Type to serve active content.
|
||||
filename, detected, err := chatfiles.PrepareStoredFile(filename, filename, data)
|
||||
if err != nil {
|
||||
@@ -6209,11 +6209,6 @@ func (api *API) postChatFile(rw http.ResponseWriter, r *http.Request) {
|
||||
Message: "Filename is required.",
|
||||
Detail: "Provide a filename in the Content-Disposition header.",
|
||||
})
|
||||
case errors.Is(err, chatfiles.ErrUnsupportedStoredFileType):
|
||||
httpapi.Write(ctx, rw, http.StatusBadRequest, codersdk.Response{
|
||||
Message: "Unsupported file type.",
|
||||
Detail: fmt.Sprintf("Allowed types: %s.", chatfiles.AllowedStoredMediaTypesString()),
|
||||
})
|
||||
default:
|
||||
httpapi.Write(ctx, rw, http.StatusBadRequest, codersdk.Response{
|
||||
Message: "Invalid file.",
|
||||
@@ -6222,6 +6217,13 @@ func (api *API) postChatFile(rw http.ResponseWriter, r *http.Request) {
|
||||
}
|
||||
return
|
||||
}
|
||||
if !chatfiles.IsAllowedPromptInputMediaType(detected) {
|
||||
httpapi.Write(ctx, rw, http.StatusBadRequest, codersdk.Response{
|
||||
Message: "Unsupported file type.",
|
||||
Detail: fmt.Sprintf("Allowed types: %s.", chatfiles.AllowedPromptInputMediaTypesString()),
|
||||
})
|
||||
return
|
||||
}
|
||||
// The compatibility check below is security-critical: it keeps exact
|
||||
// media-type matching by default while allowing application/
|
||||
// octet-stream uploads to defer to byte classification, and letting
|
||||
@@ -6370,6 +6372,12 @@ func createChatInputFromParts(
|
||||
Detail: fmt.Sprintf("Failed to retrieve file for %s[%d].", fieldName, i),
|
||||
}
|
||||
}
|
||||
if !chatfiles.IsAllowedPromptInputMediaType(chatFile.Mimetype) {
|
||||
return nil, "", nil, &codersdk.Response{
|
||||
Message: "Invalid input part.",
|
||||
Detail: fmt.Sprintf("%s[%d].file_id references a file type that cannot be used as prompt input. Allowed types: %s.", fieldName, i, chatfiles.AllowedPromptInputMediaTypesString()),
|
||||
}
|
||||
}
|
||||
content = append(content, codersdk.ChatMessageFile(part.FileID, chatFile.Mimetype, chatFile.Name))
|
||||
fileIDs = append(fileIDs, part.FileID)
|
||||
// file-reference parts carry inline code snippets, not uploaded
|
||||
|
||||
@@ -7438,6 +7438,51 @@ func TestChatMessageWithFiles(t *testing.T) {
|
||||
require.Contains(t, sdkErr.Detail, "does not exist")
|
||||
})
|
||||
|
||||
t.Run("UnsupportedPromptInputFileType", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
ctx := testutil.Context(t, testutil.WaitLong)
|
||||
client, store := newChatClientWithDatabase(t)
|
||||
firstUser := coderdtest.CreateFirstUser(t, client.Client)
|
||||
_ = createChatModelConfig(t, client)
|
||||
|
||||
chat, err := client.CreateChat(ctx, codersdk.CreateChatRequest{
|
||||
OrganizationID: firstUser.OrganizationID,
|
||||
Content: []codersdk.ChatInputPart{
|
||||
{Type: codersdk.ChatInputPartTypeText, Text: "initial message"},
|
||||
},
|
||||
})
|
||||
require.NoError(t, err)
|
||||
|
||||
//nolint:gocritic // Using AsChatd to mimic an agent-created artifact.
|
||||
chatdCtx := dbauthz.AsChatd(ctx)
|
||||
fileRow, err := store.InsertChatFile(chatdCtx, database.InsertChatFileParams{
|
||||
OwnerID: firstUser.UserID,
|
||||
OrganizationID: firstUser.OrganizationID,
|
||||
Name: "artifact.zip",
|
||||
Mimetype: "application/zip",
|
||||
Data: []byte("zip data"),
|
||||
})
|
||||
require.NoError(t, err)
|
||||
rejected, err := store.LinkChatFiles(chatdCtx, database.LinkChatFilesParams{
|
||||
ChatID: chat.ID,
|
||||
MaxFileLinks: int32(codersdk.MaxChatFileIDs),
|
||||
FileIds: []uuid.UUID{fileRow.ID},
|
||||
})
|
||||
require.NoError(t, err)
|
||||
require.Zero(t, rejected)
|
||||
|
||||
_, err = client.CreateChatMessage(ctx, chat.ID, codersdk.CreateChatMessageRequest{
|
||||
Content: []codersdk.ChatInputPart{
|
||||
{Type: codersdk.ChatInputPartTypeFile, FileID: fileRow.ID},
|
||||
},
|
||||
})
|
||||
sdkErr := requireSDKError(t, err, http.StatusBadRequest)
|
||||
require.Equal(t, "Invalid input part.", sdkErr.Message)
|
||||
require.Contains(t, sdkErr.Detail, "cannot be used as prompt input")
|
||||
require.Contains(t, sdkErr.Detail, "application/json")
|
||||
})
|
||||
|
||||
t.Run("FilesLinkedOnSend", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
@@ -10067,6 +10112,42 @@ func TestGetChatFile(t *testing.T) {
|
||||
require.Equal(t, "report.pdf", params["filename"])
|
||||
})
|
||||
|
||||
t.Run("AgentArtifactZipServedAsAttachment", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
ctx := testutil.Context(t, testutil.WaitLong)
|
||||
client, store := newChatClientWithDatabase(t)
|
||||
firstUser := coderdtest.CreateFirstUser(t, client.Client)
|
||||
|
||||
data := []byte("zip data")
|
||||
//nolint:gocritic // Using AsChatd to mimic an agent-created artifact.
|
||||
chatdCtx := dbauthz.AsChatd(ctx)
|
||||
row, err := store.InsertChatFile(chatdCtx, database.InsertChatFileParams{
|
||||
OwnerID: firstUser.UserID,
|
||||
OrganizationID: firstUser.OrganizationID,
|
||||
Name: "artifact.zip",
|
||||
Mimetype: "application/zip",
|
||||
Data: data,
|
||||
})
|
||||
require.NoError(t, err)
|
||||
|
||||
res, err := client.Request(ctx, http.MethodGet,
|
||||
fmt.Sprintf("/api/experimental/chats/files/%s", row.ID), nil)
|
||||
require.NoError(t, err)
|
||||
defer res.Body.Close()
|
||||
require.Equal(t, http.StatusOK, res.StatusCode)
|
||||
require.Equal(t, "application/zip", res.Header.Get("Content-Type"))
|
||||
require.Equal(t, "nosniff", res.Header.Get("X-Content-Type-Options"))
|
||||
|
||||
disposition, params, err := mime.ParseMediaType(res.Header.Get("Content-Disposition"))
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, "attachment", disposition)
|
||||
require.Equal(t, "artifact.zip", params["filename"])
|
||||
|
||||
got, err := io.ReadAll(res.Body)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, data, got)
|
||||
})
|
||||
|
||||
t.Run("LongFilename", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
ctx := testutil.Context(t, testutil.WaitLong)
|
||||
|
||||
@@ -12,7 +12,6 @@ import (
|
||||
"github.com/coder/coder/v2/coderd/database"
|
||||
"github.com/coder/coder/v2/coderd/database/dbmock"
|
||||
"github.com/coder/coder/v2/coderd/x/chatd/chattool"
|
||||
"github.com/coder/coder/v2/coderd/x/chatfiles"
|
||||
"github.com/coder/coder/v2/codersdk"
|
||||
)
|
||||
|
||||
@@ -102,29 +101,51 @@ func TestStoreChatAttachment_UsesDetectNameForClassification(t *testing.T) {
|
||||
require.Equal(t, "application/json", attachment.MediaType)
|
||||
}
|
||||
|
||||
func TestStoreChatAttachment_RejectsUnsupportedStoredFileTypeBeforeDBWork(t *testing.T) {
|
||||
func TestStoreChatAttachment_AllowsUnsupportedPromptInputType(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
ctrl := gomock.NewController(t)
|
||||
db := dbmock.NewMockStore(ctrl)
|
||||
tx := dbmock.NewMockStore(ctrl)
|
||||
server := &Server{db: db}
|
||||
|
||||
chatID := uuid.New()
|
||||
ownerID := uuid.New()
|
||||
workspaceID := uuid.New()
|
||||
orgID := uuid.New()
|
||||
fileID := uuid.New()
|
||||
chatSnapshot := database.Chat{
|
||||
ID: uuid.New(),
|
||||
OwnerID: uuid.New(),
|
||||
WorkspaceID: uuid.NullUUID{UUID: uuid.New(), Valid: true},
|
||||
ID: chatID,
|
||||
OwnerID: ownerID,
|
||||
WorkspaceID: uuid.NullUUID{UUID: workspaceID, Valid: true},
|
||||
}
|
||||
data := []byte(`<svg xmlns="http://www.w3.org/2000/svg"><rect/></svg>`)
|
||||
|
||||
attachment, err := server.storeChatAttachment(
|
||||
context.Background(),
|
||||
chatSnapshot,
|
||||
"evil.svg",
|
||||
"evil.svg",
|
||||
[]byte(`<svg xmlns="http://www.w3.org/2000/svg"><rect/></svg>`),
|
||||
expectStoreChatAttachmentTx(t, db, tx)
|
||||
tx.EXPECT().GetWorkspaceByID(gomock.Any(), workspaceID).Return(database.Workspace{ID: workspaceID, OrganizationID: orgID}, nil)
|
||||
tx.EXPECT().InsertChatFile(gomock.Any(), gomock.AssignableToTypeOf(database.InsertChatFileParams{})).DoAndReturn(
|
||||
func(_ context.Context, arg database.InsertChatFileParams) (database.InsertChatFileRow, error) {
|
||||
require.Equal(t, ownerID, arg.OwnerID)
|
||||
require.Equal(t, orgID, arg.OrganizationID)
|
||||
require.Equal(t, "evil.svg", arg.Name)
|
||||
require.Equal(t, "image/svg+xml", arg.Mimetype)
|
||||
require.Equal(t, data, arg.Data)
|
||||
return database.InsertChatFileRow{ID: fileID}, nil
|
||||
},
|
||||
)
|
||||
require.ErrorIs(t, err, chatfiles.ErrUnsupportedStoredFileType)
|
||||
require.ErrorContains(t, err, "image/svg+xml")
|
||||
require.Equal(t, chattool.AttachmentMetadata{}, attachment)
|
||||
tx.EXPECT().LinkChatFiles(gomock.Any(), database.LinkChatFilesParams{
|
||||
ChatID: chatID,
|
||||
MaxFileLinks: int32(codersdk.MaxChatFileIDs),
|
||||
FileIds: []uuid.UUID{fileID},
|
||||
}).Return(int32(0), nil)
|
||||
|
||||
attachment, err := server.storeChatAttachment(context.Background(), chatSnapshot, "evil.svg", "evil.svg", data)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, chattool.AttachmentMetadata{
|
||||
FileID: fileID,
|
||||
MediaType: "image/svg+xml",
|
||||
Name: "evil.svg",
|
||||
}, attachment)
|
||||
}
|
||||
|
||||
func TestStoreChatAttachment_NoWorkspace(t *testing.T) {
|
||||
|
||||
+12
-21
@@ -24,17 +24,13 @@ var (
|
||||
// after normalization.
|
||||
ErrStoredFileNameRequired = xerrors.New("stored file name is required")
|
||||
|
||||
// ErrUnsupportedStoredFileType indicates that classified file bytes do not
|
||||
// map to an allowed durable file type.
|
||||
ErrUnsupportedStoredFileType = xerrors.New("unsupported attachment type")
|
||||
|
||||
utf8BOM = []byte{0xEF, 0xBB, 0xBF}
|
||||
|
||||
// allowedStoredMediaTypes is derived from codersdk.AllChatAttachmentMediaTypes
|
||||
// allowedPromptInputMediaTypes is derived from codersdk.AllChatAttachmentMediaTypes
|
||||
// so the frontend file picker and the server enforcement share a single
|
||||
// source of truth. Do not edit this map directly; add new entries to the
|
||||
// codersdk const block instead.
|
||||
allowedStoredMediaTypes = func() map[string]struct{} {
|
||||
allowedPromptInputMediaTypes = func() map[string]struct{} {
|
||||
m := make(map[string]struct{}, len(codersdk.AllChatAttachmentMediaTypes))
|
||||
for _, t := range codersdk.AllChatAttachmentMediaTypes {
|
||||
m[string(t)] = struct{}{}
|
||||
@@ -61,16 +57,16 @@ func BaseMediaType(mediaType string) string {
|
||||
return mediaType
|
||||
}
|
||||
|
||||
// AllowedStoredMediaTypesString returns the supported durable chat file media
|
||||
// AllowedPromptInputMediaTypesString returns the supported prompt input media
|
||||
// types as a comma-separated list.
|
||||
func AllowedStoredMediaTypesString() string {
|
||||
return strings.Join(slices.Sorted(maps.Keys(allowedStoredMediaTypes)), ", ")
|
||||
func AllowedPromptInputMediaTypesString() string {
|
||||
return strings.Join(slices.Sorted(maps.Keys(allowedPromptInputMediaTypes)), ", ")
|
||||
}
|
||||
|
||||
// IsAllowedStoredMediaType reports whether the media type is supported for
|
||||
// durable chat file storage.
|
||||
func IsAllowedStoredMediaType(mediaType string) bool {
|
||||
_, ok := allowedStoredMediaTypes[BaseMediaType(mediaType)]
|
||||
// IsAllowedPromptInputMediaType reports whether the media type is supported for
|
||||
// user-provided prompt input.
|
||||
func IsAllowedPromptInputMediaType(mediaType string) bool {
|
||||
_, ok := allowedPromptInputMediaTypes[BaseMediaType(mediaType)]
|
||||
return ok
|
||||
}
|
||||
|
||||
@@ -80,7 +76,7 @@ func IsAllowedStoredMediaType(mediaType string) bool {
|
||||
// attack surface than the other media types we allow inline.
|
||||
func IsInlineRenderableStoredMediaType(mediaType string) bool {
|
||||
mediaType = BaseMediaType(mediaType)
|
||||
if !IsAllowedStoredMediaType(mediaType) {
|
||||
if !IsAllowedPromptInputMediaType(mediaType) {
|
||||
return false
|
||||
}
|
||||
return mediaType != "application/pdf"
|
||||
@@ -102,8 +98,7 @@ func NormalizeStoredFileName(name string) string {
|
||||
|
||||
// PrepareStoredFile normalizes the display name, rejects empty normalized
|
||||
// names, and classifies the file bytes using detectName when provided, so
|
||||
// callers can preserve subtype detection even when the user-facing filename is
|
||||
// overridden.
|
||||
// callers can apply the right policy for their file source.
|
||||
func PrepareStoredFile(name, detectName string, data []byte) (storedName, mediaType string, err error) {
|
||||
storedName = NormalizeStoredFileName(name)
|
||||
if storedName == "" {
|
||||
@@ -112,11 +107,7 @@ func PrepareStoredFile(name, detectName string, data []byte) (storedName, mediaT
|
||||
if strings.TrimSpace(detectName) == "" {
|
||||
detectName = storedName
|
||||
}
|
||||
mediaType = ClassifyStoredMediaType(detectName, data)
|
||||
if !IsAllowedStoredMediaType(mediaType) {
|
||||
return "", "", xerrors.Errorf("%w %q", ErrUnsupportedStoredFileType, mediaType)
|
||||
}
|
||||
return storedName, mediaType, nil
|
||||
return storedName, ClassifyStoredMediaType(detectName, data), nil
|
||||
}
|
||||
|
||||
// PrepareRecordingArtifact normalizes the recording artifact name, rejects
|
||||
|
||||
@@ -166,16 +166,17 @@ func TestPrepareStoredFile(t *testing.T) {
|
||||
require.ErrorIs(t, err, chatfiles.ErrStoredFileNameRequired)
|
||||
})
|
||||
|
||||
t.Run("RejectsUnsupportedStoredFileType", func(t *testing.T) {
|
||||
t.Run("ClassifiesUnsupportedPromptInputType", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
_, _, err := chatfiles.PrepareStoredFile(
|
||||
name, mediaType, err := chatfiles.PrepareStoredFile(
|
||||
"evil.svg",
|
||||
"evil.svg",
|
||||
[]byte(`<svg xmlns="http://www.w3.org/2000/svg"><rect/></svg>`),
|
||||
)
|
||||
require.ErrorIs(t, err, chatfiles.ErrUnsupportedStoredFileType)
|
||||
require.ErrorContains(t, err, "image/svg+xml")
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, "evil.svg", name)
|
||||
require.Equal(t, "image/svg+xml", mediaType)
|
||||
})
|
||||
|
||||
t.Run("TruncatesNamesAtRuneBoundaries", func(t *testing.T) {
|
||||
@@ -322,18 +323,18 @@ func TestIsCompatibleUploadMediaType(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestIsAllowedStoredMediaType(t *testing.T) {
|
||||
func TestIsAllowedPromptInputMediaType(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
require.True(t, chatfiles.IsAllowedStoredMediaType("text/plain; charset=utf-8"))
|
||||
require.True(t, chatfiles.IsAllowedStoredMediaType("text/markdown"))
|
||||
require.True(t, chatfiles.IsAllowedStoredMediaType("text/csv"))
|
||||
require.True(t, chatfiles.IsAllowedStoredMediaType("application/json"))
|
||||
require.True(t, chatfiles.IsAllowedStoredMediaType("application/pdf"))
|
||||
require.True(t, chatfiles.IsAllowedStoredMediaType("image/png"))
|
||||
require.False(t, chatfiles.IsAllowedStoredMediaType("image/svg+xml"))
|
||||
require.False(t, chatfiles.IsAllowedStoredMediaType("image/avif"))
|
||||
require.False(t, chatfiles.IsAllowedStoredMediaType("application/zip"))
|
||||
require.True(t, chatfiles.IsAllowedPromptInputMediaType("text/plain; charset=utf-8"))
|
||||
require.True(t, chatfiles.IsAllowedPromptInputMediaType("text/markdown"))
|
||||
require.True(t, chatfiles.IsAllowedPromptInputMediaType("text/csv"))
|
||||
require.True(t, chatfiles.IsAllowedPromptInputMediaType("application/json"))
|
||||
require.True(t, chatfiles.IsAllowedPromptInputMediaType("application/pdf"))
|
||||
require.True(t, chatfiles.IsAllowedPromptInputMediaType("image/png"))
|
||||
require.False(t, chatfiles.IsAllowedPromptInputMediaType("image/svg+xml"))
|
||||
require.False(t, chatfiles.IsAllowedPromptInputMediaType("image/avif"))
|
||||
require.False(t, chatfiles.IsAllowedPromptInputMediaType("application/zip"))
|
||||
}
|
||||
|
||||
func TestIsInlineRenderableStoredMediaType(t *testing.T) {
|
||||
@@ -344,6 +345,7 @@ func TestIsInlineRenderableStoredMediaType(t *testing.T) {
|
||||
require.True(t, chatfiles.IsInlineRenderableStoredMediaType("image/png"))
|
||||
require.False(t, chatfiles.IsInlineRenderableStoredMediaType("application/pdf"))
|
||||
require.False(t, chatfiles.IsInlineRenderableStoredMediaType("image/svg+xml"))
|
||||
require.False(t, chatfiles.IsInlineRenderableStoredMediaType("application/zip"))
|
||||
}
|
||||
|
||||
func TestHasSVGRootElement(t *testing.T) {
|
||||
|
||||
Reference in New Issue
Block a user