renjue
2665ac1dee
CI / docker (push) Failing after 10s
Resolve public_base_url at SignURL time so admin settings apply without restart. Co-authored-by: Cursor <cursoragent@cursor.com>
101 lines
2.5 KiB
Go
101 lines
2.5 KiB
Go
package imagestore
|
|
|
|
import (
|
|
"crypto/hmac"
|
|
"crypto/sha256"
|
|
"encoding/hex"
|
|
"fmt"
|
|
"os"
|
|
"path/filepath"
|
|
"strconv"
|
|
"strings"
|
|
"time"
|
|
|
|
"github.com/google/uuid"
|
|
"github.com/rose_cat707/luminary/internal/model"
|
|
"gorm.io/gorm"
|
|
)
|
|
|
|
type Store struct {
|
|
db *gorm.DB
|
|
baseDir string
|
|
signKey string
|
|
publicURL func() string
|
|
ttl time.Duration
|
|
}
|
|
|
|
func New(db *gorm.DB, baseDir string, publicURL func() string, signKey string, ttl time.Duration) *Store {
|
|
return &Store{db: db, baseDir: baseDir, publicURL: publicURL, signKey: signKey, ttl: ttl}
|
|
}
|
|
|
|
func (s *Store) baseURL() string {
|
|
if s.publicURL == nil {
|
|
return ""
|
|
}
|
|
return strings.TrimRight(s.publicURL(), "/")
|
|
}
|
|
|
|
func (s *Store) EnsureDir() error {
|
|
return os.MkdirAll(s.baseDir, 0o755)
|
|
}
|
|
|
|
func (s *Store) Save(predictionID, filename string, data []byte, mime string) (*model.StoredImage, error) {
|
|
if err := s.EnsureDir(); err != nil {
|
|
return nil, err
|
|
}
|
|
dir := filepath.Join(s.baseDir, predictionID)
|
|
if err := os.MkdirAll(dir, 0o755); err != nil {
|
|
return nil, err
|
|
}
|
|
localPath := filepath.Join(dir, filename)
|
|
if err := os.WriteFile(localPath, data, 0o644); err != nil {
|
|
return nil, err
|
|
}
|
|
img := model.StoredImage{
|
|
ID: uuid.NewString(),
|
|
PredictionID: predictionID,
|
|
Filename: filename,
|
|
LocalPath: localPath,
|
|
Mime: mime,
|
|
Size: int64(len(data)),
|
|
CreatedAt: time.Now(),
|
|
}
|
|
if err := s.db.Create(&img).Error; err != nil {
|
|
return nil, err
|
|
}
|
|
return &img, nil
|
|
}
|
|
|
|
func (s *Store) Get(id string) (*model.StoredImage, error) {
|
|
var img model.StoredImage
|
|
if err := s.db.First(&img, "id = ?", id).Error; err != nil {
|
|
return nil, err
|
|
}
|
|
return &img, nil
|
|
}
|
|
|
|
func (s *Store) ListByPrediction(predictionID string) ([]model.StoredImage, error) {
|
|
var imgs []model.StoredImage
|
|
err := s.db.Where("prediction_id = ?", predictionID).Order("created_at asc").Find(&imgs).Error
|
|
return imgs, err
|
|
}
|
|
|
|
func (s *Store) SignURL(imageID string) string {
|
|
exp := time.Now().Add(s.ttl).Unix()
|
|
sig := s.sign(imageID, exp)
|
|
return fmt.Sprintf("%s/files/%s?exp=%d&sig=%s", s.baseURL(), imageID, exp, sig)
|
|
}
|
|
|
|
func (s *Store) Verify(imageID string, exp int64, sig string) bool {
|
|
if time.Now().Unix() > exp {
|
|
return false
|
|
}
|
|
return hmac.Equal([]byte(s.sign(imageID, exp)), []byte(sig))
|
|
}
|
|
|
|
func (s *Store) sign(imageID string, exp int64) string {
|
|
mac := hmac.New(sha256.New, []byte(s.signKey))
|
|
mac.Write([]byte(imageID + "|" + strconv.FormatInt(exp, 10)))
|
|
return hex.EncodeToString(mac.Sum(nil))
|
|
}
|