Paperclip/pkg/internal/services/attachments.go

245 lines
6.1 KiB
Go

package services
import (
"fmt"
"github.com/spf13/viper"
"gorm.io/datatypes"
"math"
"mime"
"mime/multipart"
"net/http"
"path/filepath"
"strconv"
"sync"
"git.solsynth.dev/hydrogen/paperclip/pkg/internal/database"
"git.solsynth.dev/hydrogen/paperclip/pkg/internal/models"
"github.com/google/uuid"
"github.com/rs/zerolog/log"
"gorm.io/gorm"
)
const metadataCacheLimit = 512
var metadataCache sync.Map
func GetAttachmentByID(id uint) (models.Attachment, error) {
var attachment models.Attachment
if err := database.C.Where(models.Attachment{
BaseModel: models.BaseModel{ID: id},
}).Preload("Pool").Preload("Account").First(&attachment).Error; err != nil {
return attachment, err
} else {
CacheAttachment(attachment)
}
return attachment, nil
}
func GetAttachmentByRID(rid string) (models.Attachment, error) {
if val, ok := metadataCache.Load(rid); ok && val.(models.Attachment).Account.ID > 0 {
return val.(models.Attachment), nil
}
var attachment models.Attachment
if err := database.C.Where(models.Attachment{
Rid: rid,
}).Preload("Pool").Preload("Account").First(&attachment).Error; err != nil {
return attachment, err
} else {
CacheAttachment(attachment)
}
return attachment, nil
}
func GetAttachmentByHash(hash string) (models.Attachment, error) {
var attachment models.Attachment
if err := database.C.Where(models.Attachment{
HashCode: hash,
}).Preload("Pool").First(&attachment).Error; err != nil {
return attachment, err
}
return attachment, nil
}
func GetAttachmentCache(id any) (models.Attachment, bool) {
if val, ok := metadataCache.Load(id); ok && val.(models.Attachment).Account.ID > 0 {
return val.(models.Attachment), ok
}
return models.Attachment{}, false
}
func CacheAttachment(item models.Attachment) {
MaintainAttachmentCache()
metadataCache.Store(item.Rid, item)
}
func NewAttachmentMetadata(tx *gorm.DB, user models.Account, file *multipart.FileHeader, attachment models.Attachment) (models.Attachment, error) {
attachment.Uuid = uuid.NewString()
attachment.Rid = RandString(16)
attachment.Size = file.Size
attachment.Name = file.Filename
attachment.AccountID = user.ID
// If the user didn't provide file mimetype manually, we have to detect it
if len(attachment.MimeType) == 0 {
if ext := filepath.Ext(attachment.Name); len(ext) > 0 {
// Detect mimetype by file extensions
attachment.MimeType = mime.TypeByExtension(ext)
} else {
// Detect mimetype by file header
// This method as a fallback method, because this isn't pretty accurate
header, err := file.Open()
if err != nil {
return attachment, fmt.Errorf("failed to read file header: %v", err)
}
defer header.Close()
fileHeader := make([]byte, 512)
_, err = header.Read(fileHeader)
if err != nil {
return attachment, err
}
attachment.MimeType = http.DetectContentType(fileHeader)
}
}
if err := tx.Save(&attachment).Error; err != nil {
return attachment, fmt.Errorf("failed to save attachment record: %v", err)
} else {
CacheAttachment(attachment)
}
return attachment, nil
}
func NewAttachmentPlaceholder(tx *gorm.DB, user models.Account, attachment models.Attachment) (models.Attachment, error) {
attachment.Uuid = uuid.NewString()
attachment.Rid = RandString(16)
attachment.IsUploaded = false
attachment.FileChunks = datatypes.JSONMap{}
attachment.AccountID = user.ID
chunkSize := viper.GetInt64("performance.file_chunk_size")
chunkCount := math.Ceil(float64(attachment.Size) / float64(chunkSize))
for idx := 0; idx < int(chunkCount); idx++ {
cid := RandString(8)
attachment.FileChunks[cid] = idx
}
// If the user didn't provide file mimetype manually, we have to detect it
if len(attachment.MimeType) == 0 {
if ext := filepath.Ext(attachment.Name); len(ext) > 0 {
// Detect mimetype by file extensions
attachment.MimeType = mime.TypeByExtension(ext)
}
}
if err := tx.Save(&attachment).Error; err != nil {
return attachment, fmt.Errorf("failed to save attachment record: %v", err)
} else {
CacheAttachment(attachment)
}
return attachment, nil
}
func TryLinkAttachment(tx *gorm.DB, og models.Attachment, hash string) (bool, error) {
prev, err := GetAttachmentByHash(hash)
if err != nil {
return false, err
}
if prev.PoolID != nil && og.PoolID != nil && prev.PoolID != og.PoolID {
if !prev.Pool.Config.Data().AllowCrossPoolEgress || !og.Pool.Config.Data().AllowCrossPoolIngress {
// Pool config doesn't allow reference
return false, nil
}
}
prev.RefCount++
og.RefID = &prev.ID
og.Uuid = prev.Uuid
og.Destination = prev.Destination
if og.AccountID == prev.AccountID {
og.IsSelfRef = true
}
if err := tx.Save(&og).Error; err != nil {
tx.Rollback()
return true, err
} else if err = tx.Save(&prev).Error; err != nil {
tx.Rollback()
return true, err
}
CacheAttachment(prev)
CacheAttachment(og)
return true, nil
}
func UpdateAttachment(item models.Attachment) (models.Attachment, error) {
if err := database.C.Updates(&item).Error; err != nil {
return item, err
} else {
CacheAttachment(item)
}
return item, nil
}
func DeleteAttachment(item models.Attachment) error {
dat := item
tx := database.C.Begin()
if item.RefID != nil {
var refTarget models.Attachment
if err := database.C.Where(models.Attachment{
BaseModel: models.BaseModel{ID: *item.RefID},
}).First(&refTarget).Error; err == nil {
refTarget.RefCount--
if err := tx.Save(&refTarget).Error; err != nil {
tx.Rollback()
return fmt.Errorf("unable to update ref count: %v", err)
}
}
}
if err := database.C.Delete(&item).Error; err != nil {
tx.Rollback()
return err
} else {
strId := strconv.Itoa(int(item.ID))
metadataCache.Delete(strId)
metadataCache.Delete(item.Rid)
}
tx.Commit()
if dat.RefCount == 0 {
PublishDeleteFileTask(dat)
}
return nil
}
func MaintainAttachmentCache() {
var keySet []any
metadataCache.Range(func(k any, v any) bool {
keySet = append(keySet, k)
return true
})
if len(keySet) > metadataCacheLimit {
go func() {
log.Debug().Int("count", len(keySet)).Msg("Cleaning attachment metadata cache...")
for _, k := range keySet {
metadataCache.Delete(k)
}
}()
}
}