new-api/service/asset_doubao.go

package service

import (
	"bytes"
	"context"
	"crypto/hmac"
	"crypto/sha256"
	"encoding/hex"
	"fmt"
	"io"
	"net/http"
	"net/url"
	"strings"
	"time"

	"github.com/QuantumNous/new-api/common"
	"github.com/QuantumNous/new-api/model"
	"github.com/QuantumNous/new-api/setting/system_setting"
)

// DoubaoVideoAssetAdapter serves the Ark-compatible asset API on official
// Volcengine (DoubaoVideo) channels. Unlike the aiping adapter (Bearer key),
// the official endpoint authenticates every request with a Volcengine V4
// HMAC-SHA256 signature built from an AccessKey/SecretKey pair stored in
// channel_asset_credentials (PoolID holds the optional project code).
type DoubaoVideoAssetAdapter struct {
	operation map[AssetOperation]struct{}
}

func NewDoubaoVideoAssetAdapter() AssetAdapter {
	operations := []AssetOperation{
		AssetOperationAssetCreate,
		AssetOperationAssetList,
		AssetOperationAssetGet,
		AssetOperationAssetUpdate,
		AssetOperationAssetDelete,
		AssetOperationAssetGroupCreate,
		AssetOperationAssetGroupList,
		AssetOperationAssetGroupGet,
		AssetOperationAssetGroupUpdate,
		AssetOperationAssetGroupDelete,
	}
	supported := make(map[AssetOperation]struct{}, len(operations))
	for _, op := range operations {
		supported[op] = struct{}{}
	}
	return &DoubaoVideoAssetAdapter{operation: supported}
}

func (a *DoubaoVideoAssetAdapter) Name() string {
	return "doubao_video_asset"
}

func (a *DoubaoVideoAssetAdapter) Supports(operation AssetOperation) bool {
	_, ok := a.operation[operation]
	return ok
}

func (a *DoubaoVideoAssetAdapter) DoAssetRequest(ctx context.Context, channel *model.Channel, req AssetRequest) (*AssetUpstreamResponse, *AssetError) {
	if !a.Supports(req.Action.Operation) {
		return nil, newAssetError(AssetErrorOperationNotSupported, fmt.Sprintf("asset operation %s is not supported", req.Action.Operation), http.StatusBadRequest)
	}
	credential, err := model.GetChannelAssetCredential(channel.Id)
	if err != nil {
		return nil, newAssetError(AssetErrorServer, err.Error(), http.StatusInternalServerError)
	}
	if credential == nil || strings.TrimSpace(credential.AccessKey) == "" || strings.TrimSpace(credential.SecretKey) == "" {
		return nil, newAssetError(AssetErrorServer, "Volcengine asset AccessKey and SecretKey are required (configure the channel asset credential)", http.StatusBadRequest)
	}

	upstreamURL, rawQuery, err := buildDoubaoVideoAssetURL(channel, req.Action.Action, req.Version)
	if err != nil {
		return nil, newAssetError(AssetErrorServer, err.Error(), http.StatusInternalServerError)
	}
	fetchSetting := system_setting.GetFetchSetting()
	if err := common.ValidateURLWithFetchSetting(upstreamURL, fetchSetting.EnableSSRFProtection, fetchSetting.AllowPrivateIp, fetchSetting.DomainFilterMode, fetchSetting.IpFilterMode, fetchSetting.DomainList, fetchSetting.IpList, fetchSetting.AllowedPorts, fetchSetting.ApplyIPFilterForDomain); err != nil {
		return nil, newAssetError(AssetErrorServer, fmt.Sprintf("request blocked: %v", err), http.StatusForbidden)
	}

	httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, upstreamURL, bytes.NewReader(req.RawBody))
	if err != nil {
		return nil, newAssetError(AssetErrorServer, err.Error(), http.StatusInternalServerError)
	}
	signDoubaoVideoAssetRequest(httpReq, credential.AccessKey, credential.SecretKey, req.RawBody, rawQuery)
	if projectCode := strings.TrimSpace(credential.PoolID); projectCode != "" {
		httpReq.Header.Set("X-Project-Code", projectCode)
	}

	client, err := GetHttpClientWithProxy(channel.GetSetting().Proxy)
	if err != nil {
		return nil, newAssetError(AssetErrorServer, err.Error(), http.StatusInternalServerError)
	}
	if client == nil {
		client = http.DefaultClient
	}
	resp, err := client.Do(httpReq)
	if err != nil {
		return nil, newAssetError(AssetErrorUpstream, err.Error(), http.StatusBadGateway)
	}
	defer resp.Body.Close()
	data, err := io.ReadAll(resp.Body)
	if err != nil {
		return nil, newAssetError(AssetErrorUpstream, err.Error(), http.StatusBadGateway)
	}
	if resp.StatusCode < 200 || resp.StatusCode >= 300 {
		return nil, newAssetError(AssetErrorUpstream, string(data), http.StatusBadGateway)
	}
	return &AssetUpstreamResponse{StatusCode: resp.StatusCode, Header: resp.Header, Body: data}, nil
}

