package service import ( "context" "io" "net/http" "net/http/httptest" "strings" "testing" "github.com/QuantumNous/new-api/common" "github.com/QuantumNous/new-api/model" "github.com/QuantumNous/new-api/setting/system_setting" "github.com/stretchr/testify/require" ) func TestTianyiYunAssetAdapterCreateAssetMapsActionRequest(t *testing.T) { disableSSRFProtectionForTianyiYunAssetTest(t) server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { require.Equal(t, http.MethodPost, r.Method) require.Equal(t, "/v1/api/assets/upload", r.URL.Path) require.Equal(t, "Bearer key", r.Header.Get("Authorization")) body, err := io.ReadAll(r.Body) require.NoError(t, err) require.JSONEq(t, `{"url":"https://example.test/a.png","asset_type":"Image","name":"cover"}`, string(body)) _, _ = w.Write([]byte(`{"code":0,"data":{"Id":"asset-1"}}`)) })) defer server.Close() adapter := NewTianyiYunAssetAdapter() channel := &model.Channel{BaseURL: common.GetPointer(server.URL + "/v1"), Key: "key"} response, assetErr := adapter.DoAssetRequest(context.Background(), channel, AssetRequest{ Action: AssetActionSpec{Action: "CreateAsset", Operation: AssetOperationAssetCreate}, Body: map[string]any{"URL": "https://example.test/a.png", "AssetType": "Image", "Name": "cover"}, }) require.Nil(t, assetErr) require.Equal(t, http.StatusOK, response.StatusCode) require.JSONEq(t, `{"code":0,"data":{"Id":"asset-1"}}`, string(response.Body)) } func TestTianyiYunAssetAdapterGetAssetMapsActionRequest(t *testing.T) { disableSSRFProtectionForTianyiYunAssetTest(t) server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { require.Equal(t, http.MethodGet, r.Method) require.Equal(t, "/v1/api/assets/asset-1%2Fpart", r.URL.EscapedPath()) require.Equal(t, "Bearer key", r.Header.Get("Authorization")) body, err := io.ReadAll(r.Body) require.NoError(t, err) require.Empty(t, body) _, _ = w.Write([]byte(`{"code":0,"data":{"Id":"asset-1"}}`)) })) defer server.Close() adapter := NewTianyiYunAssetAdapter() channel := &model.Channel{BaseURL: common.GetPointer(server.URL), Key: "key"} response, assetErr := adapter.DoAssetRequest(context.Background(), channel, AssetRequest{ Action: AssetActionSpec{Action: "GetAsset", Operation: AssetOperationAssetGet}, Body: map[string]any{"Id": "asset-1/part"}, }) require.Nil(t, assetErr) require.Equal(t, http.StatusOK, response.StatusCode) } func TestTianyiYunAssetAdapterRejectsInvalidCreateRequest(t *testing.T) { cases := []map[string]any{ {"AssetType": "Image"}, {"URL": "data:image/png;base64,AAAA", "AssetType": "Image"}, {"URL": "aGVsbG8=", "AssetType": "Image"}, {"URL": "ftp://example.test/a.png", "AssetType": "Image"}, {"URL": "https://example.test/a.png", "AssetType": "Document"}, } for _, body := range cases { t.Run("invalid", func(t *testing.T) { adapter := NewTianyiYunAssetAdapter() response, assetErr := adapter.DoAssetRequest(context.Background(), &model.Channel{Key: "key"}, AssetRequest{ Action: AssetActionSpec{Action: "CreateAsset", Operation: AssetOperationAssetCreate}, Body: body, }) require.Nil(t, response) require.NotNil(t, assetErr) require.Equal(t, http.StatusBadRequest, assetErr.HTTPStatus) }) } } func disableSSRFProtectionForTianyiYunAssetTest(t *testing.T) { t.Helper() setting := system_setting.GetFetchSetting() old := *setting setting.EnableSSRFProtection = false t.Cleanup(func() { *setting = old }) } func TestTianyiYunAssetAdapterRejectsPrivateSourceURL(t *testing.T) { setting := system_setting.GetFetchSetting() old := *setting setting.EnableSSRFProtection = true setting.AllowPrivateIp = false setting.DomainFilterMode = false setting.IpFilterMode = false setting.DomainList = nil setting.IpList = nil setting.AllowedPorts = []string{"80", "443", "8080", "8443"} t.Cleanup(func() { *setting = old }) for _, sourceURL := range []string{"https://127.0.0.1/a.png", "http://10.0.0.1/a.png"} { t.Run(sourceURL, func(t *testing.T) { response, assetErr := NewTianyiYunAssetAdapter().DoAssetRequest(context.Background(), &model.Channel{Key: "key"}, AssetRequest{ Action: AssetActionSpec{Action: "CreateAsset", Operation: AssetOperationAssetCreate}, Body: map[string]any{"URL": sourceURL, "AssetType": "Image"}, }) require.Nil(t, response) require.NotNil(t, assetErr) require.Equal(t, http.StatusBadRequest, assetErr.HTTPStatus) }) } } func TestTianyiYunAssetAdapterSupportsOnlyCreateAndGet(t *testing.T) { adapter := NewTianyiYunAssetAdapter() require.True(t, adapter.Supports(AssetOperationAssetCreate)) require.True(t, adapter.Supports(AssetOperationAssetGet)) require.False(t, adapter.Supports(AssetOperationAssetList)) _, assetErr := adapter.DoAssetRequest(context.Background(), &model.Channel{Key: "key"}, AssetRequest{ Action: AssetActionSpec{Operation: AssetOperationAssetList}, }) require.NotNil(t, assetErr) require.Equal(t, AssetErrorOperationNotSupported, assetErr.Type) require.True(t, strings.Contains(assetErr.Message, "not supported")) }