larksuite__cli/shortcuts/common/drive_media_upload.go

// Copyright (c) 2026 Lark Technologies Pte. Ltd.
// SPDX-License-Identifier: MIT

package common

import (
	"bytes"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"net/http"
	"strings"

	larkcore "github.com/larksuite/oapi-sdk-go/v3/core"

	"github.com/larksuite/cli/errs"
	"github.com/larksuite/cli/internal/client"
	"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
	// Reader, when non-nil, is used as the upload source instead of opening
	// FilePath. Callers must set FileName and FileSize explicitly. The reader
	// is NOT closed by UploadDriveMediaAll; the caller owns its lifetime.
	// Used by the clipboard path in docs +media-insert.
	Reader io.Reader
}

type DriveMediaMultipartUploadConfig struct {
	FilePath   string
	FileName   string
	FileSize   int64
	ParentType string
	ParentNode string
	Extra      string
	// Reader mirrors DriveMediaUploadAllConfig.Reader for chunked uploads.
	Reader io.Reader
}

// Deprecated: use UploadDriveMediaAllTyped for typed error envelopes.
func UploadDriveMediaAll(runtime *RuntimeContext, cfg DriveMediaUploadAllConfig) (string, error) {
	var fileReader io.Reader
	if cfg.Reader != nil {
		fileReader = cfg.Reader
	} else {
		f, err := runtime.FileIO().Open(cfg.FilePath)
		if err != nil {
			return "", WrapInputStatError(err)
		}
		defer f.Close()
		fileReader = f
	}

	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", fileReader)

	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)
}

// UploadDriveMediaAllTyped is the typed-error counterpart of
// UploadDriveMediaAll: file-open failures surface as typed validation errors,
// transport failures as typed network errors, and API failures are classified
// via ClassifyAPIResponse so subtype / code / log_id survive on the error.
func UploadDriveMediaAllTyped(runtime *RuntimeContext, cfg DriveMediaUploadAllConfig) (string, error) {
	var fileReader io.Reader
	if cfg.Reader != nil {
		fileReader = cfg.Reader
	} else {
		f, err := runtime.FileIO().Open(cfg.FilePath)
		if err != nil {
			return "", WrapInputStatErrorTyped(err)
		}
		defer f.Close()
		fileReader = f
	}

	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", fileReader)

	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 "", prefixDriveMediaUploadProblem(client.WrapDoAPIError(err), driveMediaUploadAllAction)
	}

	data, err := runtime.ClassifyAPIResponse(apiResp)
	if err != nil {
		return "", prefixDriveMediaUploadProblem(err, driveMediaUploadAllAction)
	}
	return extractDriveMediaUploadFileTokenTyped(data, driveMediaUploadAllAction)
}

// Deprecated: use UploadDriveMediaMultipartTyped for typed error envelopes.
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, session); err != nil {
		return "", err
	}

	return finishDriveMediaMultipartUpload(runtime, session.UploadID, session.BlockNum)
}

// UploadDriveMediaMultipartTyped is the typed-error counterpart of
// UploadDriveMediaMultipart: prepare/finish failures come back typed from
// CallAPITyped, malformed session plans surface as invalid-response internal
// errors, and per-part transport/API failures are classified the same way as
// UploadDriveMediaAllTyped.
func UploadDriveMediaMultipartTyped(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.CallAPITyped("POST", "/open-apis/drive/v1/medias/upload_prepare", nil, prepareBody)
	if err != nil {
		return "", err
	}

	session, err := parseDriveMediaMultipartUploadSessionTyped(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 = uploadDriveMediaMultipartPartsTyped(runtime, cfg, session); err != nil {
		return "", err
	}

	return finishDriveMediaMultipartUploadTyped(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 {
	// Preserve any already-classified error: legacy *output.ExitError or any
	// typed errs.* error. Only un-classified errors get wrapped as network.
	var exitErr *output.ExitError
	if errors.As(err, &exitErr) {
		return err
	}
	if _, ok := errs.ProblemOf(err); ok {
		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), driveMediaUploadErrorDetail(apiResp, result["error"]))
	}

	data, _ := result["data"].(map[string]interface{})
	return data, nil
}

func driveMediaUploadErrorDetail(apiResp *larkcore.ApiResp, detail interface{}) interface{} {
	logID := ""
	if apiResp != nil {
		logID = strings.TrimSpace(apiResp.LogId())
	}
	if logID == "" {
		return detail
	}
	detailMap, ok := detail.(map[string]interface{})
	if !ok {
		if detail == nil {
			return map[string]interface{}{"log_id": logID}
		}
		return map[string]interface{}{"error": detail, "log_id": logID}
	}
	if _, exists := detailMap["log_id"]; !exists {
		detailMap["log_id"] = logID
	}
	return detailMap
}

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, cfg DriveMediaMultipartUploadConfig, session DriveMediaMultipartUploadSession) error {
	var r io.Reader
	if cfg.Reader != nil {
		r = cfg.Reader
	} else {
		f, err := runtime.FileIO().Open(cfg.FilePath)
		if err != nil {
			return WrapInputStatError(err)
		}
		defer f.Close()
		r = f
	}

	maxInt := int64(^uint(0) >> 1)
	bufferSize := session.BlockSize
	if bufferSize <= 0 || bufferSize > maxInt {
		return output.Errorf(output.ExitAPI, "api_error", "upload prepare failed: invalid block_size returned")
	}
	buffer := make([]byte, int(bufferSize))
	remaining := cfg.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(r, 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)
}

