Skip to content

feat(im): use Content-Disposition filename when downloading message r… #536

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
chenxingtong-bytedance wants to merge 1 commit into
base: main
Choose a base branch
from
+132 −20
Open

feat(im): use Content-Disposition filename when downloading message r… #536

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
22 changes: 11 additions & 11 deletions shortcuts/im/helpers_network_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -262,7 +262,7 @@ func TestDownloadIMResourceToPathSuccess(t *testing.T) {
cmdutil.TestChdir(t, t.TempDir())

target := filepath.Join("nested", "resource.bin")
_, size, err := downloadIMResourceToPath(context.Background(), runtime, "om_123", "file_123", "file", target)
_, size, err := downloadIMResourceToPath(context.Background(), runtime, "om_123", "file_123", "file", target, true)
if err != nil {
t.Fatalf("downloadIMResourceToPath() error = %v", err)
}
Expand Down Expand Up @@ -308,7 +308,7 @@ func TestDownloadIMResourceToPathImageUsesSingleRequestWithoutRange(t *testing.T

cmdutil.TestChdir(t, t.TempDir())

gotPath, size, err := downloadIMResourceToPath(context.Background(), runtime, "om_img", "img_123", "image", "image")
gotPath, size, err := downloadIMResourceToPath(context.Background(), runtime, "om_img", "img_123", "image", "image", true)
if err != nil {
t.Fatalf("downloadIMResourceToPath() error = %v", err)
}
Expand Down Expand Up @@ -342,7 +342,7 @@ func TestDownloadIMResourceToPathHTTPErrorBody(t *testing.T) {

cmdutil.TestChdir(t, t.TempDir())

_, _, err := downloadIMResourceToPath(context.Background(), runtime, "om_403", "file_403", "file", "out.bin")
_, _, err := downloadIMResourceToPath(context.Background(), runtime, "om_403", "file_403", "file", "out.bin", true)
if err == nil || !strings.Contains(err.Error(), "HTTP 403: denied") {
t.Fatalf("downloadIMResourceToPath() error = %v", err)
}
Expand Down Expand Up @@ -372,7 +372,7 @@ func TestDownloadIMResourceToPathRetriesNetworkError(t *testing.T) {

cmdutil.TestChdir(t, t.TempDir())
target := "out.bin"
_, size, err := downloadIMResourceToPath(context.Background(), runtime, "om_retry", "file_retry", "file", target)
_, size, err := downloadIMResourceToPath(context.Background(), runtime, "om_retry", "file_retry", "file", target, true)
if err != nil {
t.Fatalf("downloadIMResourceToPath() error = %v", err)
}
Expand Down Expand Up @@ -408,7 +408,7 @@ func TestDownloadIMResourceToPathRetrySecondAttemptSuccess(t *testing.T) {

cmdutil.TestChdir(t, t.TempDir())
target := "out.bin"
_, size, err := downloadIMResourceToPath(context.Background(), runtime, "om_retry2", "file_retry2", "file", target)
_, size, err := downloadIMResourceToPath(context.Background(), runtime, "om_retry2", "file_retry2", "file", target, true)
if err != nil {
t.Fatalf("downloadIMResourceToPath() error = %v", err)
}
Expand Down Expand Up @@ -444,7 +444,7 @@ func TestDownloadIMResourceToPathRetryContextCanceled(t *testing.T) {

cmdutil.TestChdir(t, t.TempDir())
target := "out.bin"
_, _, err := downloadIMResourceToPath(ctx, runtime, "om_cancel", "file_cancel", "file", target)
_, _, err := downloadIMResourceToPath(ctx, runtime, "om_cancel", "file_cancel", "file", target, true)
if err != context.Canceled {
t.Fatalf("downloadIMResourceToPath() error = %v, want context.Canceled", err)
}
Expand Down Expand Up @@ -525,7 +525,7 @@ func TestDownloadIMResourceToPathRangeDownload(t *testing.T) {

cmdutil.TestChdir(t, t.TempDir())
target := filepath.Join("nested", "resource.bin")
_, size, err := downloadIMResourceToPath(context.Background(), runtime, "om_range", "file_range", "file", target)
_, size, err := downloadIMResourceToPath(context.Background(), runtime, "om_range", "file_range", "file", target, true)
if err != nil {
t.Fatalf("downloadIMResourceToPath() error = %v", err)
}
Expand Down Expand Up @@ -567,7 +567,7 @@ func TestDownloadIMResourceToPathInvalidContentRange(t *testing.T) {
}))

cmdutil.TestChdir(t, t.TempDir())
_, _, err := downloadIMResourceToPath(context.Background(), runtime, "om_bad", "file_bad", "file", "out.bin")
_, _, err := downloadIMResourceToPath(context.Background(), runtime, "om_bad", "file_bad", "file", "out.bin", true)
if err == nil || !strings.Contains(err.Error(), "invalid Content-Range header") {
t.Fatalf("downloadIMResourceToPath() error = %v", err)
}
Expand Down Expand Up @@ -596,7 +596,7 @@ func TestDownloadIMResourceToPathRangeChunkFailureCleansOutput(t *testing.T) {
cmdutil.TestChdir(t, t.TempDir())

target := "out.bin"
_, _, err := downloadIMResourceToPath(context.Background(), runtime, "om_miderr", "file_miderr", "file", target)
_, _, err := downloadIMResourceToPath(context.Background(), runtime, "om_miderr", "file_miderr", "file", target, true)
if err == nil || !strings.Contains(err.Error(), "HTTP 500: chunk failed") {
t.Fatalf("downloadIMResourceToPath() error = %v", err)
}
Expand All @@ -622,7 +622,7 @@ func TestDownloadIMResourceToPathRangeOverflowCleansOutput(t *testing.T) {
cmdutil.TestChdir(t, t.TempDir())

target := "out.bin"
_, _, err := downloadIMResourceToPath(context.Background(), runtime, "om_overflow", "file_overflow", "file", target)
_, _, err := downloadIMResourceToPath(context.Background(), runtime, "om_overflow", "file_overflow", "file", target, true)
if err == nil || !strings.Contains(err.Error(), "chunk overflow") {
t.Fatalf("downloadIMResourceToPath() error = %v", err)
}
Expand Down Expand Up @@ -658,7 +658,7 @@ func TestDownloadIMResourceToPathRangeShortChunkSizeMismatch(t *testing.T) {

cmdutil.TestChdir(t, t.TempDir())

_, _, err := downloadIMResourceToPath(context.Background(), runtime, "om_short", "file_short", "file", "out.bin")
_, _, err := downloadIMResourceToPath(context.Background(), runtime, "om_short", "file_short", "file", "out.bin", true)
if err == nil || !strings.Contains(err.Error(), "file size mismatch") {
t.Fatalf("downloadIMResourceToPath() error = %v", err)
}
Expand Down
64 changes: 63 additions & 1 deletion shortcuts/im/helpers_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -593,7 +593,7 @@ func TestDownloadIMResourceToPathHTTPClientError(t *testing.T) {
return nil, errors.New("http client unavailable")
}))

_, _, err := downloadIMResourceToPath(context.Background(), runtime, "om_123", "img_123", "image", "out.bin")
_, _, err := downloadIMResourceToPath(context.Background(), runtime, "om_123", "img_123", "image", "out.bin", true)
if err == nil || !strings.Contains(err.Error(), "http client unavailable") {
t.Fatalf("downloadIMResourceToPath() error = %v", err)
}
Expand Down Expand Up @@ -637,6 +637,68 @@ func TestParseTotalSize(t *testing.T) {
}
}

func TestParseContentDispositionFilename(t *testing.T) {
tests := []struct {
name string
header string
want string
}{
{name: "empty header", header: "", want: ""},
{name: "no filename param", header: "attachment", want: ""},
{name: "plain filename", header: `attachment; filename="report.xlsx"`, want: "report.xlsx"},
{name: "unquoted filename", header: `attachment; filename=report.xlsx`, want: "report.xlsx"},
{name: "RFC 5987 UTF-8 encoded", header: `attachment; filename*=UTF-8''%E5%AD%A3%E5%BA%A6%E6%8A%A5%E5%91%8A.xlsx`, want: "季度报告.xlsx"},
{name: "RFC 5987 takes priority over plain", header: `attachment; filename="fallback.xlsx"; filename*=UTF-8''%E5%AD%A3%E5%BA%A6%E6%8A%A5%E5%91%8A.xlsx`, want: "季度报告.xlsx"},
{name: "path traversal stripped", header: `attachment; filename="../../etc/passwd"`, want: "passwd"},
{name: "windows path stripped", header: `attachment; filename="C:\\Windows\\evil.exe"`, want: "evil.exe"},
{name: "control char rejected", header: "attachment; filename=\"evil\x01file.txt\"", want: ""},
{name: "malformed header", header: "not/valid/mime; ===", want: ""},
{name: "whitespace trimmed", header: `attachment; filename=" report.pdf "`, want: "report.pdf"},
}

for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
if got := parseContentDispositionFilename(tt.header); got != tt.want {
t.Fatalf("parseContentDispositionFilename(%q) = %q, want %q", tt.header, got, tt.want)
}
})
}
}

func TestResolveIMResourceDownloadPath(t *testing.T) {
tests := []struct {
name string
safePath string
contentType string
contentDisposition string
userSpecifiedOutput bool
want string
}{
// safePath already has extension: always return as-is
{name: "user path with ext, no CD", safePath: "out.xlsx", contentType: "application/pdf", userSpecifiedOutput: true, want: "out.xlsx"},
{name: "user path with ext, CD present", safePath: "out.xlsx", contentDisposition: `attachment; filename="server.pdf"`, userSpecifiedOutput: true, want: "out.xlsx"},
// No --output: use CD filename when present
{name: "default path, CD filename", safePath: "file_xxx", contentDisposition: `attachment; filename="季度报告.xlsx"`, want: "季度报告.xlsx"},
{name: "default path, CD RFC5987", safePath: "file_xxx", contentDisposition: `attachment; filename*=UTF-8''%E5%AD%A3%E5%BA%A6%E6%8A%A5%E5%91%8A.xlsx`, want: "季度报告.xlsx"},
{name: "default path, no CD, MIME ext", safePath: "file_xxx", contentType: "application/pdf", want: "file_xxx.pdf"},
{name: "default path, no CD, unknown MIME", safePath: "file_xxx", contentType: "application/x-unknown", want: "file_xxx"},
{name: "default path, CD with dir component", safePath: "downloads/file_xxx", contentDisposition: `attachment; filename="report.xlsx"`, want: "downloads/report.xlsx"},
// User --output without extension: use CD filename's extension
{name: "user path no ext, CD with ext", safePath: "myfile", contentDisposition: `attachment; filename="server.pdf"`, userSpecifiedOutput: true, want: "myfile.pdf"},
{name: "user path no ext, CD no ext, MIME ext", safePath: "myfile", contentDisposition: `attachment; filename="noext"`, contentType: "image/png", userSpecifiedOutput: true, want: "myfile.png"},
{name: "user path no ext, no CD, MIME ext", safePath: "myfile", contentType: "image/jpeg", userSpecifiedOutput: true, want: "myfile.jpg"},
}

for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
got := resolveIMResourceDownloadPath(tt.safePath, tt.contentType, tt.contentDisposition, tt.userSpecifiedOutput)
if got != tt.want {
t.Fatalf("resolveIMResourceDownloadPath() = %q, want %q", got, tt.want)
}
})
}
}

