bring back connection cleanup

Author: dignifiedquire

src/index.js

@@ -12,19 +12,19 @@ const log = debug('libp2p:tcp')
 const createListener = require('./listener')
 
 module.exports = class TCP {
-  dial (ma, options, callback) {
+  dial (ma, options, cb) {
     if (isFunction(options)) {
-      callback = options
+      cb = options
       options = {}
     }
 
-    if (!callback) {
-      callback = () => {}
+    if (!cb) {
+      cb = () => {}
     }
 
     const cOpts = ma.toOptions()
     log('Connecting to %s %s', cOpts.port, cOpts.host)
-    const socket = toPull.duplex(net.connect(cOpts, callback))
+    const socket = toPull.duplex(net.connect(cOpts, cb))
 
     socket.getObservedAddrs = (cb) => {
       return cb(null, [ma])

src/listener.js

@@ -12,6 +12,7 @@ const toPull = require('stream-to-pull-stream')
 const getMultiaddr = require('./get-multiaddr')
 
 const IPFS_CODE = 421
+const CLOSE_TIMEOUT = 2000
 
 module.exports = (handler) => {
   const listener = net.createServer((socket) => {
@@ -24,7 +25,41 @@ module.exports = (handler) => {
     handler(conn)
   })
 
+  // Keep track of open connections to destroy in case of timeout
   listener.__connections = {}
+  listener.on('connection', (socket) => {
+    const key = `${socket.remoteAddress}:${socket.remotePort}`
+    listener.__connections[key] = socket
+
+    socket.on('close', () => {
+      delete listener.__connections[key]
+    })
+  })
+
+  listener._close = listener.close
+  listener.close = (options, cb) => {
+    if (typeof options === 'function') {
+      cb = options
+      options = {}
+    }
+    cb = cb || (() => {})
+    options = options || {}
+
+    let closed = false
+    listener._close(cb)
+    listener.once('close', () => {
+      closed = true
+    })
+    setTimeout(() => {
+      if (closed) return
+
+      log('unable to close graciously, destroying conns')
+      Object.keys(listener.__connections).forEach((key) => {
+        log('destroying %s', key)
+        listener.__connections[key].destroy()
+      })
+    }, options.timeout || CLOSE_TIMEOUT)
+  }
   let ipfsId
   let listeningAddr
 
@@ -41,10 +76,14 @@ module.exports = (handler) => {
     return listener._listen(lOpts.port, lOpts.host, cb)
   }
 
-  listener.getAddrs = (callback) => {
+  listener.getAddrs = (cb) => {
     const multiaddrs = []
     const address = listener.address()
 
+    if (!address) {
+      return cb(new Error('Listener is not ready yet'))
+    }
+
     // Because TCP will only return the IPv6 version
     // we need to capture from the passed multiaddr
     if (listeningAddr.toString().indexOf('ip4') !== -1) {
@@ -77,7 +116,7 @@ module.exports = (handler) => {
       multiaddrs.push(ma)
     }
 
-    callback(null, multiaddrs)
+    cb(null, multiaddrs)
   }
 
   return listener