package mail import ( "context" "encoding/json" "errors" "io" "path/filepath" "strings" "github.com/google/uuid" "github.com/jackc/pgx/v5" "github.com/ultisuite/ulti-backend/internal/mail/limits" "github.com/ultisuite/ulti-backend/internal/mail/storage" ) var ( ErrAttachmentNotFound = errors.New("attachment not found") ErrAttachmentTooLarge = limits.ErrAttachmentTooLarge ErrTooManyAttachments = limits.ErrTooManyAttachments ErrAttachmentsTotalTooLarge = limits.ErrAttachmentsTotalTooLarge ) type draftAttachmentRef struct { ID string `json:"id"` Filename string `json:"filename"` ContentType string `json:"content_type"` Size int64 `json:"size"` S3Bucket string `json:"s3_bucket"` S3Key string `json:"s3_key"` ContentID string `json:"content_id,omitempty"` IsInline bool `json:"is_inline"` } func (s *Service) ListMessageAttachments(ctx context.Context, externalID, messageID string) ([]map[string]any, error) { if _, err := s.ensureMessageOwned(ctx, externalID, messageID); err != nil { return nil, err } rows, err := s.db.Query(ctx, ` SELECT id, filename, content_type, size, content_id, is_inline, COALESCE(drive_path, '') FROM attachments WHERE message_id = $1 ORDER BY created_at ASC `, messageID) if err != nil { return nil, err } defer rows.Close() out := make([]map[string]any, 0) for rows.Next() { var id, filename, contentType, contentID, drivePath string var size int64 var isInline bool if err := rows.Scan(&id, &filename, &contentType, &size, &contentID, &isInline, &drivePath); err != nil { return nil, err } entry := map[string]any{ "id": id, "filename": filename, "content_type": contentType, "size": size, "is_inline": isInline, } if contentID != "" { entry["content_id"] = contentID } if drivePath != "" { entry["drive_path"] = drivePath } out = append(out, entry) } return out, rows.Err() } func (s *Service) MessageAttachmentCIDMap(ctx context.Context, externalID, messageID string) (map[string]string, error) { if _, err := s.ensureMessageOwned(ctx, externalID, messageID); err != nil { return nil, err } rows, err := s.db.Query(ctx, ` SELECT id, content_id, filename, is_inline FROM attachments WHERE message_id = $1 AND (content_id <> '' OR is_inline) `, messageID) if err != nil { return nil, err } defer rows.Close() mapping := make(map[string]string) for rows.Next() { var id, contentID, filename string var isInline bool if err := rows.Scan(&id, &contentID, &filename, &isInline); err != nil { return nil, err } if contentID != "" { registerCIDMapKeys(mapping, contentID, id) } if isInline && filename != "" { registerCIDMapKeys(mapping, filepath.Base(filename), id) } } return mapping, rows.Err() } func registerCIDMapKeys(mapping map[string]string, contentID, attachmentID string) { key := strings.Trim(contentID, "<> \t") if key == "" { return } mapping[key] = attachmentID mapping[strings.ToLower(key)] = attachmentID mapping["cid:"+key] = attachmentID mapping["cid:"+strings.ToLower(key)] = attachmentID } func (s *Service) UploadMessageAttachment( ctx context.Context, externalID, messageID, filename, contentType, contentID string, isInline bool, reader io.Reader, size int64, ) (string, error) { if s.storage == nil { return "", errors.New("object storage unavailable") } if err := limits.ValidateAttachmentSize(size); err != nil { return "", err } userID, err := s.ensureMessageOwned(ctx, externalID, messageID) if err != nil { return "", err } var count int var totalSize int64 err = s.db.QueryRow(ctx, ` SELECT COUNT(*)::int, COALESCE(SUM(size), 0)::bigint FROM attachments WHERE message_id = $1 `, messageID).Scan(&count, &totalSize) if err != nil { return "", err } if err := limits.ValidateAttachmentQuota(count, totalSize, size); err != nil { return "", err } objectKey := storage.MessageObjectKey(userID, messageID, filename) if err := s.storage.Put(ctx, objectKey, reader, size, contentType); err != nil { return "", err } var id string err = s.db.QueryRow(ctx, ` INSERT INTO attachments (message_id, filename, content_type, size, s3_bucket, s3_key, content_id, is_inline) VALUES ($1, $2, $3, $4, $5, $6, $7, $8) RETURNING id `, messageID, filename, contentType, size, s.storageBucket(), objectKey, contentID, isInline).Scan(&id) if err != nil { _ = s.storage.Delete(ctx, objectKey) return "", err } _, err = s.db.Exec(ctx, ` UPDATE messages SET has_attachments = true, updated_at = NOW() WHERE id = $1 `, messageID) if err != nil { return "", err } return id, nil } func (s *Service) OpenAttachment(ctx context.Context, externalID, attachmentID string) ( filename, contentType string, size int64, isInline bool, body io.ReadCloser, err error, ) { if s.storage == nil { return "", "", 0, false, nil, errors.New("object storage unavailable") } var s3Key string err = s.db.QueryRow(ctx, ` SELECT a.filename, a.content_type, a.size, a.is_inline, a.s3_key FROM attachments a JOIN messages m ON a.message_id = m.id JOIN mail_accounts ma ON m.account_id = ma.id WHERE a.id = $1 AND ma.user_id = (SELECT id FROM users WHERE external_id = $2) `, attachmentID, externalID).Scan(&filename, &contentType, &size, &isInline, &s3Key) if err != nil { if errors.Is(err, pgx.ErrNoRows) { return "", "", 0, false, nil, ErrAttachmentNotFound } return "", "", 0, false, nil, err } obj, err := s.storage.Get(ctx, s3Key) if err != nil { return "", "", 0, false, nil, err } return filename, contentType, size, isInline, obj, nil } func (s *Service) UploadDraftAttachment( ctx context.Context, externalID, draftID, filename, contentType, contentID string, isInline bool, reader io.Reader, size int64, ) (string, error) { if s.storage == nil { return "", errors.New("object storage unavailable") } if err := limits.ValidateAttachmentSize(size); err != nil { return "", err } userID, err := s.ResolveUserID(ctx, externalID) if err != nil { return "", err } var attachmentsJSON []byte err = s.db.QueryRow(ctx, ` SELECT attachments FROM outbox WHERE id = $1 AND user_id = $2 AND status = 'draft' `, draftID, userID).Scan(&attachmentsJSON) if err != nil { if errors.Is(err, pgx.ErrNoRows) { return "", ErrNotFound } return "", err } objectKey := storage.DraftObjectKey(userID, draftID, filename) if err := s.storage.Put(ctx, objectKey, reader, size, contentType); err != nil { return "", err } refs := make([]draftAttachmentRef, 0) if len(attachmentsJSON) > 0 && string(attachmentsJSON) != "[]" { _ = json.Unmarshal(attachmentsJSON, &refs) } var totalSize int64 for _, ref := range refs { totalSize += ref.Size } if err := limits.ValidateAttachmentQuota(len(refs), totalSize, size); err != nil { return "", err } attID := uuid.NewString() refs = append(refs, draftAttachmentRef{ ID: attID, Filename: filename, ContentType: contentType, Size: size, S3Bucket: s.storageBucket(), S3Key: objectKey, ContentID: contentID, IsInline: isInline, }) updated, _ := json.Marshal(refs) result, err := s.db.Exec(ctx, ` UPDATE outbox SET attachments = $1, updated_at = NOW() WHERE id = $2 AND user_id = $3 AND status = 'draft' `, updated, draftID, userID) if err != nil { _ = s.storage.Delete(ctx, objectKey) return "", err } if result.RowsAffected() == 0 { _ = s.storage.Delete(ctx, objectKey) return "", ErrNotFound } return attID, nil } func (s *Service) OpenDraftAttachment(ctx context.Context, externalID, draftID, attachmentID string) ( filename, contentType string, body io.ReadCloser, err error, ) { if s.storage == nil { return "", "", nil, errors.New("object storage unavailable") } userID, err := s.ResolveUserID(ctx, externalID) if err != nil { return "", "", nil, err } var attachmentsJSON []byte err = s.db.QueryRow(ctx, ` SELECT attachments FROM outbox WHERE id = $1 AND user_id = $2 AND status = 'draft' `, draftID, userID).Scan(&attachmentsJSON) if err != nil { if errors.Is(err, pgx.ErrNoRows) { return "", "", nil, ErrNotFound } return "", "", nil, err } var refs []draftAttachmentRef if err := json.Unmarshal(attachmentsJSON, &refs); err != nil { return "", "", nil, err } for _, ref := range refs { if ref.ID != attachmentID { continue } obj, err := s.storage.Get(ctx, ref.S3Key) if err != nil { return "", "", nil, err } return ref.Filename, ref.ContentType, obj, nil } return "", "", nil, ErrAttachmentNotFound } func (s *Service) ensureMessageOwned(ctx context.Context, externalID, messageID string) (userID string, err error) { err = s.db.QueryRow(ctx, ` SELECT u.id FROM messages m JOIN mail_accounts ma ON m.account_id = ma.id JOIN users u ON ma.user_id = u.id WHERE m.id = $1 AND u.external_id = $2 `, messageID, externalID).Scan(&userID) if err != nil { if errors.Is(err, pgx.ErrNoRows) { return "", ErrNotFound } return "", err } return userID, nil } func (s *Service) storageBucket() string { if s.attachmentsBucket != "" { return s.attachmentsBucket } return "mail-attachments" }