func TestShortcuts(t *testing.T) {
var commands []string
for _, shortcut := range Shortcuts() {
Expand Down
62 changes: 56 additions & 6 deletions shortcuts/im/im_messages_resources_download.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,7 @@ import (
"context"
"fmt"
"io"
"mime"
"net/http"
"path/filepath"
"strconv"
Expand Down Expand Up @@ -72,7 +73,8 @@ var ImMessagesResourcesDownload = common.Shortcut{
return output.ErrValidation("unsafe output path: %s", err)
}

finalPath, sizeBytes, err := downloadIMResourceToPath(ctx, runtime, messageId, fileKey, fileType, relPath)
userSpecifiedOutput := runtime.Str("output") != ""
finalPath, sizeBytes, err := downloadIMResourceToPath(ctx, runtime, messageId, fileKey, fileType, relPath, userSpecifiedOutput)
if err != nil {
return err
}
Expand Down Expand Up @@ -262,18 +264,21 @@ func initialIMResourceDownloadHeaders(fileType string) map[string]string {
}
}

func downloadIMResourceToPath(ctx context.Context, runtime *common.RuntimeContext, messageID, fileKey, fileType, outputPath string) (string, int64, error) {
func downloadIMResourceToPath(ctx context.Context, runtime *common.RuntimeContext, messageID, fileKey, fileType, outputPath string, userSpecifiedOutput bool) (string, int64, error) {
downloadResp, err := doIMResourceDownloadRequest(ctx, runtime, messageID, fileKey, fileType, initialIMResourceDownloadHeaders(fileType))
if err != nil {
return "", 0, err
}
if downloadResp == nil {
return "", 0, output.ErrNetwork("download failed: empty response")
}

if downloadResp.StatusCode >= 400 {
defer downloadResp.Body.Close()
return "", 0, downloadResponseError(downloadResp)
}

finalPath := resolveIMResourceDownloadPath(outputPath, downloadResp.Header.Get("Content-Type"))
finalPath := resolveIMResourceDownloadPath(outputPath, downloadResp.Header.Get("Content-Type"), downloadResp.Header.Get("Content-Disposition"), userSpecifiedOutput)

var (
body io.ReadCloser
Expand Down Expand Up @@ -316,18 +321,63 @@ func downloadIMResourceToPath(ctx context.Context, runtime *common.RuntimeContex
return savedPath, result.Size(), nil
}

func resolveIMResourceDownloadPath(safePath, contentType string) string {
func resolveIMResourceDownloadPath(safePath, contentType, contentDisposition string, userSpecifiedOutput bool) string {
if filepath.Ext(safePath) != "" {
return safePath
}
mimeType := strings.Split(contentType, ";")[0]
mimeType = strings.TrimSpace(mimeType)
if cdFilename := parseContentDispositionFilename(contentDisposition); cdFilename != "" {
if !userSpecifiedOutput {
// No --output flag: use the original filename from the server.
dir := filepath.Dir(safePath)
if dir == "." {
return cdFilename
}
return filepath.Join(dir, cdFilename)
}
// User specified a path without extension: append the extension from the CD filename.
if ext := filepath.Ext(cdFilename); ext != "" {
return safePath + ext
}
}
mimeType := strings.TrimSpace(strings.Split(contentType, ";")[0])
if ext, ok := imMimeToExt[mimeType]; ok {
return safePath + ext
}
return safePath
}

// parseContentDispositionFilename extracts and sanitizes the filename from a
// Content-Disposition header. It handles RFC 5987 encoded filenames (filename*)
// with priority over plain filename via the standard mime package.
// Returns an empty string if no valid filename can be extracted.
func parseContentDispositionFilename(header string) string {
if header == "" {
return ""
}
_, params, err := mime.ParseMediaType(header)
if err != nil {
return ""
}
name := strings.TrimSpace(params["filename"])
if name == "" {
return ""
}
// Strip any path component (Unix or Windows style) to prevent path traversal.
if i := strings.LastIndexAny(name, "/\\"); i >= 0 {
name = name[i+1:]
}
if name == "" || name == "." || name == ".." {
return ""
}
// Reject control characters (including null bytes).
for _, r := range name {
if r < 0x20 || r == 0x7f {
return ""
}
}
return name
}

func doIMResourceDownloadRequest(ctx context.Context, runtime *common.RuntimeContext, messageID, fileKey, fileType string, headers map[string]string) (*http.Response, error) {
query := larkcore.QueryParams{}
query.Set("type", fileType)
Expand Down
4 changes: 2 additions & 2 deletions skills/lark-im/references/lark-im-messages-resources-download.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -34,7 +34,7 @@ lark-cli im +messages-resources-download --message-id om_xxx --file-key img_v3_x
| `--message-id <id>` | Yes | Message ID (`om_xxx` format) |
| `--file-key <key>` | Yes | Resource key (`img_xxx` or `file_xxx`) |
| `--type <type>` | Yes | Resource type: `image` or `file` |
| `--output <path>` | No | Output path (relative paths only; `..` traversal is not allowed; defaults to `file_key` as the file name). File extension is automatically added based on Content-Type if not provided |
| `--output <path>` | No | Output path (relative paths only; `..` traversal is not allowed). When omitted, the server's original filename from `Content-Disposition` is used if available; otherwise defaults to `file_key`. File extension is automatically inferred from `Content-Disposition` or `Content-Type` if not provided |
| `--as <identity>` | No | Identity type: `user` (default) or `bot` |
| `--dry-run` | No | Print the request only, do not execute it |

Expand All @@ -51,7 +51,7 @@ When downloading large files, the command automatically uses **HTTP Range reques

**Benefits:**
- Reduces the impact of transient request failures during large downloads
- Automatically detects and appends correct file extension from Content-Type
- Preserves the server's original filename via `Content-Disposition` (supports RFC 5987 UTF-8 encoding); falls back to `Content-Type`-based extension inference
- Validates file size integrity after download completion

## `file_key` Sources
Expand Down
Loading