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) } }() } }