// prefixDriveMediaUploadProblem prepends the upload action to a typed error's
// message so callers see which upload step failed. Non-typed errors are
// returned unchanged.
func prefixDriveMediaUploadProblem(err error, action string) error {
	if p, ok := errs.ProblemOf(err); ok {
		p.Message = action + ": " + p.Message
	}
	return err
}

// parseDriveMediaMultipartUploadSessionTyped validates the upload_prepare
// session plan like ParseDriveMediaMultipartUploadSession, but reports a
// malformed plan as a typed invalid-response internal error.
func parseDriveMediaMultipartUploadSessionTyped(data map[string]interface{}) (DriveMediaMultipartUploadSession, error) {
	session := DriveMediaMultipartUploadSession{
		UploadID:  GetString(data, "upload_id"),
		BlockSize: int64(GetFloat(data, "block_size")),
		BlockNum:  int(GetFloat(data, "block_num")),
	}
	if session.UploadID == "" {
		return DriveMediaMultipartUploadSession{}, errs.NewInternalError(errs.SubtypeInvalidResponse, "upload prepare failed: no upload_id returned")
	}
	if session.BlockSize <= 0 {
		return DriveMediaMultipartUploadSession{}, errs.NewInternalError(errs.SubtypeInvalidResponse, "upload prepare failed: invalid block_size returned")
	}
	if session.BlockNum <= 0 {
		return DriveMediaMultipartUploadSession{}, errs.NewInternalError(errs.SubtypeInvalidResponse, "upload prepare failed: invalid block_num returned")
	}
	return session, nil
}

// extractDriveMediaUploadFileTokenTyped mirrors ExtractDriveMediaUploadFileToken
// with a typed invalid-response internal error for a missing file_token.
func extractDriveMediaUploadFileTokenTyped(data map[string]interface{}, action string) (string, error) {
	fileToken := GetString(data, "file_token")
	if fileToken == "" {
		return "", errs.NewInternalError(errs.SubtypeInvalidResponse, "%s: no file_token returned", action)
	}
	return fileToken, nil
}

// uploadDriveMediaMultipartPartsTyped mirrors uploadDriveMediaMultipartParts
// with typed errors for file-open, file-read, and per-part upload failures.
func uploadDriveMediaMultipartPartsTyped(runtime *RuntimeContext, cfg DriveMediaMultipartUploadConfig, session DriveMediaMultipartUploadSession) error {
	var r io.Reader
	if cfg.Reader != nil {
		r = cfg.Reader
	} else {
		f, err := runtime.FileIO().Open(cfg.FilePath)
		if err != nil {
			return WrapInputStatErrorTyped(err)
		}
		defer f.Close()
		r = f
	}

	maxInt := int64(^uint(0) >> 1)
	bufferSize := session.BlockSize
	if bufferSize <= 0 || bufferSize > maxInt {
		return errs.NewInternalError(errs.SubtypeInvalidResponse, "upload prepare failed: invalid block_size returned")
	}
	buffer := make([]byte, int(bufferSize))
	remaining := cfg.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(r, buffer[:int(chunkSize)])
		if readErr != nil {
			return WrapInputStatErrorTyped(readErr)
		}

		if err := uploadDriveMediaMultipartPartTyped(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 uploadDriveMediaMultipartPartTyped(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 prefixDriveMediaUploadProblem(client.WrapDoAPIError(err), driveMediaUploadPartAction)
	}

	if _, err := runtime.ClassifyAPIResponse(apiResp); err != nil {
		return prefixDriveMediaUploadProblem(err, driveMediaUploadPartAction)
	}
	return nil
}

func finishDriveMediaMultipartUploadTyped(runtime *RuntimeContext, uploadID string, blockNum int) (string, error) {
	data, err := runtime.CallAPITyped("POST", "/open-apis/drive/v1/medias/upload_finish", nil, map[string]interface{}{
		"upload_id": uploadID,
		"block_num": blockNum,
	})
	if err != nil {
		return "", err
	}
	return extractDriveMediaUploadFileTokenTyped(data, driveMediaUploadFinishAction)
}