Remove handle variable from Document class

ballerina-platform · Oct 13, 2023 · 0739c63 · 0739c63
1 parent d852fd8
commit 0739c63
Show file tree

Hide file tree

Showing 3 changed files with 11 additions and 19 deletions.
diff --git a/ballerina/modules/wssec/document.bal b/ballerina/modules/wssec/document.bal
@@ -16,15 +16,11 @@
 import ballerina/jballerina.java;
 
 isolated class Document {
-
-    private handle nativeDocumentBuilder;
 
     public isolated function init(xml xmlPayload) returns Error? {
-        handle|error documentBuilder = newDocument(xmlPayload);
+        handle|error documentBuilder = newDocument(self, xmlPayload);
         if documentBuilder is error {
             return error Error(documentBuilder.message());
-        } else {
-            self.nativeDocumentBuilder = documentBuilder;
         }
     }
 
@@ -37,6 +33,6 @@ isolated class Document {
     } external;
 }
 
-isolated function newDocument(xml xmlPayload) returns handle|error = @java:Constructor {
+isolated function newDocument(Document doc, xml xmlPayload) returns handle|error = @java:Constructor {
     'class: "org.wssec.DocumentBuilder"
 } external;
diff --git a/native/src/main/java/org/wssec/DocumentBuilder.java b/native/src/main/java/org/wssec/DocumentBuilder.java
@@ -17,9 +17,7 @@
 package org.wssec;
 
 import io.ballerina.runtime.api.creators.ValueCreator;
-import io.ballerina.runtime.api.utils.StringUtils;
 import io.ballerina.runtime.api.values.BArray;
-import io.ballerina.runtime.api.values.BHandle;
 import io.ballerina.runtime.api.values.BObject;
 import io.ballerina.runtime.api.values.BXml;
 import org.w3c.dom.Document;
@@ -35,26 +33,25 @@ public class DocumentBuilder {
 
     private final Document document;
 
-    public DocumentBuilder(BXml xmlPayload) throws Exception {
+    public DocumentBuilder(BObject documentBuilder, BXml xmlPayload) throws Exception {
         DocumentBuilderFactory factory = DocumentBuilderFactory.newInstance();
         factory.setNamespaceAware(true);
         this.document = factory.newDocumentBuilder().parse(new InputSource(new StringReader(xmlPayload.toString())));
+        documentBuilder.addNativeData(NATIVE_DOCUMENT, this.document);
     }
 
     protected DocumentBuilder(Document document) {
         this.document = document;
     }
 
     public static BArray getSignatureData(BObject document) {
-        BHandle handle = (BHandle) document.get(StringUtils.fromString(NATIVE_DOCUMENT));
-        DocumentBuilder docBuilder = (DocumentBuilder) handle.getValue();
-        return ValueCreator.createArrayValue(WsSecurityUtils.getSignatureValue(docBuilder.getNativeDocument()));
+        Document nativeDocument = (Document) document.getNativeData().get(NATIVE_DOCUMENT);
+        return ValueCreator.createArrayValue(WsSecurityUtils.getSignatureValue(nativeDocument));
     }
 
     public static BArray getEncryptedData(BObject document) {
-        BHandle handle = (BHandle) document.get(StringUtils.fromString(NATIVE_DOCUMENT));
-        DocumentBuilder docBuilder = (DocumentBuilder) handle.getValue();
-        return ValueCreator.createArrayValue(WsSecurityUtils.getEncryptedData(docBuilder.getNativeDocument()));
+        Document nativeDocument = (Document) document.getNativeData().get(NATIVE_DOCUMENT);
+        return ValueCreator.createArrayValue(WsSecurityUtils.getEncryptedData(nativeDocument));
     }
 
     protected Document getNativeDocument() {

diff --git a/native/src/main/java/org/wssec/WsSecurityHeader.java b/native/src/main/java/org/wssec/WsSecurityHeader.java
@@ -34,10 +34,9 @@ public class WsSecurityHeader {
     private final Document document;
 
     public WsSecurityHeader(BObject documentBuilder) {
-        BHandle handle = (BHandle) documentBuilder.get(StringUtils.fromString(NATIVE_DOCUMENT));
-        DocumentBuilder docBuilder = (DocumentBuilder) handle.getValue();
-        this.wsSecHeader = new WSSecHeader(docBuilder.getNativeDocument());
-        this.document = docBuilder.getNativeDocument();
+        Document document = (Document) documentBuilder.getNativeData().get(NATIVE_DOCUMENT);
+        this.wsSecHeader = new WSSecHeader(document);
+        this.document = document;
     }
 
     protected Document getDocument() {