package migration import ( "bytes" "context" "encoding/json" "fmt" "mime" "strings" "github.com/jackc/pgx/v5/pgxpool" "github.com/ultisuite/ulti-backend/internal/mail/limits" "github.com/ultisuite/ulti-backend/internal/mail/storage" ) type gmailAttachmentPart struct { Filename string MimeType string ContentID string IsInline bool Size int BodyData string AttachmentID string } func extractGmailAttachmentParts(payload gmailPayload) []gmailAttachmentPart { var out []gmailAttachmentPart var walk func(gmailPayload) walk = func(node gmailPayload) { headers := gmailHeaderMap(node.Headers) mimeType := strings.TrimSpace(node.MimeType) if strings.HasPrefix(strings.ToLower(mimeType), "multipart/") { for _, part := range node.Parts { walk(part) } return } if isGmailAttachmentPart(mimeType, headers, node.Body) { out = append(out, buildGmailAttachmentPart(mimeType, headers, node.Body)) } } walk(payload) if len(out) > limits.MaxAttachmentsPerMessage { return out[:limits.MaxAttachmentsPerMessage] } return out } func gmailHeaderMap(headers []gmailHeader) map[string]string { out := make(map[string]string, len(headers)) for _, h := range headers { key := strings.ToLower(strings.TrimSpace(h.Name)) if key != "" && out[key] == "" { out[key] = h.Value } } return out } func isGmailAttachmentPart(mimeType string, headers map[string]string, body gmailBody) bool { if body.AttachmentID != "" { return true } disposition := strings.ToLower(headers["content-disposition"]) if strings.HasPrefix(disposition, "attachment") { return true } contentID := normalizeGmailContentID(headers["content-id"]) filename := gmailPartFilename(headers) if strings.HasPrefix(disposition, "inline") && (filename != "" || contentID != "") { return true } if contentID != "" && body.Data != "" { return true } lower := strings.ToLower(strings.TrimSpace(mimeType)) if lower == "text/plain" || lower == "text/html" || lower == "" { return false } return filename != "" && body.Data != "" } func buildGmailAttachmentPart(mimeType string, headers map[string]string, body gmailBody) gmailAttachmentPart { disposition := strings.ToLower(headers["content-disposition"]) contentID := normalizeGmailContentID(headers["content-id"]) filename := gmailPartFilename(headers) isInline := strings.HasPrefix(disposition, "inline") || (contentID != "" && !strings.HasPrefix(disposition, "attachment")) if filename == "" { filename = inlineGmailAttachmentFilename(contentID, mimeType) } size := body.Size if size <= 0 && body.Data != "" { size = len(body.Data) } return gmailAttachmentPart{ Filename: filename, MimeType: normalizeAttachmentMimeType(mimeType), ContentID: contentID, IsInline: isInline, Size: size, BodyData: body.Data, AttachmentID: body.AttachmentID, } } func gmailPartFilename(headers map[string]string) string { if name := strings.TrimSpace(headers["filename"]); name != "" { return name } disposition := headers["content-disposition"] if disposition == "" { return "" } _, params, err := mime.ParseMediaType(disposition) if err != nil { return "" } return strings.TrimSpace(params["filename"]) } func normalizeGmailContentID(raw string) string { return strings.Trim(raw, "<> \t") } func inlineGmailAttachmentFilename(contentID, mimeType string) string { base := "inline" if contentID != "" { base = strings.Map(func(r rune) rune { switch r { case '<', '>', '/', '\\', ':', '"', '\'', '?', '*': return '_' default: return r } }, contentID) } ext := extensionFromMimeType(mimeType) if ext != "" && !strings.HasSuffix(strings.ToLower(base), ext) { return base + ext } return base } func extensionFromMimeType(mimeType string) string { switch strings.ToLower(strings.TrimSpace(strings.Split(mimeType, ";")[0])) { case "image/jpeg", "image/jpg": return ".jpg" case "image/png": return ".png" case "image/gif": return ".gif" case "image/webp": return ".webp" case "application/pdf": return ".pdf" default: return "" } } func normalizeAttachmentMimeType(mimeType string) string { mimeType = strings.TrimSpace(mimeType) if mimeType == "" { return "application/octet-stream" } mediaType, _, err := mime.ParseMediaType(mimeType) if err != nil || mediaType == "" { return mimeType } return mediaType } func (g *GmailImporter) storeGmailAttachments( ctx context.Context, userID, messageID, gmailID, accessToken string, payload gmailPayload, messageExisted bool, ) error { if g.storage == nil { return nil } parts := extractGmailAttachmentParts(payload) if len(parts) == 0 { return nil } var existingCount int var existingTotal int64 if err := g.db.QueryRow(ctx, ` SELECT COUNT(*)::int, COALESCE(SUM(size), 0)::bigint FROM attachments WHERE message_id = $1::uuid `, messageID).Scan(&existingCount, &existingTotal); err != nil { return err } stored := 0 for _, part := range parts { if messageExisted && gmailAttachmentExists(ctx, g.db, messageID, part) { continue } data, err := g.loadGmailAttachmentData(ctx, accessToken, gmailID, part) if err != nil { return err } if len(data) == 0 { continue } if err := limits.ValidateAttachmentSize(int64(len(data))); err != nil { continue } if err := limits.ValidateAttachmentQuota(existingCount, existingTotal, int64(len(data))); err != nil { break } objectKey := storage.MessageObjectKey(userID, messageID, part.Filename) if err := g.storage.Put(ctx, objectKey, bytes.NewReader(data), int64(len(data)), part.MimeType); err != nil { return err } bucket := g.attachBucket if bucket == "" { bucket = "mail-attachments" } _, err = g.db.Exec(ctx, ` INSERT INTO attachments (message_id, filename, content_type, size, s3_bucket, s3_key, content_id, is_inline, virus_scan_status) VALUES ($1::uuid, $2, $3, $4, $5, $6, $7, $8, 'skipped') `, messageID, part.Filename, part.MimeType, len(data), bucket, objectKey, part.ContentID, part.IsInline) if err != nil { _ = g.storage.Delete(ctx, objectKey) return err } existingCount++ existingTotal += int64(len(data)) stored++ } if stored > 0 { _, err := g.db.Exec(ctx, `UPDATE messages SET has_attachments = true, updated_at = NOW() WHERE id = $1::uuid`, messageID) return err } return nil } func (g *GmailImporter) loadGmailAttachmentData(ctx context.Context, accessToken, gmailID string, part gmailAttachmentPart) ([]byte, error) { if part.AttachmentID != "" { url := fmt.Sprintf( "https://gmail.googleapis.com/gmail/v1/users/me/messages/%s/attachments/%s", gmailID, part.AttachmentID, ) raw, err := g.apiGet(ctx, url, accessToken) if err != nil { return nil, err } var parsed struct { Data string `json:"data"` } if err := json.Unmarshal(raw, &parsed); err != nil { return nil, err } if parsed.Data == "" { return nil, nil } return []byte(decodeGmailBody(parsed.Data)), nil } if part.BodyData == "" { return nil, nil } return []byte(decodeGmailBody(part.BodyData)), nil } func gmailAttachmentExists(ctx context.Context, db *pgxpool.Pool, messageID string, part gmailAttachmentPart) bool { var count int if part.ContentID != "" { _ = db.QueryRow(ctx, ` SELECT COUNT(*) FROM attachments WHERE message_id = $1::uuid AND (content_id = $2 OR filename = $3) `, messageID, part.ContentID, part.Filename).Scan(&count) return count > 0 } _ = db.QueryRow(ctx, ` SELECT COUNT(*) FROM attachments WHERE message_id = $1::uuid AND filename = $2 `, messageID, part.Filename).Scan(&count) return count > 0 }