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

package runner

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

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

const (
	objectStorageRequestTimeout = 120 * time.Second
	objectStorageReadLimit      = 256 << 20
	objectStorageSignedURLTTL   = 15 * time.Minute
)

var sharedObjectStorageHTTPClient = &http.Client{
	Timeout: objectStorageRequestTimeout,
	Transport: &http.Transport{
		MaxIdleConns:        256,
		MaxIdleConnsPerHost: 256,
		MaxConnsPerHost:     256,
		IdleConnTimeout:     90 * time.Second,
	},
}

type objectStorageAdapter struct {
	channel store.FileStorageChannel
	client  *http.Client
	now     func() time.Time
}

type FileStorageChannelTestResult struct {
	Provider        string `json:"provider"`
	PutSucceeded    bool   `json:"putSucceeded"`
	HeadSucceeded   bool   `json:"headSucceeded"`
	DeleteSucceeded bool   `json:"deleteSucceeded"`
	DurationMS      int64  `json:"durationMs"`
}

// TestFileStorageChannel performs an isolated write, metadata read and cleanup
// against one object-storage channel. The random probe avoids overwriting a
// content-addressed business object and no object key or credential is exposed.
func (s *Service) TestFileStorageChannel(ctx context.Context, channel store.FileStorageChannel) (FileStorageChannelTestResult, error) {
	startedAt := time.Now()
	result := FileStorageChannelTestResult{Provider: strings.ToLower(strings.TrimSpace(channel.Provider))}
	adapter, err := newObjectStorageAdapter(channel)
	if err != nil {
		return result, err
	}
	payload := FileUploadPayload{
		Bytes:       []byte("easyai-storage-probe:" + uuid.NewString()),
		ContentType: "application/octet-stream",
		FileName:    "probe.bin",
		Scene:       store.FileStorageSceneUpload,
		Source:      "admin-connection-test",
	}
	upload, err := adapter.put(ctx, payload)
	if err != nil {
		return result, err
	}
	result.PutSucceeded = true
	objectKey := stringFromAny(upload["objectKey"])
	if objectKey == "" {
		return result, storageClientError("storage_write_failed", "object storage probe did not return an object reference", 0, false)
	}
	if err := adapter.head(ctx, objectKey); err != nil {
		_ = adapter.delete(context.WithoutCancel(ctx), objectKey)
		return result, err
	}
	result.HeadSucceeded = true
	if err := adapter.delete(ctx, objectKey); err != nil {
		return result, err
	}
	result.DeleteSucceeded = true
	result.DurationMS = time.Since(startedAt).Milliseconds()
	return result, nil
}

func newObjectStorageAdapter(channel store.FileStorageChannel) (*objectStorageAdapter, error) {
	provider := strings.ToLower(strings.TrimSpace(channel.Provider))
	if provider != "aliyun_oss" && provider != "s3" {
		return nil, &clients.ClientError{Code: "storage_config_invalid", Message: "unsupported object storage provider", Retryable: false}
	}
	if objectStorageConfigString(channel.Config, "endpoint") == "" || objectStorageConfigString(channel.Config, "bucket") == "" {
		return nil, &clients.ClientError{Code: "storage_config_invalid", Message: "object storage endpoint and bucket are required", Retryable: false}
	}
	if strings.TrimSpace(channel.AccessKeyID) == "" || strings.TrimSpace(channel.AccessKeySecret) == "" {
		return nil, &clients.ClientError{Code: "storage_auth_failed", Message: "object storage credentials are not configured", Retryable: false}
	}
	return &objectStorageAdapter{
		channel: channel,
		client:  sharedObjectStorageHTTPClient,
		now:     time.Now,
	}, nil
}

