Merge pull request #1336 from nats-io/fix-1320

[BUG] Fix UNSUBs after disconnect can cause auth violations

Author: scottf

src/main/java/io/nats/client/impl/NatsConnection.java

@@ -1059,7 +1059,7 @@ void sendUnsub(NatsSubscription sub, int after) {
         if (after > 0) {
             bab.append(SP).append(after);
         }
-        queueInternalOutgoing(new ProtocolMessage(bab));
+        queueOutgoing(new ProtocolMessage(bab, true));
     }
 
     // Assumes the null/empty checks were handled elsewhere
@@ -1107,8 +1107,11 @@ void sendSubscriptionMessage(String sid, String subject, String queueName, boole
         }
         bab.append(SP).append(sid);
 
-        ProtocolMessage subMsg = new ProtocolMessage(bab);
-
+        // setting this to filter on stop.
+        // if it's an "internal" message, it won't be filtered
+        // if it's a normal message, the subscription will already be registered
+        // and therefore will be re-subscribed after a stop anyway
+        ProtocolMessage subMsg = new ProtocolMessage(bab, true);
         if (treatAsInternal) {
             queueInternalOutgoing(subMsg);
         } else {
@@ -1499,7 +1502,7 @@ void sendConnect(NatsUri nuri) throws IOException {
             ByteArrayBuilder bab =
                 new ByteArrayBuilder(OP_CONNECT_SP_LEN + connectOptions.limit(), UTF_8)
                     .append(CONNECT_SP_BYTES).append(connectOptions);
-            queueInternalOutgoing(new ProtocolMessage(bab));
+            queueInternalOutgoing(new ProtocolMessage(bab, false));
         } catch (Exception exp) {
             throw new IOException("Error sending connect string", exp);
         }
@@ -1527,7 +1530,7 @@ public Duration RTT() throws IOException {
         pongQueue.add(pongFuture);
         try {
             long time = NatsSystemClock.nanoTime();
-            writer.queueInternalMessage(new ProtocolMessage(OP_PING_BYTES));
+            writer.queueInternalMessage(new ProtocolMessage(PING_PROTO));
             pongFuture.get(timeout, TimeUnit.MILLISECONDS);
             return Duration.ofNanos(NatsSystemClock.nanoTime() - time);
         }

src/main/java/io/nats/client/impl/NatsConnectionWriter.java

@@ -30,7 +30,8 @@
 import java.util.concurrent.locks.ReentrantLock;
 
 import static io.nats.client.support.BuilderBase.bufferAllocSize;
-import static io.nats.client.support.NatsConstants.*;
+import static io.nats.client.support.NatsConstants.CR;
+import static io.nats.client.support.NatsConstants.LF;
 
 class NatsConnectionWriter implements Runnable {
     private static final int BUFFER_BLOCK_SIZE = 256;
@@ -105,10 +106,7 @@ Future<Boolean> stop() {
             try {
                 this.outgoing.pause();
                 this.reconnectOutgoing.pause();
-                // Clear old ping/pong requests
-                this.outgoing.filter((msg) ->
-                    msg.isProtocol() &&
-                        (msg.getProtocolBab().equals(OP_PING_BYTES) || msg.getProtocolBab().equals(OP_PONG_BYTES)));
+                this.outgoing.filter(NatsMessage::isProtocolFilterOnStop);
             }
             finally {
                 this.startStopLock.unlock();
@@ -151,7 +149,7 @@ void sendMessageBatch(NatsMessage msg, DataPort dataPort, StatisticsCollector st
                 sendBuffer[sendPosition++] = CR;
                 sendBuffer[sendPosition++] = LF;
 
-                if (!msg.isProtocol()) { // because a protocol message does not have headers
+                if (!msg.isProtocol()) { // because a protocol message does not have headers or data
                     sendPosition += msg.copyNotEmptyHeaders(sendPosition, sendBuffer);
 
                     byte[] bytes = msg.getData(); // guaranteed to not be null

src/main/java/io/nats/client/impl/NatsMessage.java

@@ -107,6 +107,10 @@ boolean isProtocol() {
         return false; // overridden in NatsMessage.ProtocolMessage
     }
 
+    boolean isProtocolFilterOnStop() {
+        return false; // overridden in NatsMessage.ProtocolMessage
+    }
+
     private static final Headers EMPTY_READ_ONLY = new Headers(null, true, null);
 
     protected void calculate() {

src/main/java/io/nats/client/impl/ProtocolMessage.java

@@ -19,30 +19,39 @@
 // Protocol message is a special version of a NatsPublishableMessage extends NatsMessage
 // ----------------------------------------------------------------------------------------------------
 class ProtocolMessage extends NatsPublishableMessage {
+    final boolean filterOnStop;
 
-    ProtocolMessage(ByteArrayBuilder babProtocol) {
+    ProtocolMessage(ByteArrayBuilder babProtocol, boolean filterOnStop) {
         super(false);
         protocolBab = babProtocol;
         sizeInBytes = controlLineLength = protocolBab.length() + 2; // CRLF, protocol doesn't have data
+        this.filterOnStop = filterOnStop;
     }
 
     ProtocolMessage(byte[] protocol) {
         super(false);
         protocolBab = new ByteArrayBuilder(protocol);
         sizeInBytes = controlLineLength = protocolBab.length() + 2; // CRLF, protocol doesn't have data
+        this.filterOnStop = true;
     }
 
     ProtocolMessage(ProtocolMessage pm) {
         super(false);
         protocolBab = pm.protocolBab;
         sizeInBytes = controlLineLength = pm.sizeInBytes;
+        filterOnStop = pm.filterOnStop;
     }
 
     @Override
     boolean isProtocol() {
         return true;
     }
 
+    @Override
+    boolean isProtocolFilterOnStop() {
+        return filterOnStop;
+    }
+
     @Override
     int copyNotEmptyHeaders(int destPosition, byte[] dest) {
         return 0; // until a protocol messages gets headers, might as well shortcut this.