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"

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

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

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

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