Resolve NoiseHandshake TODOs and add tests for message length checks

Author: jchambers

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

@@ -161,6 +161,8 @@ public class NoiseHandshake {
 
   private int currentPreSharedKey;
 
+  static final int MAX_NOISE_MESSAGE_SIZE = 65_535;
+
   private static final byte[] EMPTY_BYTE_ARRAY = new byte[0];
   private static final ByteBuffer EMPTY_BYTE_BUFFER = ByteBuffer.wrap(EMPTY_BYTE_ARRAY);
 
@@ -643,8 +645,9 @@ static int getPayloadLength(final HandshakePattern handshakePattern,
    * @see <a href="https://noiseprotocol.org/noise.html#payload-security-properties">The Noise Protocol Framework - Payload security properties</a>
    */
   public byte[] writeMessage(@Nullable final byte[] payload) {
-    // TODO Verify that message size is within bounds
     final int payloadLength = payload != null ? payload.length : 0;
+    checkOutboundMessageSize(payloadLength);
+
     final byte[] message = new byte[getOutboundMessageLength(payloadLength)];
 
     try {
@@ -690,7 +693,11 @@ public int writeMessage(@Nullable final byte[] payload,
                           final byte[] message,
                           final int messageOffset) throws ShortBufferException {
 
-    // TODO Check message buffer length, or just let plumbing deeper down complain?
+    checkOutboundMessageSize(payloadLength);
+
+    if (message.length - messageOffset < getOutboundMessageLength(payloadLength)) {
+      throw new ShortBufferException("Message array after offset is not large enough to hold handshake message");
+    }
 
     if (!isExpectingWrite()) {
       throw new IllegalStateException("Handshake not currently expecting to write a message");
@@ -770,7 +777,10 @@ public int writeMessage(@Nullable final byte[] payload,
    * @see <a href="https://noiseprotocol.org/noise.html#payload-security-properties">The Noise Protocol Framework - Payload security properties</a>
    */
   public ByteBuffer writeMessage(@Nullable final ByteBuffer payload) {
-    final ByteBuffer message = ByteBuffer.allocate(getOutboundMessageLength(payload != null ? payload.remaining() : 0));
+    final int payloadLength = payload != null ? payload.remaining() : 0;
+    checkOutboundMessageSize(payloadLength);
+
+    final ByteBuffer message = ByteBuffer.allocate(getOutboundMessageLength(payloadLength));
 
     try {
       writeMessage(payload, message);
@@ -813,7 +823,12 @@ public ByteBuffer writeMessage(@Nullable final ByteBuffer payload) {
   public int writeMessage(@Nullable final ByteBuffer payload,
                           final ByteBuffer message) throws ShortBufferException {
 
-    // TODO Check message buffer length, or just let plumbing deeper down complain?
+    final int payloadLength = payload != null ? payload.remaining() : 0;
+    checkOutboundMessageSize(payloadLength);
+
+    if (message.remaining() < getOutboundMessageLength(payloadLength)) {
+      throw new ShortBufferException("Message buffer is not large enough to hold handshake message");
+    }
 
     if (!isExpectingWrite()) {
       throw new IllegalStateException("Handshake not currently expecting to write a message");
@@ -869,6 +884,12 @@ public int writeMessage(@Nullable final ByteBuffer payload,
     return bytesWritten;
   }
 
+  private void checkOutboundMessageSize(final int payloadLength) {
+    if (getOutboundMessageLength(payloadLength) > MAX_NOISE_MESSAGE_SIZE) {
+      throw new IllegalArgumentException("Message containing payload would be larger than maximum allowed Noise message size");
+    }
+  }
+
   /**
    * Reads the next handshake message, advancing this handshake's internal state.
    *
@@ -881,6 +902,8 @@ public int writeMessage(@Nullable final ByteBuffer payload,
    * @throws IllegalArgumentException if the given message is too short to contain the expected handshake message
    */
   public byte[] readMessage(final byte[] message) throws AEADBadTagException {
+    checkInboundMessageSize(message.length);
+
     final byte[] payload = new byte[getPayloadLength(message.length)];
 
     try {
@@ -921,6 +944,12 @@ public int readMessage(final byte[] message,
                          final byte[] payload,
                          final int payloadOffset) throws ShortBufferException, AEADBadTagException {
 
+    checkInboundMessageSize(messageLength);
+
+    if (payload.length - payloadOffset < getPayloadLength(messageLength)) {
+      throw new ShortBufferException("Payload array after offset is not large enough to hold payload");
+    }
+
     if (!isExpectingRead()) {
       throw new IllegalStateException("Handshake not currently expecting to read a message");
     }
@@ -987,6 +1016,8 @@ public int readMessage(final byte[] message,
    * @throws IllegalArgumentException if the given message is too short to contain the expected handshake message
    */
   public ByteBuffer readMessage(final ByteBuffer message) throws AEADBadTagException {
+    checkInboundMessageSize(message.remaining());
+
     final ByteBuffer payload = ByteBuffer.allocate(getPayloadLength(message.remaining()));
 
     try {
@@ -1025,6 +1056,12 @@ public ByteBuffer readMessage(final ByteBuffer message) throws AEADBadTagExcepti
   public int readMessage(final ByteBuffer message,
                          final ByteBuffer payload) throws ShortBufferException, AEADBadTagException {
 
+    checkInboundMessageSize(message.remaining());
+
+    if (payload.remaining() < getPayloadLength(message.remaining())) {
+      throw new ShortBufferException("Payload buffer is not large enough to hold payload");
+    }
+
     if (!isExpectingRead()) {
       throw new IllegalStateException("Handshake not currently expecting to read a message");
     }
@@ -1077,6 +1114,12 @@ public int readMessage(final ByteBuffer message,
     return decryptAndHash(message, payload);
   }
 
+  private void checkInboundMessageSize(final int messageSize) {
+    if (messageSize > MAX_NOISE_MESSAGE_SIZE) {
+      throw new IllegalArgumentException("Message is larger than maximum allowed Noise message size");
+    }
+  }
+
   private void handleMixKeyToken(final HandshakePattern.Token token) {
     switch (token) {
       case EE -> {

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

@@ -16,6 +16,7 @@
 
 import javax.annotation.Nullable;
 import javax.crypto.AEADBadTagException;
+import javax.crypto.ShortBufferException;
 import java.io.IOException;
 import java.io.InputStream;
 import java.nio.ByteBuffer;
@@ -191,6 +192,91 @@ void getPayloadLength() throws NoSuchPatternException {
         () -> NoiseHandshake.getPayloadLength(handshakePattern, 0, publicKeyLength, 55));
   }
 
+  @Test
+  void writeMessageOversize() throws NoSuchAlgorithmException {
+    final NoiseKeyAgreement keyAgreement = NoiseKeyAgreement.getInstance("25519");
+
+    final NoiseHandshake handshake =
+        NoiseHandshakeBuilder.forIKInitiator(keyAgreement.generateKeyPair(), keyAgreement.generateKeyPair().getPublic())
+            .setComponentsFromProtocolName("Noise_IK_25519_AESGCM_SHA256")
+            .build();
+
+    // 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 payloadLength = NoiseHandshake.MAX_NOISE_MESSAGE_SIZE - 1;
+    final int messageLength = handshake.getOutboundMessageLength(payloadLength);
+
+    assertTrue(messageLength > NoiseHandshake.MAX_NOISE_MESSAGE_SIZE);
+
+    assertThrows(IllegalArgumentException.class,
+        () -> handshake.writeMessage(new byte[payloadLength]));
+
+    assertThrows(IllegalArgumentException.class,
+        () -> handshake.writeMessage(new byte[payloadLength], 0, payloadLength, new byte[messageLength], 0));
+
+    assertThrows(IllegalArgumentException.class,
+        () -> handshake.writeMessage(ByteBuffer.allocate(payloadLength)));
+
+    assertThrows(IllegalArgumentException.class,
+        () -> handshake.writeMessage(ByteBuffer.allocate(payloadLength), ByteBuffer.allocate(messageLength)));
+  }
+
+  @Test
+  void writeMessageShortBuffer() throws NoSuchAlgorithmException {
+    final NoiseHandshake handshake =
+        NoiseHandshakeBuilder.forNNInitiator()
+            .setComponentsFromProtocolName("Noise_NN_25519_AESGCM_SHA256")
+            .build();
+
+    final byte[] payload = new byte[32];
+    final byte[] message = new byte[payload.length - 1];
+
+    assertThrows(ShortBufferException.class, () ->
+        handshake.writeMessage(payload, 0, payload.length, message, 0));
+
+    assertThrows(ShortBufferException.class, () ->
+        handshake.writeMessage(ByteBuffer.wrap(payload), ByteBuffer.wrap(message)));
+  }
+
+  @Test
+  void readMessageOversize() throws NoSuchAlgorithmException {
+    final NoiseHandshake handshake =
+        NoiseHandshakeBuilder.forNNResponder()
+            .setComponentsFromProtocolName("Noise_NN_25519_AESGCM_SHA256")
+            .build();
+
+    final int messageLength = NoiseHandshake.MAX_NOISE_MESSAGE_SIZE + 1;
+
+    assertThrows(IllegalArgumentException.class, () ->
+        handshake.readMessage(new byte[messageLength]));
+
+    assertThrows(IllegalArgumentException.class, () ->
+        handshake.readMessage(new byte[messageLength], 0, messageLength, new byte[messageLength], 0));
+
+    assertThrows(IllegalArgumentException.class, () ->
+        handshake.readMessage(ByteBuffer.allocate(messageLength)));
+
+    assertThrows(IllegalArgumentException.class, () ->
+        handshake.readMessage(ByteBuffer.allocate(messageLength), ByteBuffer.allocate(messageLength)));
+  }
+
+  @Test
+  void readMessageShortBuffer() throws NoSuchAlgorithmException {
+    final NoiseHandshake handshake =
+        NoiseHandshakeBuilder.forNNResponder()
+            .setComponentsFromProtocolName("Noise_NN_25519_AESGCM_SHA256")
+            .build();
+
+    final byte[] message = new byte[128];
+    final int payloadLength = handshake.getPayloadLength(message.length);
+
+    assertThrows(ShortBufferException.class, () ->
+        handshake.readMessage(message, 0, message.length, new byte[payloadLength - 1], 0));
+
+    assertThrows(ShortBufferException.class, () ->
+        handshake.readMessage(ByteBuffer.wrap(message), ByteBuffer.allocate(payloadLength - 1)));
+  }
+
   @ParameterizedTest
   @MethodSource("cacophonyTestVectors")
   void cacophonyTestsWithNewByteArray(final CacophonyTestVector testVector) throws AEADBadTagException {