func (a *objectStorageAdapter) put(ctx context.Context, payload FileUploadPayload) (map[string]any, error) {
	objectKey := a.objectKey(payload)
	requestURL, err := a.objectURL(objectKey)
	if err != nil {
		return nil, err
	}
	contentType := strings.TrimSpace(payload.ContentType)
	if contentType == "" {
		contentType = "application/octet-stream"
	}
	req, err := http.NewRequestWithContext(ctx, http.MethodPut, requestURL, bytes.NewReader(payload.Bytes))
	if err != nil {
		return nil, storageClientError("storage_config_invalid", err.Error(), 0, false)
	}
	req.Header.Set("Content-Type", contentType)
	if err := a.sign(req, sha256Hex(payload.Bytes), a.now().UTC()); err != nil {
		return nil, err
	}
	resp, err := a.client.Do(req)
	if err != nil {
		return nil, storageClientError("storage_write_failed", err.Error(), 0, true)
	}
	defer resp.Body.Close()
	responseBody, readErr := io.ReadAll(io.LimitReader(resp.Body, 64<<10))
	if readErr != nil {
		return nil, storageClientError("storage_write_failed", readErr.Error(), resp.StatusCode, storageStatusRetryable(resp.StatusCode))
	}
	if resp.StatusCode < 200 || resp.StatusCode >= 300 {
		return nil, objectStorageHTTPError("storage_write_failed", resp.StatusCode, responseBody)
	}
	accessScope := objectStorageAccessScope(a.channel, payload.Scene)
	publicURL := ""
	if accessScope == "public" {
		publicURL = a.publicURL(objectKey)
	}
	urlExpiresAt := ""
	if publicURL == "" {
		publicURL, err = a.presignGet(objectKey, objectStorageSignedURLTTL)
		if err != nil {
			return nil, err
		}
		urlExpiresAt = a.now().UTC().Add(objectStorageSignedURLTTL).Format(time.RFC3339)
	}
	digest := sha256.Sum256(payload.Bytes)
	result := map[string]any{
		"url":         publicURL,
		"fileName":    path.Base(objectKey),
		"objectKey":   objectKey,
		"contentType": contentType,
		"size":        len(payload.Bytes),
		"sha256":      hex.EncodeToString(digest[:]),
		"accessScope": accessScope,
		"storageChannel": map[string]any{
			"id":         a.channel.ID,
			"channelKey": a.channel.ChannelKey,
			"name":       a.channel.Name,
			"provider":   a.channel.Provider,
		},
	}
	if urlExpiresAt != "" {
		result["urlExpiresAt"] = urlExpiresAt
	}
	return result, nil
}

func (a *objectStorageAdapter) get(ctx context.Context, objectKey string) ([]byte, error) {
	requestURL, err := a.objectURL(objectKey)
	if err != nil {
		return nil, err
	}
	req, err := http.NewRequestWithContext(ctx, http.MethodGet, requestURL, nil)
	if err != nil {
		return nil, storageClientError("storage_config_invalid", err.Error(), 0, false)
	}
	if err := a.sign(req, sha256Hex(nil), a.now().UTC()); err != nil {
		return nil, err
	}
	resp, err := a.client.Do(req)
	if err != nil {
		return nil, storageClientError("storage_read_failed", err.Error(), 0, true)
	}
	defer resp.Body.Close()
	if resp.StatusCode < 200 || resp.StatusCode >= 300 {
		body, _ := io.ReadAll(io.LimitReader(resp.Body, 64<<10))
		return nil, objectStorageHTTPError("storage_read_failed", resp.StatusCode, body)
	}
	payload, err := io.ReadAll(io.LimitReader(resp.Body, objectStorageReadLimit+1))
	if err != nil {
		return nil, storageClientError("storage_read_failed", err.Error(), resp.StatusCode, true)
	}
	if len(payload) > objectStorageReadLimit {
		return nil, storageClientError("storage_read_failed", "stored object exceeds the read limit", resp.StatusCode, false)
	}
	return payload, nil
}

func (a *objectStorageAdapter) head(ctx context.Context, objectKey string) error {
	return a.emptyObjectRequest(ctx, http.MethodHead, objectKey)
}

func (a *objectStorageAdapter) delete(ctx context.Context, objectKey string) error {
	return a.emptyObjectRequest(ctx, http.MethodDelete, objectKey)
}

func (a *objectStorageAdapter) emptyObjectRequest(ctx context.Context, method string, objectKey string) error {
	requestURL, err := a.objectURL(objectKey)
	if err != nil {
		return err
	}
	req, err := http.NewRequestWithContext(ctx, method, requestURL, nil)
	if err != nil {
		return storageClientError("storage_config_invalid", err.Error(), 0, false)
	}
	if err := a.sign(req, sha256Hex(nil), a.now().UTC()); err != nil {
		return err
	}
	resp, err := a.client.Do(req)
	if err != nil {
		return storageClientError("storage_read_failed", err.Error(), 0, true)
	}
	defer resp.Body.Close()
	_, _ = io.Copy(io.Discard, io.LimitReader(resp.Body, 64<<10))
	if resp.StatusCode < 200 || resp.StatusCode >= 300 {
		return objectStorageHTTPError("storage_read_failed", resp.StatusCode, nil)
	}
	return nil
}

