Merge pull request #90 from hwdavr/enum-support

Revert "Code optimisation by removing suffix of USR"
rockbruno · Jan 25, 2020 · 672f7f5 · 672f7f5
2 parents c161764 + 1a27574
commit 672f7f5
Showing 1 changed file with 12 additions and 23 deletions.
diff --git a/swiftshield-Sources/AutomaticSwiftShield.swift b/swiftshield-Sources/AutomaticSwiftShield.swift
@@ -81,13 +81,6 @@ class AutomaticSwiftShield: Protector {
                 }
                 let name = data.name
                 let usr = data.usr
-                let kind = data.kind
-
-                if let type = self.sourceKit.referenceType(kind: kind), type == .enum && name.hasSuffix("CodingKeys") {
-                    obfuscationData.excludedEnums.insert(self.usrWithoutSuffix(usr))
-                    return
-                }
-
                 obfuscationData.usrDict.insert(usr)
                 if dict.getString(.receiverId) == nil {
                     obfuscationData.usrRelationDict[usr] = variant
@@ -137,8 +130,7 @@ extension AutomaticSwiftShield {
     private func getNameData(from dict: SourceKitdResponse.Dictionary,
                              obfuscationData: ObfuscationData) -> (name: String,
                                                                    usr: String,
-                                                                   obfuscatedName: String,
-                                                                   kind: String)? {
+                                                                   obfuscatedName: String)? {
         let kind = dict.getUID(.kindId).asString
         guard sourceKit.declarationType(for: kind) != nil else {
             return nil
@@ -150,18 +142,9 @@ extension AutomaticSwiftShield {
         guard let protected = obfuscationData.obfuscationDict[name] else {
             let newName = String.random(length: self.protectedClassNameSize, excluding: obfuscationData.allObfuscatedNames)
             obfuscationData.obfuscationDict[name] = newName
-            return (name, usr, newName, kind)
+            return (name, usr, newName)
         }
-        return (name, usr, protected, kind)
-    }
-
-    private func usrWithoutSuffix(_ usr: String) -> String {
-        let lastIndex = usr.lastIndex(of: "_")
-        if let index = lastIndex {
-            let usrPrefix = usr[..<index]
-            return String(usrPrefix)
-        }
-        return usr
+        return (name, usr, protected)
     }
 
     func findReferencesInIndexed(obfuscationData: AutomaticObfuscationData) {
@@ -182,11 +165,17 @@ extension AutomaticSwiftShield {
                     return
                 }
 
+                if type == .enum && name.hasSuffix("CodingKeys") {
+                    obfuscationData.excludedEnums.insert(usr)
+                    return
+                }
 
                 if type == .enumelement {
-                    let usrPrefix = self.usrWithoutSuffix(usr)
-                    if obfuscationData.excludedEnums.contains(usrPrefix) {
-                        return
+                    for exclusion in obfuscationData.excludedEnums {
+                        // Enum element belongs to excluded enum
+                        if usr.hasPrefix(exclusion) {
+                            return 
+                        }
                     }
                 }