Fix major memory safety issue where MyHTML support class was destroye…

…d too quickly
PopFlamingo · Jan 20, 2019 · 607f2c1 · 607f2c1
1 parent dd9a25c
commit 607f2c1
Show file tree

Hide file tree

Showing 3 changed files with 35 additions and 4 deletions.
diff --git a/Sources/MyHTML/MyHTML.swift b/Sources/MyHTML/MyHTML.swift
@@ -21,8 +21,12 @@ public class MyHTML {
 
     }
 
+    init(raw: OpaquePointer) {
+        self.raw = raw
+    }
+
     deinit {
-        myhtml_destroy(raw)
+        assert(myhtml_destroy(raw) == nil, "Unsuccessful destroy")
     }
 
 

diff --git a/Sources/MyHTML/Tree.swift b/Sources/MyHTML/Tree.swift
@@ -1,12 +1,13 @@
 import CMyHTML
 
 public class Tree {
-    var raw: OpaquePointer
+
     public init(context: MyHTML, html: String) throws {
         guard let raw = myhtml_tree_create() else {
             throw Error.cannotCreateBaseStructure
         }
         self.raw = raw
+        self.context = context
         let status = myhtml_tree_init(raw, context.raw)
         guard status == MyHTML_STATUS_OK.rawValue else {
             throw Error.statusError(rawValue: status)
@@ -22,9 +23,18 @@ public class Tree {
     }
 
     deinit {
-        myhtml_tree_destroy(raw)
+        assert(myhtml_tree_destroy(raw) == nil, "Unsuccessful destroy")
     }
 
+
+    var raw: OpaquePointer
+
+    /// The context is a private variable that keeps MyHTML ref count > 0
+    /// as long as it is needed, this way, no early call to its C free
+    /// functions is done. Test `testRightTimeContextDeinit` attempts
+    /// to catch cases where this wouldn't be done correctly;
+    private var context: MyHTML
+
     public var htmlNode: Node? {
         guard let rawNode = myhtml_tree_get_node_html(raw) else { return nil }
         return Node(raw: rawNode)

diff --git a/Tests/MyHTMLTests/MyHTMLTests.swift b/Tests/MyHTMLTests/MyHTMLTests.swift
@@ -99,6 +99,22 @@ final class MyHTMLTests: XCTestCase {
         XCTAssertEqual(tree.getChildNodes(named: "img")[0].attributes[1].value,
                        nil)
     }
+
+
+    func testRightTimeContextDeinit() throws {
+        func create() throws -> (Tree, OpaquePointer) {
+            let fooHtml = try MyHTML(options: .single, threadCount: 1)
+            return try (Tree(context: fooHtml, html: sampleCodeA), fooHtml.raw)
+        }
+
+        let (tree, htmlPointer) = try create()
+        let htmlPointerThroughC: OpaquePointer? = myhtml_tree_get_myhtml(tree.raw)
+        XCTAssertNotNil(htmlPointerThroughC)
+        XCTAssertEqual(htmlPointer, htmlPointerThroughC)
+        let rawTree = myhtml_tree_create()
+        myhtml_tree_init(rawTree, htmlPointerThroughC) // Should fail if early free
+        myhtml_tree_destroy(rawTree)
+    }
 
     static var allTests = [
         ("testAttributeContainsValue", testAttributeContainsValue),
@@ -110,7 +126,8 @@ final class MyHTMLTests: XCTestCase {
         ("testGetNodeTextContent", testGetNodeTextContent),
         ("testNodeIdentity", testNodeIdentity),
         ("testAttributesCount", testAttributesCount),
-        ("testAttributesValues", testAttributesValues)
+        ("testAttributesValues", testAttributesValues),
+        ("testRightTimeContextDeinit", testRightTimeContextDeinit)
     ]
 
     let sampleCodeA = """