easyai-ai-gateway/apps/api/internal/runner/object_storage_test.go

package runner

import (
	"context"
	"io"
	"net/http"
	"net/http/httptest"
	"net/url"
	"strings"
	"sync/atomic"
	"testing"
	"time"

	"github.com/easyai/easyai-ai-gateway/apps/api/internal/clients"
	"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
)

type objectStorageMetricEvent struct {
	event    string
	provider string
	bytes    int64
}

type objectStorageMetricsRecorder struct {
	events []objectStorageMetricEvent
}

func (m *objectStorageMetricsRecorder) ObserveBillingEvent(string) {}

func (m *objectStorageMetricsRecorder) ObserveObjectStorage(event string, provider string, bytes int64, _ time.Duration) {
	m.events = append(m.events, objectStorageMetricEvent{event: event, provider: provider, bytes: bytes})
}

func TestS3ObjectStorageUsesSigV4AndDeterministicObjectKey(t *testing.T) {
	payload := []byte("same-media-payload")
	var calls atomic.Int64
	var requestPath string
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		calls.Add(1)
		requestPath = r.URL.Path
		if !strings.HasPrefix(r.Header.Get("Authorization"), "AWS4-HMAC-SHA256 Credential=access-id/") {
			t.Errorf("missing SigV4 authorization: %q", r.Header.Get("Authorization"))
		}
		if r.Header.Get("x-amz-content-sha256") != sha256Hex(payload) {
			t.Errorf("payload hash=%q", r.Header.Get("x-amz-content-sha256"))
		}
		body, _ := io.ReadAll(r.Body)
		if string(body) != string(payload) {
			t.Errorf("payload=%q", body)
		}
		w.WriteHeader(http.StatusOK)
	}))
	defer server.Close()

	channel := testObjectStorageChannel("s3", server.URL, "s3-primary")
	adapter, err := newObjectStorageAdapter(channel)
	if err != nil {
		t.Fatal(err)
	}
	adapter.now = func() time.Time { return time.Date(2026, time.August, 4, 1, 2, 3, 0, time.UTC) }
	result, err := adapter.put(t.Context(), FileUploadPayload{Bytes: payload, ContentType: "image/png", FileName: "ignored.jpg", Scene: store.FileStorageSceneImageResult})
	if err != nil {
		t.Fatal(err)
	}
	wantSuffix := "/bucket/media/image_result/2026/08/04/" + sha256Hex(payload) + ".png"
	if requestPath != wantSuffix {
		t.Fatalf("request path=%q, want %q", requestPath, wantSuffix)
	}
	if stringFromAny(result["objectKey"]) != strings.TrimPrefix(wantSuffix, "/bucket/") {
		t.Fatalf("object key=%q", result["objectKey"])
	}
	if calls.Load() != 1 {
		t.Fatalf("calls=%d", calls.Load())
	}
}

func TestRequestAssetUsesPrivateSignedURLEvenWhenPublicBaseURLExists(t *testing.T) {
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
		w.WriteHeader(http.StatusOK)
	}))
	defer server.Close()
	channel := testObjectStorageChannel("s3", server.URL, "s3-private")
	channel.Config["publicBaseUrl"] = "https://cdn.example"
	channel.Config["accessScope"] = "public"
	adapter, err := newObjectStorageAdapter(channel)
	if err != nil {
		t.Fatal(err)
	}
	adapter.now = func() time.Time { return time.Date(2026, time.August, 4, 0, 0, 0, 0, time.UTC) }

	result, err := adapter.put(t.Context(), FileUploadPayload{
		Bytes: []byte("private input"), ContentType: "image/png", FileName: "input.png", Scene: store.FileStorageSceneRequestAsset,
	})
	if err != nil {
		t.Fatal(err)
	}
	if result["accessScope"] != "private" {
		t.Fatalf("accessScope=%v, want private", result["accessScope"])
	}
	gotURL := stringFromAny(result["url"])
	if strings.HasPrefix(gotURL, "https://cdn.example/") || !strings.Contains(gotURL, "X-Amz-Signature=") {
		t.Fatalf("request asset URL must be signed and private: %s", gotURL)
	}
}

