Add message length checks and tests for NoiseTransportImpl

Author: jchambers

src/main/java/com/eatthepath/noise/NoiseTransportImpl.java

@@ -9,6 +9,8 @@ class NoiseTransportImpl implements NoiseTransport {
   private final CipherState readerState;
   private final CipherState writerState;
 
+  static final int MAX_NOISE_MESSAGE_SIZE = 65_535;
+
   NoiseTransportImpl(final CipherState readerState, final CipherState writerState) {
     this.readerState = readerState;
     this.writerState = writerState;
@@ -26,16 +28,26 @@ public int getCiphertextLength(final int plaintextLength) {
 
   @Override
   public ByteBuffer readMessage(final ByteBuffer ciphertext) throws AEADBadTagException {
+    checkInboundMessageSize(ciphertext.remaining());
+
     return readerState.decrypt(null, ciphertext);
   }
 
   @Override
   public int readMessage(final ByteBuffer ciphertext, final ByteBuffer plaintext) throws ShortBufferException, AEADBadTagException {
+    checkInboundMessageSize(ciphertext.remaining());
+
+    if (plaintext.remaining() < getPlaintextLength(ciphertext.remaining())) {
+      throw new ShortBufferException("Plaintext buffer does not have enough remaining capacity to hold plaintext");
+    }
+
     return readerState.decrypt(null, ciphertext, plaintext);
   }
 
   @Override
   public byte[] readMessage(final byte[] ciphertext) throws AEADBadTagException {
+    checkInboundMessageSize(ciphertext.length);
+
     return readerState.decrypt(null, ciphertext);
   }
 
@@ -46,23 +58,45 @@ public int readMessage(final byte[] ciphertext,
                          final byte[] plaintext,
                          final int plaintextOffset) throws ShortBufferException, AEADBadTagException {
 
+    checkInboundMessageSize(ciphertextLength);
+
+    if (plaintext.length - plaintextOffset < getPlaintextLength(ciphertextLength)) {
+      throw new ShortBufferException("Plaintext array after offset is not large enough to hold plaintext");
+    }
+
     return readerState.decrypt(null, 0, 0,
         ciphertext, ciphertextOffset, ciphertextLength,
         plaintext, plaintextOffset);
   }
 
+  private void checkInboundMessageSize(final int ciphertextLength) {
+    if (ciphertextLength > MAX_NOISE_MESSAGE_SIZE) {
+      throw new IllegalArgumentException("Message is larger than maximum allowed Noise transport message size");
+    }
+  }
+
   @Override
   public ByteBuffer writeMessage(final ByteBuffer plaintext) {
+    checkOutboundMessageSize(plaintext.remaining());
+
     return writerState.encrypt(null, plaintext);
   }
 
   @Override
   public int writeMessage(final ByteBuffer plaintext, final ByteBuffer ciphertext) throws ShortBufferException {
+    checkOutboundMessageSize(plaintext.remaining());
+
+    if (ciphertext.remaining() < getCiphertextLength(plaintext.remaining())) {
+      throw new ShortBufferException("Ciphertext buffer does not have enough remaining capacity to hold ciphertext");
+    }
+
     return writerState.encrypt(null, plaintext, ciphertext);
   }
 
   @Override
   public byte[] writeMessage(final byte[] plaintext) {
+    checkOutboundMessageSize(plaintext.length);
+
     return writerState.encrypt(null, plaintext);
   }
 
@@ -73,11 +107,23 @@ public int writeMessage(final byte[] plaintext,
                           final byte[] ciphertext,
                           final int ciphertextOffset) throws ShortBufferException {
 
+    checkOutboundMessageSize(plaintextLength);
+
+    if (ciphertext.length - ciphertextOffset < getCiphertextLength(plaintextLength)) {
+      throw new ShortBufferException("Ciphertext array after offset is not large enough to hold ciphertext");
+    }
+
     return writerState.encrypt(null, 0, 0,
         plaintext, plaintextOffset, plaintextLength,
         ciphertext, ciphertextOffset);
   }
 
+  void checkOutboundMessageSize(final int plaintextLength) {
+    if (getCiphertextLength(plaintextLength) > MAX_NOISE_MESSAGE_SIZE) {
+      throw new IllegalArgumentException("Ciphertext would be larger than maximum allowed Noise transport message size");
+    }
+  }
+
   @Override
   public void rekeyReader() {
     readerState.rekey();

src/test/java/com/eatthepath/noise/NoiseTransportImplTest.java

@@ -0,0 +1,109 @@
+package com.eatthepath.noise;
+
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+
+import javax.crypto.AEADBadTagException;
+import javax.crypto.ShortBufferException;
+import java.nio.ByteBuffer;
+import java.security.NoSuchAlgorithmException;
+
+import static org.junit.jupiter.api.Assertions.*;
+
+class NoiseTransportImplTest {
+
+  private NoiseTransportImpl noiseTransport;
+
+  @BeforeEach
+  void setUp() throws NoSuchAlgorithmException, AEADBadTagException {
+    final NoiseHandshake initiatorHandshake =
+        NoiseHandshakeBuilder.forNNInitiator()
+            .setComponentsFromProtocolName("Noise_NN_25519_AESGCM_SHA256")
+            .build();
+
+    final NoiseHandshake responderHandshake =
+        NoiseHandshakeBuilder.forNNResponder()
+            .setComponentsFromProtocolName("Noise_NN_25519_AESGCM_SHA256")
+            .build();
+
+    responderHandshake.readMessage(initiatorHandshake.writeMessage((byte[]) null));
+    initiatorHandshake.readMessage(responderHandshake.writeMessage((byte[]) null));
+
+    noiseTransport = (NoiseTransportImpl) initiatorHandshake.toTransport();
+  }
+
+  @Test
+  void getPlaintextLength() {
+    final int ciphertextLength = 77;
+    assertEquals(ciphertextLength - 16, noiseTransport.getPlaintextLength(ciphertextLength));
+  }
+
+  @Test
+  void getCiphertextLength() {
+    final int plaintextLength = 83;
+    assertEquals(plaintextLength + 16, noiseTransport.getCiphertextLength(plaintextLength));
+  }
+
+  @Test
+  void writeMessageOversize() {
+    // We want to make sure we're testing the size of the resulting message (which may include key material and AEAD
+    // tags) rather than the length of just the payload
+    final int plaintextLength = NoiseTransportImpl.MAX_NOISE_MESSAGE_SIZE - 1;
+    final int messageLength = noiseTransport.getCiphertextLength(plaintextLength);
+
+    assertTrue(messageLength > NoiseTransportImpl.MAX_NOISE_MESSAGE_SIZE);
+
+    assertThrows(IllegalArgumentException.class,
+        () -> noiseTransport.writeMessage(new byte[plaintextLength]));
+
+    assertThrows(IllegalArgumentException.class,
+        () -> noiseTransport.writeMessage(new byte[plaintextLength], 0, plaintextLength, new byte[messageLength], 0));
+
+    assertThrows(IllegalArgumentException.class,
+        () -> noiseTransport.writeMessage(ByteBuffer.allocate(plaintextLength)));
+
+    assertThrows(IllegalArgumentException.class,
+        () -> noiseTransport.writeMessage(ByteBuffer.allocate(plaintextLength), ByteBuffer.allocate(messageLength)));
+  }
+
+  @Test
+  void writeMessageShortBuffer() {
+    final byte[] plaintext = new byte[32];
+    final byte[] message = new byte[noiseTransport.getCiphertextLength(plaintext.length) - 1];
+
+    assertThrows(ShortBufferException.class, () ->
+        noiseTransport.writeMessage(plaintext, 0, plaintext.length, message, 0));
+
+    assertThrows(ShortBufferException.class, () ->
+        noiseTransport.writeMessage(ByteBuffer.wrap(plaintext), ByteBuffer.wrap(message)));
+  }
+
+  @Test
+  void readMessageOversize() throws NoSuchAlgorithmException {
+    final int messageLength = NoiseTransportImpl.MAX_NOISE_MESSAGE_SIZE + 1;
+
+    assertThrows(IllegalArgumentException.class, () ->
+        noiseTransport.readMessage(new byte[messageLength]));
+
+    assertThrows(IllegalArgumentException.class, () ->
+        noiseTransport.readMessage(new byte[messageLength], 0, messageLength, new byte[messageLength], 0));
+
+    assertThrows(IllegalArgumentException.class, () ->
+        noiseTransport.readMessage(ByteBuffer.allocate(messageLength)));
+
+    assertThrows(IllegalArgumentException.class, () ->
+        noiseTransport.readMessage(ByteBuffer.allocate(messageLength), ByteBuffer.allocate(messageLength)));
+  }
+
+  @Test
+  void readMessageShortBuffer() {
+    final byte[] message = new byte[128];
+    final int plaintextLength = noiseTransport.getPlaintextLength(message.length);
+
+    assertThrows(ShortBufferException.class, () ->
+        noiseTransport.readMessage(message, 0, message.length, new byte[plaintextLength - 1], 0));
+
+    assertThrows(ShortBufferException.class, () ->
+        noiseTransport.readMessage(ByteBuffer.wrap(message), ByteBuffer.allocate(plaintextLength - 1)));
+  }
+}