Fix relayed keepalive and direct-upgrade handoff

lib/connect.js

@@ -11,7 +11,7 @@ const Holepuncher = require('./holepuncher')
 const Sleeper = require('./sleeper')
 const {
   clearRelayTimeout,
-  closeRelayConnection,
+  confirmDirectUpgrade,
   destroyRelayConnection
 } = require('./relay-connection')
 const { FIREWALL, ERROR } = require('./constants')
@@ -102,6 +102,7 @@ module.exports = function connect(dht, publicKey, opts = {}) {
     relaySocket: null,
     relayClient: null,
     relayPaired: false,
+    validUpgrade: true,
     relayKeepAlive: opts.relayKeepAlive || 5000
   }
 
@@ -135,9 +136,11 @@ module.exports = function connect(dht, publicKey, opts = {}) {
     // Check if the traffic originated from the socket on which we're expecting relay traffic. If so,
     // we haven't hole punched yet and the other side is just sending us traffic through the relay.
     if (c.relaySocket && isRelay(c.relaySocket, socket, port, host)) {
+      c.validUpgrade = false
       return false
     }
 
+    c.validUpgrade = true
     if (c.onsocket) {
       c.onsocket(socket, port, host)
     } else {
@@ -496,12 +499,9 @@ async function connectThroughNode(c, address, socket) {
     if (c.rawStream === null) return // Already hole punched
 
     if (c.rawStream.connected) {
-      // Once this fires, changeRemote has moved the raw stream onto the direct path.
-      c.rawStream.once('remote-changed', () => closeRelayConnection(c))
-
-      const remoteChanging = c.rawStream.changeRemote(socket, c.connect.payload.udx.id, port, host)
-
-      if (remoteChanging) remoteChanging.catch(safetyCatch)
+      const rawStream = c.rawStream
+      const remoteChanging = rawStream.changeRemote(socket, c.connect.payload.udx.id, port, host)
+      confirmDirectUpgrade(c, rawStream, remoteChanging)
     } else {
       // cache the relay addrs for a future reconnect, we prefer the remote one so they
       // can give us the correct ones from their pov

lib/relay-connection.js

@@ -1,5 +1,9 @@
+const b4a = require('b4a')
+const safetyCatch = require('safety-catch')
+
 exports.clearRelayTimeout = clearRelayTimeout
 exports.closeRelayConnection = closeRelayConnection
+exports.confirmDirectUpgrade = confirmDirectUpgrade
 exports.destroyRelayConnection = destroyRelayConnection
 
 function clearRelayTimeout(connectionOrHandshake) {
@@ -13,6 +17,42 @@ function closeRelayConnection(connectionOrHandshake) {
   if (socket) socket.end()
 }
 
+function confirmDirectUpgrade(connectionOrHandshake, rawStream, remoteChanging) {
+  const cleanup = () => {
+    rawStream.off('data', ondirect)
+    rawStream.off('message', ondirect)
+    rawStream.off('close', cleanup)
+  }
+
+  function ondirect() {
+    if (!connectionOrHandshake.validUpgrade) {
+      // reset, aka assume from direct
+      connectionOrHandshake.validUpgrade = true
+      return
+    }
+
+    cleanup()
+    closeRelayConnection(connectionOrHandshake)
+  }
+
+  const confirm = () => {
+    rawStream.on('data', ondirect)
+    rawStream.on('message', ondirect)
+    rawStream.once('close', cleanup)
+
+    // Use a raw UDX message to make the peer observe the direct path without
+    // writing application data into the secret stream.
+    if (rawStream.trySend) rawStream.trySend(b4a.alloc(0))
+  }
+
+  if (!remoteChanging) {
+    confirm()
+    return
+  }
+
+  remoteChanging.then(confirm).catch(safetyCatch)
+}
+
 function destroyRelayConnection(connectionOrHandshake) {
   const socket = resetRelayConnection(connectionOrHandshake)
 

lib/server.js

@@ -11,7 +11,7 @@ const SecurePayload = require('./secure-payload')
 const Holepuncher = require('./holepuncher')
 const {
   clearRelayTimeout,
-  closeRelayConnection,
+  confirmDirectUpgrade,
   destroyRelayConnection
 } = require('./relay-connection')
 const { ALREADY_LISTENING, NODE_DESTROYED, KEYPAIR_ALREADY_USED } = require('./errors')
@@ -235,7 +235,8 @@ module.exports = class Server extends EventEmitter {
       relayToken: null,
       relaySocket: null,
       relayClient: null,
-      relayPaired: false
+      relayPaired: false,
+      validUpgrade: true
     }
 
     this._holepunches[id] = hs
@@ -291,10 +292,12 @@ module.exports = class Server extends EventEmitter {
           // Check if the traffic originated from the socket on which we're expecting relay traffic. If so,
           // we haven't hole punched yet and the other side is just sending us traffic through the relay.
           if (hs.relaySocket && isRelay(hs.relaySocket, socket, port, host)) {
+            hs.validUpgrade = false
             return false
           }
 
           hs.onsocket(socket, port, host)
+          hs.validUpgrade = true
           return false
         }
       })
@@ -328,12 +331,9 @@ module.exports = class Server extends EventEmitter {
         hs.rawStream.removeListener('close', onrawstreamclose)
 
         if (hs.rawStream.connected) {
-          // Once this fires, changeRemote has moved the raw stream onto the direct path.
-          hs.rawStream.once('remote-changed', () => closeRelayConnection(hs))
-
-          const remoteChanging = hs.rawStream.changeRemote(socket, remotePayload.udx.id, port, host)
-
-          if (remoteChanging) remoteChanging.catch(safetyCatch)
+          const rawStream = hs.rawStream
+          const remoteChanging = rawStream.changeRemote(socket, remotePayload.udx.id, port, host)
+          confirmDirectUpgrade(hs, rawStream, remoteChanging)
         } else {
           hs.rawStream.connect(socket, remotePayload.udx.id, port, host)
           hs.encryptedSocket = this.createSecretStream(false, hs.rawStream, {
@@ -664,7 +664,10 @@ module.exports = class Server extends EventEmitter {
           .on('close', () => destroyRelayConnection(hs))
           .connect(socket, remoteId, remotePort, remoteHost)
 
-        hs.encryptedSocket = this.createSecretStream(false, hs.rawStream, { handshake: h })
+        hs.encryptedSocket = this.createSecretStream(false, hs.rawStream, {
+          handshake: h,
+          keepAlive: this.dht.connectionKeepAlive
+        })
 
         this.onconnection(hs.encryptedSocket)
       })