new-api/service/asset_compatible.go

package service

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

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

type CompatibleAssetAdapter struct {
	name      string
	operation map[AssetOperation]struct{}
}

func NewCompatibleAssetAdapter(name string, operations []AssetOperation) AssetAdapter {
	supported := make(map[AssetOperation]struct{}, len(operations))
	for _, op := range operations {
		supported[op] = struct{}{}
	}
	return &CompatibleAssetAdapter{name: name, operation: supported}
}

func (a *CompatibleAssetAdapter) Name() string {
	return a.name
}

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

func (a *CompatibleAssetAdapter) 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)
	}
	upstreamURL, err := buildCompatibleAssetURL(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)
	}
	httpReq.Header.Set("Content-Type", "application/json")
	httpReq.Header.Set("Accept", "application/json")
	httpReq.Header.Set("Authorization", "Bearer "+strings.TrimSpace(channel.Key))
	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 buildCompatibleAssetURL(channel *model.Channel, action string, version string) (string, error) {
	baseURL := strings.TrimSpace(channel.GetBaseURL())
	if baseURL == "" {
		baseURL = "https://ark.cn-beijing.volcengineapi.com"
	}
	u, err := url.Parse(baseURL)
	if err != nil {
		return "", err
	}
	u.Path = strings.TrimRight(u.Path, "/") + "/api/v1/multimodal/sd/assets"
	q := u.Query()
	q.Set("Action", action)
	if strings.TrimSpace(version) == "" {
		version = "2024-01-01"
	}
	q.Set("Version", version)
	u.RawQuery = q.Encode()
	return u.String(), nil
}