From ab407902726138fda5bdf9a59c8296a9c9a79a3b Mon Sep 17 00:00:00 2001
From: Bruno Rocha <rochab@spotify.com>
Date: Mon, 18 May 2020 17:24:37 +0200
Subject: [PATCH] Fix usage of emojis

---
 SOURCEKITISSUES.md                            |  1 -
 .../Obfuscator/SourceKitObfuscator.swift      |  9 +--
 Sources/SwiftShieldCore/StringHelpers.swift   |  7 +++
 Sources/swiftshield/main.swift                |  2 +-
 Tests/SwiftShieldTests/FeatureTests.swift     | 60 +++++++++++++++++++
 5 files changed, 73 insertions(+), 6 deletions(-)

diff --git a/SOURCEKITISSUES.md b/SOURCEKITISSUES.md
index 8e81f66..3d85712 100644
--- a/SOURCEKITISSUES.md
+++ b/SOURCEKITISSUES.md
@@ -20,7 +20,6 @@ These are problems that SourceKit has that are unrelated to a specific feature,
 **Note: You can use the `--ignore-targets` argument to completely disable the obfuscation of specific targets.**
 
 - (SR-9020)](https://bugs.swift.org/browse/SR-9020) Legacy KeyPaths that include types (like `#keyPath(Foo.bar)`) will not get indexed.
-- Any file that has an emoji will break the obfuscation process. This may not be a SourceKit bug itself, but something that we have to treat on our side.
 - `@objc optional` protocol methods don't have their references indexed.
 
 # Additional important information
diff --git a/Sources/SwiftShieldCore/Obfuscator/SourceKitObfuscator.swift b/Sources/SwiftShieldCore/Obfuscator/SourceKitObfuscator.swift
index fb6fe64..4355110 100644
--- a/Sources/SwiftShieldCore/Obfuscator/SourceKitObfuscator.swift
+++ b/Sources/SwiftShieldCore/Obfuscator/SourceKitObfuscator.swift
@@ -199,27 +199,28 @@ extension SourceKitObfuscator {
                 // Avoid duplicates.
                 currentReferenceIndex += 1
             }
+            let currentCharacter = charArray[currentCharIndex]
             if line == reference.line, column == reference.column {
                 previousReference = reference
                 let originalName = reference.name
                 let obfuscatedName = obfuscate(name: originalName)
-                let wasInternalKeyword = charArray[currentCharIndex] == "`"
+                let wasInternalKeyword = currentCharacter == "`"
                 for i in 1 ..< (originalName.count + (wasInternalKeyword ? 2 : 0)) {
                     charArray[currentCharIndex + i] = ""
                 }
                 charArray[currentCharIndex] = obfuscatedName
                 currentReferenceIndex += 1
                 currentCharIndex += originalName.count
-                column += originalName.count
+                column += originalName.utf8Count
                 if wasInternalKeyword {
                     charArray[currentCharIndex] = ""
                 }
-            } else if charArray[currentCharIndex] == "\n" {
+            } else if currentCharacter == "\n" {
                 line += 1
                 column = 1
                 currentCharIndex += 1
             } else {
-                column += 1
+                column += currentCharacter.utf8Count
                 currentCharIndex += 1
             }
         }
diff --git a/Sources/SwiftShieldCore/StringHelpers.swift b/Sources/SwiftShieldCore/StringHelpers.swift
index 50c9a79..9494299 100644
--- a/Sources/SwiftShieldCore/StringHelpers.swift
+++ b/Sources/SwiftShieldCore/StringHelpers.swift
@@ -68,3 +68,10 @@ extension NSTextCheckingResult {
         return groupStartIndex ..< groupEndIndex
     }
 }
+
+extension String {
+    /// Considers emoji scalars when counting.
+    var utf8Count: Int {
+        return utf8.count
+    }
+}
diff --git a/Sources/swiftshield/main.swift b/Sources/swiftshield/main.swift
index 065bb11..dea3c06 100644
--- a/Sources/swiftshield/main.swift
+++ b/Sources/swiftshield/main.swift
@@ -4,7 +4,7 @@ import SwiftShieldCore
 
 struct Swiftshield: ParsableCommand {
     static var configuration = CommandConfiguration(
-        abstract: "SwiftShield 4.0.0",
+        abstract: "SwiftShield 4.0.1",
         subcommands: [Obfuscate.self, Deobfuscate.self]
     )
 }
diff --git a/Tests/SwiftShieldTests/FeatureTests.swift b/Tests/SwiftShieldTests/FeatureTests.swift
index 3f2acb4..14194c9 100644
--- a/Tests/SwiftShieldTests/FeatureTests.swift
+++ b/Tests/SwiftShieldTests/FeatureTests.swift
@@ -188,4 +188,64 @@ final class FeatureTests: XCTestCase {
         //}
         """)
     }
+
+    func test_files_withEmojis() throws {
+        let (obfs, store, delegate) = baseTestData()
+        let module = try testModule(withContents: """
+        enum JSON {
+            static func parse(_ a: String) -> String { return a }
+        }
+
+        extension String {
+            func unobfuscate() -> String { return self }
+        }
+
+        func l3️⃣og(_ a: String) -> String { return JSON.parse("") }
+
+        var paramsString = "foo"
+
+        struct Logger {
+            func log() {
+                log("Hello 👨‍👩‍👧‍👧 3 A̛͚̖ 3️⃣ response up message 📲: \\(JSON.parse(paramsString.unobfuscate()).description) 🇹🇩👫👨‍👩‍👧‍👧👨‍👨‍👦"); log("")
+            }
+
+            func log(_ a: String) {
+                _ = l3️⃣og("foo".unobfuscate())
+            }
+        }
+        """)
+        store.obfuscationDictionary["JSON"] = "OBS1"
+        store.obfuscationDictionary["parse"] = "OBS2"
+        store.obfuscationDictionary["unobfuscate"] = "OBS3"
+        store.obfuscationDictionary["log"] = "OBS4"
+        store.obfuscationDictionary["Logger"] = "OBS5"
+        store.obfuscationDictionary["l3️⃣og"] = "OBS6"
+
+        try obfs.registerModuleForObfuscation(module)
+        try obfs.obfuscate()
+
+        XCTAssertEqual(delegate.receivedContent[modifiableFilePath], """
+        enum OBS1 {
+            static func OBS2(_ a: String) -> String { return a }
+        }
+
+        extension String {
+            func OBS3() -> String { return self }
+        }
+
+        func OBS6(_ a: String) -> String { return OBS1.OBS2("") }
+
+        var paramsString = "foo"
+
+        struct OBS5 {
+            func OBS4() {
+                OBS4("Hello 👨‍👩‍👧‍👧 3 A̛͚̖ 3️⃣ response up message 📲: \\(OBS1.OBS2(paramsString.OBS3()).description) 🇹🇩👫👨‍👩‍👧‍👧👨‍👨‍👦"); OBS4("")
+            }
+
+            func OBS4(_ a: String) {
+                _ = OBS6("foo".OBS3())
+            }
+        }
+        """)
+    }
 }