func TestPrivateSignedURLUsesConfiguredTTL(t *testing.T) {
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
		w.WriteHeader(http.StatusOK)
	}))
	defer server.Close()
	channel := testObjectStorageChannel("s3", server.URL, "s3-private-ttl")
	channel.Config["signedUrlExpiresSeconds"] = 3600
	adapter, err := newObjectStorageAdapter(channel)
	if err != nil {
		t.Fatal(err)
	}
	adapter.now = func() time.Time { return time.Date(2026, time.August, 4, 0, 0, 0, 0, time.UTC) }

	result, err := adapter.put(t.Context(), FileUploadPayload{
		Bytes: []byte("private input"), ContentType: "image/png", FileName: "input.png", Scene: store.FileStorageSceneRequestAsset,
	})
	if err != nil {
		t.Fatal(err)
	}
	parsed, err := url.Parse(stringFromAny(result["url"]))
	if err != nil {
		t.Fatal(err)
	}
	if got := parsed.Query().Get("X-Amz-Expires"); got != "3600" {
		t.Fatalf("X-Amz-Expires=%q, want 3600", got)
	}
	if got := stringFromAny(result["urlExpiresAt"]); got != "2026-08-04T01:00:00Z" {
		t.Fatalf("urlExpiresAt=%q", got)
	}
}

func TestTemporaryObjectEnsuresS3LifecycleAndAddsTagging(t *testing.T) {
	var methods []string
	var lifecyclePayload string
	var tagging string
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if _, lifecycle := r.URL.Query()["lifecycle"]; lifecycle {
			methods = append(methods, r.Method+" lifecycle")
			if r.Method == http.MethodGet {
				w.WriteHeader(http.StatusNotFound)
				return
			}
			body, _ := io.ReadAll(r.Body)
			lifecyclePayload = string(body)
			w.WriteHeader(http.StatusOK)
			return
		}
		methods = append(methods, r.Method+" object")
		tagging = r.Header.Get("x-amz-tagging")
		w.WriteHeader(http.StatusOK)
	}))
	defer server.Close()

	channel := testObjectStorageChannel("s3", server.URL, "s3-expiring")
	channel.Config["temporaryFileExpirePolicy"] = "1m"
	adapter, err := newObjectStorageAdapter(channel)
	if err != nil {
		t.Fatal(err)
	}
	adapter.now = func() time.Time { return time.Date(2026, time.August, 4, 0, 0, 0, 0, time.UTC) }
	result, err := adapter.put(t.Context(), FileUploadPayload{
		Bytes: []byte("temporary image"), ContentType: "image/png", Scene: store.FileStorageSceneImageResult,
	})
	if err != nil {
		t.Fatal(err)
	}
	if got := strings.Join(methods, ","); got != "GET lifecycle,PUT lifecycle,PUT object" {
		t.Fatalf("methods=%s", got)
	}
	for _, fragment := range []string{"DeleteTempFiles-1d", "DeleteTempFiles-1m", "DeleteTempFiles-3m", "DeleteTempFiles-6m", "<Value>1m</Value>"} {
		if !strings.Contains(lifecyclePayload, fragment) {
			t.Fatalf("lifecycle payload missing %q: %s", fragment, lifecyclePayload)
		}
	}
	values, err := url.ParseQuery(tagging)
	if err != nil {
		t.Fatal(err)
	}
	if values.Get("TempFile") != "true" || values.Get("ExpiresIn") != "1m" || values.Get("ExpiresAt") != "1788393600" {
		t.Fatalf("unexpected tagging: %q", tagging)
	}
	if got := stringFromAny(result["objectExpiresAt"]); got != "2026-09-03T00:00:00Z" {
		t.Fatalf("objectExpiresAt=%q", got)
	}
}

