Refactor CrtResource to enforce more invariants (#55)

awslabs · May 14, 2019 · 6ac6a87 · 6ac6a87
1 parent b8912c4
commit 6ac6a87
Show file tree

Hide file tree

Showing 11 changed files with 107 additions and 57 deletions.
diff --git a/src/main/java/software/amazon/awssdk/crt/CrtResource.java b/src/main/java/software/amazon/awssdk/crt/CrtResource.java
@@ -16,14 +16,17 @@
 
 import java.util.Collection;
 import java.util.Collections;
+import java.util.LinkedList;
 import java.util.concurrent.ConcurrentHashMap;
 
 /**
  * This wraps a native pointer to an AWS Common Runtime resource. It also ensures
  * that the first time a resource is referenced, the CRT will be loaded and bound.
  */
-public class CrtResource {
+public class CrtResource implements AutoCloseable {
     private static final ConcurrentHashMap<Long, String> NATIVE_RESOURCES = new ConcurrentHashMap<>();
+    private static final long NULL = 0;
+    private final LinkedList<CrtResource> ownedSubResources = new LinkedList<>();
     private long ptr;
 
     static {
@@ -42,12 +45,42 @@ public static Collection<String> getAllocatedNativeResources() {
     public CrtResource() {
     }
 
+    /**
+     * Marks a SubResource as owned by this Resource. This will cause the SubResource to be closed after this Resource
+     * is closed.
+     * @param subresource The owned subresource
+     * @return The original subresource.
+     */
+    public <T extends CrtResource> T own(T subresource) {
+        if (!isNull()) {
+            throw new IllegalStateException("Parent Resource already created and acquired, can't add sub-resources after the fact.");
+        }
+        ownedSubResources.push(subresource);
+        return subresource;
+    }
+
     protected void acquire(long _ptr) {
-        NATIVE_RESOURCES.put(_ptr, this.getClass().getCanonicalName());
+        if (!isNull()) {
+            throw new IllegalStateException("Can't acquire >1 Native Pointer");
+        }
+
+        if (_ptr == NULL) {
+            throw new IllegalStateException("Can't acquire NULL Pointer");
+        }
+
+        String lastValue = NATIVE_RESOURCES.put(_ptr, this.getClass().getCanonicalName());
+
+        if (lastValue != null) {
+            throw new IllegalStateException("Acquired two CrtResources to the same Native Resource! Class: " + lastValue);
+        }
+
         ptr = _ptr;
     }
 
     protected long release() {
+        if (isNull()) {
+            throw new IllegalStateException("Already Released Resource!");
+        }
         NATIVE_RESOURCES.remove(ptr);
         long addr = ptr;
         ptr = 0;
@@ -57,4 +90,20 @@ protected long release() {
     public long native_ptr() {
         return ptr;
     }
+
+    public boolean isNull() {
+        return (ptr == NULL);
+    }
+
+    @Override
+    public void close() {
+        if (!isNull()) {
+            throw new IllegalStateException("Closing sub-resources before releasing parent Resource may cause use-after-free bugs!");
+        }
+
+        while(ownedSubResources.size() > 0) {
+            CrtResource r = ownedSubResources.pop();
+            r.close();
+        }
+    }
 }
diff --git a/src/main/java/software/amazon/awssdk/crt/io/ClientBootstrap.java b/src/main/java/software/amazon/awssdk/crt/io/ClientBootstrap.java
@@ -14,29 +14,26 @@
  */
 package software.amazon.awssdk.crt.io;
 
-import software.amazon.awssdk.crt.CrtRuntimeException;
 import software.amazon.awssdk.crt.CrtResource;
-
-import java.io.Closeable;
+import software.amazon.awssdk.crt.CrtRuntimeException;
 
 /**
  * This class wraps the aws_client_bootstrap from aws-c-io to provide
  * a client context for all protocol stacks in the AWS Common Runtime.
  */
-public final class ClientBootstrap extends CrtResource implements Closeable {
+public final class ClientBootstrap extends CrtResource {
     private final HostResolver hostResolver;
     private final EventLoopGroup elg;
-    private final boolean ownResources;
 
     /**
      * Creates a new ClientBootstrap. Most applications will only ever need one instance of this.
      * @param numThreads The number of Threads to use in the EventLoop
      * @throws CrtRuntimeException If the system is unable to allocate space for a native client bootstrap object
      */
     public ClientBootstrap(int numThreads) throws CrtRuntimeException {
-        this.elg = new EventLoopGroup(numThreads);
-        this.hostResolver = new HostResolver(elg);
-        this.ownResources = true;
+        this.elg = own(new EventLoopGroup(numThreads));
+        this.hostResolver = own(new HostResolver(elg));
+        acquire(clientBootstrapNew(elg.native_ptr(), hostResolver.native_ptr()));
     }
 
     /**
@@ -49,19 +46,15 @@ public ClientBootstrap(int numThreads) throws CrtRuntimeException {
     public ClientBootstrap(EventLoopGroup elg, HostResolver hr) throws CrtRuntimeException {
         this.hostResolver = hr;
         this.elg = elg;
-        this.ownResources = false;
         acquire(clientBootstrapNew(elg.native_ptr(), hostResolver.native_ptr()));
     }
 
     @Override
     public void close() {
-        if (native_ptr() != 0) {
+        if (!isNull()) {
             clientBootstrapDestroy(release());
         }
-        if (ownResources) {
-            hostResolver.close();
-            elg.close();
-        }
+        super.close();
     }
 
     /*******************************************************************************

diff --git a/src/main/java/software/amazon/awssdk/crt/io/EventLoopGroup.java b/src/main/java/software/amazon/awssdk/crt/io/EventLoopGroup.java
@@ -14,17 +14,15 @@
  */
 package software.amazon.awssdk.crt.io;
 
-import software.amazon.awssdk.crt.CrtRuntimeException;
 import software.amazon.awssdk.crt.CrtResource;
-
-import java.io.Closeable;
+import software.amazon.awssdk.crt.CrtRuntimeException;
 
 /**
  * This class wraps the aws_event_loop_group from aws-c-io to provide
  * access to an event loop for the MQTT protocol stack in the AWS Common
  * Runtime.
  */
-public final class EventLoopGroup extends CrtResource implements Closeable {
+public final class EventLoopGroup extends CrtResource {
 
     /**
      * Creates a new event loop group for the I/O subsystem to use to run blocking I/O requests
@@ -41,9 +39,10 @@ public EventLoopGroup(int numThreads) throws CrtRuntimeException {
      */
     @Override
     public void close() {
-        if (native_ptr() != 0) {
+        if (!isNull()) {
             eventLoopGroupDestroy(release());
         }
+        super.close();
     }
 
     /*******************************************************************************

diff --git a/src/main/java/software/amazon/awssdk/crt/io/HostResolver.java b/src/main/java/software/amazon/awssdk/crt/io/HostResolver.java
@@ -1,11 +1,23 @@
+/*
+ * Copyright 2010-2019 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License").
+ * You may not use this file except in compliance with the License.
+ * A copy of the License is located at
+ *
+ *  http://aws.amazon.com/apache2.0
+ *
+ * or in the "license" file accompanying this file. This file is distributed
+ * on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
+ * express or implied. See the License for the specific language governing
+ * permissions and limitations under the License.
+ */
 package software.amazon.awssdk.crt.io;
 
 import software.amazon.awssdk.crt.CrtResource;
 import software.amazon.awssdk.crt.CrtRuntimeException;
 
-import java.io.Closeable;
-
-public class HostResolver extends CrtResource implements Closeable {
+public class HostResolver extends CrtResource {
     private final static int DEFAULT_MAX_ENTRIES = 8;
     private final EventLoopGroup elg;
     private final int maxEntries;
@@ -22,9 +34,10 @@ public HostResolver(EventLoopGroup elg, int maxEntries) throws CrtRuntimeExcepti
 
     @Override
     public void close() {
-        if (native_ptr() != 0) {
+        if (!isNull()) {
             hostResolverRelease(release());
         }
+        super.close();
     }
 
     /*******************************************************************************

diff --git a/src/main/java/software/amazon/awssdk/crt/io/SocketOptions.java b/src/main/java/software/amazon/awssdk/crt/io/SocketOptions.java
@@ -14,16 +14,14 @@
  */
 package software.amazon.awssdk.crt.io;
 
-import software.amazon.awssdk.crt.CrtRuntimeException;
 import software.amazon.awssdk.crt.CrtResource;
-
-import java.io.Closeable;
+import software.amazon.awssdk.crt.CrtRuntimeException;
 
 /**
  * This class wraps the aws_socket_options from aws-c-io to provide
  * access to TCP/UDP socket configuration in the AWS Common Runtime.
  */
-public final class SocketOptions extends CrtResource implements Closeable {
+public final class SocketOptions extends CrtResource {
 
     /**
      * Socket communications domain
@@ -90,9 +88,10 @@ public SocketOptions() throws CrtRuntimeException {
      */
     @Override
     public void close() {
-        if (native_ptr() != 0) {
+        if (!isNull()) {
             socketOptionsDestroy(release());
         }
+        super.close();
     }
 
     /**

diff --git a/src/main/java/software/amazon/awssdk/crt/io/TlsContext.java b/src/main/java/software/amazon/awssdk/crt/io/TlsContext.java
@@ -14,16 +14,14 @@
  */
 package software.amazon.awssdk.crt.io;
 
-import software.amazon.awssdk.crt.CrtRuntimeException;
 import software.amazon.awssdk.crt.CrtResource;
-
-import java.io.Closeable;
+import software.amazon.awssdk.crt.CrtRuntimeException;
 
 /**
  * This class wraps the aws_tls_context from aws-c-io to provide
  * access to TLS configuration contexts in the AWS Common Runtime.
  */
-public final class TlsContext extends CrtResource implements Closeable {
+public final class TlsContext extends CrtResource {
 
     /**
      * Creates a new TlsContext. There are significant native resources consumed to create a TlsContext, so most
@@ -41,9 +39,10 @@ public TlsContext(TlsContextOptions options) throws CrtRuntimeException {
      */
     @Override
     public void close() {
-        if (native_ptr() != 0) {
+        if (!isNull()) {
             tlsContextDestroy(release());
         }
+        super.close();
     }
 
     /*******************************************************************************

diff --git a/src/main/java/software/amazon/awssdk/crt/io/TlsContextOptions.java b/src/main/java/software/amazon/awssdk/crt/io/TlsContextOptions.java
@@ -14,16 +14,14 @@
  */
 package software.amazon.awssdk.crt.io;
 
-import software.amazon.awssdk.crt.CrtRuntimeException;
 import software.amazon.awssdk.crt.CrtResource;
-
-import java.io.Closeable;
+import software.amazon.awssdk.crt.CrtRuntimeException;
 
 /**
  * This class wraps the aws_tls_connection_options from aws-c-io to provide
  * access to TLS configuration contexts in the AWS Common Runtime.
  */
-public final class TlsContextOptions extends CrtResource implements Closeable {
+public final class TlsContextOptions extends CrtResource {
 
     public enum TlsVersions {
         /**
@@ -70,9 +68,10 @@ public TlsContextOptions() throws CrtRuntimeException {
      */
     @Override
     public void close() {
-        if (native_ptr() != 0) {
+        if (!isNull()) {
             tlsContextOptionsDestroy(release());
         }
+        super.close();
     }
 
     /**

diff --git a/src/main/java/software/amazon/awssdk/crt/mqtt/MqttClient.java b/src/main/java/software/amazon/awssdk/crt/mqtt/MqttClient.java
@@ -15,21 +15,19 @@
  */
 package software.amazon.awssdk.crt.mqtt;
 
-import software.amazon.awssdk.crt.io.ClientBootstrap;
 import software.amazon.awssdk.crt.CrtResource;
 import software.amazon.awssdk.crt.CrtRuntimeException;
+import software.amazon.awssdk.crt.io.ClientBootstrap;
 import software.amazon.awssdk.crt.io.TlsContext;
 
-import java.io.Closeable;
-
 /**
  * This class wraps aws-c-mqtt to provide the basic MQTT pub/sub functionalities
  * via the AWS Common Runtime
  * 
  * One MqttClient class is needed per application. It can create any number of connections to
  * any number of MQTT endpoints
  */
-public class MqttClient extends CrtResource implements Closeable {
+public class MqttClient extends CrtResource {
     private final ClientBootstrap bootstrap;
     private final TlsContext tlsContext;
 
@@ -67,9 +65,10 @@ public MqttClient(ClientBootstrap clientBootstrap, TlsContext tlsContext) throws
      */
     @Override
     public void close() {
-        if (native_ptr() != 0) {
+        if (!isNull()) {
             mqttClientDestroy(release());
         }
+        super.close();
     }
 
     /**