func buildDoubaoVideoAssetURL(channel *model.Channel, action string, version string) (string, string, error) {
	baseURL := strings.TrimSpace(channel.GetBaseURL())
	if baseURL == "" {
		baseURL = "https://ark.cn-beijing.volces.com"
	}
	u, err := url.Parse(baseURL)
	if err != nil {
		return "", "", err
	}
	u.Path = strings.TrimRight(u.Path, "/") + "/openApi/portrait"
	if strings.TrimSpace(version) == "" {
		version = "2024-01-01"
	}
	// Action < Version lexicographically, so this raw query is already the
	// sorted canonical form required by the V4 signature.
	rawQuery := "Action=" + url.QueryEscape(action) + "&Version=" + url.QueryEscape(version)
	u.RawQuery = rawQuery
	return u.String(), rawQuery, nil
}

// signDoubaoVideoAssetRequest applies the Volcengine V4 HMAC-SHA256 signature
// to the request: canonical request over (method, path, sorted query,
// content-type/host/x-content-sha256/x-date headers, body hash) -> string to
// sign with the {date}/{region}/ark/request scope -> chained HMAC keys.
func signDoubaoVideoAssetRequest(req *http.Request, accessKey string, secretKey string, body []byte, rawQuery string) {
	now := time.Now().UTC()
	xDate := now.Format("20060102T150405Z")
	shortDate := now.Format("20060102")
	payloadHash := sha256.Sum256(body)
	payloadHashHex := hex.EncodeToString(payloadHash[:])

	canonicalHeaders := fmt.Sprintf(
		"content-type:application/json\nhost:%s\nx-content-sha256:%s\nx-date:%s\n",
		req.URL.Host, payloadHashHex, xDate,
	)
	signedHeaders := "content-type;host;x-content-sha256;x-date"
	canonicalRequest := strings.Join([]string{
		req.Method,
		req.URL.EscapedPath(),
		rawQuery,
		canonicalHeaders,
		signedHeaders,
		payloadHashHex,
	}, "\n")

	scope := shortDate + "/cn-beijing/ark/request"
	canonicalHash := sha256.Sum256([]byte(canonicalRequest))
	stringToSign := strings.Join([]string{
		"HMAC-SHA256",
		xDate,
		scope,
		hex.EncodeToString(canonicalHash[:]),
	}, "\n")

	signingKey := volcengineSigningKey(secretKey, shortDate)
	signature := hmac.New(sha256.New, signingKey)
	signature.Write([]byte(stringToSign))

	req.Header.Set("Content-Type", "application/json")
	req.Header.Set("Accept", "application/json")
	req.Header.Set("X-Date", xDate)
	req.Header.Set("X-Content-Sha256", payloadHashHex)
	req.Header.Set("Authorization", fmt.Sprintf(
		"HMAC-SHA256 Credential=%s/%s, SignedHeaders=%s, Signature=%s",
		accessKey, scope, signedHeaders, hex.EncodeToString(signature.Sum(nil)),
	))
}

func volcengineSigningKey(secretKey string, shortDate string) []byte {
	key := hmacSHA256([]byte(secretKey), shortDate)
	key = hmacSHA256(key, "cn-beijing")
	key = hmacSHA256(key, "ark")
	return hmacSHA256(key, "request")
}

func hmacSHA256(key []byte, message string) []byte {
	mac := hmac.New(sha256.New, key)
	mac.Write([]byte(message))
	return mac.Sum(nil)
}