refactor: iterate login response tokens directly

src/connection.ts

@@ -31,6 +31,7 @@ import RpcRequestPayload from './rpcrequest-payload';
 import SqlBatchPayload from './sqlbatch-payload';
 import MessageIO from './message-io';
 import { Parser as TokenStreamParser } from './token/token-stream-parser';
+import StreamParser from './token/stream-parser';
 import { Transaction, ISOLATION_LEVEL, assertValidIsolationLevel } from './transaction';
 import { ConnectionError, RequestError } from './errors';
 import { connectInParallel, connectInSequence } from './connector';
@@ -2212,6 +2213,33 @@ class Connection extends EventEmitter {
     return new TokenStreamParser(message, this.debug, handler, this.config.options);
   }
 
+  /**
+   * Parse the tokens in the given message and dispatch each token to the
+   * given handler, until the message ends or the given abort promise
+   * rejects.
+   *
+   * @private
+   */
+  async processTokens(message: Message, handler: TokenHandler, signalAborted: Promise<never>) {
+    const tokens = StreamParser.parseTokens(message, this.debug, this.config.options);
+
+    while (true) {
+      const result = await Promise.race([tokens.next(), signalAborted]);
+
+      if (result.done) {
+        break;
+      }
+
+      const token = result.value;
+      if (token === undefined) {
+        continue;
+      }
+
+      this.debug.token(token);
+      handler[token.handlerName as keyof TokenHandler](token as any);
+    }
+  }
+
   async wrapWithTls(socket: net.Socket, signal: AbortSignal): Promise<tls.TLSSocket> {
     signal.throwIfAborted();
 
@@ -3474,11 +3502,7 @@ class Connection extends EventEmitter {
       ]);
 
       const handler = new Login7TokenHandler(this);
-      const tokenStreamParser = this.createTokenStreamParser(message, handler);
-      await Promise.race([
-        once(tokenStreamParser, 'end'),
-        signalAborted
-      ]);
+      await this.processTokens(message, handler, signalAborted);
 
       if (handler.loginAckReceived) {
         return handler.routingData;
@@ -3504,11 +3528,7 @@ class Connection extends EventEmitter {
         ]);
 
         const handler = new Login7TokenHandler(this);
-        const tokenStreamParser = this.createTokenStreamParser(message, handler);
-        await Promise.race([
-          once(tokenStreamParser, 'end'),
-          signalAborted
-        ]);
+        await this.processTokens(message, handler, signalAborted);
 
         if (handler.loginAckReceived) {
           return handler.routingData;
@@ -3550,11 +3570,7 @@ class Connection extends EventEmitter {
       ]);
 
       const handler = new Login7TokenHandler(this);
-      const tokenStreamParser = this.createTokenStreamParser(message, handler);
-      await Promise.race([
-        once(tokenStreamParser, 'end'),
-        signalAborted
-      ]);
+      await this.processTokens(message, handler, signalAborted);
 
       if (handler.loginAckReceived) {
         return handler.routingData;
@@ -3648,11 +3664,7 @@ class Connection extends EventEmitter {
         signalAborted
       ]);
 
-      const tokenStreamParser = this.createTokenStreamParser(message, new InitialSqlTokenHandler(this));
-      await Promise.race([
-        once(tokenStreamParser, 'end'),
-        signalAborted
-      ]);
+      await this.processTokens(message, new InitialSqlTokenHandler(this), signalAborted);
     });
   }
 }

test/unit/connection-failure-test.ts

@@ -429,6 +429,73 @@ describe('Connection failure handling', function() {
     });
   });
 
+  it('should fail cleanly when the Login7 response contains an invalid token', function(done) {
+    server.on('connection', async (connection) => {
+      const debug = new Debug();
+      const incomingMessageStream = new IncomingMessageStream(debug);
+      const outgoingMessageStream = new OutgoingMessageStream(debug, { packetSize: 4 * 1024 });
+
+      connection.pipe(incomingMessageStream);
+      outgoingMessageStream.pipe(connection);
+
+      try {
+        const messageIterator = incomingMessageStream[Symbol.asyncIterator]();
+
+        // PRELOGIN
+        {
+          const { value: message } = await messageIterator.next();
+          assert.strictEqual(message.type, 0x12);
+
+          const chunks: Buffer[] = [];
+          for await (const data of message) {
+            chunks.push(data);
+          }
+
+          const responsePayload = new PreloginPayload({ encrypt: false, version: { major: 1, minor: 2, build: 3, subbuild: 0 } });
+          const responseMessage = new Message({ type: 0x12 });
+          responseMessage.end(responsePayload.data);
+          outgoingMessageStream.write(responseMessage);
+        }
+
+        // LOGIN7
+        {
+          const { value: message } = await messageIterator.next();
+          assert.strictEqual(message.type, 0x10);
+
+          const chunks: Buffer[] = [];
+          for await (const data of message) {
+            chunks.push(data);
+          }
+
+          // Respond with data that is not a valid token stream.
+          // `0x00` is not a valid token type.
+          const responseMessage = new Message({ type: 0x04 });
+          responseMessage.end(Buffer.from([0x00, 0x01, 0x02, 0x03]));
+          outgoingMessageStream.write(responseMessage);
+        }
+      } catch (err) {
+        console.log(err);
+      }
+    });
+
+    const connection = new Connection({
+      server: (server.address() as net.AddressInfo).address,
+      options: {
+        port: (server.address() as net.AddressInfo).port,
+        encrypt: false
+      }
+    });
+
+    connection.connect((err) => {
+      connection.close();
+
+      assert.instanceOf(err, Error);
+      assert.match(err!.message, /Unknown type/);
+
+      done();
+    });
+  });
+
   it('should fail correctly when the connection is aborted after the initial SQL message is sent', function(done) {
     server.on('connection', async (connection) => {
       const debug = new Debug();