diff --git a/shortcuts/common/drive_media_upload.go b/shortcuts/common/drive_media_upload.go new file mode 100644 index 00000000..fda69b97 --- /dev/null +++ b/shortcuts/common/drive_media_upload.go @@ -0,0 +1,245 @@ +// Copyright (c) 2026 Lark Technologies Pte. Ltd. +// SPDX-License-Identifier: MIT + +package common + +import ( + "bytes" + "encoding/json" + "errors" + "fmt" + "io" + "net/http" + + "github.com/larksuite/cli/internal/validate" + "github.com/larksuite/cli/internal/vfs" + larkcore "github.com/larksuite/oapi-sdk-go/v3/core" + + "github.com/larksuite/cli/internal/output" +) + +const MaxDriveMediaUploadSinglePartSize int64 = 20 * 1024 * 1024 // 20MB + +const ( + driveMediaUploadAllAction = "upload media failed" + driveMediaUploadPartAction = "upload media part failed" + driveMediaUploadFinishAction = "upload media finish failed" +) + +type DriveMediaMultipartUploadSession struct { + UploadID string + BlockSize int64 + BlockNum int +} + +type DriveMediaUploadAllConfig struct { + FilePath string + FileName string + FileSize int64 + ParentType string + ParentNode *string + Extra string +} + +type DriveMediaMultipartUploadConfig struct { + FilePath string + FileName string + FileSize int64 + ParentType string + ParentNode string + Extra string +} + +func UploadDriveMediaAll(runtime *RuntimeContext, cfg DriveMediaUploadAllConfig) (string, error) { + safeFilePath, err := validate.SafeInputPath(cfg.FilePath) + if err != nil { + return "", output.ErrValidation("invalid file path: %s", err) + } + f, err := vfs.Open(safeFilePath) + if err != nil { + return "", output.ErrValidation("cannot read file: %s", err) + } + defer f.Close() + + fd := larkcore.NewFormdata() + fd.AddField("file_name", cfg.FileName) + fd.AddField("parent_type", cfg.ParentType) + fd.AddField("size", fmt.Sprintf("%d", cfg.FileSize)) + if cfg.ParentNode != nil { + fd.AddField("parent_node", *cfg.ParentNode) + } + if cfg.Extra != "" { + fd.AddField("extra", cfg.Extra) + } + fd.AddFile("file", f) + + apiResp, err := runtime.DoAPI(&larkcore.ApiReq{ + HttpMethod: http.MethodPost, + ApiPath: "/open-apis/drive/v1/medias/upload_all", + Body: fd, + }, larkcore.WithFileUpload()) + if err != nil { + return "", WrapDriveMediaUploadRequestError(err, driveMediaUploadAllAction) + } + + data, err := ParseDriveMediaUploadResponse(apiResp, driveMediaUploadAllAction) + if err != nil { + return "", err + } + return ExtractDriveMediaUploadFileToken(data, driveMediaUploadAllAction) +} + +func UploadDriveMediaMultipart(runtime *RuntimeContext, cfg DriveMediaMultipartUploadConfig) (string, error) { + // upload_prepare expects parent_node to be present even when the caller wants + // the service default/root behavior, so multipart callers pass an explicit + // string instead of relying on field omission like upload_all does. + prepareBody := map[string]interface{}{ + "file_name": cfg.FileName, + "parent_type": cfg.ParentType, + "parent_node": cfg.ParentNode, + "size": cfg.FileSize, + } + if cfg.Extra != "" { + prepareBody["extra"] = cfg.Extra + } + + data, err := runtime.CallAPI("POST", "/open-apis/drive/v1/medias/upload_prepare", nil, prepareBody) + if err != nil { + return "", err + } + + session, err := ParseDriveMediaMultipartUploadSession(data) + if err != nil { + return "", err + } + fmt.Fprintf(runtime.IO().ErrOut, "Multipart upload initialized: %d chunks x %s\n", session.BlockNum, FormatSize(session.BlockSize)) + + if err = uploadDriveMediaMultipartParts(runtime, cfg.FilePath, cfg.FileSize, session); err != nil { + return "", err + } + + return finishDriveMediaMultipartUpload(runtime, session.UploadID, session.BlockNum) +} + +func ParseDriveMediaMultipartUploadSession(data map[string]interface{}) (DriveMediaMultipartUploadSession, error) { + // The backend chooses both chunk size and chunk count. Validate them once so + // the streaming loop can follow the returned plan without re-checking shape. + session := DriveMediaMultipartUploadSession{ + UploadID: GetString(data, "upload_id"), + BlockSize: int64(GetFloat(data, "block_size")), + BlockNum: int(GetFloat(data, "block_num")), + } + if session.UploadID == "" { + return DriveMediaMultipartUploadSession{}, output.Errorf(output.ExitAPI, "api_error", "upload prepare failed: no upload_id returned") + } + if session.BlockSize <= 0 { + return DriveMediaMultipartUploadSession{}, output.Errorf(output.ExitAPI, "api_error", "upload prepare failed: invalid block_size returned") + } + if session.BlockNum <= 0 { + return DriveMediaMultipartUploadSession{}, output.Errorf(output.ExitAPI, "api_error", "upload prepare failed: invalid block_num returned") + } + return session, nil +} + +func WrapDriveMediaUploadRequestError(err error, action string) error { + var exitErr *output.ExitError + if errors.As(err, &exitErr) { + return err + } + return output.ErrNetwork("%s: %v", action, err) +} + +func ParseDriveMediaUploadResponse(apiResp *larkcore.ApiResp, action string) (map[string]interface{}, error) { + var result map[string]interface{} + if err := json.Unmarshal(apiResp.RawBody, &result); err != nil { + return nil, output.Errorf(output.ExitAPI, "api_error", "%s: invalid response JSON: %v", action, err) + } + + if larkCode := int(GetFloat(result, "code")); larkCode != 0 { + msg, _ := result["msg"].(string) + return nil, output.ErrAPI(larkCode, fmt.Sprintf("%s: [%d] %s", action, larkCode, msg), result["error"]) + } + + data, _ := result["data"].(map[string]interface{}) + return data, nil +} + +func ExtractDriveMediaUploadFileToken(data map[string]interface{}, action string) (string, error) { + fileToken := GetString(data, "file_token") + if fileToken == "" { + return "", output.Errorf(output.ExitAPI, "api_error", "%s: no file_token returned", action) + } + return fileToken, nil +} + +func uploadDriveMediaMultipartParts(runtime *RuntimeContext, filePath string, fileSize int64, session DriveMediaMultipartUploadSession) error { + safeFilePath, err := validate.SafeInputPath(filePath) + if err != nil { + return output.ErrValidation("invalid file path: %s", err) + } + f, err := vfs.Open(safeFilePath) + if err != nil { + return output.ErrValidation("cannot read file: %s", err) + } + defer f.Close() + + maxInt := int64(^uint(0) >> 1) + if session.BlockSize > maxInt { + return output.Errorf(output.ExitAPI, "api_error", "upload prepare failed: invalid block_size returned") + } + + buffer := make([]byte, int(session.BlockSize)) + remaining := fileSize + // Follow the server-declared block plan exactly; upload_finish expects the + // same block count returned by upload_prepare. + for seq := 0; seq < session.BlockNum; seq++ { + chunkSize := session.BlockSize + if remaining > 0 && chunkSize > remaining { + chunkSize = remaining + } + + n, readErr := io.ReadFull(f, buffer[:int(chunkSize)]) + if readErr != nil { + return output.ErrValidation("cannot read file: %s", readErr) + } + + if err = uploadDriveMediaMultipartPart(runtime, session.UploadID, seq, buffer[:n]); err != nil { + return err + } + fmt.Fprintf(runtime.IO().ErrOut, " Block %d/%d uploaded (%s)\n", seq+1, session.BlockNum, FormatSize(int64(n))) + remaining -= int64(n) + } + + return nil +} + +func uploadDriveMediaMultipartPart(runtime *RuntimeContext, uploadID string, seq int, chunk []byte) error { + fd := larkcore.NewFormdata() + fd.AddField("upload_id", uploadID) + fd.AddField("seq", fmt.Sprintf("%d", seq)) + fd.AddField("size", fmt.Sprintf("%d", len(chunk))) + fd.AddFile("file", bytes.NewReader(chunk)) + + apiResp, err := runtime.DoAPI(&larkcore.ApiReq{ + HttpMethod: http.MethodPost, + ApiPath: "/open-apis/drive/v1/medias/upload_part", + Body: fd, + }, larkcore.WithFileUpload()) + if err != nil { + return WrapDriveMediaUploadRequestError(err, driveMediaUploadPartAction) + } + + _, err = ParseDriveMediaUploadResponse(apiResp, driveMediaUploadPartAction) + return err +} + +func finishDriveMediaMultipartUpload(runtime *RuntimeContext, uploadID string, blockNum int) (string, error) { + data, err := runtime.CallAPI("POST", "/open-apis/drive/v1/medias/upload_finish", nil, map[string]interface{}{ + "upload_id": uploadID, + "block_num": blockNum, + }) + if err != nil { + return "", err + } + return ExtractDriveMediaUploadFileToken(data, driveMediaUploadFinishAction) +} diff --git a/shortcuts/common/drive_media_upload_test.go b/shortcuts/common/drive_media_upload_test.go new file mode 100644 index 00000000..025a6508 --- /dev/null +++ b/shortcuts/common/drive_media_upload_test.go @@ -0,0 +1,520 @@ +// Copyright (c) 2026 Lark Technologies Pte. Ltd. +// SPDX-License-Identifier: MIT + +package common + +import ( + "bytes" + "context" + "encoding/json" + "errors" + "fmt" + "io" + "mime" + "mime/multipart" + "os" + "strings" + "sync/atomic" + "testing" + + larkcore "github.com/larksuite/oapi-sdk-go/v3/core" + + "github.com/larksuite/cli/internal/cmdutil" + "github.com/larksuite/cli/internal/core" + "github.com/larksuite/cli/internal/httpmock" + "github.com/larksuite/cli/internal/output" +) + +var commonDriveMediaUploadTestSeq atomic.Int64 + +func TestUploadDriveMediaAllBuildsMultipartBody(t *testing.T) { + tests := []struct { + name string + parentNode *string + wantParentNode string + wantParentSet bool + }{ + { + name: "includes parent_node when provided", + parentNode: strPtr("blk_parent"), + wantParentNode: "blk_parent", + wantParentSet: true, + }, + { + name: "omits parent_node when not provided", + parentNode: nil, + wantParentSet: false, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + runtime, reg := newDriveMediaUploadTestRuntime(t) + withDriveMediaUploadWorkingDir(t, t.TempDir()) + + uploadStub := &httpmock.Stub{ + Method: "POST", + URL: "/open-apis/drive/v1/medias/upload_all", + Body: map[string]interface{}{ + "code": 0, + "data": map[string]interface{}{"file_token": "file_all_123"}, + }, + } + reg.Register(uploadStub) + + filePath := writeDriveMediaUploadTestFile(t, "small.bin", 3) + fileToken, err := UploadDriveMediaAll(runtime, DriveMediaUploadAllConfig{ + FilePath: filePath, + FileName: "small.bin", + FileSize: 3, + ParentType: "docx_file", + ParentNode: tt.parentNode, + Extra: `{"drive_route_token":"doxcn123"}`, + }) + if err != nil { + t.Fatalf("UploadDriveMediaAll() error: %v", err) + } + if fileToken != "file_all_123" { + t.Fatalf("fileToken = %q, want %q", fileToken, "file_all_123") + } + + body := decodeCapturedDriveMediaMultipartBody(t, uploadStub) + if got := body.Fields["file_name"]; got != "small.bin" { + t.Fatalf("file_name = %q, want %q", got, "small.bin") + } + if got := body.Fields["parent_type"]; got != "docx_file" { + t.Fatalf("parent_type = %q, want %q", got, "docx_file") + } + if got := body.Fields["size"]; got != "3" { + t.Fatalf("size = %q, want %q", got, "3") + } + if got := body.Fields["extra"]; got != `{"drive_route_token":"doxcn123"}` { + t.Fatalf("extra = %q, want drive route token payload", got) + } + if got := len(body.Files["file"]); got != 3 { + t.Fatalf("file size = %d, want %d", got, 3) + } + + gotParentNode, hasParentNode := body.Fields["parent_node"] + if hasParentNode != tt.wantParentSet { + t.Fatalf("parent_node present = %v, want %v", hasParentNode, tt.wantParentSet) + } + if hasParentNode && gotParentNode != tt.wantParentNode { + t.Fatalf("parent_node = %q, want %q", gotParentNode, tt.wantParentNode) + } + }) + } +} + +func TestUploadDriveMediaMultipartBuildsPreparePartsAndFinish(t *testing.T) { + runtime, reg := newDriveMediaUploadTestRuntime(t) + withDriveMediaUploadWorkingDir(t, t.TempDir()) + + prepareStub := &httpmock.Stub{ + Method: "POST", + URL: "/open-apis/drive/v1/medias/upload_prepare", + Body: map[string]interface{}{ + "code": 0, + "data": map[string]interface{}{ + "upload_id": "upload_123", + "block_size": float64(4 * 1024 * 1024), + "block_num": float64(6), + }, + }, + } + reg.Register(prepareStub) + + partStubs := make([]*httpmock.Stub, 0, 6) + for i := 0; i < 6; i++ { + stub := &httpmock.Stub{ + Method: "POST", + URL: "/open-apis/drive/v1/medias/upload_part", + Body: map[string]interface{}{ + "code": 0, + "msg": "ok", + }, + } + partStubs = append(partStubs, stub) + reg.Register(stub) + } + + finishStub := &httpmock.Stub{ + Method: "POST", + URL: "/open-apis/drive/v1/medias/upload_finish", + Body: map[string]interface{}{ + "code": 0, + "data": map[string]interface{}{"file_token": "file_multi_123"}, + }, + } + reg.Register(finishStub) + + filePath := writeDriveMediaUploadSizedFile(t, "large.bin", MaxDriveMediaUploadSinglePartSize+1) + fileToken, err := UploadDriveMediaMultipart(runtime, DriveMediaMultipartUploadConfig{ + FilePath: filePath, + FileName: "large.bin", + FileSize: MaxDriveMediaUploadSinglePartSize + 1, + ParentType: "ccm_import_open", + ParentNode: "", + Extra: `{"obj_type":"sheet","file_extension":"xlsx"}`, + }) + if err != nil { + t.Fatalf("UploadDriveMediaMultipart() error: %v", err) + } + if fileToken != "file_multi_123" { + t.Fatalf("fileToken = %q, want %q", fileToken, "file_multi_123") + } + + prepareBody := decodeCapturedDriveMediaJSONBody(t, prepareStub) + if got, _ := prepareBody["parent_type"].(string); got != "ccm_import_open" { + t.Fatalf("prepare parent_type = %q, want %q", got, "ccm_import_open") + } + if got, _ := prepareBody["parent_node"].(string); got != "" { + t.Fatalf("prepare parent_node = %q, want empty string", got) + } + if got, _ := prepareBody["size"].(float64); got != float64(MaxDriveMediaUploadSinglePartSize+1) { + t.Fatalf("prepare size = %v, want %d", got, MaxDriveMediaUploadSinglePartSize+1) + } + + firstPart := decodeCapturedDriveMediaMultipartBody(t, partStubs[0]) + if got := firstPart.Fields["upload_id"]; got != "upload_123" { + t.Fatalf("first part upload_id = %q, want %q", got, "upload_123") + } + if got := firstPart.Fields["seq"]; got != "0" { + t.Fatalf("first part seq = %q, want %q", got, "0") + } + if got := firstPart.Fields["size"]; got != "4194304" { + t.Fatalf("first part size = %q, want %q", got, "4194304") + } + + lastPart := decodeCapturedDriveMediaMultipartBody(t, partStubs[len(partStubs)-1]) + if got := lastPart.Fields["seq"]; got != "5" { + t.Fatalf("last part seq = %q, want %q", got, "5") + } + if got := lastPart.Fields["size"]; got != "1" { + t.Fatalf("last part size = %q, want %q", got, "1") + } + if got := len(lastPart.Files["file"]); got != 1 { + t.Fatalf("last part file size = %d, want %d", got, 1) + } + + finishBody := decodeCapturedDriveMediaJSONBody(t, finishStub) + if got, _ := finishBody["upload_id"].(string); got != "upload_123" { + t.Fatalf("finish upload_id = %q, want %q", got, "upload_123") + } + if got, _ := finishBody["block_num"].(float64); got != 6 { + t.Fatalf("finish block_num = %v, want %d", got, 6) + } +} + +func TestParseDriveMediaMultipartUploadSessionValidatesResponseFields(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + data map[string]interface{} + wantText string + }{ + { + name: "missing upload id", + data: map[string]interface{}{ + "block_size": 4 * 1024 * 1024, + "block_num": 6, + }, + wantText: "upload prepare failed: no upload_id returned", + }, + { + name: "missing block size", + data: map[string]interface{}{ + "upload_id": "upload_123", + "block_num": 6, + }, + wantText: "upload prepare failed: invalid block_size returned", + }, + { + name: "missing block num", + data: map[string]interface{}{ + "upload_id": "upload_123", + "block_size": 4 * 1024 * 1024, + }, + wantText: "upload prepare failed: invalid block_num returned", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + _, err := ParseDriveMediaMultipartUploadSession(tt.data) + if err == nil || !strings.Contains(err.Error(), tt.wantText) { + t.Fatalf("err = %v, want substring %q", err, tt.wantText) + } + }) + } +} + +func TestUploadDriveMediaMultipartPartAPIFailure(t *testing.T) { + runtime, reg := newDriveMediaUploadTestRuntime(t) + withDriveMediaUploadWorkingDir(t, t.TempDir()) + reg.Register(&httpmock.Stub{ + Method: "POST", + URL: "/open-apis/drive/v1/medias/upload_prepare", + Body: map[string]interface{}{ + "code": 0, + "data": map[string]interface{}{ + "upload_id": "upload_123", + "block_size": float64(4 * 1024 * 1024), + "block_num": float64(6), + }, + }, + }) + reg.Register(&httpmock.Stub{ + Method: "POST", + URL: "/open-apis/drive/v1/medias/upload_part", + Body: map[string]interface{}{ + "code": 999, + "msg": "chunk rejected", + }, + }) + + filePath := writeDriveMediaUploadSizedFile(t, "large.bin", MaxDriveMediaUploadSinglePartSize+1) + _, err := UploadDriveMediaMultipart(runtime, DriveMediaMultipartUploadConfig{ + FilePath: filePath, + FileName: "large.bin", + FileSize: MaxDriveMediaUploadSinglePartSize + 1, + ParentType: "ccm_import_open", + ParentNode: "", + }) + if err == nil { + t.Fatal("expected error, got nil") + } + if !strings.Contains(err.Error(), "upload media part failed: [999] chunk rejected") { + t.Fatalf("unexpected error: %v", err) + } +} + +func TestUploadDriveMediaMultipartFinishRequiresFileToken(t *testing.T) { + runtime, reg := newDriveMediaUploadTestRuntime(t) + withDriveMediaUploadWorkingDir(t, t.TempDir()) + reg.Register(&httpmock.Stub{ + Method: "POST", + URL: "/open-apis/drive/v1/medias/upload_prepare", + Body: map[string]interface{}{ + "code": 0, + "data": map[string]interface{}{ + "upload_id": "upload_123", + "block_size": float64(4 * 1024 * 1024), + "block_num": float64(6), + }, + }, + }) + for i := 0; i < 6; i++ { + reg.Register(&httpmock.Stub{ + Method: "POST", + URL: "/open-apis/drive/v1/medias/upload_part", + Body: map[string]interface{}{ + "code": 0, + "msg": "ok", + }, + }) + } + reg.Register(&httpmock.Stub{ + Method: "POST", + URL: "/open-apis/drive/v1/medias/upload_finish", + Body: map[string]interface{}{ + "code": 0, + "data": map[string]interface{}{}, + }, + }) + + filePath := writeDriveMediaUploadSizedFile(t, "large.bin", MaxDriveMediaUploadSinglePartSize+1) + _, err := UploadDriveMediaMultipart(runtime, DriveMediaMultipartUploadConfig{ + FilePath: filePath, + FileName: "large.bin", + FileSize: MaxDriveMediaUploadSinglePartSize + 1, + ParentType: "ccm_import_open", + ParentNode: "", + }) + if err == nil { + t.Fatal("expected error, got nil") + } + if !strings.Contains(err.Error(), "upload media finish failed: no file_token returned") { + t.Fatalf("unexpected error: %v", err) + } +} + +func TestParseDriveMediaUploadResponseErrors(t *testing.T) { + t.Parallel() + + t.Run("invalid json", func(t *testing.T) { + t.Parallel() + + _, err := ParseDriveMediaUploadResponse(&larkcore.ApiResp{RawBody: []byte("{")}, "upload media failed") + if err == nil || !strings.Contains(err.Error(), "invalid response JSON") { + t.Fatalf("expected invalid JSON error, got %v", err) + } + }) + + t.Run("api code error", func(t *testing.T) { + t.Parallel() + + _, err := ParseDriveMediaUploadResponse(&larkcore.ApiResp{RawBody: []byte(`{"code":999,"msg":"boom","error":{"detail":"x"}}`)}, "upload media failed") + if err == nil || !strings.Contains(err.Error(), "upload media failed: [999] boom") { + t.Fatalf("expected API error, got %v", err) + } + }) +} + +func TestExtractDriveMediaUploadFileTokenRequiresToken(t *testing.T) { + t.Parallel() + + _, err := ExtractDriveMediaUploadFileToken(map[string]interface{}{}, "upload media failed") + if err == nil || !strings.Contains(err.Error(), "upload media failed: no file_token returned") { + t.Fatalf("err = %v, want missing file_token error", err) + } +} + +func TestWrapDriveMediaUploadRequestError(t *testing.T) { + t.Parallel() + + t.Run("preserves exit error", func(t *testing.T) { + t.Parallel() + + original := output.ErrValidation("bad input") + got := WrapDriveMediaUploadRequestError(original, "upload media failed") + if got != original { + t.Fatalf("expected same exit error pointer, got %v", got) + } + }) + + t.Run("wraps generic error as network", func(t *testing.T) { + t.Parallel() + + got := WrapDriveMediaUploadRequestError(io.EOF, "upload media failed") + var exitErr *output.ExitError + if !errors.As(got, &exitErr) { + t.Fatalf("expected ExitError, got %T", got) + } + if exitErr.Code != output.ExitNetwork { + t.Fatalf("exit code = %d, want %d", exitErr.Code, output.ExitNetwork) + } + if !strings.Contains(got.Error(), "upload media failed") { + t.Fatalf("unexpected error: %v", got) + } + }) +} + +type capturedDriveMediaMultipartBody struct { + Fields map[string]string + Files map[string][]byte +} + +func newDriveMediaUploadTestRuntime(t *testing.T) (*RuntimeContext, *httpmock.Registry) { + t.Helper() + + cfg := &core.CliConfig{ + AppID: fmt.Sprintf("common-drive-media-test-%d", commonDriveMediaUploadTestSeq.Add(1)), AppSecret: "test-secret", Brand: core.BrandFeishu, + } + f, _, _, reg := cmdutil.TestFactory(t, cfg) + runtime := &RuntimeContext{ + ctx: context.Background(), + Config: cfg, + Factory: f, + resolvedAs: core.AsBot, + } + return runtime, reg +} + +func withDriveMediaUploadWorkingDir(t *testing.T, dir string) { + t.Helper() + + cwd, err := os.Getwd() + if err != nil { + t.Fatalf("Getwd() error: %v", err) + } + if err := os.Chdir(dir); err != nil { + t.Fatalf("Chdir(%q) error: %v", dir, err) + } + t.Cleanup(func() { + if err := os.Chdir(cwd); err != nil { + t.Fatalf("restore cwd error: %v", err) + } + }) +} + +func writeDriveMediaUploadTestFile(t *testing.T, name string, size int) string { + t.Helper() + + if err := os.WriteFile(name, bytes.Repeat([]byte("a"), size), 0644); err != nil { + t.Fatalf("WriteFile(%q) error: %v", name, err) + } + return name +} + +func writeDriveMediaUploadSizedFile(t *testing.T, name string, size int64) string { + t.Helper() + + fh, err := os.Create(name) + if err != nil { + t.Fatalf("Create(%q) error: %v", name, err) + } + if err := fh.Truncate(size); err != nil { + t.Fatalf("Truncate(%q) error: %v", name, err) + } + if err := fh.Close(); err != nil { + t.Fatalf("Close(%q) error: %v", name, err) + } + return name +} + +func decodeCapturedDriveMediaJSONBody(t *testing.T, stub *httpmock.Stub) map[string]interface{} { + t.Helper() + + var body map[string]interface{} + if err := json.Unmarshal(stub.CapturedBody, &body); err != nil { + t.Fatalf("decode captured JSON body: %v", err) + } + return body +} + +func decodeCapturedDriveMediaMultipartBody(t *testing.T, stub *httpmock.Stub) capturedDriveMediaMultipartBody { + t.Helper() + + contentType := stub.CapturedHeaders.Get("Content-Type") + mediaType, params, err := mime.ParseMediaType(contentType) + if err != nil { + t.Fatalf("parse multipart content type: %v", err) + } + if mediaType != "multipart/form-data" { + t.Fatalf("content type = %q, want multipart/form-data", mediaType) + } + + reader := multipart.NewReader(bytes.NewReader(stub.CapturedBody), params["boundary"]) + body := capturedDriveMediaMultipartBody{ + Fields: map[string]string{}, + Files: map[string][]byte{}, + } + for { + part, err := reader.NextPart() + if err == io.EOF { + break + } + if err != nil { + t.Fatalf("read multipart part: %v", err) + } + + data, err := io.ReadAll(part) + if err != nil { + t.Fatalf("read multipart data: %v", err) + } + if part.FileName() != "" { + body.Files[part.FormName()] = data + continue + } + body.Fields[part.FormName()] = string(data) + } + return body +} + +func strPtr(s string) *string { + return &s +} diff --git a/shortcuts/doc/doc_media_insert.go b/shortcuts/doc/doc_media_insert.go index 8242080e..635ed8ee 100644 --- a/shortcuts/doc/doc_media_insert.go +++ b/shortcuts/doc/doc_media_insert.go @@ -5,23 +5,15 @@ package doc import ( "context" - "encoding/json" - "errors" "fmt" - "net/http" "path/filepath" - larkcore "github.com/larksuite/oapi-sdk-go/v3/core" - "github.com/larksuite/cli/internal/output" - "github.com/larksuite/cli/internal/util" "github.com/larksuite/cli/internal/validate" "github.com/larksuite/cli/internal/vfs" "github.com/larksuite/cli/shortcuts/common" ) -const maxFileSize = 20 * 1024 * 1024 // 20MB - var alignMap = map[string]int{ "left": 1, "center": 2, @@ -36,7 +28,7 @@ var DocMediaInsert = common.Shortcut{ Scopes: []string{"docs:document.media:upload", "docx:document:write_only", "docx:document:readonly"}, AuthTypes: []string{"user", "bot"}, Flags: []common.Flag{ - {Name: "file", Desc: "local file path (max 20MB)", Required: true}, + {Name: "file", Desc: "local file path (files > 20MB use multipart upload automatically)", Required: true}, {Name: "doc", Desc: "document URL or document_id", Required: true}, {Name: "type", Default: "image", Desc: "type: image | file"}, {Name: "align", Desc: "alignment: left | center | right"}, @@ -86,16 +78,9 @@ var DocMediaInsert = common.Shortcut{ Desc(fmt.Sprintf("[%d] Get document root block", stepBase)). POST("/open-apis/docx/v1/documents/:document_id/blocks/:document_id/children"). Desc(fmt.Sprintf("[%d] Create empty block at document end", stepBase+1)). - Body(createBlockData). - POST("/open-apis/drive/v1/medias/upload_all"). - Desc(fmt.Sprintf("[%d] Upload local file (multipart/form-data)", stepBase+2)). - Body(map[string]interface{}{ - "file_name": filepath.Base(filePath), - "parent_type": parentType, - "parent_node": "", - "file": "@" + filePath, - }). - PATCH("/open-apis/docx/v1/documents/:document_id/blocks/batch_update"). + Body(createBlockData) + appendDocMediaInsertUploadDryRun(d, filePath, parentType, stepBase+2) + d.PATCH("/open-apis/docx/v1/documents/:document_id/blocks/batch_update"). Desc(fmt.Sprintf("[%d] Bind uploaded file token to the new block", stepBase+3)). Body(batchUpdateData) @@ -112,7 +97,6 @@ var DocMediaInsert = common.Shortcut{ if pathErr != nil { return output.ErrValidation("unsafe file path: %s", pathErr) } - filePath = safeFilePath documentID, err := resolveDocxDocumentID(runtime, docInput) if err != nil { @@ -120,16 +104,19 @@ var DocMediaInsert = common.Shortcut{ } // Validate file - stat, err := vfs.Stat(filePath) + stat, err := vfs.Stat(safeFilePath) if err != nil { return output.ErrValidation("file not found: %s", filePath) } - if stat.Size() > maxFileSize { - return output.ErrValidation("file %.1fMB exceeds 20MB limit", float64(stat.Size())/1024/1024) + if !stat.Mode().IsRegular() { + return output.ErrValidation("file must be a regular file: %s", filePath) } fileName := filepath.Base(filePath) fmt.Fprintf(runtime.IO().ErrOut, "Inserting: %s -> document %s\n", fileName, common.MaskToken(documentID)) + if stat.Size() > common.MaxDriveMediaUploadSinglePartSize { + fmt.Fprintf(runtime.IO().ErrOut, "File exceeds 20MB, using multipart upload\n") + } // Step 1: Get document root block to find where to insert rootData, err := runtime.CallAPI("GET", @@ -166,7 +153,8 @@ var DocMediaInsert = common.Shortcut{ fmt.Fprintf(runtime.IO().ErrOut, "Resolved file block targets: upload=%s replace=%s\n", uploadParentNode, replaceBlockID) } - // Rollback helper + // The placeholder block is created before any upload starts, so failures in + // later steps should try to remove it instead of leaving an empty artifact. rollback := func() error { fmt.Fprintf(runtime.IO().ErrOut, "Rolling back: deleting block %s\n", blockId) _, err := runtime.CallAPI("DELETE", @@ -185,7 +173,7 @@ var DocMediaInsert = common.Shortcut{ } // Step 3: Upload media file - fileToken, err := uploadMediaFile(ctx, runtime, filePath, fileName, mediaType, uploadParentNode, documentID) + fileToken, err := uploadDocMediaFile(runtime, filePath, fileName, stat.Size(), parentTypeForMediaType(mediaType), uploadParentNode, documentID) if err != nil { return withRollbackWarning(err) } @@ -346,6 +334,8 @@ func extractCreatedBlockTargets(createData map[string]interface{}, mediaType str return blockID, uploadParentNode, replaceBlockID } + // File blocks are wrapped: the created top-level block owns a nested child + // that is both the upload target and the replace_file target. nestedChildren, _ := child["children"].([]interface{}) if len(nestedChildren) == 0 { return blockID, uploadParentNode, replaceBlockID @@ -357,66 +347,43 @@ func extractCreatedBlockTargets(createData map[string]interface{}, mediaType str return blockID, uploadParentNode, replaceBlockID } -// uploadMediaFile uploads a file to Feishu drive as media. -func uploadMediaFile(ctx context.Context, runtime *common.RuntimeContext, filePath, fileName, mediaType, parentNode, docId string) (string, error) { - f, err := vfs.Open(filePath) - if err != nil { - return "", err - } - defer f.Close() - - stat, err := f.Stat() - if err != nil { - return "", output.Errorf(output.ExitInternal, "internal_error", "failed to stat file: %v", err) - } - fileSize := stat.Size() - - parentType := parentTypeForMediaType(mediaType) - - // Build SDK Formdata - fd := larkcore.NewFormdata() - fd.AddField("file_name", fileName) - fd.AddField("parent_type", parentType) - fd.AddField("parent_node", parentNode) - fd.AddField("size", fmt.Sprintf("%d", fileSize)) - if docId != "" { - extra, err := buildDriveRouteExtra(docId) - if err != nil { - return "", err - } - fd.AddField("extra", extra) - } - fd.AddFile("file", f) - - apiResp, err := runtime.DoAPI(&larkcore.ApiReq{ - HttpMethod: http.MethodPost, - ApiPath: "/open-apis/drive/v1/medias/upload_all", - Body: fd, - }, larkcore.WithFileUpload()) - if err != nil { - var exitErr *output.ExitError - if errors.As(err, &exitErr) { - return "", err - } - return "", output.ErrNetwork("file upload failed: %v", err) - } - - var result map[string]interface{} - if err := json.Unmarshal(apiResp.RawBody, &result); err != nil { - return "", output.Errorf(output.ExitAPI, "api_error", "file upload failed: invalid response JSON: %v", err) - } - - code, _ := util.ToFloat64(result["code"]) - if code != 0 { - msg, _ := result["msg"].(string) - return "", output.ErrAPI(int(code), fmt.Sprintf("file upload failed: [%d] %s", int(code), msg), result["error"]) - } - - data, _ := result["data"].(map[string]interface{}) - fileToken, _ := data["file_token"].(string) - if fileToken == "" { - return "", output.Errorf(output.ExitAPI, "api_error", "file upload failed: no file_token returned") +func appendDocMediaInsertUploadDryRun(d *common.DryRunAPI, filePath, parentType string, step int) { + // The upload step runs only after the empty placeholder block is created, so + // dry-run can refer to that future block ID only symbolically. For large + // files, keep multipart internals as substeps of the single user-facing + // "upload file" step. + if docMediaShouldUseMultipart(filePath) { + d.POST("/open-apis/drive/v1/medias/upload_prepare"). + Desc(fmt.Sprintf("[%da] Initialize multipart upload", step)). + Body(map[string]interface{}{ + "file_name": filepath.Base(filePath), + "parent_type": parentType, + "parent_node": "", + "size": "", + }). + POST("/open-apis/drive/v1/medias/upload_part"). + Desc(fmt.Sprintf("[%db] Upload file parts (repeated)", step)). + Body(map[string]interface{}{ + "upload_id": "", + "seq": "", + "size": "", + "file": "", + }). + POST("/open-apis/drive/v1/medias/upload_finish"). + Desc(fmt.Sprintf("[%dc] Finalize multipart upload and get file_token", step)). + Body(map[string]interface{}{ + "upload_id": "", + "block_num": "", + }) + return } - return fileToken, nil + d.POST("/open-apis/drive/v1/medias/upload_all"). + Desc(fmt.Sprintf("[%d] Upload local file (multipart/form-data)", step)). + Body(map[string]interface{}{ + "file_name": filepath.Base(filePath), + "parent_type": parentType, + "parent_node": "", + "file": "@" + filePath, + }) } diff --git a/shortcuts/doc/doc_media_test.go b/shortcuts/doc/doc_media_test.go index eab88a73..ee810f67 100644 --- a/shortcuts/doc/doc_media_test.go +++ b/shortcuts/doc/doc_media_test.go @@ -5,6 +5,8 @@ package doc import ( "bytes" + "context" + "encoding/json" "net/http" "os" "path/filepath" @@ -19,10 +21,6 @@ import ( "github.com/larksuite/cli/shortcuts/common" ) -func docsTestConfig() *core.CliConfig { - return docsTestConfigWithAppID("docs-test-app") -} - func docsTestConfigWithAppID(appID string) *core.CliConfig { return &core.CliConfig{ AppID: appID, AppSecret: "test-secret", Brand: core.BrandFeishu, @@ -59,7 +57,7 @@ func withDocsWorkingDir(t *testing.T, dir string) { } func TestDocMediaInsertRejectsOldDocURL(t *testing.T) { - f, _, _, _ := cmdutil.TestFactory(t, docsTestConfig()) + f, _, _, _ := cmdutil.TestFactory(t, docsTestConfigWithAppID("docs-test-app")) err := mountAndRunDocs(t, DocMediaInsert, []string{ "+media-insert", @@ -77,7 +75,7 @@ func TestDocMediaInsertRejectsOldDocURL(t *testing.T) { } func TestDocMediaInsertDryRunWikiAddsResolveStep(t *testing.T) { - f, stdout, _, _ := cmdutil.TestFactory(t, docsTestConfig()) + f, stdout, _, _ := cmdutil.TestFactory(t, docsTestConfigWithAppID("docs-test-app")) err := mountAndRunDocs(t, DocMediaInsert, []string{ "+media-insert", @@ -99,6 +97,98 @@ func TestDocMediaInsertDryRunWikiAddsResolveStep(t *testing.T) { } } +func TestDocMediaUploadDryRunUsesMultipartForLargeFile(t *testing.T) { + tmpDir := t.TempDir() + withDocsWorkingDir(t, tmpDir) + writeSizedDocTestFile(t, "large.bin", common.MaxDriveMediaUploadSinglePartSize+1) + + cmd := &cobra.Command{Use: "docs +media-upload"} + cmd.Flags().String("file", "", "") + cmd.Flags().String("parent-type", "", "") + cmd.Flags().String("parent-node", "", "") + cmd.Flags().String("doc-id", "", "") + if err := cmd.Flags().Set("file", "./large.bin"); err != nil { + t.Fatalf("set --file: %v", err) + } + if err := cmd.Flags().Set("parent-type", "docx_file"); err != nil { + t.Fatalf("set --parent-type: %v", err) + } + if err := cmd.Flags().Set("parent-node", "blk_parent"); err != nil { + t.Fatalf("set --parent-node: %v", err) + } + + dry := decodeDocDryRun(t, MediaUpload.DryRun(context.Background(), common.TestNewRuntimeContext(cmd, nil))) + if dry.Description != "chunked media upload (files > 20MB)" { + t.Fatalf("dry-run description = %q", dry.Description) + } + if len(dry.API) != 3 { + t.Fatalf("expected 3 API calls, got %d", len(dry.API)) + } + if dry.API[0].URL != "/open-apis/drive/v1/medias/upload_prepare" { + t.Fatalf("first URL = %q, want upload_prepare", dry.API[0].URL) + } + if dry.API[1].URL != "/open-apis/drive/v1/medias/upload_part" { + t.Fatalf("second URL = %q, want upload_part", dry.API[1].URL) + } + if dry.API[2].URL != "/open-apis/drive/v1/medias/upload_finish" { + t.Fatalf("third URL = %q, want upload_finish", dry.API[2].URL) + } + if got, _ := dry.API[0].Body["parent_node"].(string); got != "blk_parent" { + t.Fatalf("prepare parent_node = %q, want %q", got, "blk_parent") + } +} + +func TestDocMediaInsertDryRunUsesMultipartForLargeFile(t *testing.T) { + tmpDir := t.TempDir() + withDocsWorkingDir(t, tmpDir) + writeSizedDocTestFile(t, "large.bin", common.MaxDriveMediaUploadSinglePartSize+1) + + cmd := &cobra.Command{Use: "docs +media-insert"} + cmd.Flags().String("file", "", "") + cmd.Flags().String("doc", "", "") + cmd.Flags().String("type", "", "") + cmd.Flags().String("align", "", "") + cmd.Flags().String("caption", "", "") + if err := cmd.Flags().Set("doc", "doxcnDryRunLarge"); err != nil { + t.Fatalf("set --doc: %v", err) + } + if err := cmd.Flags().Set("file", "./large.bin"); err != nil { + t.Fatalf("set --file: %v", err) + } + + dry := decodeDocDryRun(t, DocMediaInsert.DryRun(context.Background(), common.TestNewRuntimeContext(cmd, nil))) + if dry.Description != "4-step orchestration: query root → create block → upload file → bind to block (auto-rollback on failure)" { + t.Fatalf("dry-run description = %q", dry.Description) + } + if len(dry.API) != 6 { + t.Fatalf("expected 6 API calls, got %d", len(dry.API)) + } + if dry.API[2].URL != "/open-apis/drive/v1/medias/upload_prepare" { + t.Fatalf("third URL = %q, want upload_prepare", dry.API[2].URL) + } + if dry.API[3].URL != "/open-apis/drive/v1/medias/upload_part" { + t.Fatalf("fourth URL = %q, want upload_part", dry.API[3].URL) + } + if dry.API[4].URL != "/open-apis/drive/v1/medias/upload_finish" { + t.Fatalf("fifth URL = %q, want upload_finish", dry.API[4].URL) + } + if dry.API[5].URL != "/open-apis/docx/v1/documents/doxcnDryRunLarge/blocks/batch_update" { + t.Fatalf("last URL = %q, want batch_update", dry.API[5].URL) + } + if !strings.Contains(dry.API[2].Desc, "[3a]") { + t.Fatalf("upload_prepare desc = %q, want [3a] step marker", dry.API[2].Desc) + } + if !strings.Contains(dry.API[3].Desc, "[3b]") { + t.Fatalf("upload_part desc = %q, want [3b] step marker", dry.API[3].Desc) + } + if !strings.Contains(dry.API[4].Desc, "[3c]") { + t.Fatalf("upload_finish desc = %q, want [3c] step marker", dry.API[4].Desc) + } + if !strings.Contains(dry.API[5].Desc, "[4]") { + t.Fatalf("batch_update desc = %q, want [4] step marker", dry.API[5].Desc) + } +} + func TestDocMediaInsertExecuteResolvesWikiBeforeFileCheck(t *testing.T) { f, _, stderr, reg := cmdutil.TestFactory(t, docsTestConfigWithAppID("docs-insert-exec-app")) reg.Register(&httpmock.Stub{ @@ -194,3 +284,42 @@ func TestDocMediaDownloadRejectsHTTPErrorBeforeWrite(t *testing.T) { t.Fatalf("download target should not be created, statErr=%v", statErr) } } + +type docDryRunOutput struct { + Description string `json:"description"` + API []struct { + Desc string `json:"desc"` + URL string `json:"url"` + Body map[string]interface{} `json:"body"` + } `json:"api"` +} + +func writeSizedDocTestFile(t *testing.T, name string, size int64) { + t.Helper() + + fh, err := os.Create(name) + if err != nil { + t.Fatalf("Create(%q) error: %v", name, err) + } + if err := fh.Truncate(size); err != nil { + t.Fatalf("Truncate(%q) error: %v", name, err) + } + if err := fh.Close(); err != nil { + t.Fatalf("Close(%q) error: %v", name, err) + } +} + +func decodeDocDryRun(t *testing.T, dryAPI *common.DryRunAPI) docDryRunOutput { + t.Helper() + + raw, err := json.Marshal(dryAPI) + if err != nil { + t.Fatalf("marshal dry-run output: %v", err) + } + + var dry docDryRunOutput + if err := json.Unmarshal(raw, &dry); err != nil { + t.Fatalf("decode dry-run output: %v", err) + } + return dry +} diff --git a/shortcuts/doc/doc_media_upload.go b/shortcuts/doc/doc_media_upload.go index 39db9300..5b072793 100644 --- a/shortcuts/doc/doc_media_upload.go +++ b/shortcuts/doc/doc_media_upload.go @@ -5,16 +5,10 @@ package doc import ( "context" - "encoding/json" - "errors" "fmt" - "net/http" "path/filepath" - larkcore "github.com/larksuite/oapi-sdk-go/v3/core" - "github.com/larksuite/cli/internal/output" - "github.com/larksuite/cli/internal/util" "github.com/larksuite/cli/internal/validate" "github.com/larksuite/cli/internal/vfs" "github.com/larksuite/cli/shortcuts/common" @@ -28,7 +22,7 @@ var MediaUpload = common.Shortcut{ Scopes: []string{"docs:document.media:upload"}, AuthTypes: []string{"user", "bot"}, Flags: []common.Flag{ - {Name: "file", Desc: "local file path (max 20MB)", Required: true}, + {Name: "file", Desc: "local file path (files > 20MB use multipart upload automatically)", Required: true}, {Name: "parent-type", Desc: "parent type: docx_image | docx_file", Required: true}, {Name: "parent-node", Desc: "parent node ID (block_id)", Required: true}, {Name: "doc-id", Desc: "document ID (for drive_route_token)"}, @@ -42,13 +36,41 @@ var MediaUpload = common.Shortcut{ "file_name": filepath.Base(filePath), "parent_type": parentType, "parent_node": parentNode, - "file": "@" + filePath, } if docId != "" { body["extra"] = fmt.Sprintf(`{"drive_route_token":"%s"}`, docId) } - return common.NewDryRunAPI(). - Desc("multipart/form-data upload"). + dry := common.NewDryRunAPI() + if docMediaShouldUseMultipart(filePath) { + prepareBody := map[string]interface{}{ + "file_name": filepath.Base(filePath), + "parent_type": parentType, + "parent_node": parentNode, + "size": "", + } + if extra, ok := body["extra"]; ok { + prepareBody["extra"] = extra + } + dry.Desc("chunked media upload (files > 20MB)"). + POST("/open-apis/drive/v1/medias/upload_prepare"). + Body(prepareBody). + POST("/open-apis/drive/v1/medias/upload_part"). + Body(map[string]interface{}{ + "upload_id": "", + "seq": "", + "size": "", + "file": "", + }). + POST("/open-apis/drive/v1/medias/upload_finish"). + Body(map[string]interface{}{ + "upload_id": "", + "block_num": "", + }) + return dry + } + + body["file"] = "@" + filePath + return dry.Desc("multipart/form-data upload"). POST("/open-apis/drive/v1/medias/upload_all"). Body(body) }, @@ -62,69 +84,25 @@ var MediaUpload = common.Shortcut{ if pathErr != nil { return output.ErrValidation("unsafe file path: %s", pathErr) } - filePath = safeFilePath // Validate file - stat, err := vfs.Stat(filePath) + stat, err := vfs.Stat(safeFilePath) if err != nil { return output.ErrValidation("file not found: %s", filePath) } - if stat.Size() > maxFileSize { - return output.ErrValidation("file %.1fMB exceeds 20MB limit", float64(stat.Size())/1024/1024) + if !stat.Mode().IsRegular() { + return output.ErrValidation("file must be a regular file: %s", filePath) } fileName := filepath.Base(filePath) fmt.Fprintf(runtime.IO().ErrOut, "Uploading: %s (%d bytes)\n", fileName, stat.Size()) - - f, err := vfs.Open(filePath) - if err != nil { - return output.ErrValidation("cannot open file: %v", err) + if stat.Size() > common.MaxDriveMediaUploadSinglePartSize { + fmt.Fprintf(runtime.IO().ErrOut, "File exceeds 20MB, using multipart upload\n") } - defer f.Close() - // Build SDK Formdata - fd := larkcore.NewFormdata() - fd.AddField("file_name", fileName) - fd.AddField("parent_type", parentType) - fd.AddField("parent_node", parentNode) - fd.AddField("size", fmt.Sprintf("%d", stat.Size())) - if docId != "" { - extra, err := buildDriveRouteExtra(docId) - if err != nil { - return err - } - fd.AddField("extra", extra) - } - fd.AddFile("file", f) - - apiResp, err := runtime.DoAPI(&larkcore.ApiReq{ - HttpMethod: http.MethodPost, - ApiPath: "/open-apis/drive/v1/medias/upload_all", - Body: fd, - }, larkcore.WithFileUpload()) + fileToken, err := uploadDocMediaFile(runtime, filePath, fileName, stat.Size(), parentType, parentNode, docId) if err != nil { - var exitErr *output.ExitError - if errors.As(err, &exitErr) { - return err - } - return output.ErrNetwork("upload failed: %v", err) - } - - var result map[string]interface{} - if err := json.Unmarshal(apiResp.RawBody, &result); err != nil { - return output.Errorf(output.ExitAPI, "api_error", "upload failed: invalid response JSON: %v", err) - } - - code, _ := util.ToFloat64(result["code"]) - if code != 0 { - msg, _ := result["msg"].(string) - return output.ErrAPI(int(code), fmt.Sprintf("upload failed: [%d] %s", int(code), msg), result["error"]) - } - - data, _ := result["data"].(map[string]interface{}) - fileToken, _ := data["file_token"].(string) - if fileToken == "" { - return output.Errorf(output.ExitAPI, "api_error", "upload failed: no file_token returned") + return err } runtime.Out(map[string]interface{}{ @@ -135,3 +113,49 @@ var MediaUpload = common.Shortcut{ return nil }, } + +func uploadDocMediaFile(runtime *common.RuntimeContext, filePath, fileName string, fileSize int64, parentType, parentNode, docID string) (string, error) { + var extra string + if docID != "" { + var err error + extra, err = buildDriveRouteExtra(docID) + if err != nil { + return "", err + } + } + + // Doc media uploads share the generic Drive media transport. The doc-specific + // routing only shows up in parent_type/parent_node and optional route extra. + if fileSize <= common.MaxDriveMediaUploadSinglePartSize { + return common.UploadDriveMediaAll(runtime, common.DriveMediaUploadAllConfig{ + FilePath: filePath, + FileName: fileName, + FileSize: fileSize, + ParentType: parentType, + ParentNode: &parentNode, + Extra: extra, + }) + } + return common.UploadDriveMediaMultipart(runtime, common.DriveMediaMultipartUploadConfig{ + FilePath: filePath, + FileName: fileName, + FileSize: fileSize, + ParentType: parentType, + ParentNode: parentNode, + Extra: extra, + }) +} + +func docMediaShouldUseMultipart(filePath string) bool { + // Dry-run uses local stat as a best-effort planning hint. Execute re-validates + // the file before choosing the actual upload path. + safeFilePath, err := validate.SafeInputPath(filePath) + if err != nil { + return false + } + info, err := vfs.Stat(safeFilePath) + if err != nil { + return false + } + return info.Mode().IsRegular() && info.Size() > common.MaxDriveMediaUploadSinglePartSize +} diff --git a/shortcuts/drive/drive_export_test.go b/shortcuts/drive/drive_export_test.go index 46ba0bfa..fed45ed1 100644 --- a/shortcuts/drive/drive_export_test.go +++ b/shortcuts/drive/drive_export_test.go @@ -65,7 +65,6 @@ func TestValidateDriveExportSpec(t *testing.T) { func TestDriveExportMarkdownWritesFile(t *testing.T) { f, stdout, _, reg := cmdutil.TestFactory(t, driveTestConfig()) - registerDriveBotTokenStub(reg) reg.Register(&httpmock.Stub{ Method: "GET", URL: "/open-apis/docs/v1/content", @@ -117,7 +116,6 @@ func TestDriveExportMarkdownWritesFile(t *testing.T) { func TestDriveExportAsyncSuccess(t *testing.T) { f, stdout, _, reg := cmdutil.TestFactory(t, driveTestConfig()) - registerDriveBotTokenStub(reg) reg.Register(&httpmock.Stub{ Method: "POST", URL: "/open-apis/drive/v1/export_tasks", @@ -188,7 +186,6 @@ func TestDriveExportAsyncSuccess(t *testing.T) { func TestDriveExportReadyDownloadFailureIncludesRecoveryHint(t *testing.T) { f, _, _, reg := cmdutil.TestFactory(t, driveTestConfig()) - registerDriveBotTokenStub(reg) reg.Register(&httpmock.Stub{ Method: "POST", URL: "/open-apis/drive/v1/export_tasks", @@ -267,7 +264,6 @@ func TestDriveExportReadyDownloadFailureIncludesRecoveryHint(t *testing.T) { func TestDriveExportTimeoutReturnsFollowUpCommand(t *testing.T) { f, stdout, _, reg := cmdutil.TestFactory(t, driveTestConfig()) - registerDriveBotTokenStub(reg) reg.Register(&httpmock.Stub{ Method: "POST", URL: "/open-apis/drive/v1/export_tasks", @@ -333,7 +329,6 @@ func TestDriveExportTimeoutReturnsFollowUpCommand(t *testing.T) { func TestDriveExportPollErrorsReturnLastErrorWithRecoveryHint(t *testing.T) { f, stdout, _, reg := cmdutil.TestFactory(t, driveTestConfig()) - registerDriveBotTokenStub(reg) reg.Register(&httpmock.Stub{ Method: "POST", URL: "/open-apis/drive/v1/export_tasks", @@ -389,7 +384,6 @@ func TestDriveExportPollErrorsReturnLastErrorWithRecoveryHint(t *testing.T) { func TestDriveExportDownloadUsesProvidedFileName(t *testing.T) { f, stdout, _, reg := cmdutil.TestFactory(t, driveTestConfig()) - registerDriveBotTokenStub(reg) reg.Register(&httpmock.Stub{ Method: "GET", URL: "/open-apis/drive/v1/export_tasks/file/box_789/download", @@ -424,7 +418,6 @@ func TestDriveExportDownloadUsesProvidedFileName(t *testing.T) { func TestDriveExportDownloadRejectsOverwriteWithoutFlag(t *testing.T) { f, _, _, reg := cmdutil.TestFactory(t, driveTestConfig()) - registerDriveBotTokenStub(reg) reg.Register(&httpmock.Stub{ Method: "GET", URL: "/open-apis/drive/v1/export_tasks/file/box_dup/download", @@ -480,7 +473,6 @@ func TestSaveContentToOutputDirRejectsOverwriteWithoutFlag(t *testing.T) { func TestDriveTaskResultExportIncludesReadyFlags(t *testing.T) { f, stdout, _, reg := cmdutil.TestFactory(t, driveTestConfig()) - registerDriveBotTokenStub(reg) reg.Register(&httpmock.Stub{ Method: "GET", URL: "/open-apis/drive/v1/export_tasks/tk_export", diff --git a/shortcuts/drive/drive_import.go b/shortcuts/drive/drive_import.go index f2aed91e..9c25b2af 100644 --- a/shortcuts/drive/drive_import.go +++ b/shortcuts/drive/drive_import.go @@ -7,7 +7,6 @@ import ( "context" "fmt" "path/filepath" - "strings" "github.com/larksuite/cli/internal/vfs" @@ -147,9 +146,8 @@ func preflightDriveImportFile(spec *driveImportSpec) (int64, error) { if err != nil { return 0, output.ErrValidation("unsafe file path: %s", err) } - spec.FilePath = safeFilePath - info, err := vfs.Stat(spec.FilePath) + info, err := vfs.Stat(safeFilePath) if err != nil { return 0, output.ErrValidation("cannot read file: %s", err) } @@ -168,7 +166,7 @@ func appendDriveImportUploadDryRun(dry *common.DryRunAPI, spec driveImportSpec, extra = fmt.Sprintf(`{"obj_type":"%s","file_extension":"%s"}`, spec.DocType, spec.FileExtension()) } - if fileSize > maxDriveUploadFileSize { + if fileSize > common.MaxDriveMediaUploadSinglePartSize { dry.POST("/open-apis/drive/v1/medias/upload_prepare"). Desc("[1a] Initialize multipart upload"). Body(map[string]interface{}{ diff --git a/shortcuts/drive/drive_import_common.go b/shortcuts/drive/drive_import_common.go index f3b61c47..35c14772 100644 --- a/shortcuts/drive/drive_import_common.go +++ b/shortcuts/drive/drive_import_common.go @@ -4,21 +4,15 @@ package drive import ( - "bytes" "context" "encoding/json" - "errors" "fmt" - "io" - "net/http" "path/filepath" "strings" "time" "github.com/larksuite/cli/internal/vfs" - larkcore "github.com/larksuite/oapi-sdk-go/v3/core" - "github.com/larksuite/cli/internal/output" "github.com/larksuite/cli/internal/validate" "github.com/larksuite/cli/shortcuts/common" @@ -37,12 +31,6 @@ const ( driveImport800MBFileSizeLimit int64 = 800 * 1024 * 1024 ) -type driveMultipartUploadSession struct { - UploadID string - BlockSize int - BlockNum int -} - // driveImportExtToDocTypes defines which source file extensions can be imported // into which Drive-native document types. var driveImportExtToDocTypes = map[string][]string{ @@ -106,163 +94,41 @@ func uploadMediaForImport(ctx context.Context, runtime *common.RuntimeContext, f if err = validateDriveImportFileSize(filePath, docType, fileSize); err != nil { return "", err } - fileSizeValue, err := driveUploadSizeValue(fileSize) - if err != nil { - return "", err - } extra, err := buildImportMediaExtra(filePath, docType) if err != nil { return "", err } - if fileSize <= maxDriveUploadFileSize { + if fileSize <= common.MaxDriveMediaUploadSinglePartSize { fmt.Fprintf(runtime.IO().ErrOut, "Uploading media for import: %s (%s)\n", fileName, common.FormatSize(fileSize)) - return uploadMediaForImportAll(runtime, filePath, fileName, fileSizeValue, extra) + // upload_all for import works without parent_node; omitting it preserves + // the existing root-level import staging behavior. + return common.UploadDriveMediaAll(runtime, common.DriveMediaUploadAllConfig{ + FilePath: filePath, + FileName: fileName, + FileSize: fileSize, + ParentType: "ccm_import_open", + Extra: extra, + }) } fmt.Fprintf(runtime.IO().ErrOut, "Uploading media for import via multipart upload: %s (%s)\n", fileName, common.FormatSize(fileSize)) - return uploadMediaForImportMultipart(runtime, filePath, fileName, fileSizeValue, extra) -} - -func uploadMediaForImportAll(runtime *common.RuntimeContext, filePath, fileName string, fileSize int, extra string) (string, error) { - f, err := vfs.Open(filePath) - if err != nil { - return "", output.ErrValidation("cannot read file: %s", err) - } - defer f.Close() - - fd := larkcore.NewFormdata() - fd.AddField("file_name", fileName) - fd.AddField("parent_type", "ccm_import_open") - fd.AddField("size", fmt.Sprintf("%d", fileSize)) - fd.AddField("extra", extra) - fd.AddFile("file", f) - - apiResp, err := runtime.DoAPI(&larkcore.ApiReq{ - HttpMethod: http.MethodPost, - ApiPath: "/open-apis/drive/v1/medias/upload_all", - Body: fd, - }, larkcore.WithFileUpload()) - if err != nil { - return "", wrapDriveUploadRequestError(err, "upload media failed") - } - - data, err := parseDriveUploadResponse(apiResp, "upload media failed") - if err != nil { - return "", err - } - return extractDriveUploadFileToken(data, "upload media failed") -} - -func uploadMediaForImportMultipart(runtime *common.RuntimeContext, filePath, fileName string, fileSize int, extra string) (string, error) { - session, err := prepareMediaImportUpload(runtime, fileName, fileSize, extra) - if err != nil { - fmt.Fprintf(runtime.IO().ErrOut, "Multipart upload prepare failed: %s\n", err) - return "", err - } - - totalBlocks := session.BlockNum - fmt.Fprintf(runtime.IO().ErrOut, "Multipart upload initialized: %d chunks x %s\n", totalBlocks, common.FormatSize(int64(session.BlockSize))) - - f, err := vfs.Open(filePath) - if err != nil { - return "", output.ErrValidation("cannot read file: %s", err) - } - defer f.Close() - - buffer := make([]byte, session.BlockSize) - remaining := fileSize - uploadedBlocks := 0 - for remaining > 0 { - chunkSize := session.BlockSize - if chunkSize > remaining { - chunkSize = remaining - } - - n, readErr := io.ReadFull(f, buffer[:chunkSize]) - if readErr != nil { - return "", output.ErrValidation("cannot read file: %s", readErr) - } - - if err = uploadMediaImportPart(runtime, session.UploadID, uploadedBlocks, buffer[:n]); err != nil { - fmt.Fprintf(runtime.IO().ErrOut, "Multipart upload part failed: %s\n", err) - return "", err - } - - remaining -= n - uploadedBlocks++ - } - - if session.BlockNum > 0 && session.BlockNum != uploadedBlocks { - return "", output.Errorf(output.ExitAPI, "api_error", "upload prepare mismatch: expected %d blocks, uploaded %d", session.BlockNum, uploadedBlocks) - } - - return finishMediaImportUpload(runtime, session.UploadID, uploadedBlocks) -} - -func prepareMediaImportUpload(runtime *common.RuntimeContext, fileName string, fileSize int, extra string) (driveMultipartUploadSession, error) { - data, err := runtime.CallAPI("POST", "/open-apis/drive/v1/medias/upload_prepare", nil, map[string]interface{}{ - "file_name": fileName, - "parent_type": "ccm_import_open", // For media import uploads, parent_type must be ccm_import_open. - "size": fileSize, - "extra": extra, - "parent_node": "", // For media import uploads, parent_node must be an explicit empty string; unlike medias/upload_all, this field cannot be omitted. - }) - if err != nil { - return driveMultipartUploadSession{}, err - } - - session := driveMultipartUploadSession{ - UploadID: common.GetString(data, "upload_id"), - BlockSize: int(common.GetFloat(data, "block_size")), - BlockNum: int(common.GetFloat(data, "block_num")), - } - if session.UploadID == "" { - return driveMultipartUploadSession{}, output.Errorf(output.ExitAPI, "api_error", "upload prepare failed: no upload_id returned") - } - if session.BlockSize <= 0 { - return driveMultipartUploadSession{}, output.Errorf(output.ExitAPI, "api_error", "upload prepare failed: invalid block_size returned") - } - if session.BlockNum <= 0 { - return driveMultipartUploadSession{}, output.Errorf(output.ExitAPI, "api_error", "upload prepare failed: invalid block_num returned") - } - return session, nil -} - -func uploadMediaImportPart(runtime *common.RuntimeContext, uploadID string, seq int, chunk []byte) error { - fd := larkcore.NewFormdata() - fd.AddField("upload_id", uploadID) - fd.AddField("seq", seq) - fd.AddField("size", len(chunk)) - fd.AddFile("file", bytes.NewReader(chunk)) - - apiResp, err := runtime.DoAPI(&larkcore.ApiReq{ - HttpMethod: http.MethodPost, - ApiPath: "/open-apis/drive/v1/medias/upload_part", - Body: fd, - }, larkcore.WithFileUpload()) - if err != nil { - return wrapDriveUploadRequestError(err, "upload media part failed") - } - - _, err = parseDriveUploadResponse(apiResp, "upload media part failed") - return err -} - -func finishMediaImportUpload(runtime *common.RuntimeContext, uploadID string, blockNum int) (string, error) { - data, err := runtime.CallAPI("POST", "/open-apis/drive/v1/medias/upload_finish", nil, map[string]interface{}{ - "upload_id": uploadID, - "block_num": blockNum, + // upload_prepare is stricter than upload_all here and expects parent_node to + // be sent explicitly, even when import uses the implicit root staging area. + return common.UploadDriveMediaMultipart(runtime, common.DriveMediaMultipartUploadConfig{ + FilePath: filePath, + FileName: fileName, + FileSize: fileSize, + ParentType: "ccm_import_open", + ParentNode: "", + Extra: extra, }) - if err != nil { - fmt.Fprintf(runtime.IO().ErrOut, "Multipart upload finish failed: %s\n", err) - return "", err - } - return extractDriveUploadFileToken(data, "upload media finish failed") } func buildImportMediaExtra(filePath, docType string) (string, error) { + // The import media endpoint uses extra to decide both the target native type + // and how to interpret the uploaded source file. extraBytes, err := json.Marshal(map[string]string{ "obj_type": docType, "file_extension": strings.TrimPrefix(strings.ToLower(filepath.Ext(filePath)), "."), @@ -318,45 +184,6 @@ func validateDriveImportFileSize(filePath, docType string, fileSize int64) error ) } -func driveUploadSizeValue(fileSize int64) (int, error) { - maxInt := int64(^uint(0) >> 1) - if fileSize > maxInt { - return 0, output.ErrValidation("file %s is too large to upload", common.FormatSize(fileSize)) - } - return int(fileSize), nil -} - -func wrapDriveUploadRequestError(err error, action string) error { - var exitErr *output.ExitError - if errors.As(err, &exitErr) { - return err - } - return output.ErrNetwork("%s: %v", action, err) -} - -func parseDriveUploadResponse(apiResp *larkcore.ApiResp, action string) (map[string]interface{}, error) { - var result map[string]interface{} - if err := json.Unmarshal(apiResp.RawBody, &result); err != nil { - return nil, output.Errorf(output.ExitAPI, "api_error", "%s: invalid response JSON: %v", action, err) - } - - if larkCode := int(common.GetFloat(result, "code")); larkCode != 0 { - msg, _ := result["msg"].(string) - return nil, output.ErrAPI(larkCode, fmt.Sprintf("%s: [%d] %s", action, larkCode, msg), result["error"]) - } - - data, _ := result["data"].(map[string]interface{}) - return data, nil -} - -func extractDriveUploadFileToken(data map[string]interface{}, action string) (string, error) { - fileToken := common.GetString(data, "file_token") - if fileToken == "" { - return "", output.Errorf(output.ExitAPI, "api_error", "%s: no file_token returned", action) - } - return fileToken, nil -} - // validateDriveImportSpec enforces the CLI-level compatibility rules before any // upload or import request is sent to the backend. func validateDriveImportSpec(spec driveImportSpec) error { diff --git a/shortcuts/drive/drive_import_common_test.go b/shortcuts/drive/drive_import_common_test.go index 85b97763..c546134c 100644 --- a/shortcuts/drive/drive_import_common_test.go +++ b/shortcuts/drive/drive_import_common_test.go @@ -5,20 +5,12 @@ package drive import ( "bytes" - "encoding/json" - "errors" - "io" - "mime" - "mime/multipart" "os" "strings" "testing" - larkcore "github.com/larksuite/oapi-sdk-go/v3/core" - "github.com/larksuite/cli/internal/cmdutil" "github.com/larksuite/cli/internal/httpmock" - "github.com/larksuite/cli/internal/output" ) func TestValidateDriveImportSpecRejectsMismatchedType(t *testing.T) { @@ -144,7 +136,6 @@ func TestDriveImportStatusPendingWithoutToken(t *testing.T) { func TestDriveImportTimeoutReturnsFollowUpCommand(t *testing.T) { f, stdout, _, reg := cmdutil.TestFactory(t, driveTestConfig()) - registerDriveBotTokenStub(reg) reg.Register(&httpmock.Stub{ Method: "POST", URL: "/open-apis/drive/v1/medias/upload_all", @@ -207,295 +198,6 @@ func TestDriveImportTimeoutReturnsFollowUpCommand(t *testing.T) { } } -func TestDriveImportUsesMultipartUploadForLargeFile(t *testing.T) { - f, stdout, _, reg := cmdutil.TestFactory(t, driveTestConfig()) - registerDriveBotTokenStub(reg) - - prepareStub := &httpmock.Stub{ - Method: "POST", - URL: "/open-apis/drive/v1/medias/upload_prepare", - Body: map[string]interface{}{ - "code": 0, - "data": map[string]interface{}{ - "upload_id": "upload_123", - "block_size": 4 * 1024 * 1024, - "block_num": 6, - }, - }, - } - reg.Register(prepareStub) - - partStubs := make([]*httpmock.Stub, 0, 6) - for i := 0; i < 6; i++ { - stub := &httpmock.Stub{ - Method: "POST", - URL: "/open-apis/drive/v1/medias/upload_part", - Body: map[string]interface{}{ - "code": 0, - "msg": "ok", - }, - } - partStubs = append(partStubs, stub) - reg.Register(stub) - } - - finishStub := &httpmock.Stub{ - Method: "POST", - URL: "/open-apis/drive/v1/medias/upload_finish", - Body: map[string]interface{}{ - "code": 0, - "data": map[string]interface{}{ - "file_token": "file_123", - }, - }, - } - reg.Register(finishStub) - reg.Register(&httpmock.Stub{ - Method: "POST", - URL: "/open-apis/drive/v1/import_tasks", - Body: map[string]interface{}{ - "code": 0, - "data": map[string]interface{}{"ticket": "tk_import"}, - }, - }) - reg.Register(&httpmock.Stub{ - Method: "GET", - URL: "/open-apis/drive/v1/import_tasks/tk_import", - Body: map[string]interface{}{ - "code": 0, - "data": map[string]interface{}{ - "result": map[string]interface{}{ - "type": "sheet", - "job_status": 0, - "token": "sheet_123", - "url": "https://example.com/sheets/sheet_123", - }, - }, - }, - }) - - tmpDir := t.TempDir() - withDriveWorkingDir(t, tmpDir) - writeSizedDriveImportFile(t, "large.xlsx", int64(maxDriveUploadFileSize)+1) - - err := mountAndRunDrive(t, DriveImport, []string{ - "+import", - "--file", "large.xlsx", - "--type", "sheet", - "--as", "bot", - }, f, stdout) - if err != nil { - t.Fatalf("unexpected error: %v", err) - } - - if !bytes.Contains(stdout.Bytes(), []byte(`"token": "sheet_123"`)) { - t.Fatalf("stdout missing imported token: %s", stdout.String()) - } - - prepareBody := decodeCapturedJSONBody(t, prepareStub) - if got, _ := prepareBody["parent_type"].(string); got != "ccm_import_open" { - t.Fatalf("prepare parent_type = %q, want %q", got, "ccm_import_open") - } - if got, _ := prepareBody["file_name"].(string); got != "large.xlsx" { - t.Fatalf("prepare file_name = %q, want %q", got, "large.xlsx") - } - if got, _ := prepareBody["size"].(float64); got != float64(maxDriveUploadFileSize+1) { - t.Fatalf("prepare size = %v, want %d", got, maxDriveUploadFileSize+1) - } - - firstPart := decodeCapturedMultipartBody(t, partStubs[0]) - if got := firstPart.Fields["upload_id"]; got != "upload_123" { - t.Fatalf("first part upload_id = %q, want %q", got, "upload_123") - } - if got := firstPart.Fields["seq"]; got != "0" { - t.Fatalf("first part seq = %q, want %q", got, "0") - } - if got := firstPart.Fields["size"]; got != "4194304" { - t.Fatalf("first part size = %q, want %q", got, "4194304") - } - if got := len(firstPart.Files["file"]); got != 4*1024*1024 { - t.Fatalf("first part file size = %d, want %d", got, 4*1024*1024) - } - - lastPart := decodeCapturedMultipartBody(t, partStubs[len(partStubs)-1]) - if got := lastPart.Fields["seq"]; got != "5" { - t.Fatalf("last part seq = %q, want %q", got, "5") - } - if got := lastPart.Fields["size"]; got != "1" { - t.Fatalf("last part size = %q, want %q", got, "1") - } - if got := len(lastPart.Files["file"]); got != 1 { - t.Fatalf("last part file size = %d, want %d", got, 1) - } - - finishBody := decodeCapturedJSONBody(t, finishStub) - if got, _ := finishBody["upload_id"].(string); got != "upload_123" { - t.Fatalf("finish upload_id = %q, want %q", got, "upload_123") - } - if got, _ := finishBody["block_num"].(float64); got != 6 { - t.Fatalf("finish block_num = %v, want %d", got, 6) - } -} - -func TestDriveImportMultipartPrepareValidatesResponseFields(t *testing.T) { - tests := []struct { - name string - data map[string]interface{} - wantText string - }{ - { - name: "missing upload id", - data: map[string]interface{}{ - "block_size": 4 * 1024 * 1024, - "block_num": 6, - }, - wantText: "upload prepare failed: no upload_id returned", - }, - { - name: "missing block size", - data: map[string]interface{}{ - "upload_id": "upload_123", - "block_num": 6, - }, - wantText: "upload prepare failed: invalid block_size returned", - }, - { - name: "missing block num", - data: map[string]interface{}{ - "upload_id": "upload_123", - "block_size": 4 * 1024 * 1024, - }, - wantText: "upload prepare failed: invalid block_num returned", - }, - } - - for _, tt := range tests { - t.Run(tt.name, func(t *testing.T) { - f, _, _, reg := cmdutil.TestFactory(t, driveTestConfig()) - registerDriveBotTokenStub(reg) - reg.Register(&httpmock.Stub{ - Method: "POST", - URL: "/open-apis/drive/v1/medias/upload_prepare", - Body: map[string]interface{}{ - "code": 0, - "data": tt.data, - }, - }) - - tmpDir := t.TempDir() - withDriveWorkingDir(t, tmpDir) - writeSizedDriveImportFile(t, "large.xlsx", int64(maxDriveUploadFileSize)+1) - - err := mountAndRunDrive(t, DriveImport, []string{ - "+import", - "--file", "large.xlsx", - "--type", "sheet", - "--as", "bot", - }, f, nil) - if err == nil { - t.Fatal("expected error, got nil") - } - if !strings.Contains(err.Error(), tt.wantText) { - t.Fatalf("error = %v, want substring %q", err, tt.wantText) - } - }) - } -} - -func TestDriveImportMultipartUploadPartAPIFailure(t *testing.T) { - f, _, _, reg := cmdutil.TestFactory(t, driveTestConfig()) - registerDriveBotTokenStub(reg) - reg.Register(&httpmock.Stub{ - Method: "POST", - URL: "/open-apis/drive/v1/medias/upload_prepare", - Body: map[string]interface{}{ - "code": 0, - "data": map[string]interface{}{ - "upload_id": "upload_123", - "block_size": 4 * 1024 * 1024, - "block_num": 6, - }, - }, - }) - reg.Register(&httpmock.Stub{ - Method: "POST", - URL: "/open-apis/drive/v1/medias/upload_part", - Body: map[string]interface{}{ - "code": 999, - "msg": "chunk rejected", - }, - }) - - tmpDir := t.TempDir() - withDriveWorkingDir(t, tmpDir) - writeSizedDriveImportFile(t, "large.xlsx", int64(maxDriveUploadFileSize)+1) - - err := mountAndRunDrive(t, DriveImport, []string{ - "+import", - "--file", "large.xlsx", - "--type", "sheet", - "--as", "bot", - }, f, nil) - if err == nil { - t.Fatal("expected error, got nil") - } - if !strings.Contains(err.Error(), "upload media part failed: [999] chunk rejected") { - t.Fatalf("unexpected error: %v", err) - } -} - -func TestDriveImportMultipartFinishRequiresFileToken(t *testing.T) { - f, _, _, reg := cmdutil.TestFactory(t, driveTestConfig()) - registerDriveBotTokenStub(reg) - reg.Register(&httpmock.Stub{ - Method: "POST", - URL: "/open-apis/drive/v1/medias/upload_prepare", - Body: map[string]interface{}{ - "code": 0, - "data": map[string]interface{}{ - "upload_id": "upload_123", - "block_size": 4 * 1024 * 1024, - "block_num": 6, - }, - }, - }) - for i := 0; i < 6; i++ { - reg.Register(&httpmock.Stub{ - Method: "POST", - URL: "/open-apis/drive/v1/medias/upload_part", - Body: map[string]interface{}{ - "code": 0, - "msg": "ok", - }, - }) - } - reg.Register(&httpmock.Stub{ - Method: "POST", - URL: "/open-apis/drive/v1/medias/upload_finish", - Body: map[string]interface{}{ - "code": 0, - "data": map[string]interface{}{}, - }, - }) - - tmpDir := t.TempDir() - withDriveWorkingDir(t, tmpDir) - writeSizedDriveImportFile(t, "large.xlsx", int64(maxDriveUploadFileSize)+1) - - err := mountAndRunDrive(t, DriveImport, []string{ - "+import", - "--file", "large.xlsx", - "--type", "sheet", - "--as", "bot", - }, f, nil) - if err == nil { - t.Fatal("expected error, got nil") - } - if !strings.Contains(err.Error(), "upload media finish failed: no file_token returned") { - t.Fatalf("unexpected error: %v", err) - } -} - func TestDriveImportRejectsOversizedFileByImportLimit(t *testing.T) { f, _, _, _ := cmdutil.TestFactory(t, driveTestConfig()) @@ -517,73 +219,6 @@ func TestDriveImportRejectsOversizedFileByImportLimit(t *testing.T) { } } -func TestParseDriveUploadResponseErrors(t *testing.T) { - t.Parallel() - - t.Run("invalid json", func(t *testing.T) { - t.Parallel() - - _, err := parseDriveUploadResponse(&larkcore.ApiResp{RawBody: []byte("{")}, "upload media failed") - if err == nil || !strings.Contains(err.Error(), "invalid response JSON") { - t.Fatalf("expected invalid JSON error, got %v", err) - } - }) - - t.Run("api code error", func(t *testing.T) { - t.Parallel() - - _, err := parseDriveUploadResponse(&larkcore.ApiResp{RawBody: []byte(`{"code":999,"msg":"boom","error":{"detail":"x"}}`)}, "upload media failed") - if err == nil || !strings.Contains(err.Error(), "upload media failed: [999] boom") { - t.Fatalf("expected API error, got %v", err) - } - }) -} - -func TestWrapDriveUploadRequestError(t *testing.T) { - t.Parallel() - - t.Run("preserves exit error", func(t *testing.T) { - t.Parallel() - - original := output.ErrValidation("bad input") - got := wrapDriveUploadRequestError(original, "upload media failed") - if got != original { - t.Fatalf("expected same exit error pointer, got %v", got) - } - }) - - t.Run("wraps generic error as network", func(t *testing.T) { - t.Parallel() - - got := wrapDriveUploadRequestError(io.EOF, "upload media failed") - var exitErr *output.ExitError - if !errors.As(got, &exitErr) { - t.Fatalf("expected ExitError, got %T", got) - } - if exitErr.Code != output.ExitNetwork { - t.Fatalf("exit code = %d, want %d", exitErr.Code, output.ExitNetwork) - } - if !strings.Contains(got.Error(), "upload media failed") { - t.Fatalf("unexpected error: %v", got) - } - }) -} - -type capturedMultipartBody struct { - Fields map[string]string - Files map[string][]byte -} - -func decodeCapturedJSONBody(t *testing.T, stub *httpmock.Stub) map[string]interface{} { - t.Helper() - - var body map[string]interface{} - if err := json.Unmarshal(stub.CapturedBody, &body); err != nil { - t.Fatalf("decode captured JSON body: %v", err) - } - return body -} - func writeSizedDriveImportFile(t *testing.T, name string, size int64) { t.Helper() @@ -598,42 +233,3 @@ func writeSizedDriveImportFile(t *testing.T, name string, size int64) { t.Fatalf("Close(%q) error: %v", name, err) } } - -func decodeCapturedMultipartBody(t *testing.T, stub *httpmock.Stub) capturedMultipartBody { - t.Helper() - - contentType := stub.CapturedHeaders.Get("Content-Type") - mediaType, params, err := mime.ParseMediaType(contentType) - if err != nil { - t.Fatalf("parse multipart content type: %v", err) - } - if mediaType != "multipart/form-data" { - t.Fatalf("content type = %q, want multipart/form-data", mediaType) - } - - reader := multipart.NewReader(bytes.NewReader(stub.CapturedBody), params["boundary"]) - body := capturedMultipartBody{ - Fields: map[string]string{}, - Files: map[string][]byte{}, - } - for { - part, err := reader.NextPart() - if err == io.EOF { - break - } - if err != nil { - t.Fatalf("read multipart part: %v", err) - } - - data, err := io.ReadAll(part) - if err != nil { - t.Fatalf("read multipart data: %v", err) - } - if part.FileName() != "" { - body.Files[part.FormName()] = data - continue - } - body.Fields[part.FormName()] = string(data) - } - return body -} diff --git a/shortcuts/drive/drive_import_test.go b/shortcuts/drive/drive_import_test.go index 91301a1d..9caef243 100644 --- a/shortcuts/drive/drive_import_test.go +++ b/shortcuts/drive/drive_import_test.go @@ -132,7 +132,7 @@ func TestDriveImportDryRunShowsMultipartUploadForLargeFile(t *testing.T) { if err != nil { t.Fatalf("Create() error: %v", err) } - if err := fh.Truncate(int64(maxDriveUploadFileSize) + 1); err != nil { + if err := fh.Truncate(common.MaxDriveMediaUploadSinglePartSize + 1); err != nil { t.Fatalf("Truncate() error: %v", err) } if err := fh.Close(); err != nil { diff --git a/shortcuts/drive/drive_io_test.go b/shortcuts/drive/drive_io_test.go index f9e91ae3..da2db5b2 100644 --- a/shortcuts/drive/drive_io_test.go +++ b/shortcuts/drive/drive_io_test.go @@ -18,9 +18,6 @@ import ( "github.com/larksuite/cli/shortcuts/common" ) -// registerDriveBotTokenStub is a no-op. TAT is now managed by CredentialProvider, not SDK. -func registerDriveBotTokenStub(_ *httpmock.Registry) {} - func driveTestConfig() *core.CliConfig { return &core.CliConfig{ AppID: "drive-test-app", AppSecret: "test-secret", Brand: core.BrandFeishu, @@ -62,7 +59,6 @@ func TestDriveUploadLargeFileUsesMultipart(t *testing.T) { AppID: "drive-upload-test-app", AppSecret: "test-secret", Brand: core.BrandFeishu, } f, stdout, _, reg := cmdutil.TestFactory(t, uploadTestConfig) - registerDriveBotTokenStub(reg) // Step 1: upload_prepare reg.Register(&httpmock.Stub{ @@ -72,7 +68,7 @@ func TestDriveUploadLargeFileUsesMultipart(t *testing.T) { "code": 0, "msg": "ok", "data": map[string]interface{}{ "upload_id": "test-upload-id", - "block_size": float64(maxDriveUploadFileSize), + "block_size": float64(common.MaxDriveMediaUploadSinglePartSize), "block_num": float64(2), }, }, @@ -116,7 +112,7 @@ func TestDriveUploadLargeFileUsesMultipart(t *testing.T) { if err != nil { t.Fatalf("Create() error: %v", err) } - if err := fh.Truncate(maxDriveUploadFileSize + 1); err != nil { + if err := fh.Truncate(common.MaxDriveMediaUploadSinglePartSize + 1); err != nil { t.Fatalf("Truncate() error: %v", err) } if err := fh.Close(); err != nil { @@ -141,7 +137,6 @@ func TestDriveUploadSmallFile(t *testing.T) { AppID: "drive-upload-small-test", AppSecret: "test-secret", Brand: core.BrandFeishu, } f, stdout, _, reg := cmdutil.TestFactory(t, uploadTestConfig) - registerDriveBotTokenStub(reg) reg.Register(&httpmock.Stub{ Method: "POST", @@ -181,7 +176,6 @@ func TestDriveUploadSmallFileAPIError(t *testing.T) { AppID: "drive-upload-small-err", AppSecret: "test-secret", Brand: core.BrandFeishu, } f, stdout, _, reg := cmdutil.TestFactory(t, uploadTestConfig) - registerDriveBotTokenStub(reg) reg.Register(&httpmock.Stub{ Method: "POST", @@ -218,7 +212,6 @@ func TestDriveUploadSmallFileNoToken(t *testing.T) { AppID: "drive-upload-small-notoken", AppSecret: "test-secret", Brand: core.BrandFeishu, } f, stdout, _, reg := cmdutil.TestFactory(t, uploadTestConfig) - registerDriveBotTokenStub(reg) reg.Register(&httpmock.Stub{ Method: "POST", @@ -256,7 +249,6 @@ func TestDriveUploadSmallFileInvalidJSON(t *testing.T) { AppID: "drive-upload-small-json", AppSecret: "test-secret", Brand: core.BrandFeishu, } f, stdout, _, reg := cmdutil.TestFactory(t, uploadTestConfig) - registerDriveBotTokenStub(reg) reg.Register(&httpmock.Stub{ Method: "POST", @@ -291,7 +283,6 @@ func TestDriveUploadPrepareInvalidResponse(t *testing.T) { AppID: "drive-upload-prepare-bad", AppSecret: "test-secret", Brand: core.BrandFeishu, } f, stdout, _, reg := cmdutil.TestFactory(t, uploadTestConfig) - registerDriveBotTokenStub(reg) reg.Register(&httpmock.Stub{ Method: "POST", @@ -317,7 +308,7 @@ func TestDriveUploadPrepareInvalidResponse(t *testing.T) { if err != nil { t.Fatalf("Create() error: %v", err) } - if err := fh.Truncate(maxDriveUploadFileSize + 1); err != nil { + if err := fh.Truncate(common.MaxDriveMediaUploadSinglePartSize + 1); err != nil { t.Fatalf("Truncate() error: %v", err) } fh.Close() @@ -338,7 +329,6 @@ func TestDriveUploadPartAPIError(t *testing.T) { AppID: "drive-upload-part-err", AppSecret: "test-secret", Brand: core.BrandFeishu, } f, stdout, _, reg := cmdutil.TestFactory(t, uploadTestConfig) - registerDriveBotTokenStub(reg) reg.Register(&httpmock.Stub{ Method: "POST", @@ -347,7 +337,7 @@ func TestDriveUploadPartAPIError(t *testing.T) { "code": 0, "msg": "ok", "data": map[string]interface{}{ "upload_id": "test-upload-id", - "block_size": float64(maxDriveUploadFileSize), + "block_size": float64(common.MaxDriveMediaUploadSinglePartSize), "block_num": float64(2), }, }, @@ -380,7 +370,7 @@ func TestDriveUploadPartAPIError(t *testing.T) { if err != nil { t.Fatalf("Create() error: %v", err) } - if err := fh.Truncate(maxDriveUploadFileSize + 1); err != nil { + if err := fh.Truncate(common.MaxDriveMediaUploadSinglePartSize + 1); err != nil { t.Fatalf("Truncate() error: %v", err) } fh.Close() @@ -401,7 +391,6 @@ func TestDriveUploadPartInvalidJSON(t *testing.T) { AppID: "drive-upload-part-json", AppSecret: "test-secret", Brand: core.BrandFeishu, } f, stdout, _, reg := cmdutil.TestFactory(t, uploadTestConfig) - registerDriveBotTokenStub(reg) reg.Register(&httpmock.Stub{ Method: "POST", @@ -410,7 +399,7 @@ func TestDriveUploadPartInvalidJSON(t *testing.T) { "code": 0, "msg": "ok", "data": map[string]interface{}{ "upload_id": "test-upload-id", - "block_size": float64(maxDriveUploadFileSize + 1), + "block_size": float64(common.MaxDriveMediaUploadSinglePartSize + 1), "block_num": float64(1), }, }, @@ -433,7 +422,7 @@ func TestDriveUploadPartInvalidJSON(t *testing.T) { if err != nil { t.Fatalf("Create() error: %v", err) } - if err := fh.Truncate(maxDriveUploadFileSize + 1); err != nil { + if err := fh.Truncate(common.MaxDriveMediaUploadSinglePartSize + 1); err != nil { t.Fatalf("Truncate() error: %v", err) } fh.Close() @@ -454,7 +443,6 @@ func TestDriveUploadFinishNoToken(t *testing.T) { AppID: "drive-upload-finish-notoken", AppSecret: "test-secret", Brand: core.BrandFeishu, } f, stdout, _, reg := cmdutil.TestFactory(t, uploadTestConfig) - registerDriveBotTokenStub(reg) reg.Register(&httpmock.Stub{ Method: "POST", @@ -463,7 +451,7 @@ func TestDriveUploadFinishNoToken(t *testing.T) { "code": 0, "msg": "ok", "data": map[string]interface{}{ "upload_id": "test-upload-id", - "block_size": float64(maxDriveUploadFileSize + 1), + "block_size": float64(common.MaxDriveMediaUploadSinglePartSize + 1), "block_num": float64(1), }, }, @@ -495,7 +483,7 @@ func TestDriveUploadFinishNoToken(t *testing.T) { if err != nil { t.Fatalf("Create() error: %v", err) } - if err := fh.Truncate(maxDriveUploadFileSize + 1); err != nil { + if err := fh.Truncate(common.MaxDriveMediaUploadSinglePartSize + 1); err != nil { t.Fatalf("Truncate() error: %v", err) } fh.Close() @@ -516,7 +504,6 @@ func TestDriveUploadWithCustomName(t *testing.T) { AppID: "drive-upload-name-test", AppSecret: "test-secret", Brand: core.BrandFeishu, } f, stdout, _, reg := cmdutil.TestFactory(t, uploadTestConfig) - registerDriveBotTokenStub(reg) reg.Register(&httpmock.Stub{ Method: "POST", diff --git a/shortcuts/drive/drive_move_common_test.go b/shortcuts/drive/drive_move_common_test.go index 8221ada9..09327115 100644 --- a/shortcuts/drive/drive_move_common_test.go +++ b/shortcuts/drive/drive_move_common_test.go @@ -104,7 +104,6 @@ func TestDriveMoveDryRunFolderIncludesTaskCheckParams(t *testing.T) { func TestDriveMoveFolderSuccessUsesTaskCheckHelper(t *testing.T) { f, stdout, _, reg := cmdutil.TestFactory(t, driveTestConfig()) - registerDriveBotTokenStub(reg) reg.Register(&httpmock.Stub{ Method: "POST", URL: "/open-apis/drive/v1/files/fld_src/move", @@ -148,7 +147,6 @@ func TestDriveMoveFolderSuccessUsesTaskCheckHelper(t *testing.T) { func TestDriveMoveFolderTimeoutReturnsFollowUpCommand(t *testing.T) { f, stdout, _, reg := cmdutil.TestFactory(t, driveTestConfig()) - registerDriveBotTokenStub(reg) reg.Register(&httpmock.Stub{ Method: "POST", URL: "/open-apis/drive/v1/files/fld_src/move", diff --git a/shortcuts/drive/drive_move_test.go b/shortcuts/drive/drive_move_test.go index 184b9e5e..5b5ee906 100644 --- a/shortcuts/drive/drive_move_test.go +++ b/shortcuts/drive/drive_move_test.go @@ -13,7 +13,6 @@ import ( func TestDriveMoveUsesRootFolderWhenFolderTokenMissing(t *testing.T) { f, stdout, _, reg := cmdutil.TestFactory(t, driveTestConfig()) - registerDriveBotTokenStub(reg) reg.Register(&httpmock.Stub{ Method: "GET", URL: "/open-apis/drive/explorer/v2/root_folder/meta", @@ -52,7 +51,6 @@ func TestDriveMoveUsesRootFolderWhenFolderTokenMissing(t *testing.T) { func TestDriveMoveRootFolderLookupRequiresToken(t *testing.T) { f, _, _, reg := cmdutil.TestFactory(t, driveTestConfig()) - registerDriveBotTokenStub(reg) reg.Register(&httpmock.Stub{ Method: "GET", URL: "/open-apis/drive/explorer/v2/root_folder/meta", diff --git a/shortcuts/drive/drive_task_result_test.go b/shortcuts/drive/drive_task_result_test.go index cb11ec75..f31ae818 100644 --- a/shortcuts/drive/drive_task_result_test.go +++ b/shortcuts/drive/drive_task_result_test.go @@ -127,7 +127,6 @@ func TestDriveTaskResultDryRunExportIncludesTokenParam(t *testing.T) { func TestDriveTaskResultImportIncludesReadyFlags(t *testing.T) { f, stdout, _, reg := cmdutil.TestFactory(t, driveTestConfig()) - registerDriveBotTokenStub(reg) reg.Register(&httpmock.Stub{ Method: "GET", URL: "/open-apis/drive/v1/import_tasks/tk_import", @@ -161,7 +160,6 @@ func TestDriveTaskResultImportIncludesReadyFlags(t *testing.T) { func TestDriveTaskResultTaskCheckIncludesReadyFlags(t *testing.T) { f, stdout, _, reg := cmdutil.TestFactory(t, driveTestConfig()) - registerDriveBotTokenStub(reg) reg.Register(&httpmock.Stub{ Method: "GET", URL: "/open-apis/drive/v1/files/task_check", diff --git a/shortcuts/drive/drive_upload.go b/shortcuts/drive/drive_upload.go index 2846f604..a828f88c 100644 --- a/shortcuts/drive/drive_upload.go +++ b/shortcuts/drive/drive_upload.go @@ -21,8 +21,6 @@ import ( "github.com/larksuite/cli/shortcuts/common" ) -const maxDriveUploadFileSize = 20 * 1024 * 1024 // 20MB - var DriveUpload = common.Shortcut{ Service: "drive", Command: "+upload", @@ -78,7 +76,7 @@ var DriveUpload = common.Shortcut{ fmt.Fprintf(runtime.IO().ErrOut, "Uploading: %s (%s)\n", fileName, common.FormatSize(fileSize)) var fileToken string - if fileSize > maxDriveUploadFileSize { + if fileSize > common.MaxDriveMediaUploadSinglePartSize { fmt.Fprintf(runtime.IO().ErrOut, "File exceeds 20MB, using multipart upload\n") fileToken, err = uploadFileMultipart(ctx, runtime, filePath, fileName, folderToken, fileSize) } else { diff --git a/skills/lark-doc/references/lark-doc-media-insert.md b/skills/lark-doc/references/lark-doc-media-insert.md index 767745f4..301de79e 100644 --- a/skills/lark-doc/references/lark-doc-media-insert.md +++ b/skills/lark-doc/references/lark-doc-media-insert.md @@ -30,7 +30,7 @@ lark-cli docs +media-insert --doc doxcnXXX --file ./arch.png --align center --ca | 参数 | 必填 | 说明 | |------|------|------| | `--doc ` | 是 | 文档 ID 或 docx URL（仅支持 `/docx/` 形式自动提取；**不支持 `/wiki/...` URL 自动提取**） | -| `--file ` | 是 | 本地文件路径（最大 20MB） | +| `--file ` | 是 | 本地文件路径（文件大于 20MB 时自动切换分片上传） | | `--type ` | 否 | `image`（默认）或 `file` | | `--align ` | 否 | 仅图片：`left` / `center`（默认）/ `right` | | `--caption ` | 否 | 仅图片：图片描述 |