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

package runner

import (
	"bytes"
	"context"
	"crypto/hmac"
	"crypto/sha1"
	"encoding/base64"
	"fmt"
	"io"
	"net/http"
	"net/url"
	"path"
	"strings"
	"time"

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

const directOSSUploadTimeout = 120 * time.Second

type directOSSUploader struct {
	endpoint        string
	bucket          string
	accessKeyID     string
	accessKeySecret string
	publicBaseURL   string
	objectPrefix    string
	client          *http.Client
	now             func() time.Time
}

func newDirectOSSUploader(cfg config.Config) *directOSSUploader {
	if !cfg.MediaOSSDirectEnabled {
		return nil
	}
	return &directOSSUploader{
		endpoint:        strings.TrimRight(cfg.MediaOSSEndpoint, "/"),
		bucket:          strings.TrimSpace(cfg.MediaOSSBucket),
		accessKeyID:     strings.TrimSpace(cfg.MediaOSSAccessKeyID),
		accessKeySecret: strings.TrimSpace(cfg.MediaOSSAccessKeySecret),
		publicBaseURL:   strings.TrimRight(cfg.MediaOSSPublicBaseURL, "/"),
		objectPrefix:    strings.Trim(cfg.MediaOSSObjectPrefix, "/"),
		client: &http.Client{
			Timeout: directOSSUploadTimeout,
			Transport: &http.Transport{
				MaxIdleConns:        256,
				MaxIdleConnsPerHost: 256,
				MaxConnsPerHost:     256,
				IdleConnTimeout:     90 * time.Second,
			},
		},
		now: time.Now,
	}
}

func directOSSScene(scene string) bool {
	switch strings.TrimSpace(scene) {
	case store.FileStorageSceneRequestAsset, store.FileStorageSceneImageResult:
		return true
	default:
		return false
	}
}

func (u *directOSSUploader) upload(ctx context.Context, payload FileUploadPayload) (map[string]any, error) {
	if u == nil {
		return nil, &clients.ClientError{Code: "upload_config_failed", Message: "direct OSS uploader is not configured", Retryable: false}
	}
	objectKey := u.objectKey(payload)
	escapedKey := escapeOSSObjectKey(objectKey)
	uploadURL := u.endpoint + "/" + escapedKey
	contentType := strings.TrimSpace(payload.ContentType)
	if contentType == "" {
		contentType = "application/octet-stream"
	}

	var lastErr error
	for attempt := 0; attempt < 3; attempt++ {
		if attempt > 0 {
			if err := sleepWithContext(ctx, time.Duration(attempt*attempt)*200*time.Millisecond); err != nil {
				return nil, err
			}
		}
		status, err := u.put(ctx, uploadURL, escapedKey, contentType, payload.Bytes)
		if err == nil && status >= 200 && status < 300 {
			publicURL := u.publicBaseURL + "/" + escapedKey
			return map[string]any{
				"url":      publicURL,
				"fileName": objectKey,
				"storageChannel": map[string]any{
					"channelKey": "environment-direct-oss",
					"name":       "Environment Direct OSS",
					"provider":   "aliyun_oss_direct",
				},
				"assetStorage": map[string]any{
					"scene":    payload.Scene,
					"source":   firstNonEmptyString(payload.Source, "ai-gateway"),
					"strategy": "direct_aliyun_oss",
				},
			}, nil
		}
		if err != nil {
			lastErr = err
			continue
		}
		lastErr = fmt.Errorf("HTTP %d", status)
		if status != http.StatusTooManyRequests && status < 500 {
			break
		}
	}
	message := "direct OSS upload failed"
	if lastErr != nil {
		message += ": " + lastErr.Error()
	}
	return nil, &clients.ClientError{Code: "upload_failed", Message: message, Retryable: true}
}

func (u *directOSSUploader) put(ctx context.Context, uploadURL string, escapedKey string, contentType string, payload []byte) (int, error) {
	date := u.now().UTC().Format(http.TimeFormat)
	canonicalResource := "/" + u.bucket + "/" + escapedKey
	stringToSign := "PUT\n\n" + contentType + "\n" + date + "\n" + canonicalResource
	mac := hmac.New(sha1.New, []byte(u.accessKeySecret))
	_, _ = mac.Write([]byte(stringToSign))
	signature := base64.StdEncoding.EncodeToString(mac.Sum(nil))

	req, err := http.NewRequestWithContext(ctx, http.MethodPut, uploadURL, bytes.NewReader(payload))
	if err != nil {
		return 0, err
	}
	req.Header.Set("Authorization", "OSS "+u.accessKeyID+":"+signature)
	req.Header.Set("Content-Type", contentType)
	req.Header.Set("Date", date)
	resp, err := u.client.Do(req)
	if err != nil {
		return 0, err
	}
	defer resp.Body.Close()
	_, _ = io.Copy(io.Discard, io.LimitReader(resp.Body, 64<<10))
	return resp.StatusCode, nil
}

func (u *directOSSUploader) objectKey(payload FileUploadPayload) string {
	now := u.now().UTC()
	extension := uploadFileExtension(payload.ContentType, path.Ext(payload.FileName))
	baseName := strings.TrimSuffix(path.Base(payload.FileName), path.Ext(payload.FileName))
	baseName = sanitizeGeneratedAssetNamePart(baseName)
	if baseName == "" {
		baseName = "gateway-media"
	}
	if len(baseName) > 48 {
		baseName = baseName[:48]
	}
	return fmt.Sprintf(
		"%s/%s/%04d/%02d/%02d/%s-%s%s",
		u.objectPrefix,
		strings.TrimSpace(payload.Scene),
		now.Year(),
		now.Month(),
		now.Day(),
		baseName,
		randomHexSuffix(8),
		extension,
	)
}

func escapeOSSObjectKey(objectKey string) string {
	parts := strings.Split(objectKey, "/")
	for index, part := range parts {
		parts[index] = url.PathEscape(part)
	}
	return strings.Join(parts, "/")
}