Additional fixes

rockbruno · May 23, 2020 · 23bc23e · 23bc23e
1 parent 8e36f56
commit 23bc23e
Show file tree

Hide file tree

Showing 4 changed files with 20 additions and 33 deletions.
diff --git a/Sources/SwiftShieldCore/Obfuscator/SourceKitObfuscator.swift b/Sources/SwiftShieldCore/Obfuscator/SourceKitObfuscator.swift
@@ -95,22 +95,13 @@ extension SourceKitObfuscator {
             }
         }
 
-        if kind == .property, dict.parent != nil, let parentKind: SKUID = dict.parent[keys.kind] {
-            let parentUSR: String
-            if parentKind.declarationType(onlyObfuscable: false) == .extension {
-                // The type's USR of an extension is hidden inside the entities.
-                guard let entities: SKResponseArray = dict.parent[keys.entities] else {
-                    throw logger.fatalError(forMessage: "Parent of \(usr) is an extension with no entities!")
-                }
-                guard let pUSR: String = entities[0][keys.usr] else {
-                    throw logger.fatalError(forMessage: "Parent of \(usr) is an extension with no USR!")
-                }
-                parentUSR = pUSR
-            } else {
-                guard let pUSR: String = dict.parent[keys.usr] else {
-                    throw logger.fatalError(forMessage: "Parent of \(usr) is has no USR!")
-                }
-                parentUSR = pUSR
+        if kind == .property,
+           dict.parent != nil,
+           let parentKind: SKUID = dict.parent[keys.kind],
+           parentKind.declarationType() == .object
+        {
+            guard let parentUSR: String = dict.parent[keys.usr] else {
+                throw logger.fatalError(forMessage: "Parent of \(usr) is has no USR!")
             }
             let codableUSRs: Set<String> = ["s:s7Codablea", "s:SE", "s:Se"]
             if try inheritsFromAnyUSR(
@@ -154,7 +145,7 @@ extension SourceKitObfuscator {
         var referenceArray = [Reference]()
         index.response.recurseEntities { [unowned self] dict in
             guard let kindId: SKUID = dict[self.keys.kind],
-                kindId.referenceType() != nil,
+                kindId.referenceType() != nil || kindId.declarationType() != nil,
                 let rawName: String = dict[self.keys.name],
                 let usr: String = dict[self.keys.usr],
                 self.dataStore.processedUsrs.contains(usr),
@@ -341,7 +332,7 @@ extension SKResponseDictionary {
         guard let kindId: SKUID = self[sourcekitd.keys.kind] else {
             return false
         }
-        let type = kindId.referenceType()
+        let type = kindId.referenceType() ?? kindId.declarationType()
         guard type == .method || type == .property else {
             return false
         }

diff --git a/Sources/SwiftShieldCore/SourceKit/SourceKit.swift b/Sources/SwiftShieldCore/SourceKit/SourceKit.swift
@@ -773,14 +773,14 @@ public struct SKUID: CustomStringConvertible {
     }
 
     func referenceType() -> SourceKit.DeclarationType? {
-        declarationType() ?? declarationType(firstSuffix: ".ref.")
+        declarationType(firstSuffix: ".ref.")
     }
 
-    func declarationType(onlyObfuscable: Bool = true) -> SourceKit.DeclarationType? {
-        declarationType(firstSuffix: ".decl.", onlyObfuscable: onlyObfuscable)
+    func declarationType() -> SourceKit.DeclarationType? {
+        declarationType(firstSuffix: ".decl.")
     }
 
-    func declarationType(firstSuffix: String, onlyObfuscable: Bool = true) -> SourceKit.DeclarationType? {
+    func declarationType(firstSuffix: String) -> SourceKit.DeclarationType? {
         let kind = description
         let prefix = "source.lang.swift" + firstSuffix
         guard kind.hasPrefix(prefix) else {
@@ -809,14 +809,6 @@ public struct SKUID: CustomStringConvertible {
         case "enumelement":
             return .enumelement
         default:
-            break
-        }
-        guard onlyObfuscable == false else {
-            return nil
-        }
-        if kindSuffix.hasPrefix("extension.") {
-            return .extension
-        } else {
             return nil
         }
     }

diff --git a/Sources/SwiftShieldCore/SourceKit/SourceKitDeclarationType.swift b/Sources/SwiftShieldCore/SourceKit/SourceKitDeclarationType.swift
@@ -6,6 +6,5 @@ extension SourceKit {
         case method
         case `enum`
         case enumelement
-        case `extension`
     }
 }
diff --git a/Tests/SwiftShieldTests/FeatureTests.swift b/Tests/SwiftShieldTests/FeatureTests.swift
@@ -350,7 +350,9 @@ final class FeatureTests: XCTestCase {
             let prop1: String
         }
 
-        class BarClass: Codable {
+        protocol SomeProt {}
+
+        class BarClass: Codable, SomeProt {
             let prop1: String
         }
         """)
@@ -365,6 +367,7 @@ final class FeatureTests: XCTestCase {
         store.obfuscationDictionary["BarExternal"] = "OBSEX"
         store.obfuscationDictionary["BarClass"] = "OBSOBJC"
         store.obfuscationDictionary["CodableProtocolInAnotherFile"] = "EXCOD"
+        store.obfuscationDictionary["SomeProt"] = "OBSSOMEPROT"
         store.obfuscationDictionary["prop1"] = "OBS1"
         store.obfuscationDictionary["prop2"] = "OBS2"
         store.obfuscationDictionary["prop3"] = "OBS3"
@@ -406,7 +409,9 @@ final class FeatureTests: XCTestCase {
             let prop1: String
         }
 
-        class OBSOBJC: Codable {
+        protocol OBSSOMEPROT {}
+
+        class OBSOBJC: Codable, OBSSOMEPROT {
             let prop1: String
         }
         """)