fix: Clean message body

Mail/Views/AI Writer/AIModel.swift

@@ -151,7 +151,7 @@ extension AIModel {
             replyingString = replyingBody.value
         } else if let value = replyingBody.value {
             let splitReply = await MessageBodyUtils.splitBodyAndQuote(messageBody: value)
-            replyingString = try await SwiftSoupUtils.extractText(from: splitReply.messageBody)
+            replyingString = try await SwiftSoupUtils(from: splitReply.messageBody).extractText()
         }
 
         guard let replyingString else { return }
@@ -219,32 +219,32 @@ extension AIModel {
 // MARK: - Insert result
 
 extension AIModel {
-    func didTapInsert() {
+    func didTapInsert() async {
         let shouldReplaceBody = shouldOverrideBody()
         guard !shouldReplaceBody || UserDefaults.shared.doNotShowAIReplaceMessageAgain else {
             isShowingReplaceBodyAlert = true
             return
         }
-        splitPropositionAndInsert(shouldReplaceBody: shouldReplaceBody)
+        await splitPropositionAndInsert(shouldReplaceBody: shouldReplaceBody)
     }
 
-    func splitPropositionAndInsert(shouldReplaceBody: Bool) {
+    func splitPropositionAndInsert(shouldReplaceBody: Bool) async {
         let (subject, body) = splitSubjectAndBody()
         if let subject, !subject.isEmpty && shouldOverrideSubject() {
             isShowingReplaceSubjectAlert = AIProposition(subject: subject, body: body, shouldReplaceContent: shouldReplaceBody)
         } else {
-            insertProposition(subject: subject, body: body, shouldReplaceBody: shouldReplaceBody)
+            await insertProposition(subject: subject, body: body, shouldReplaceBody: shouldReplaceBody)
         }
     }
 
-    func insertProposition(subject: String?, body: String, shouldReplaceBody: Bool) {
+    func insertProposition(subject: String?, body: String, shouldReplaceBody: Bool) async {
         matomo.track(
             eventWithCategory: .aiWriter,
             action: .data,
             name: shouldReplaceBody ? "replaceProposition" : "insertProposition"
         )
 
-        draftContentManager.replaceContent(subject: subject, body: body)
+        await draftContentManager.replaceContent(subject: subject, body: body)
         withAnimation {
             isShowingProposition = false
         }

Mail/Views/AI Writer/Proposition/AIPropositionView.swift

@@ -97,7 +97,9 @@ struct AIPropositionView: View {
                             AIProgressView()
                         case .standard, .error:
                             Button {
-                                aiModel.didTapInsert()
+                                Task {
+                                    await aiModel.didTapInsert()
+                                }
                             } label: {
                                 Label { Text(MailResourcesStrings.Localizable.aiButtonInsert) } icon: {
                                     IKIcon(MailResourcesAsset.plus)
@@ -123,16 +125,20 @@ struct AIPropositionView: View {
             }
             .customAlert(isPresented: $aiModel.isShowingReplaceBodyAlert) {
                 ReplaceMessageBodyView {
-                    aiModel.splitPropositionAndInsert(shouldReplaceBody: true)
+                    Task {
+                        await aiModel.splitPropositionAndInsert(shouldReplaceBody: true)
+                    }
                 }
             }
             .customAlert(item: $aiModel.isShowingReplaceSubjectAlert) { proposition in
                 ReplaceMessageSubjectView(subject: proposition.subject) { shouldReplaceSubject in
-                    aiModel.insertProposition(
-                        subject: shouldReplaceSubject ? proposition.subject : nil,
-                        body: proposition.body,
-                        shouldReplaceBody: proposition.shouldReplaceContent
-                    )
+                    Task {
+                        await aiModel.insertProposition(
+                            subject: shouldReplaceSubject ? proposition.subject : nil,
+                            body: proposition.body,
+                            shouldReplaceBody: proposition.shouldReplaceContent
+                        )
+                    }
                 }
             }
             .ikButtonPrimaryStyle(MailResourcesAsset.aiColor.swiftUIColor)

Mail/Views/Thread/WebView/WebViewModel+SwiftSoup.swift

@@ -33,10 +33,10 @@ extension WebViewModel {
     }
 
     func loadHTMLString(value: String?, blockRemoteContent: Bool) async -> LoadResult {
-        guard let rawHtml = value else { return .errorEmptyInputValue }
+        guard let rawHTML = value else { return .errorEmptyInputValue }
 
         do {
-            guard let safeDocument = MessageWebViewUtils.cleanHTMLContent(rawHTML: rawHtml)
+            guard let safeDocument = try? await SwiftSoupUtils(from: rawHTML).cleanCompleteDocument()
             else { return .errorCleanHTMLContent }
 
             try updateHeadContent(of: safeDocument)

MailCore/Cache/DraftContentManager.swift

@@ -137,18 +137,19 @@ public class DraftContentManager: ObservableObject {
         return try await loadReplyingMessage(messageReply.message, replyMode: messageReply.replyMode).body?.freezeIfNeeded()
     }
 
-    public func replaceContent(subject: String? = nil, body: String) {
-        guard let liveDraft = try? getLiveDraft() else { return }
-        guard let parsedMessage = try? SwiftSoup.parse(liveDraft.body) else { return }
+    public func replaceContent(subject: String? = nil, body: String) async {
+        guard let draft = try? getFrozenDraft() else { return }
+        guard let parsedMessage = try? await SwiftSoup.parse(draft.body) else { return }
 
         var extractedElements = ""
         for itemToExtract in Draft.appendedHTMLElements {
-            if let element = try? SwiftSoupUtils.extractHTML(from: parsedMessage, ".\(itemToExtract)") {
+            if let element = try? await SwiftSoupUtils(document: parsedMessage).extractHTML(".\(itemToExtract)") {
                 extractedElements.append(element)
             }
         }
 
         let realm = mailboxManager.getRealm()
+        guard let liveDraft = draft.thaw() else { return }
         try? realm.write {
             if let subject {
                 liveDraft.subject = subject
@@ -166,6 +167,10 @@ public class DraftContentManager: ObservableObject {
         return liveDraft
     }
 
+    private func getFrozenDraft() throws -> Draft {
+        return try getLiveDraft().freezeIfNeeded()
+    }
+
     private func writeCompleteDraft(
         completeBody: String,
         signature: Signature,
@@ -313,7 +318,7 @@ public class DraftContentManager: ObservableObject {
 
     private func loadReplyingMessageAndFormat(_ message: Message, replyMode: ReplyMode) async throws -> String {
         let replyingMessage = try await loadReplyingMessage(message, replyMode: replyMode)
-        return Draft.replyingBody(message: replyingMessage, replyMode: replyMode)
+        return await Draft.replyingBody(message: replyingMessage.freezeIfNeeded(), replyMode: replyMode)
     }
 
     private func loadReplyingAttachments(message: Message, replyMode: ReplyMode) async throws -> [Attachment] {

MailCore/Models/Draft.swift

@@ -225,7 +225,7 @@ public final class Draft: Object, Codable, Identifiable {
         return Draft(to: [recipient.detached()])
     }
 
-    public static func replyingBody(message: Message, replyMode: ReplyMode) -> String {
+    public static func replyingBody(message: Message, replyMode: ReplyMode) async -> String {
         let unsafeQuote: String
         switch replyMode {
         case .reply, .replyAll:
@@ -234,7 +234,7 @@ public final class Draft: Object, Codable, Identifiable {
             unsafeQuote = Constants.forwardQuote(message: message)
         }
 
-        let quote = (try? MessageWebViewUtils.cleanHTMLContent(rawHTML: unsafeQuote)?.outerHtml()) ?? ""
+        let quote = await (try? SwiftSoupUtils(from: unsafeQuote).cleanCompleteDocument().outerHtml()) ?? ""
 
         return "<br><br>" + quote
     }

MailCore/Utils/MessageWebViewUtils.swift

@@ -47,23 +47,4 @@ public enum MessageWebViewUtils {
 
         return resources
     }
-
-    public static func cleanHTMLContent(rawHTML: String) -> Document? {
-        do {
-            let dirtyDocument = try SwiftSoup.parse(rawHTML)
-            let cleanedDocument = try SwiftSoup.Cleaner(headWhitelist: .headWhitelist, bodyWhitelist: .extendedBodyWhitelist)
-                .clean(dirtyDocument)
-
-            // We need to remove the tag <meta http-equiv="refresh" content="x">
-            let metaRefreshTags = try cleanedDocument.select("meta[http-equiv='refresh']")
-            for metaRefreshTag in metaRefreshTags {
-                try metaRefreshTag.parent()?.removeChild(metaRefreshTag)
-            }
-
-            return cleanedDocument
-        } catch {
-            DDLogError("An error occurred while parsing body \(error)")
-            return nil
-        }
-    }
 }

MailCore/Utils/NotificationsHelper.swift

@@ -201,10 +201,10 @@ public enum NotificationsHelper {
         }
 
         do {
-            let basicHtml = try SwiftSoup.clean(body, Whitelist.basic())!
-            let parsedBody = try await SwiftSoup.parse(basicHtml)
+            let cleanedDocument = try await SwiftSoupUtils(from: body).cleanBody()
+            guard let extractedBody = cleanedDocument.body() else { return message.preview }
 
-            let rawText = try parsedBody.text(trimAndNormaliseWhitespace: false)
+            let rawText = try extractedBody.text(trimAndNormaliseWhitespace: false)
             return rawText.trimmingCharacters(in: .whitespacesAndNewlines)
         } catch {
             return message.preview

MailCore/Utils/SwiftSoupUtils.swift

@@ -18,17 +18,43 @@
 
 import SwiftSoup
 
-public enum SwiftSoupUtils {
-    public static func extractHTML(from document: Document, _ cssQuery: String) throws -> String {
-        guard let foundElement = try document.select(cssQuery).first() else {
-            throw SwiftSoupError.elementNotFound
+public struct SwiftSoupUtils {
+    private let document: Document
+
+    public init(document: Document) {
+        self.document = document
+    }
+
+    public init(from html: String) throws {
+        document = try SwiftSoup.parse(html)
+    }
+
+    public func cleanBody() async throws -> Document {
+        let cleanedDocument = try SwiftSoup.Cleaner(headWhitelist: nil, bodyWhitelist: .extendedBodyWhitelist).clean(document)
+        return cleanedDocument
+    }
+
+    public func cleanCompleteDocument() async throws -> Document {
+        let cleanedDocument = try SwiftSoup.Cleaner(headWhitelist: .headWhitelist, bodyWhitelist: .extendedBodyWhitelist)
+            .clean(document)
+
+        // We need to remove the tag <meta http-equiv="refresh" content="x">
+        let metaRefreshTags = try await cleanedDocument.select("meta[http-equiv='refresh']")
+        for metaRefreshTag in metaRefreshTags {
+            try metaRefreshTag.parent()?.removeChild(metaRefreshTag)
         }
+
+        return cleanedDocument
+    }
+
+    public func extractHTML(_ cssQuery: String) async throws -> String {
+        guard let foundElement = try await document.select(cssQuery).first() else { throw SwiftSoupError.elementNotFound }
+
         let htmlContent = try foundElement.outerHtml()
         return htmlContent
     }
 
-    public static func extractText(from html: String) async throws -> String? {
-        let document = try await SwiftSoup.parse(html)
+    public func extractText() async throws -> String? {
         return try document.body()?.text()
     }
 }