From 85217346c0a0c6713ca56cae6131f05e3bc097d2 Mon Sep 17 00:00:00 2001
From: Zachary Boyd <zachary@zacharyboyd.nyc>
Date: Sun, 26 Feb 2017 16:07:47 -0500
Subject: [PATCH] wraps connections in a domain

---
 src/SOCKSServer.js | 66 ++++++++++++++++++++++++++++------------------
 1 file changed, 41 insertions(+), 25 deletions(-)

diff --git a/src/SOCKSServer.js b/src/SOCKSServer.js
index ef19e2b..71835bf 100644
--- a/src/SOCKSServer.js
+++ b/src/SOCKSServer.js
@@ -1,23 +1,34 @@
 const socks = require('socksv5');
 const SOCKS5Server = socks.Server;
+const domain = require('domain');
 
 class SOCKSServer extends SOCKS5Server{
 	constructor(tor_pool, logger) {
-		super((info, accept, deny) => {
+		let handleConnection = (info, accept, deny) => {
+			let d = domain.create();
+
 			let inbound_socket = accept(true);
+			d.add(inbound_socket);
 			var outbound_socket;
 			let buffer = [];
 
 			let socks_port = (tor_pool.next().socks_port);
 			logger && logger.info(`[socks]: ${info.srcAddr}:${info.srcPort} → 127.0.0.1:${socks_port} → ${info.dstAddr}:${info.dstPort}`)
 
-			let onClose = () => {
+			let onClose = (error) => {
 				inbound_socket && inbound_socket.end();
 				outbound_socket && outbound_socket.end();
 
 				inbound_socket = outbound_socket = buffer = void(0);
+
+				if (error)
+					this.logger.error(`[socks]: an error occured: ${error.message}`)
+
+				d.exit();
 			};
 
+			d.on('error', onClose);
+
 			let onInboundData = (data) => buffer.push(data)
 			
 			if (!inbound_socket) return;
@@ -26,33 +37,38 @@ class SOCKSServer extends SOCKS5Server{
 			inbound_socket.on('data', onInboundData);	
 			inbound_socket.on('error', onClose);
 
-			socks.connect({
-				host: info.dstAddr,
-				port: info.dstPort,
-				proxyHost: '127.0.0.1',
-				proxyPort: socks_port,
-				localDNS: false,
-				auths: [ socks.auth.None() ]
-			}, ($outbound_socket) => {
-				outbound_socket = $outbound_socket;
-				outbound_socket && outbound_socket.on('close', onClose);
+			d.run(() => {
+				socks.connect({
+					host: info.dstAddr,
+					port: info.dstPort,
+					proxyHost: '127.0.0.1',
+					proxyPort: socks_port,
+					localDNS: false,
+					auths: [ socks.auth.None() ]
+				}, ($outbound_socket) => {
+					outbound_socket = $outbound_socket;
+					d.add(outbound_socket);
+					outbound_socket && outbound_socket.on('close', onClose);
 
-				inbound_socket && inbound_socket.removeListener('data', onInboundData);
-				inbound_socket &&  inbound_socket.on('data', (data) => {
-					outbound_socket && outbound_socket.write(data);
-				});
+					inbound_socket && inbound_socket.removeListener('data', onInboundData);
+					inbound_socket &&  inbound_socket.on('data', (data) => {
+						outbound_socket && outbound_socket.write(data);
+					});
 
-				outbound_socket && outbound_socket.on('data', (data) => {
-					inbound_socket && inbound_socket.write(data);
-				});
+					outbound_socket && outbound_socket.on('data', (data) => {
+						inbound_socket && inbound_socket.write(data);
+					});
 
-				outbound_socket && outbound_socket.on('error', onClose);
+					outbound_socket && outbound_socket.on('error', onClose);
 
-				while (buffer && buffer.length && outbound_socket) {
-					outbound_socket.write(buffer.shift());
-				}
-			})
-		});
+					while (buffer && buffer.length && outbound_socket) {
+						outbound_socket.write(buffer.shift());
+					}
+				})
+			});
+		};
+
+		super(handleConnection);
 
 		this.logger = logger;