From 2a355eea9487c7dfcd09cffd757facbd8d0a5138 Mon Sep 17 00:00:00 2001
From: daniele margutti <hello@danielemargutti.com>
Date: Sun, 26 Jan 2020 19:57:27 +0100
Subject: [PATCH] #104 Added escapeString options in XMLParsingOptions of
 StyleXML. Applied by default.

---
 .../SwiftRichString/Style/StyleGroup.swift    |  3 +-
 .../SwiftRichString/Support/Extensions.swift  | 48 +++++++++++++++++++
 .../Support/XMLStringBuilder.swift            |  9 +++-
 3 files changed, 58 insertions(+), 2 deletions(-)

diff --git a/Sources/SwiftRichString/Style/StyleGroup.swift b/Sources/SwiftRichString/Style/StyleGroup.swift
index 326527a..1547a2f 100644
--- a/Sources/SwiftRichString/Style/StyleGroup.swift
+++ b/Sources/SwiftRichString/Style/StyleGroup.swift
@@ -52,7 +52,8 @@ public class StyleXML: StyleProtocol {
     public var baseStyle: StyleProtocol?
     
     /// XML Parsing options.
-    public var xmlParsingOptions: XMLParsingOptions = []
+    /// By default `escapeString` is applied.
+    public var xmlParsingOptions: XMLParsingOptions = [.escapeString]
     
     /// Image provider is called to provide custom image when `StyleXML` encounter a `img` tag image.
     /// If not implemented the image is automatically searched inside any bundled `xcassets`.
diff --git a/Sources/SwiftRichString/Support/Extensions.swift b/Sources/SwiftRichString/Support/Extensions.swift
index 4f8fdb5..8a9be4b 100644
--- a/Sources/SwiftRichString/Support/Extensions.swift
+++ b/Sources/SwiftRichString/Support/Extensions.swift
@@ -35,6 +35,54 @@ import AppKit
 import UIKit
 #endif
 
+extension String {
+    
+    // Current implementation by @alexaubry in
+    // https://github.com/alexaubry/HTMLString
+    public func escapeWithUnicodeEntities() -> String {
+        var copy = self
+        copy.addUnicodeEntities()
+        return copy
+    }
+    
+    internal mutating func addUnicodeEntities() {
+        var position: String.Index? = startIndex
+        let requiredEscapes: Set<Character> = ["!", "\"", "$", "%", "&", "'", "+", ",", "<", "=", ">", "@", "[", "]", "`", "{", "}"]
+        
+        while let cursorPosition = position {
+            guard cursorPosition != endIndex else { break }
+            let character = self[cursorPosition]
+            
+            if requiredEscapes.contains(character) {
+                // One of the required escapes for security reasons
+                let escape = "&#\(character.asciiValue!);" // required escapes can can only be ASCII
+                position = positionAfterReplacingCharacter(at: cursorPosition, with: escape)
+            } else {
+                // Not a required escape, no need to replace the character
+                position = index(cursorPosition, offsetBy: 1, limitedBy: endIndex)
+            }
+        }
+    }
+    
+    /// Replaces the character at the given position with the escape and returns the new position.
+    fileprivate mutating func positionAfterReplacingCharacter(at position: String.Index, with escape: String) -> String.Index? {
+        let nextIndex = index(position, offsetBy: 1)
+        
+        if let fittingPosition = index(position, offsetBy: escape.count, limitedBy: endIndex) {
+            // Check if we can fit the whole escape in the receiver
+            replaceSubrange(position ..< nextIndex, with: escape)
+            return fittingPosition
+        } else {
+            // If we can't, remove the character and insert the escape to make it fit.
+            remove(at: position)
+            insert(contentsOf: escape, at: position)
+            return index(position, offsetBy: escape.count, limitedBy: endIndex)
+        }
+    }
+
+    
+}
+
 extension NSNumber {
 	
 	internal static func from(float: Float?) -> NSNumber? {
diff --git a/Sources/SwiftRichString/Support/XMLStringBuilder.swift b/Sources/SwiftRichString/Support/XMLStringBuilder.swift
index 786b119..4ace995 100644
--- a/Sources/SwiftRichString/Support/XMLStringBuilder.swift
+++ b/Sources/SwiftRichString/Support/XMLStringBuilder.swift
@@ -52,6 +52,12 @@ public struct XMLParsingOptions: OptionSet {
     /// recommended that you include a root node yourself and pass this option.
     public static let doNotWrapXML = XMLParsingOptions(rawValue: 1)
     
+    /// Perform string escaping to replace all characters which is not supported by NSXMLParser
+    /// into the specified encoding with decimal entity.
+    /// For example if your string contains '&' character parser will break the style.
+    /// This option is active by default.
+    public static let escapeString = XMLParsingOptions(rawValue: 2)
+    
 }
 
 // MARK: - XMLStringBuilder
@@ -104,7 +110,8 @@ public class XMLStringBuilder: NSObject, XMLParserDelegate {
     public init(styleXML: StyleXML, string: String) {
         self.styleXML = styleXML
 
-        let xml = (styleXML.xmlParsingOptions.contains(.doNotWrapXML) ? string : "<\(XMLStringBuilder.topTag)>\(string)</\(XMLStringBuilder.topTag)>")
+        let xmlString = (styleXML.xmlParsingOptions.contains(.escapeString) ? string.escapeWithUnicodeEntities() : string)
+        let xml = (styleXML.xmlParsingOptions.contains(.doNotWrapXML) ? xmlString : "<\(XMLStringBuilder.topTag)>\(xmlString)</\(XMLStringBuilder.topTag)>")
         guard let data = xml.data(using: String.Encoding.utf8) else {
             fatalError("Unable to convert to UTF8")
         }