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 TianyiYunAssetAdapter struct{} func NewTianyiYunAssetAdapter() AssetAdapter { return &TianyiYunAssetAdapter{} } func (a *TianyiYunAssetAdapter) Name() string { return "tianyiyun_asset" } func (a *TianyiYunAssetAdapter) Supports(operation AssetOperation) bool { return operation == AssetOperationAssetCreate || operation == AssetOperationAssetGet } func (a *TianyiYunAssetAdapter) 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) } if channel == nil { return nil, newAssetError(AssetErrorInvalidRequest, "channel is required", http.StatusBadRequest) } upstreamURL, body, method, assetErr := buildTianyiYunAssetRequest(channel, req) if assetErr != nil { return nil, assetErr } 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, method, upstreamURL, bytes.NewReader(body)) if err != nil { return nil, newAssetError(AssetErrorServer, err.Error(), http.StatusInternalServerError) } httpReq.Header.Set("Accept", "application/json") httpReq.Header.Set("Authorization", "Bearer "+strings.TrimSpace(channel.Key)) if method == http.MethodPost { httpReq.Header.Set("Content-Type", "application/json") } 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 < http.StatusOK || resp.StatusCode >= http.StatusMultipleChoices { return nil, newAssetError(AssetErrorUpstream, string(data), http.StatusBadGateway) } return &AssetUpstreamResponse{StatusCode: resp.StatusCode, Header: resp.Header, Body: data}, nil } func buildTianyiYunAssetRequest(channel *model.Channel, req AssetRequest) (string, []byte, string, *AssetError) { baseURL, err := tianyiYunAssetBaseURL(channel) if err != nil { return "", nil, "", newAssetError(AssetErrorInvalidRequest, err.Error(), http.StatusBadRequest) } switch req.Action.Operation { case AssetOperationAssetCreate: payload, assetErr := buildTianyiYunAssetCreatePayload(req.Body) if assetErr != nil { return "", nil, "", assetErr } data, err := common.Marshal(payload) if err != nil { return "", nil, "", newAssetError(AssetErrorServer, err.Error(), http.StatusInternalServerError) } return baseURL + "/api/assets/upload", data, http.MethodPost, nil case AssetOperationAssetGet: id := tianyiYunAssetField(req.Body, "Id", "id") if id == "" { return "", nil, "", newAssetError(AssetErrorInvalidRequest, "asset id is required", http.StatusBadRequest) } return baseURL + "/api/assets/" + url.PathEscape(id), nil, http.MethodGet, nil default: return "", nil, "", newAssetError(AssetErrorOperationNotSupported, fmt.Sprintf("asset operation %s is not supported", req.Action.Operation), http.StatusBadRequest) } } func tianyiYunAssetBaseURL(channel *model.Channel) (string, error) { baseURL := strings.TrimSpace(channel.GetBaseURL()) if baseURL == "" { baseURL = "https://ai.ctaigw.cn/v1" } u, err := url.Parse(baseURL) if err != nil || u.Scheme == "" || u.Host == "" { return "", fmt.Errorf("invalid TianyiYun asset base URL") } u.Path = strings.TrimRight(u.Path, "/") if !strings.HasSuffix(u.Path, "/v1") { u.Path += "/v1" } u.RawQuery = "" u.Fragment = "" return strings.TrimRight(u.String(), "/"), nil } func buildTianyiYunAssetCreatePayload(body map[string]any) (map[string]string, *AssetError) { sourceURL := tianyiYunAssetField(body, "URL", "url") if sourceURL == "" { return nil, newAssetError(AssetErrorInvalidRequest, "asset URL is required", http.StatusBadRequest) } parsedURL, err := url.Parse(sourceURL) if err != nil || (parsedURL.Scheme != "http" && parsedURL.Scheme != "https") || parsedURL.Host == "" { return nil, newAssetError(AssetErrorInvalidRequest, "asset URL must use http or https", http.StatusBadRequest) } assetType := tianyiYunAssetField(body, "AssetType", "asset_type") if assetType != "Image" && assetType != "Video" && assetType != "Audio" { return nil, newAssetError(AssetErrorInvalidRequest, "asset type must be Image, Video, or Audio", http.StatusBadRequest) } fetchSetting := system_setting.GetFetchSetting() if err := common.ValidateURLWithFetchSetting(sourceURL, fetchSetting.EnableSSRFProtection, fetchSetting.AllowPrivateIp, fetchSetting.DomainFilterMode, fetchSetting.IpFilterMode, fetchSetting.DomainList, fetchSetting.IpList, fetchSetting.AllowedPorts, fetchSetting.ApplyIPFilterForDomain); err != nil { return nil, newAssetError(AssetErrorInvalidRequest, fmt.Sprintf("asset URL blocked: %v", err), http.StatusBadRequest) } payload := map[string]string{"url": sourceURL, "asset_type": assetType} if name := tianyiYunAssetField(body, "Name", "name"); name != "" { payload["name"] = name } return payload, nil } func tianyiYunAssetField(body map[string]any, upperKey string, lowerKey string) string { if body == nil { return "" } if value, ok := body[upperKey]; ok { if text, ok := value.(string); ok { return strings.TrimSpace(text) } } if value, ok := body[lowerKey]; ok { if text, ok := value.(string); ok { return strings.TrimSpace(text) } } return "" }