func (a *objectStorageAdapter) presignGet(objectKey string, ttl time.Duration) (string, error) {
	requestURL, err := a.objectURL(objectKey)
	if err != nil {
		return "", err
	}
	if ttl <= 0 {
		ttl = objectStorageSignedURLTTL
	}
	if strings.EqualFold(a.channel.Provider, "s3") {
		return a.presignS3Get(requestURL, ttl, a.now().UTC())
	}
	return a.presignOSSGet(requestURL, objectKey, ttl, a.now().UTC())
}

func (a *objectStorageAdapter) objectKey(payload FileUploadPayload) string {
	now := a.now().UTC()
	digest := sha256.Sum256(payload.Bytes)
	extension := uploadFileExtension(payload.ContentType, path.Ext(payload.FileName))
	prefix := strings.Trim(objectStorageConfigString(a.channel.Config, "objectPrefix"), "/")
	parts := make([]string, 0, 7)
	if prefix != "" {
		parts = append(parts, prefix)
	}
	parts = append(parts,
		firstNonEmptyString(strings.TrimSpace(payload.Scene), store.FileStorageSceneUpload),
		fmt.Sprintf("%04d", now.Year()),
		fmt.Sprintf("%02d", now.Month()),
		fmt.Sprintf("%02d", now.Day()),
		hex.EncodeToString(digest[:])+extension,
	)
	return strings.Join(parts, "/")
}

func (a *objectStorageAdapter) objectURL(objectKey string) (string, error) {
	endpoint, err := url.Parse(strings.TrimRight(objectStorageConfigString(a.channel.Config, "endpoint"), "/"))
	if err != nil || endpoint.Scheme == "" || endpoint.Host == "" || endpoint.User != nil {
		return "", storageClientError("storage_config_invalid", "invalid object storage endpoint", 0, false)
	}
	bucket := objectStorageConfigString(a.channel.Config, "bucket")
	if strings.EqualFold(a.channel.Provider, "s3") {
		if objectStorageConfigBool(a.channel.Config, "forcePathStyle") {
			endpoint.Path = strings.TrimRight(endpoint.Path, "/") + "/" + bucket + "/" + strings.TrimLeft(objectKey, "/")
		} else {
			endpoint.Host = bucket + "." + endpoint.Host
			endpoint.Path = strings.TrimRight(endpoint.Path, "/") + "/" + strings.TrimLeft(objectKey, "/")
		}
	} else {
		hostname := endpoint.Hostname()
		forcePathStyle := objectStorageConfigBool(a.channel.Config, "forcePathStyle")
		if !forcePathStyle && net.ParseIP(hostname) == nil && hostname != "localhost" && !strings.HasPrefix(strings.ToLower(hostname), strings.ToLower(bucket)+".") {
			if port := endpoint.Port(); port != "" {
				endpoint.Host = bucket + "." + hostname + ":" + port
			} else {
				endpoint.Host = bucket + "." + endpoint.Host
			}
		}
		endpoint.Path = strings.TrimRight(endpoint.Path, "/") + "/" + strings.TrimLeft(objectKey, "/")
	}
	return endpoint.String(), nil
}

func (a *objectStorageAdapter) publicURL(objectKey string) string {
	baseURL := strings.TrimRight(objectStorageConfigString(a.channel.Config, "publicBaseUrl"), "/")
	if baseURL == "" {
		baseURL = strings.TrimRight(objectStorageConfigString(a.channel.Config, "publicBaseURL"), "/")
	}
	if baseURL == "" {
		return ""
	}
	return baseURL + "/" + escapeObjectKey(objectKey)
}

func (a *objectStorageAdapter) sign(req *http.Request, payloadHash string, now time.Time) error {
	if strings.EqualFold(a.channel.Provider, "s3") {
		a.signS3Request(req, payloadHash, now)
		return nil
	}
	a.signOSSRequest(req, now)
	return nil
}

