new-api/service/asset_tianyiyun_test.go

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"))
}