wx_service/internal/common/upload/service/upload_service.go

package service

import (
	"crypto/hmac"
	"crypto/rand"
	"crypto/sha1"
	"encoding/base64"
	"encoding/hex"
	"encoding/json"
	"errors"
	"fmt"
	"net/http"
	"path"
	"regexp"
	"strings"
	"time"

	"wx_service/config"
)

var (
	ErrUploadNotConfigured  = errors.New("oss upload is not configured")
	ErrCallbackUnauthorized = errors.New("upload callback unauthorized")
	ErrCallbackInvalidHeader = errors.New("upload callback authorization header is invalid")
)

type UploadService struct {
	cfg config.OSSConfig
}

func NewUploadService(cfg config.OSSConfig) *UploadService {
	return &UploadService{cfg: cfg}
}

type UploadToken struct {
	Token     string `json:"token"`
	Key       string `json:"key"`
	UploadURL string `json:"upload_url"`
	ExpireAt  int64  `json:"expire"`
	CDNDomain string `json:"cdn_domain"`
}

var extPattern = regexp.MustCompile(`^\.[a-z0-9]{1,10}$`)

func (s *UploadService) CreateUploadToken(miniProgramID uint, userID uint, filename string) (UploadToken, error) {
	if s.cfg.AccessKey == "" || s.cfg.SecretKey == "" || s.cfg.Bucket == "" {
		return UploadToken{}, ErrUploadNotConfigured
	}

	expireSeconds := s.cfg.TokenExpireSeconds
	if expireSeconds <= 0 {
		expireSeconds = 300
	}
	expireAt := time.Now().Add(time.Duration(expireSeconds) * time.Second).Unix()

	ext := strings.ToLower(path.Ext(filename))
	if !extPattern.MatchString(ext) {
		ext = ""
	}

	randomHex, err := randomHex(16)
	if err != nil {
		return UploadToken{}, fmt.Errorf("generate random key: %w", err)
	}

	keyPrefix := strings.Trim(s.cfg.KeyPrefix, "/")
	key := fmt.Sprintf("%s/mp_%d/user_%d/%s/%s%s",
		keyPrefix,
		miniProgramID,
		userID,
		time.Now().Format("20060102"),
		randomHex,
		ext,
	)

	putPolicy := map[string]interface{}{
		"scope":      fmt.Sprintf("%s:%s", s.cfg.Bucket, key),
		"deadline":   expireAt,
		"returnBody": `{"key":"$(key)","hash":"$(etag)","fsize":$(fsize),"mimeType":"$(mimeType)"}`,
	}
	if callbackURL := strings.TrimSpace(s.cfg.CallbackURL); callbackURL != "" {
		putPolicy["callbackUrl"] = callbackURL
		callbackBody := strings.TrimSpace(s.cfg.CallbackBody)
		if callbackBody == "" {
			callbackBody = "key=$(key)&hash=$(etag)&fsize=$(fsize)&mimeType=$(mimeType)"
		}
		callbackBodyType := strings.TrimSpace(s.cfg.CallbackBodyType)
		if callbackBodyType == "" {
			callbackBodyType = "application/x-www-form-urlencoded"
		}
		putPolicy["callbackBody"] = callbackBody
		putPolicy["callbackBodyType"] = callbackBodyType
	}

	policyJSON, err := json.Marshal(putPolicy)
	if err != nil {
		return UploadToken{}, fmt.Errorf("marshal put policy: %w", err)
	}

	encodedPolicy := urlSafeBase64NoPad(policyJSON)
	sign := hmacSHA1([]byte(s.cfg.SecretKey), []byte(encodedPolicy))
	encodedSign := urlSafeBase64NoPad(sign)

	token := fmt.Sprintf("%s:%s:%s", s.cfg.AccessKey, encodedSign, encodedPolicy)

	return UploadToken{
		Token:     token,
		Key:       key,
		UploadURL: s.cfg.UploadURL,
		ExpireAt:  expireAt,
		CDNDomain: s.cfg.CDNDomain,
	}, nil
}

func (s *UploadService) VerifyCallbackSignature(req *http.Request, rawBody []byte) error {
	if s.cfg.AccessKey == "" || s.cfg.SecretKey == "" {
		return ErrUploadNotConfigured
	}

	authHeader := strings.TrimSpace(req.Header.Get("Authorization"))
	if authHeader == "" {
		return ErrCallbackInvalidHeader
	}

	parts := strings.SplitN(authHeader, " ", 2)
	if len(parts) != 2 {
		return ErrCallbackInvalidHeader
	}
	scheme := strings.TrimSpace(parts[0])
	if !strings.EqualFold(scheme, "QBox") {
		return ErrCallbackInvalidHeader
	}

	token := strings.TrimSpace(parts[1])
	tokenParts := strings.SplitN(token, ":", 2)
	if len(tokenParts) != 2 {
		return ErrCallbackInvalidHeader
	}
	accessKey := strings.TrimSpace(tokenParts[0])
	providedSign := strings.TrimSpace(tokenParts[1])
	if accessKey == "" || providedSign == "" {
		return ErrCallbackInvalidHeader
	}
	if accessKey != s.cfg.AccessKey {
		return ErrCallbackUnauthorized
	}

	signing := req.URL.Path
	if req.URL.RawQuery != "" {
		signing += "?" + req.URL.RawQuery
	}
	signing += "\n"
	signing += string(rawBody)

	expected := urlSafeBase64NoPad(hmacSHA1([]byte(s.cfg.SecretKey), []byte(signing)))
	if !hmac.Equal([]byte(providedSign), []byte(expected)) {
		return ErrCallbackUnauthorized
	}
	return nil
}

func hmacSHA1(secret []byte, data []byte) []byte {
	mac := hmac.New(sha1.New, secret)
	mac.Write(data)
	return mac.Sum(nil)
}

func urlSafeBase64NoPad(data []byte) string {
	return base64.URLEncoding.WithPadding(base64.NoPadding).EncodeToString(data)
}

func randomHex(nBytes int) (string, error) {
	if nBytes <= 0 {
		return "", fmt.Errorf("invalid random bytes length")
	}
	buf := make([]byte, nBytes)
	if _, err := rand.Read(buf); err != nil {
		return "", err
	}
	return hex.EncodeToString(buf), nil
}