feat: add auth and idle timeouts

Author: DASPRiD

README.md

@@ -118,6 +118,12 @@ const {server, closeOpenConnections} = await startServer({
     //    format: winston.format.simple(),
     //    transports: new winston.transport.Console(),
     //}),
+    
+    // Authentication timeout, defaults to 1 second.
+    //authTimeout: 1_000,
+
+    // Idle timeout, defaults to 30 second.
+    //idleTimeout: 30_000,
 });
 
 // This takes care of gracefully shutting down the server on CTRL+C

src/BufferedStream.ts

@@ -40,6 +40,11 @@ class BufferedStream<T extends Socket> implements Stream {
     }
 
     public async readExact(length : number) : Promise<Buffer> {
+        /* istanbul ignore next */
+        if (this.closed) {
+            throw new StreamClosedError('Socket was closed');
+        }
+
         /* istanbul ignore next */
         if (this.readResolver) {
             throw new Error('Only one read can be performed at a time');
@@ -98,16 +103,20 @@ class BufferedStream<T extends Socket> implements Stream {
     public writeUint8(value : number) : void {
         const content = Buffer.allocUnsafe(1);
         content.writeUInt8(value, 0);
-        this.socket.write(content);
+        this.writeAll(content);
     }
 
     public writeUInt32LE(value : number) : void {
         const content = Buffer.allocUnsafe(4);
         content.writeUInt32LE(value, 0);
-        this.socket.write(content);
+        this.writeAll(content);
     }
 
     public writeAll(buffer : Buffer) : void {
+        if (this.closed) {
+            throw new StreamClosedError('Socket was closed');
+        }
+
         this.socket.write(buffer);
     }
 
@@ -119,6 +128,10 @@ class BufferedStream<T extends Socket> implements Stream {
     public isClosed() : boolean {
         return this.closed;
     }
+
+    public setTimeout(timeout : number, callback : () => void) : void {
+        this.socket.setTimeout(timeout, callback);
+    }
 }
 
 export default BufferedStream;

src/Stream.ts

@@ -9,6 +9,7 @@ type Stream = {
     writeAll : (buffer : Buffer) => void;
     close : () => void;
     isClosed : () => boolean;
+    setTimeout : (timeout : number, callback : () => void) => void;
 };
 
 export default Stream;

src/client.ts

@@ -20,26 +20,59 @@ export enum Response {
     pong = 0,
 }
 
-const handleClient = async (stream : Stream, processor : Processor, logger ?: Logger) : Promise<void> => {
-    const credentialsLength = await stream.readUint8();
-    const credentials = (await stream.readExact(credentialsLength)).toString('ascii');
+const handleClient = async (
+    stream : Stream,
+    processor : Processor,
+    logger ?: Logger,
+    authTimeout = 1_000,
+    idleTimeout = 30_000,
+) : Promise<void> => {
+    // We handle the auth timeout via race instead of a socket timeout. This forces the client to complete
+    // authentication in a fixed amount of time. Otherwise, a malicious client could send individual bytes to keep the
+    // connection open for up 256 times the socket timeout.
+    let timeout;
+
+    const clientId = await Promise.race([
+        new Promise<null>(resolve => {
+            timeout = setTimeout(() => {
+                resolve(null);
+            }, authTimeout);
+        }),
+        (async () => {
+            const credentialsLength = await stream.readUint8();
+            const credentials = (await stream.readExact(credentialsLength)).toString('ascii');
+
+            if (!credentials.includes(':')) {
+                logger?.info('Malformed credentials');
+                stream.writeUint8(Response.invalidAuth);
+                return null;
+            }
 
-    if (!credentials.includes(':')) {
-        logger?.info('Malformed credentials');
-        stream.writeUint8(Response.invalidAuth);
-        return;
-    }
+            const [clientId, secret] = credentials.split(':', 2);
+            const authResult = await processor.authenticate(clientId, secret);
+
+            if (!authResult) {
+                logger?.info(`Unknown client ID "${clientId}" or invalid secret`);
+                stream.writeUint8(Response.invalidAuth);
+                return null;
+            }
+
+            stream.writeUint8(Response.validAuth);
+            return clientId;
+        })(),
+    ]);
 
-    const [clientId, secret] = credentials.split(':', 2);
-    const authResult = await processor.authenticate(clientId, secret);
+    clearTimeout(timeout);
 
-    if (!authResult) {
-        logger?.info(`Unknown client ID "${clientId}" or invalid secret`);
-        stream.writeUint8(Response.invalidAuth);
+    if (!clientId) {
         return;
     }
 
-    stream.writeUint8(Response.validAuth);
+    // At this point we can be certain that it's not a random client anymore, so further inactivity is handled via
+    // socket timeouts.
+    stream.setTimeout(idleTimeout, () => {
+        stream.close();
+    });
 
     while (!stream.isClosed()) {
         const commandCode = await stream.readUint8();

src/server.ts

@@ -11,6 +11,8 @@ export type ServerOptions = {
     tls : TlsOptions;
     port ?: number;
     hostname ?: string;
+    authTimeout ?: number;
+    idleTimeout ?: number;
     processor : Processor;
     logger ?: Logger;
 };
@@ -27,7 +29,7 @@ export const startServer = async (options : ServerOptions) : Promise<StartServer
         const stream = new BufferedStream(socket);
         streams.add(stream);
 
-        handleClient(stream, options.processor, options.logger)
+        handleClient(stream, options.processor, options.logger, options.authTimeout, options.idleTimeout)
             .then(() => {
                 stream.close();
                 streams.delete(stream);

test/client.test.ts

@@ -10,4 +10,39 @@ describe('Client', () => {
             });
         }).rejects.toBeInstanceOf(StreamClosedError);
     });
+
+    it('should time out when exceeding auth timeout', async () => {
+        let timeout;
+
+        await expect(async () => {
+            await testServer({}, async stream => {
+                stream.writeUint8(255);
+                stream.writeUint8(0);
+
+                await new Promise(resolve => {
+                    timeout = setTimeout(resolve, 75);
+                });
+
+                stream.writeUint8(1);
+            }, false, 50);
+        }).rejects.toBeInstanceOf(StreamClosedError);
+
+        clearTimeout(timeout);
+    });
+
+    it('should time out when exceeding idle timeout', async () => {
+        let timeout;
+
+        await expect(async () => {
+            await testServer({}, async stream => {
+                await new Promise(resolve => {
+                    timeout = setTimeout(resolve, 75);
+                });
+
+                stream.writeUint8(1);
+            }, true, undefined, 50);
+        }).rejects.toBeInstanceOf(StreamClosedError);
+
+        clearTimeout(timeout);
+    });
 });

test/server-tester.ts

@@ -16,9 +16,17 @@ type Tester = (socket : Stream) => Promise<void>;
 
 const startServerResults = new Set<StartServerResult>();
 
-export const testServer = async (processor : Partial<Processor>, tester : Tester, autoAuth = true) : Promise<void> => {
+export const testServer = async (
+    processor : Partial<Processor>,
+    tester : Tester,
+    autoAuth = true,
+    authTimeout ?: number,
+    idleTimeout ?: number,
+) : Promise<void> => {
     const startServerResult = await startServer({
         tls: tlsOptions,
+        authTimeout,
+        idleTimeout,
         processor: {
             authenticate: processor.authenticate ?? (() => true),
             checkUid: processor.checkUid