new-api/controller/doubao_asset.go

package controller

import (
	"context"
	"fmt"
	"io"
	"net/http"
	"strings"

	"github.com/QuantumNous/new-api/common"
	"github.com/QuantumNous/new-api/constant"
	"github.com/QuantumNous/new-api/logger"
	"github.com/QuantumNous/new-api/model"
	"github.com/QuantumNous/new-api/service"

	"github.com/gin-gonic/gin"
)

func assetProxyError(c *gin.Context, status int, errType, message string) {
	c.JSON(status, gin.H{
		"error": gin.H{
			"message": message,
			"type":    errType,
		},
	})
}

var doubaoAssetAutoGroups = service.GetUserAutoGroup

func effectiveDoubaoAssetGroup(c *gin.Context) string {
	if group := strings.TrimSpace(common.GetContextKeyString(c, constant.ContextKeyUsingGroup)); group != "" {
		return group
	}
	return strings.TrimSpace(common.GetContextKeyString(c, constant.ContextKeyTokenGroup))
}

func concreteDoubaoAssetGroupsForRequest(c *gin.Context, autoGroups func(string) []string) []string {
	group := effectiveDoubaoAssetGroup(c)
	if group != "auto" {
		if group == "" {
			return nil
		}
		return []string{group}
	}

	userGroup := common.GetContextKeyString(c, constant.ContextKeyUserGroup)
	groups := autoGroups(userGroup)
	concreteGroups := make([]string, 0, len(groups))
	for _, candidate := range groups {
		candidate = strings.TrimSpace(candidate)
		if candidate == "" || candidate == "auto" {
			continue
		}
		concreteGroups = append(concreteGroups, candidate)
	}
	return concreteGroups
}

func DoubaoAssetProxy(c *gin.Context) {
	actionRaw := strings.TrimSpace(c.Query("Action"))
	if actionRaw == "" {
		assetProxyError(c, http.StatusBadRequest, service.AssetErrorInvalidRequest, "Action query parameter is required")
		return
	}
	action, ok := service.ParseAssetAction(actionRaw)
	if !ok {
		assetProxyError(c, http.StatusBadRequest, service.AssetErrorInvalidRequest, fmt.Sprintf("unsupported asset Action: %s", actionRaw))
		return
	}

	version := strings.TrimSpace(c.Query("Version"))
	if version == "" {
		version = "2024-01-01"
	}

	rawBody, err := io.ReadAll(c.Request.Body)
	if err != nil {
		assetProxyError(c, http.StatusBadRequest, service.AssetErrorInvalidRequest, fmt.Sprintf("failed to read request body: %v", err))
		return
	}
	body := map[string]any{}
	if len(strings.TrimSpace(string(rawBody))) > 0 {
		if err := common.Unmarshal(rawBody, &body); err != nil {
			assetProxyError(c, http.StatusBadRequest, service.AssetErrorInvalidRequest, fmt.Sprintf("invalid JSON body: %v", err))
			return
		}
	}

	userID := c.GetInt("id")
	var lastAssetErr *service.AssetError
	for _, group := range concreteDoubaoAssetGroupsForRequest(c, doubaoAssetAutoGroups) {
		channel, adapter, assetErr := service.ResolveAssetChannelForOperation(userID, group, action.Operation)
		if assetErr != nil {
			lastAssetErr = assetErr
			logger.LogError(c.Request.Context(), fmt.Sprintf("Failed to resolve asset channel for group %s action %s: %s", group, action.Action, assetErr.Message))
			if assetErr.Type == service.AssetErrorChannelNotFound {
				continue
			}
			break
		}
		if channel == nil || adapter == nil {
			continue
		}
		assetRequest := service.AssetRequest{
			Action:  action,
			Version: version,
			Body:    body,
			RawBody: rawBody,
		}
		if service.IsChinaMobileAssetChannel(channel) {
			if service.IsAssetGroupOperation(action.Operation) {
				assetProxyError(c, http.StatusForbidden, service.AssetErrorOperationNotSupported, "China Mobile asset group APIs are managed by the platform")
				return
			}
			groupID, assetErr := service.GetOrCreateChinaMobileUserAssetGroup(c.Request.Context(), userID, channel, func(ctx context.Context, userID int, channel *model.Channel) (string, *service.AssetError) {
				return service.CreateChinaMobileUserAssetGroup(ctx, userID, adapter, channel)
			})
			if assetErr != nil {
				assetProxyError(c, assetErr.HTTPStatus, assetErr.Type, assetErr.Message)
				return
			}
			service.ScopeChinaMobileAssetRequest(&assetRequest, groupID)
			if action.Operation == service.AssetOperationAssetGet || action.Operation == service.AssetOperationAssetUpdate || action.Operation == service.AssetOperationAssetDelete {
				if assetErr := service.RequireChinaMobileAssetOwnership(c.Request.Context(), adapter, channel, assetRequest, groupID); assetErr != nil {
					assetProxyError(c, assetErr.HTTPStatus, assetErr.Type, assetErr.Message)
					return
				}
			}
		}
		resp, assetErr := adapter.DoAssetRequest(c.Request.Context(), channel, assetRequest)
		if assetErr != nil {
			logger.LogError(c.Request.Context(), fmt.Sprintf("Failed to proxy asset action %s via channel %d type %d: %s", action.Action, channel.Id, channel.Type, assetErr.Message))
			assetProxyError(c, assetErr.HTTPStatus, assetErr.Type, assetErr.Message)
			return
		}
		copyDoubaoAssetResponseHeaders(c, resp.Header)
		c.Data(resp.StatusCode, "application/json", resp.Body)
		return
	}
	if lastAssetErr != nil {
		assetProxyError(c, lastAssetErr.HTTPStatus, lastAssetErr.Type, lastAssetErr.Message)
		return
	}
	assetProxyError(c, http.StatusBadGateway, service.AssetErrorChannelNotFound, "no available asset channel supports requested operation")
}

func copyDoubaoAssetResponseHeaders(c *gin.Context, headers http.Header) {
	for key, values := range headers {
		switch http.CanonicalHeaderKey(key) {
		case "Connection", "Content-Length", "Keep-Alive", "Proxy-Authenticate", "Proxy-Authorization", "Te", "Trailer", "Transfer-Encoding", "Upgrade":
			continue
		}
		for _, value := range values {
			c.Writer.Header().Add(key, value)
		}
	}
}