doc_ai_backed/pkg/oss/policy.go

package oss

import (
	"context"
	"crypto/hmac"
	"crypto/sha1"
	"encoding/base64"
	"encoding/json"
	"fmt"
	"io"
	"path/filepath"
	"strings"
	"time"

	"gitea.com/bitwsd/document_ai/config"
	"gitea.com/bitwsd/document_ai/pkg/log"
	"github.com/aliyun/aliyun-oss-go-sdk/oss"
)

func GetGMTISO8601(expireEnd int64) string {
	return time.Unix(expireEnd, 0).UTC().Format("2006-01-02T15:04:05Z")
}

func GetPolicyToken() (string, error) {
	now := time.Now().Unix()
	expireEnd := now + ExpireTime
	tokenExpire := GetGMTISO8601(expireEnd)

	conf := ConfigStruct{
		Expiration: tokenExpire,
	}

	// Add file prefix restriction
	// config.Conditions = append(config.Conditions, []interface{}{"starts-with", "$key", FormulaDir})

	// Add file size restriction (1KB to 10MB)
	minSize := int64(1024)
	maxSize := int64(3 * 1024 * 1024)
	conf.Conditions = append(conf.Conditions, []interface{}{"content-length-range", minSize, maxSize})

	result, err := json.Marshal(conf)
	if err != nil {
		return "", fmt.Errorf("marshal config error: %w", err)
	}

	encodedResult := base64.StdEncoding.EncodeToString(result)
	h := hmac.New(sha1.New, []byte(config.GlobalConfig.Aliyun.OSS.AccessKeySecret))
	io.WriteString(h, encodedResult)
	signedStr := base64.StdEncoding.EncodeToString(h.Sum(nil))
	policyToken := PolicyToken{
		AccessKeyId: config.GlobalConfig.Aliyun.OSS.AccessKeyID,
		Host:        config.GlobalConfig.Aliyun.OSS.Endpoint,
		Signature:   signedStr,
		Policy:      encodedResult,
		Directory:   FormulaDir,
	}

	response, err := json.Marshal(policyToken)
	if err != nil {
		return "", fmt.Errorf("marshal policy token error: %w", err)
	}

	return string(response), nil
}

func GetPolicyURL(ctx context.Context, path string) (string, error) {
	// Create OSS client
	client, err := oss.New(config.GlobalConfig.Aliyun.OSS.Endpoint, config.GlobalConfig.Aliyun.OSS.AccessKeyID, config.GlobalConfig.Aliyun.OSS.AccessKeySecret)
	if err != nil {
		log.Error(ctx, "func", "GetPolicyURL", "msg", "create oss client failed", "error", err)
		return "", err
	}

	// Get bucket instance
	bucket, err := client.Bucket(config.GlobalConfig.Aliyun.OSS.BucketName)
	if err != nil {
		log.Error(ctx, "func", "GetPolicyURL", "msg", "get bucket failed", "error", err)
		return "", err
	}

	// Set options for the signed URL
	var contentType string
	ext := filepath.Ext(path)
	switch ext {
	case ".jpg", ".jpeg":
		contentType = "image/jpeg"
	case ".png":
		contentType = "image/png"
	case ".gif":
		contentType = "image/gif"
	case ".bmp":
		contentType = "image/bmp"
	case ".webp":
		contentType = "image/webp"
	case ".tiff":
		contentType = "image/tiff"
	case ".svg":
		contentType = "image/svg+xml"
	default:
		return "", fmt.Errorf("unsupported file type: %s", ext)
	}

	options := []oss.Option{
		oss.ContentType(contentType),
	}

	// Generate signed URL valid for 10 minutes
	signedURL, err := bucket.SignURL(path, oss.HTTPPut, ExpireTime, options...)
	if err != nil {
		log.Error(ctx, "func", "GetPolicyURL", "msg", "sign url failed", "error", err)
		return "", err
	}

	// http 转 https
	signedURL = strings.Replace(signedURL, "http://", "https://", 1)

	return signedURL, nil
}

// DownloadFile downloads a file from OSS and returns the reader, caller should close the reader
func DownloadFile(ctx context.Context, ossPath string) (io.ReadCloser, error) {
	endpoint := config.GlobalConfig.Aliyun.OSS.InnerEndpoint
	if config.GlobalConfig.Server.IsDebug() {
		endpoint = config.GlobalConfig.Aliyun.OSS.Endpoint
	}

	// Create OSS client
	client, err := oss.New(endpoint, config.GlobalConfig.Aliyun.OSS.AccessKeyID, config.GlobalConfig.Aliyun.OSS.AccessKeySecret)
	if err != nil {
		log.Error(ctx, "func", "DownloadFile", "msg", "create oss client failed", "error", err)
		return nil, err
	}

	// Get bucket instance
	bucket, err := client.Bucket(config.GlobalConfig.Aliyun.OSS.BucketName)
	if err != nil {
		log.Error(ctx, "func", "DownloadFile", "msg", "get bucket failed", "error", err)
		return nil, err
	}

	// Download the file
	reader, err := bucket.GetObject(ossPath)
	if err != nil {
		log.Error(ctx, "func", "DownloadFile", "msg", "download file failed", "ossPath", ossPath, "error", err)
		return nil, err
	}

	return reader, nil
}

func GetDownloadURL(ctx context.Context, ossPath string) (string, error) {
	endpoint := config.GlobalConfig.Aliyun.OSS.Endpoint

	client, err := oss.New(endpoint, config.GlobalConfig.Aliyun.OSS.AccessKeyID, config.GlobalConfig.Aliyun.OSS.AccessKeySecret)
	if err != nil {
		log.Error(ctx, "func", "GetDownloadURL", "msg", "create oss client failed", "error", err)
		return "", err
	}

	bucket, err := client.Bucket(config.GlobalConfig.Aliyun.OSS.BucketName)
	if err != nil {
		log.Error(ctx, "func", "GetDownloadURL", "msg", "get bucket failed", "error", err)
		return "", err
	}

	signURL, err := bucket.SignURL(ossPath, oss.HTTPGet, 60)
	if err != nil {
		log.Error(ctx, "func", "GetDownloadURL", "msg", "get object failed", "error", err)
		return "", err
	}

	return signURL, nil
}