Wait for connection status on add peer command (#4448)

iron-fish · Dec 14, 2023 · 39da0f7 · 39da0f7
1 parent d6e44e6
commit 39da0f7
Show file tree

Hide file tree

Showing 3 changed files with 145 additions and 7 deletions.
diff --git a/ironfish-cli/src/commands/peers/add.ts b/ironfish-cli/src/commands/peers/add.ts
@@ -53,6 +53,11 @@ export class AddCommand extends IronfishCommand {
 
     if (response.content.added) {
       this.log(`Successfully added peer ${request.host}:${request.port}`)
+    } else if (response.content.error !== undefined) {
+      this.log(
+        `Failed to add peer ${request.host}:${request.port} because: ${response.content.error}`,
+      )
+      this.exit(0)
     } else {
       this.log(`Could not add peer ${request.host}:${request.port}`)
       this.exit(0)

diff --git a/ironfish/src/rpc/routes/peer/addPeer.test.ts b/ironfish/src/rpc/routes/peer/addPeer.test.ts
@@ -2,28 +2,128 @@
  * License, v. 2.0. If a copy of the MPL was not distributed with this
  * file, You can obtain one at https://mozilla.org/MPL/2.0/. */
 import { formatWebSocketAddress } from '../../../network'
+import { WebSocketConnection } from '../../../network/peers/connections'
+import { mockIdentity } from '../../../network/testUtilities'
 import { createRouteTest } from '../../../testUtilities/routeTest'
 
+jest.mock('ws')
+
 describe('Route peer/addPeer', () => {
   const routeTest = createRouteTest()
 
   it('should add a peer with a correct address and port', async () => {
     const request = { host: 'testhost', port: 9037 }
+    const identity = mockIdentity('peer')
+
+    const req = await routeTest.client.request('peer/addPeer', request).waitForRoute()
+
+    const matchingPeers = routeTest.peerNetwork.peerManager.peers.filter(
+      (p) => formatWebSocketAddress(p.wsAddress) === 'ws://testhost:9037',
+    )
+
+    expect(matchingPeers.length).toBe(1)
+    expect(routeTest.peerNetwork.peerManager.peerCandidates.has(identity)).toBe(false)
+
+    const peer = matchingPeers[0]
+
+    let connection: WebSocketConnection
+    if (peer.state.type === 'CONNECTING' && peer.state.connections.webSocket) {
+      connection = peer.state.connections.webSocket
+    } else {
+      throw new Error('Peer should be CONNECTING with a WS connection')
+    }
+
+    connection.setState({
+      type: 'CONNECTED',
+      identity,
+    })
 
-    const response = await routeTest.client.request('peer/addPeer', request).waitForEnd()
+    const response = await req.waitForEnd()
 
+    expect(response.content).toMatchObject({
+      added: true,
+    })
     expect(
-      routeTest.node.peerNetwork.peerManager.peerCandidates.has('ws://testhost:9037'),
-    ).toBe(true)
+      routeTest.peerNetwork.peerManager
+        .getConnectedPeers()
+        .filter((p) => p.state.identity === identity),
+    ).toHaveLength(1)
+    expect(routeTest.peerNetwork.peerManager.peerCandidates.has(identity)).toBe(true)
+  })
+
+  it('should return false if the peer closes without an error', async () => {
+    const request = { host: 'testhost', port: 9037 }
+    const identity = mockIdentity('peer')
+
+    const req = await routeTest.client.request('peer/addPeer', request).waitForRoute()
 
-    const matchingPeers = routeTest.node.peerNetwork.peerManager.peers.filter(
+    const matchingPeers = routeTest.peerNetwork.peerManager.peers.filter(
       (p) => formatWebSocketAddress(p.wsAddress) === 'ws://testhost:9037',
     )
 
     expect(matchingPeers.length).toBe(1)
+    expect(routeTest.peerNetwork.peerManager.peerCandidates.has(identity)).toBe(false)
+
+    const peer = matchingPeers[0]
+
+    let connection: WebSocketConnection
+    if (peer.state.type === 'CONNECTING' && peer.state.connections.webSocket) {
+      connection = peer.state.connections.webSocket
+    } else {
+      throw new Error('Peer should be CONNECTING with a WS connection')
+    }
+
+    connection.close()
+
+    const response = await req.waitForEnd()
 
     expect(response.content).toMatchObject({
-      added: true,
+      added: false,
+      error: undefined,
+    })
+    expect(
+      routeTest.peerNetwork.peerManager
+        .getConnectedPeers()
+        .filter((p) => p.state.identity === identity),
+    ).toHaveLength(0)
+    expect(routeTest.peerNetwork.peerManager.peerCandidates.has(identity)).toBe(false)
+  })
+
+  it('should return false if the peer closes with an error', async () => {
+    const request = { host: 'testhost', port: 9037 }
+    const identity = mockIdentity('peer')
+
+    const req = await routeTest.client.request('peer/addPeer', request).waitForRoute()
+
+    const matchingPeers = routeTest.peerNetwork.peerManager.peers.filter(
+      (p) => formatWebSocketAddress(p.wsAddress) === 'ws://testhost:9037',
+    )
+
+    expect(matchingPeers.length).toBe(1)
+    expect(routeTest.peerNetwork.peerManager.peerCandidates.has(identity)).toBe(false)
+
+    const peer = matchingPeers[0]
+
+    let connection: WebSocketConnection
+    if (peer.state.type === 'CONNECTING' && peer.state.connections.webSocket) {
+      connection = peer.state.connections.webSocket
+    } else {
+      throw new Error('Peer should be CONNECTING with a WS connection')
+    }
+
+    connection.close(new Error('foo'))
+
+    const response = await req.waitForEnd()
+
+    expect(response.content).toMatchObject({
+      added: false,
+      error: 'foo',
     })
+    expect(
+      routeTest.peerNetwork.peerManager
+        .getConnectedPeers()
+        .filter((p) => p.state.identity === identity),
+    ).toHaveLength(0)
+    expect(routeTest.peerNetwork.peerManager.peerCandidates.has(identity)).toBe(false)
   })
 })
diff --git a/ironfish/src/rpc/routes/peer/addPeer.ts b/ironfish/src/rpc/routes/peer/addPeer.ts
@@ -4,7 +4,10 @@
 import * as yup from 'yup'
 import { Assert } from '../../../assert'
 import { DEFAULT_WEBSOCKET_PORT } from '../../../fileStores/config'
+import { Peer } from '../../../network'
+import { PeerState } from '../../../network/peers/peer'
 import { FullNode } from '../../../node'
+import { ErrorUtils } from '../../../utils'
 import { ApiNamespace } from '../namespaces'
 import { routes } from '../router'
 
@@ -16,6 +19,7 @@ export type AddPeerRequest = {
 
 export type AddPeerResponse = {
   added: boolean
+  error?: string
 }
 
 export const AddPeerRequestSchema: yup.ObjectSchema<AddPeerRequest> = yup
@@ -41,14 +45,43 @@ routes.register<typeof AddPeerRequestSchema, AddPeerResponse>(
 
     const peerManager = node.peerNetwork.peerManager
     const { host, port, whitelist } = request.data
-
     const peer = peerManager.connectToWebSocketAddress({
       host,
       port: port || DEFAULT_WEBSOCKET_PORT,
       whitelist: !!whitelist,
       forceConnect: true,
     })
 
-    request.end({ added: peer !== undefined })
+    if (peer === undefined) {
+      request.end({ added: false })
+      return
+    }
+
+    const onPeerStateChange = ({
+      peer,
+      state,
+      prevState,
+    }: {
+      peer: Peer
+      state: PeerState
+      prevState: PeerState
+    }) => {
+      if (prevState.type !== 'CONNECTED' && state.type === 'CONNECTED') {
+        request.end({ added: true })
+        peer.onStateChanged.off(onPeerStateChange)
+      } else if (prevState.type !== 'DISCONNECTED' && state.type === 'DISCONNECTED') {
+        request.end({
+          added: false,
+          error: peer.error ? ErrorUtils.renderError(peer.error) : undefined,
+        })
+        peer.onStateChanged.off(onPeerStateChange)
+      }
+    }
+
+    peer.onStateChanged.on(onPeerStateChange)
+
+    request.onClose.once(() => {
+      peer.onStateChanged.off(onPeerStateChange)
+    })
   },
 )