func TestMergeObjectStorageLifecyclePreservesUnrelatedRules(t *testing.T) {
	current := []byte(`<?xml version="1.0"?><LifecycleConfiguration><Rule><ID>KeepArchive</ID><Prefix>archive/</Prefix><Status>Enabled</Status><Expiration><Days>365</Days></Expiration></Rule></LifecycleConfiguration>`)
	merged, changed, err := mergeObjectStorageLifecycle("aliyun_oss", current)
	if err != nil {
		t.Fatal(err)
	}
	if !changed {
		t.Fatal("lifecycle merge unexpectedly reported no change")
	}
	value := string(merged)
	if !strings.Contains(value, "<ID>KeepArchive</ID>") || !strings.Contains(value, "<Prefix>archive/</Prefix>") {
		t.Fatalf("unrelated lifecycle rule was not preserved: %s", value)
	}
	for _, definition := range objectStorageLifecycleDefinitions {
		if !strings.Contains(value, "<ID>"+definition.ID+"</ID>") {
			t.Fatalf("managed lifecycle rule %s missing", definition.ID)
		}
	}
}

func TestAliyunOSSUsesBucketVirtualHostForRegionalEndpoint(t *testing.T) {
	channel := testObjectStorageChannel("aliyun_oss", "https://oss-cn-hangzhou.aliyuncs.com", "oss-primary")
	channel.Config["forcePathStyle"] = false
	adapter, err := newObjectStorageAdapter(channel)
	if err != nil {
		t.Fatal(err)
	}
	value, err := adapter.objectURL("media/request_asset/object.png")
	if err != nil {
		t.Fatal(err)
	}
	if value != "https://bucket.oss-cn-hangzhou.aliyuncs.com/media/request_asset/object.png" {
		t.Fatalf("OSS object URL=%q", value)
	}
}

func TestAliyunOSSRequestUsesSignatureV1(t *testing.T) {
	channel := testObjectStorageChannel("aliyun_oss", "https://oss-cn-hangzhou.aliyuncs.com", "oss-primary")
	channel.Config["forcePathStyle"] = false
	adapter, err := newObjectStorageAdapter(channel)
	if err != nil {
		t.Fatal(err)
	}
	request, err := http.NewRequest(http.MethodPut, "https://bucket.oss-cn-hangzhou.aliyuncs.com/media/image_result/2026/08/04/OBJECT.png", nil)
	if err != nil {
		t.Fatal(err)
	}
	request.Header.Set("Content-Type", "image/png")
	adapter.signOSSRequest(request, time.Date(2026, time.August, 4, 1, 2, 3, 0, time.UTC))
	if got, want := request.Header.Get("Authorization"), "OSS access-id:fohlh2SG7k7D+cmX48b7Keqa1Ok="; got != want {
		t.Fatalf("authorization=%q, want %q", got, want)
	}
	if got := request.Header.Get("Date"); got != "Tue, 04 Aug 2026 01:02:03 GMT" {
		t.Fatalf("date=%q", got)
	}
}

func TestAliyunOSSSignatureIncludesLifecycleTaggingHeader(t *testing.T) {
	channel := testObjectStorageChannel("aliyun_oss", "https://oss-cn-hangzhou.aliyuncs.com", "oss-tagged")
	adapter, err := newObjectStorageAdapter(channel)
	if err != nil {
		t.Fatal(err)
	}
	request, err := http.NewRequest(http.MethodPut, "https://bucket.oss-cn-hangzhou.aliyuncs.com/media/image_result/object.png", nil)
	if err != nil {
		t.Fatal(err)
	}
	request.Header.Set("Content-Type", "image/png")
	request.Header.Set("x-oss-tagging", "ExpiresAt=1788393600&ExpiresIn=1m&TempFile=true")
	now := time.Date(2026, time.August, 4, 1, 2, 3, 0, time.UTC)
	adapter.signOSSRequest(request, now)
	stringToSign := "PUT\n\nimage/png\nTue, 04 Aug 2026 01:02:03 GMT\n" +
		"x-oss-tagging:ExpiresAt=1788393600&ExpiresIn=1m&TempFile=true\n" +
		"/bucket/media/image_result/object.png"
	want := "OSS access-id:" + hmacSHA1Base64([]byte("access-secret"), stringToSign)
	if got := request.Header.Get("Authorization"); got != want {
		t.Fatalf("authorization=%q, want %q", got, want)
	}
}

func TestTemporaryRetentionDoesNotApplyToPersistentUploadScene(t *testing.T) {
	config := map[string]any{"temporaryFileExpirePolicy": "1d"}
	retention := objectStorageRetentionFor(config, store.FileStorageSceneUpload, time.Now())
	if retention.Enabled || retention.Policy != objectStorageExpirationNever {
		t.Fatalf("persistent upload unexpectedly expires: %+v", retention)
	}
}

