Merge pull request #209 from Happy0/memory_leak

Fixes memory leak when decrypting scuttlebutt incoming messages
ConsenSysMesh · Apr 7, 2019 · b67ae4f · b67ae4f
2 parents 64478ad + 6f92b79
commit b67ae4f
Show file tree

Hide file tree

Showing 3 changed files with 52 additions and 20 deletions.
diff --git a/crypto/src/main/java/net/consensys/cava/crypto/sodium/SecretBox.java b/crypto/src/main/java/net/consensys/cava/crypto/sodium/SecretBox.java
@@ -193,7 +193,7 @@ public byte[] bytesArray() {
   /**
    * A SecretBox nonce.
    */
-  public static final class Nonce {
+  public static final class Nonce implements Destroyable {
     final Allocated value;
 
     private Nonce(Pointer ptr, int length) {
@@ -230,6 +230,11 @@ public static Nonce fromBytes(byte[] bytes) {
       return Sodium.dup(bytes, Nonce::new);
     }
 
+    @Override
+    public void destroy() {
+      this.value.destroy();
+    }
+
     /**
      * Obtain the length of the nonce in bytes (24).
      *

diff --git a/...shake/src/main/java/net/consensys/cava/scuttlebutt/handshake/SecureScuttlebuttStream.java b/...shake/src/main/java/net/consensys/cava/scuttlebutt/handshake/SecureScuttlebuttStream.java
@@ -16,6 +16,7 @@
 import net.consensys.cava.bytes.MutableBytes;
 import net.consensys.cava.crypto.sodium.SHA256Hash;
 import net.consensys.cava.crypto.sodium.SecretBox;
+import net.consensys.cava.crypto.sodium.SodiumException;
 
 import java.util.ArrayList;
 import java.util.List;
@@ -100,29 +101,50 @@ private Bytes decrypt(Bytes message, SecretBox.Key key, MutableBytes nonce, bool
   }
 
   private Bytes decryptMessage(Bytes message, SecretBox.Key key, MutableBytes nonce) {
-    if (message.size() < 34) {
-      return null;
-    }
-    MutableBytes snapshotNonce = nonce.mutableCopy();
-    SecretBox.Nonce headerNonce = SecretBox.Nonce.fromBytes(snapshotNonce);
-    SecretBox.Nonce bodyNonce = SecretBox.Nonce.fromBytes(snapshotNonce.increment());
-    Bytes decryptedHeader = SecretBox.decrypt(message.slice(0, 34), key, headerNonce);
 
-    if (decryptedHeader == null) {
-      throw new StreamException("Failed to decrypt message header");
-    }
+    SecretBox.Nonce headerNonce = null;
+    SecretBox.Nonce bodyNonce = null;
+
+    try {
+      MutableBytes snapshotNonce = nonce.mutableCopy();
+      headerNonce = SecretBox.Nonce.fromBytes(snapshotNonce);
+      bodyNonce = SecretBox.Nonce.fromBytes(snapshotNonce.increment());
+
+      if (message.size() < 34) {
+        return null;
+      }
+
+      Bytes decryptedHeader = SecretBox.decrypt(message.slice(0, 34), key, headerNonce);
 
-    int bodySize = ((decryptedHeader.get(0) & 0xFF) << 8) + (decryptedHeader.get(1) & 0xFF);
-    if (message.size() < bodySize + 34) {
-      return null;
+      if (decryptedHeader == null) {
+        throw new StreamException("Failed to decrypt message header");
+      }
+
+      int bodySize = ((decryptedHeader.get(0) & 0xFF) << 8) + (decryptedHeader.get(1) & 0xFF);
+      if (message.size() < bodySize + 34) {
+        return null;
+      }
+      Bytes body = message.slice(34, bodySize);
+      Bytes decryptedBody = SecretBox.decrypt(Bytes.concatenate(decryptedHeader.slice(2), body), key, bodyNonce);
+      if (decryptedBody == null) {
+        throw new StreamException("Failed to decrypt message");
+      }
+      nonce.increment().increment();
+
+      return decryptedBody;
+    } catch (SodiumException | OutOfMemoryError ex) {
+      throw new StreamException(ex);
+    } finally {
+      // These may be null if there was an error while trying to construct them
+      destroyIfNonNull(headerNonce);
+      destroyIfNonNull(bodyNonce);
     }
-    Bytes body = message.slice(34, bodySize);
-    Bytes decryptedBody = SecretBox.decrypt(Bytes.concatenate(decryptedHeader.slice(2), body), key, bodyNonce);
-    if (decryptedBody == null) {
-      throw new StreamException("Failed to decrypt message");
+  }
+
+  private void destroyIfNonNull(SecretBox.Nonce nonce) {
+    if (nonce != null) {
+      nonce.destroy();
     }
-    nonce.increment().increment();
-    return decryptedBody;
   }
 
   private Bytes encrypt(Bytes message, SecretBox.Key clientToServerKey, MutableBytes clientToServerNonce) {

diff --git a/...utt-handshake/src/main/java/net/consensys/cava/scuttlebutt/handshake/StreamException.java b/...utt-handshake/src/main/java/net/consensys/cava/scuttlebutt/handshake/StreamException.java
@@ -17,4 +17,9 @@ public final class StreamException extends RuntimeException {
   StreamException(String message) {
     super(message);
   }
+
+  StreamException(Throwable ex) {
+    super(ex);
+  }
+
 }