func (a *objectStorageAdapter) signOSSRequest(req *http.Request, now time.Time) {
	date := now.UTC().Format(http.TimeFormat)
	contentType := req.Header.Get("Content-Type")
	canonicalHeaders := ""
	if token := strings.TrimSpace(a.channel.SessionToken); token != "" {
		req.Header.Set("x-oss-security-token", token)
		canonicalHeaders = "x-oss-security-token:" + token + "\n"
	}
	canonicalResource := "/" + objectStorageConfigString(a.channel.Config, "bucket") + req.URL.EscapedPath()
	stringToSign := req.Method + "\n\n" + contentType + "\n" + date + "\n" + canonicalHeaders + canonicalResource
	signature := hmacSHA1Base64([]byte(a.channel.AccessKeySecret), stringToSign)
	req.Header.Set("Authorization", "OSS "+a.channel.AccessKeyID+":"+signature)
	req.Header.Set("Date", date)
}

func (a *objectStorageAdapter) presignOSSGet(requestURL string, objectKey string, ttl time.Duration, now time.Time) (string, error) {
	parsed, err := url.Parse(requestURL)
	if err != nil {
		return "", storageClientError("storage_config_invalid", err.Error(), 0, false)
	}
	expires := strconv.FormatInt(now.Add(ttl).Unix(), 10)
	canonicalResource := "/" + objectStorageConfigString(a.channel.Config, "bucket") + "/" + strings.TrimLeft(objectKey, "/")
	signature := hmacSHA1Base64([]byte(a.channel.AccessKeySecret), "GET\n\n\n"+expires+"\n"+canonicalResource)
	query := parsed.Query()
	query.Set("OSSAccessKeyId", a.channel.AccessKeyID)
	query.Set("Expires", expires)
	query.Set("Signature", signature)
	if token := strings.TrimSpace(a.channel.SessionToken); token != "" {
		query.Set("security-token", token)
	}
	parsed.RawQuery = query.Encode()
	return parsed.String(), nil
}

func (a *objectStorageAdapter) signS3Request(req *http.Request, payloadHash string, now time.Time) {
	region := firstNonEmptyString(objectStorageConfigString(a.channel.Config, "region"), "us-east-1")
	amzDate := now.Format("20060102T150405Z")
	date := now.Format("20060102")
	req.Header.Set("x-amz-date", amzDate)
	req.Header.Set("x-amz-content-sha256", payloadHash)
	if token := strings.TrimSpace(a.channel.SessionToken); token != "" {
		req.Header.Set("x-amz-security-token", token)
	}
	canonicalHeaders, signedHeaders := s3CanonicalHeaders(req)
	canonicalRequest := strings.Join([]string{
		req.Method,
		s3CanonicalURI(req.URL),
		s3CanonicalQuery(req.URL.Query()),
		canonicalHeaders,
		signedHeaders,
		payloadHash,
	}, "\n")
	scope := date + "/" + region + "/s3/aws4_request"
	stringToSign := "AWS4-HMAC-SHA256\n" + amzDate + "\n" + scope + "\n" + sha256Hex([]byte(canonicalRequest))
	signature := hex.EncodeToString(s3SigningHMAC(a.channel.AccessKeySecret, date, region, stringToSign))
	req.Header.Set("Authorization", "AWS4-HMAC-SHA256 Credential="+a.channel.AccessKeyID+"/"+scope+", SignedHeaders="+signedHeaders+", Signature="+signature)
}

func (a *objectStorageAdapter) presignS3Get(requestURL string, ttl time.Duration, now time.Time) (string, error) {
	parsed, err := url.Parse(requestURL)
	if err != nil {
		return "", storageClientError("storage_config_invalid", err.Error(), 0, false)
	}
	region := firstNonEmptyString(objectStorageConfigString(a.channel.Config, "region"), "us-east-1")
	amzDate := now.Format("20060102T150405Z")
	date := now.Format("20060102")
	scope := date + "/" + region + "/s3/aws4_request"
	seconds := int64(ttl / time.Second)
	if seconds < 1 {
		seconds = 1
	}
	if seconds > 7*24*60*60 {
		seconds = 7 * 24 * 60 * 60
	}
	query := parsed.Query()
	query.Set("X-Amz-Algorithm", "AWS4-HMAC-SHA256")
	query.Set("X-Amz-Credential", a.channel.AccessKeyID+"/"+scope)
	query.Set("X-Amz-Date", amzDate)
	query.Set("X-Amz-Expires", strconv.FormatInt(seconds, 10))
	query.Set("X-Amz-SignedHeaders", "host")
	if token := strings.TrimSpace(a.channel.SessionToken); token != "" {
		query.Set("X-Amz-Security-Token", token)
	}
	canonicalRequest := "GET\n" + s3CanonicalURI(parsed) + "\n" + s3CanonicalQuery(query) + "\nhost:" + strings.ToLower(parsed.Host) + "\n\nhost\nUNSIGNED-PAYLOAD"
	stringToSign := "AWS4-HMAC-SHA256\n" + amzDate + "\n" + scope + "\n" + sha256Hex([]byte(canonicalRequest))
	query.Set("X-Amz-Signature", hex.EncodeToString(s3SigningHMAC(a.channel.AccessKeySecret, date, region, stringToSign)))
	parsed.RawQuery = query.Encode()
	return parsed.String(), nil
}