func TestObjectStorageRetriesCurrentChannelThenFailsOver(t *testing.T) {
	var primaryCalls atomic.Int64
	primary := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
		primaryCalls.Add(1)
		http.Error(w, "temporary", http.StatusServiceUnavailable)
	}))
	defer primary.Close()
	var secondaryCalls atomic.Int64
	secondary := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
		secondaryCalls.Add(1)
		w.WriteHeader(http.StatusOK)
	}))
	defer secondary.Close()

	first := testObjectStorageChannel("s3", primary.URL, "s3-primary")
	first.RetryPolicy = map[string]any{"enabled": true, "maxRetries": 1, "backoffSeconds": []any{0.001}}
	second := testObjectStorageChannel("aliyun_oss", secondary.URL, "oss-secondary")
	second.RetryPolicy = map[string]any{"enabled": false}
	metrics := &objectStorageMetricsRecorder{}
	service := &Service{billingMetrics: metrics}
	result, err := service.uploadFileWithFailover(t.Context(), FileUploadPayload{Bytes: []byte("payload"), ContentType: "image/png", Scene: store.FileStorageSceneRequestAsset}, []store.FileStorageChannel{first, second})
	if err != nil {
		t.Fatal(err)
	}
	channel, _ := result["storageChannel"].(map[string]any)
	if stringFromAny(channel["channelKey"]) != "oss-secondary" {
		t.Fatalf("unexpected winning channel: %+v", channel)
	}
	if primaryCalls.Load() != 2 || secondaryCalls.Load() != 1 {
		t.Fatalf("calls primary=%d secondary=%d", primaryCalls.Load(), secondaryCalls.Load())
	}
	assertObjectStorageMetricEvent(t, metrics.events, "retry", "s3", 0)
	assertObjectStorageMetricEvent(t, metrics.events, "failover", "s3", 0)
	assertObjectStorageMetricEvent(t, metrics.events, "write_success", "aliyun_oss", int64(len("payload")))
}

func TestFileStorageChannelConnectionProbeWritesHeadsAndDeletes(t *testing.T) {
	var methods []string
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		methods = append(methods, r.Method)
		switch r.Method {
		case http.MethodPut, http.MethodHead, http.MethodDelete:
			w.WriteHeader(http.StatusOK)
		default:
			http.Error(w, "unexpected method", http.StatusMethodNotAllowed)
		}
	}))
	defer server.Close()

	result, err := (&Service{}).TestFileStorageChannel(t.Context(), testObjectStorageChannel("s3", server.URL, "probe"))
	if err != nil {
		t.Fatal(err)
	}
	if !result.PutSucceeded || !result.HeadSucceeded || !result.DeleteSucceeded || result.Provider != "s3" {
		t.Fatalf("unexpected probe result: %+v", result)
	}
	if strings.Join(methods, ",") != "PUT,HEAD,DELETE" {
		t.Fatalf("probe methods=%v", methods)
	}
}

func TestObjectStorageReadUsesChannelRetryPolicy(t *testing.T) {
	var calls atomic.Int64
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if calls.Add(1) == 1 {
			http.Error(w, "temporary", http.StatusServiceUnavailable)
			return
		}
		_, _ = w.Write([]byte("stored-payload"))
	}))
	defer server.Close()
	channel := testObjectStorageChannel("s3", server.URL, "s3-read")
	channel.RetryPolicy = map[string]any{"enabled": true, "maxRetries": 1, "backoffSeconds": []any{0.001}}
	adapter, err := newObjectStorageAdapter(channel)
	if err != nil {
		t.Fatal(err)
	}
	payload, err := readObjectStorageWithRetries(t.Context(), adapter, "media/request_asset/object.png")
	if err != nil || string(payload) != "stored-payload" || calls.Load() != 2 {
		t.Fatalf("payload=%q calls=%d err=%v", payload, calls.Load(), err)
	}
}

