new-api/controller/doubao_asset.go

package controller

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

	"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/QuantumNous/new-api/setting/system_setting"

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

const defaultDoubaoAssetBaseURL = "https://ark.cn-beijing.volcengineapi.com"

var blockedAssetActions = map[string]struct{}{
	"createassetgroup": {},
	"getassetgroup":    {},
	"listassetgroups":  {},
	"updateassetgroup": {},
	"deleteassetgroup": {},
}

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

func buildDoubaoAssetURL(channel *model.Channel, action string, version string) (string, error) {
	baseURL := defaultDoubaoAssetBaseURL
	if channel != nil && channel.BaseURL != nil {
		if configuredBaseURL := strings.TrimSpace(*channel.BaseURL); configuredBaseURL != "" {
			baseURL = configuredBaseURL
		}
	}

	u, err := url.Parse(baseURL)
	if err != nil {
		return "", err
	}
	if u.Scheme == "" || u.Host == "" {
		return "", fmt.Errorf("invalid Doubao asset base URL: %s", baseURL)
	}

	u.Path = strings.TrimRight(u.Path, "/") + "/api/v1/multimodal/sd/assets"
	u.RawQuery = ""
	u.Fragment = ""

	query := u.Query()
	query.Set("Action", action)
	if strings.TrimSpace(version) == "" {
		version = "2024-01-01"
	}
	query.Set("Version", version)
	u.RawQuery = query.Encode()

	return u.String(), nil
}

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 resolveDoubaoAssetChannelForRequest(c *gin.Context) (*model.Channel, string, error) {
	userId := c.GetInt("id")
	var lastErr error
	for _, group := range concreteDoubaoAssetGroupsForRequest(c, service.GetUserAutoGroup) {
		channel, err := service.ResolveDoubaoAssetChannel(userId, group)
		if err != nil {
			lastErr = err
			logger.LogError(c.Request.Context(), fmt.Sprintf("Failed to resolve Doubao asset channel for group %s: %v", group, err))
			continue
		}
		if channel != nil {
			return channel, group, nil
		}
	}
	if lastErr != nil {
		return nil, "", lastErr
	}
	return nil, "", nil
}

func DoubaoAssetProxy(c *gin.Context) {
	action := strings.TrimSpace(c.Query("Action"))
	if action == "" {
		assetProxyError(c, http.StatusBadRequest, "invalid_request_error", "Action query parameter is required")
		return
	}
	if _, ok := blockedAssetActions[strings.ToLower(action)]; ok {
		assetProxyError(c, http.StatusBadRequest, "invalid_request_error", fmt.Sprintf("Asset group API (%s) is not supported", action))
		return
	}

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

	channel, _, err := resolveDoubaoAssetChannelForRequest(c)
	if err != nil {
		assetProxyError(c, http.StatusBadGateway, "server_error", fmt.Sprintf("Failed to resolve Doubao asset channel: %v", err))
		return
	}
	if channel == nil {
		assetProxyError(c, http.StatusBadGateway, "server_error", "Failed to resolve Doubao asset channel")
		return
	}
	if strings.TrimSpace(channel.Key) == "" {
		assetProxyError(c, http.StatusBadGateway, "server_error", "Doubao asset channel key is missing")
		return
	}

	upstreamURL, err := buildDoubaoAssetURL(channel, action, version)
	if err != nil {
		assetProxyError(c, http.StatusBadGateway, "server_error", fmt.Sprintf("Failed to build upstream URL: %v", err))
		return
	}

	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 {
		logger.LogError(c.Request.Context(), fmt.Sprintf("Doubao asset URL blocked: %v", err))
		assetProxyError(c, http.StatusForbidden, "server_error", fmt.Sprintf("request blocked: %v", err))
		return
	}

	client, err := service.GetHttpClientWithProxy(channel.GetSetting().Proxy)
	if err != nil {
		assetProxyError(c, http.StatusBadGateway, "server_error", fmt.Sprintf("Failed to create proxy client: %v", err))
		return
	}
	if client == nil {
		assetProxyError(c, http.StatusBadGateway, "server_error", "Failed to create proxy client")
		return
	}

	ctx, cancel := context.WithTimeout(c.Request.Context(), 60*time.Second)
	defer cancel()

	req, err := http.NewRequestWithContext(ctx, http.MethodPost, upstreamURL, c.Request.Body)
	if err != nil {
		assetProxyError(c, http.StatusBadGateway, "server_error", fmt.Sprintf("Failed to create upstream request: %v", err))
		return
	}
	req.Header.Set("Content-Type", "application/json")
	req.Header.Set("Accept", "application/json")
	req.Header.Set("Authorization", "Bearer "+strings.TrimSpace(channel.Key))

	resp, err := client.Do(req)
	if err != nil {
		logger.LogError(c.Request.Context(), fmt.Sprintf("Failed to proxy Doubao asset request to %s: %s", upstreamURL, err.Error()))
		assetProxyError(c, http.StatusBadGateway, "server_error", fmt.Sprintf("Failed to proxy Doubao asset request: %v", err))
		return
	}
	defer resp.Body.Close()

	for key, values := range resp.Header {
		for _, value := range values {
			c.Writer.Header().Add(key, value)
		}
	}
	c.Writer.WriteHeader(resp.StatusCode)
	if _, err = io.Copy(c.Writer, resp.Body); err != nil {
		logger.LogError(c.Request.Context(), fmt.Sprintf("Failed to copy Doubao asset upstream response: %s", err.Error()))
	}
}