func s3CanonicalHeaders(req *http.Request) (string, string) {
	headers := map[string]string{"host": strings.ToLower(req.URL.Host)}
	for key, values := range req.Header {
		lower := strings.ToLower(strings.TrimSpace(key))
		if lower == "content-type" || strings.HasPrefix(lower, "x-amz-") {
			headers[lower] = strings.Join(values, ",")
		}
	}
	keys := make([]string, 0, len(headers))
	for key := range headers {
		keys = append(keys, key)
	}
	sort.Strings(keys)
	var canonical strings.Builder
	for _, key := range keys {
		canonical.WriteString(key)
		canonical.WriteByte(':')
		canonical.WriteString(strings.Join(strings.Fields(headers[key]), " "))
		canonical.WriteByte('\n')
	}
	return canonical.String(), strings.Join(keys, ";")
}

func s3CanonicalURI(value *url.URL) string {
	uri := value.EscapedPath()
	if uri == "" {
		return "/"
	}
	return uri
}

func s3CanonicalQuery(values url.Values) string {
	return strings.ReplaceAll(values.Encode(), "+", "%20")
}

func s3SigningHMAC(secret string, date string, region string, stringToSign string) []byte {
	dateKey := hmacSHA256([]byte("AWS4"+secret), date)
	regionKey := hmacSHA256(dateKey, region)
	serviceKey := hmacSHA256(regionKey, "s3")
	signingKey := hmacSHA256(serviceKey, "aws4_request")
	return hmacSHA256(signingKey, stringToSign)
}

func hmacSHA256(key []byte, value string) []byte {
	mac := hmac.New(sha256.New, key)
	_, _ = mac.Write([]byte(value))
	return mac.Sum(nil)
}

func hmacSHA1Base64(key []byte, value string) string {
	mac := hmac.New(sha1.New, key)
	_, _ = mac.Write([]byte(value))
	return base64.StdEncoding.EncodeToString(mac.Sum(nil))
}

func sha256Hex(value []byte) string {
	digest := sha256.Sum256(value)
	return hex.EncodeToString(digest[:])
}

func objectStorageConfigString(config map[string]any, key string) string {
	if config == nil {
		return ""
	}
	value, _ := config[key].(string)
	return strings.TrimSpace(value)
}

func objectStorageConfigBool(config map[string]any, key string) bool {
	if config == nil {
		return false
	}
	value, _ := config[key].(bool)
	return value
}

func objectStorageAccessScope(channel store.FileStorageChannel, scene string) string {
	if scope := strings.ToLower(objectStorageConfigString(channel.Config, "accessScope")); scope == "public" || scope == "private" {
		return scope
	}
	if strings.TrimSpace(scene) == store.FileStorageSceneRequestAsset {
		return "private"
	}
	if strings.TrimSpace(objectStorageConfigString(channel.Config, "publicBaseUrl")) != "" || strings.TrimSpace(objectStorageConfigString(channel.Config, "publicBaseURL")) != "" {
		return "public"
	}
	return "private"
}

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

func storageStatusRetryable(status int) bool {
	return status == http.StatusRequestTimeout || status == http.StatusTooManyRequests || status >= 500
}

func objectStorageHTTPError(code string, status int, body []byte) error {
	if status == http.StatusUnauthorized || status == http.StatusForbidden {
		code = "storage_auth_failed"
	}
	message := http.StatusText(status)
	if message == "" {
		message = "object storage request failed"
	}
	// The provider response may contain object names, endpoints or credential
	// diagnostics. Keep it out of the public error chain; channel health stores
	// only the stable status description as well.
	return storageClientError(code, message, status, storageStatusRetryable(status))
}

func storageClientError(code string, message string, status int, retryable bool) error {
	return &clients.ClientError{Code: code, Message: message, StatusCode: status, Retryable: retryable}
}