func TestObjectStorageAuthFailureSkipsSameChannelRetry(t *testing.T) {
	var primaryCalls atomic.Int64
	primary := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
		primaryCalls.Add(1)
		http.Error(w, "forbidden details", http.StatusForbidden)
	}))
	defer primary.Close()
	secondary := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { w.WriteHeader(http.StatusOK) }))
	defer secondary.Close()

	first := testObjectStorageChannel("aliyun_oss", primary.URL, "oss-primary")
	first.RetryPolicy = map[string]any{"enabled": true, "maxRetries": 3, "backoffSeconds": []any{0.001}}
	second := testObjectStorageChannel("s3", secondary.URL, "s3-secondary")
	service := &Service{}
	if _, err := service.uploadFileWithFailover(t.Context(), FileUploadPayload{Bytes: []byte("payload"), ContentType: "image/png", Scene: store.FileStorageSceneRequestAsset}, []store.FileStorageChannel{first, second}); err != nil {
		t.Fatal(err)
	}
	if primaryCalls.Load() != 1 {
		t.Fatalf("auth failure retried %d times", primaryCalls.Load())
	}
}

func TestRequestLevelUploadFailureDoesNotSwitchChannels(t *testing.T) {
	var primaryCalls atomic.Int64
	primary := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
		primaryCalls.Add(1)
		http.Error(w, "unsupported media", http.StatusUnsupportedMediaType)
	}))
	defer primary.Close()
	var secondaryCalls atomic.Int64
	secondary := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
		secondaryCalls.Add(1)
		w.WriteHeader(http.StatusOK)
	}))
	defer secondary.Close()

	first := store.FileStorageChannel{
		ChannelKey: "server-main-primary", Provider: "server_main_openapi", UploadURL: primary.URL,
		APIKey: "test-key", RetryPolicy: map[string]any{"enabled": false},
	}
	second := testObjectStorageChannel("s3", secondary.URL, "s3-secondary")
	_, err := (&Service{}).uploadFileWithFailover(t.Context(), FileUploadPayload{
		Bytes: []byte("payload"), ContentType: "application/x-unsupported", Scene: store.FileStorageSceneUpload,
	}, []store.FileStorageChannel{first, second})
	if err == nil || clients.ErrorCode(err) != "upload_failed" {
		t.Fatalf("unexpected request-level error: %v", err)
	}
	if primaryCalls.Load() != 1 || secondaryCalls.Load() != 0 {
		t.Fatalf("request-level failure switched channels: primary=%d secondary=%d", primaryCalls.Load(), secondaryCalls.Load())
	}
}

func TestAllObjectStorageChannelsFailWithStableError(t *testing.T) {
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { http.Error(w, "failure", http.StatusBadGateway) }))
	defer server.Close()
	channel := testObjectStorageChannel("s3", server.URL, "only")
	channel.RetryPolicy = map[string]any{"enabled": false}
	_, err := (&Service{}).uploadFileWithFailover(context.Background(), FileUploadPayload{Bytes: []byte("payload"), Scene: store.FileStorageSceneUpload}, []store.FileStorageChannel{channel})
	if clients.ErrorCode(err) != "storage_write_failed" || strings.Contains(err.Error(), "failure") {
		t.Fatalf("unexpected public storage error: %v", err)
	}
}

func testObjectStorageChannel(provider string, endpoint string, channelKey string) store.FileStorageChannel {
	return store.FileStorageChannel{
		ChannelKey:      channelKey,
		Name:            channelKey,
		Provider:        provider,
		AccessKeyID:     "access-id",
		AccessKeySecret: "access-secret",
		Priority:        100,
		Config: map[string]any{
			"endpoint":       endpoint,
			"bucket":         "bucket",
			"region":         "cn-test-1",
			"objectPrefix":   "media",
			"publicBaseUrl":  "https://cdn.example.com",
			"forcePathStyle": true,
		},
	}
}

func assertObjectStorageMetricEvent(t *testing.T, events []objectStorageMetricEvent, event string, provider string, bytes int64) {
	t.Helper()
	for _, item := range events {
		if item.event == event && item.provider == provider && item.bytes == bytes {
			return
		}
	}
	t.Fatalf("missing metric event %s/%s/%d in %+v", event, provider, bytes, events)
}