Merge branch 'connection-timeout' of https://github.com/BenHall/docke…

…r-modem into BenHall-connection-timeout
apocas · Sep 20, 2019 · 36b3522 · 36b3522
2 parents 4ea0bbb + 513d8a1
commit 36b3522
Showing 1 changed file with 21 additions and 0 deletions.
diff --git a/lib/modem.js b/lib/modem.js
@@ -67,6 +67,7 @@ var Modem = function(options) {
   this.cert = opts.cert;
   this.ca = opts.ca;
   this.timeout = opts.timeout;
+  this.connectionTimeout = opts.connectionTimeout;
   this.checkServerIdentity = opts.checkServerIdentity;
   this.agent = opts.agent;
   this.headers = opts.headers || {};
@@ -188,6 +189,8 @@ Modem.prototype.dial = function(options, callback) {
 
 Modem.prototype.buildRequest = function(options, context, data, callback) {
   var self = this;
+  var connectionTimeoutTimer;
+
   var opts = self.protocol === 'ssh' ? Object.assign(options, {
     agent: ssh({'host': self.host, 'port': self.port}),
     protocol: 'http'
@@ -200,6 +203,13 @@ Modem.prototype.buildRequest = function(options, context, data, callback) {
     depth: null
   }));
 
+  if (self.connectionTimeout) {
+    connectionTimeoutTimer = setTimeout(function() {
+      debug('Connection Timeout of %s ms exceeded', self.connectionTimeout);
+      req.abort();
+    }, self.connectionTimeout);
+  }
+
   if (self.timeout) {
     req.on('socket', function(socket) {
       socket.setTimeout(self.timeout);
@@ -211,12 +221,22 @@ Modem.prototype.buildRequest = function(options, context, data, callback) {
   }
 
   if (context.hijack === true) {
+    clearTimeout(connectionTimeoutTimer);
     req.on('upgrade', function(res, sock, head) {
       return callback(null, sock);
     });
   }
 
+  req.on('connect', function() {
+    clearTimeout(connectionTimeoutTimer);
+  });
+
+  req.on('disconnect', function() {
+    clearTimeout(connectionTimeoutTimer);
+  });
+
   req.on('response', function(res) {
+    clearTimeout(connectionTimeoutTimer);
     if (context.isStream === true) {
       self.buildPayload(null, context.isStream, context.statusCodes, context.openStdin, req, res, null, callback);
     } else {
@@ -238,6 +258,7 @@ Modem.prototype.buildRequest = function(options, context, data, callback) {
   });
 
   req.on('error', function(error) {
+    clearTimeout(connectionTimeoutTimer);
     self.buildPayload(error, context.isStream, context.statusCodes, false, {}, {}, null, callback);
   });