fix: close streams when connection is closed (#214)

This adds tests for closeRead and closeWrite on the muxer streams.

Also:

- Connection.close() will now close its internal streams to avoid lingering streams.

Supersedes #90

BREAKING CHANGE: This adds closeWrite and closeRead checks in the tests, which will cause test failures for muxers that don't implement those

packages/libp2p-connection/package.json

@@ -156,6 +156,7 @@
   },
   "dependencies": {
     "@libp2p/interfaces": "^1.3.0",
+    "@libp2p/logger": "^1.1.0",
     "@multiformats/multiaddr": "^10.1.5",
     "err-code": "^3.0.1"
   },

packages/libp2p-connection/src/index.ts

@@ -4,6 +4,9 @@ import { OPEN, CLOSING, CLOSED } from '@libp2p/interfaces/connection/status'
 import { symbol } from '@libp2p/interfaces/connection'
 import type { Connection, ConnectionStat, Metadata, ProtocolStream, Stream } from '@libp2p/interfaces/connection'
 import type { PeerId } from '@libp2p/interfaces/peer-id'
+import { logger } from '@libp2p/logger'
+
+const log = logger('libp2p:connection')
 
 interface ConnectionInit {
   remoteAddr: Multiaddr
@@ -150,6 +153,17 @@ export class ConnectionImpl implements Connection {
 
     this.stat.status = CLOSING
 
+    // close all streams - this can throw if we're not multiplexed
+    try {
+      await Promise.all(
+        this.streams.map(async s => await s.close().catch(err => {
+          log.error(err)
+        }))
+      )
+    } catch (err) {
+      log.error(err)
+    }
+
     // Close raw connection
     this._closing = true
     await this._close()

packages/libp2p-connection/test/compliance.spec.ts

@@ -35,12 +35,13 @@ describe('compliance tests', () => {
           const id = `${streamId++}`
           const stream: Stream = {
             ...pair(),
-            close: () => {
-              void stream.sink(async function * () {}())
-                .then(() => {
-                  connection.removeStream(stream.id)
-                })
-                .catch()
+            close: async () => {
+              await stream.sink(async function * () {}())
+              connection.removeStream(stream.id)
+            },
+            closeRead: async () => {},
+            closeWrite: async () => {
+              await stream.sink(async function * () {}())
             },
             id,
             abort: () => {},

packages/libp2p-connection/test/index.spec.ts

@@ -54,8 +54,12 @@ describe('connection tests', () => {
         const id = `${streamId++}`
         const stream: Stream = {
           ...pair<Uint8Array>(),
-          close: () => {
-            void stream.sink(async function * () {}()).catch()
+          close: async () => {
+            await stream.sink(async function * () {}()).catch()
+          },
+          closeRead: async () => {},
+          closeWrite: async () => {
+            await stream.sink(async function * () {}())
           },
           id,
           abort: () => {},

packages/libp2p-interface-compliance-tests/src/mocks/connection.ts

@@ -169,7 +169,9 @@ export function mockConnection (maConn: MultiaddrConnection, opts: MockConnectio
 export function mockStream (stream: Duplex<Uint8Array>): Stream {
   return {
     ...stream,
-    close: () => {},
+    close: async () => {},
+    closeRead: async () => {},
+    closeWrite: async () => {},
     abort: () => {},
     reset: () => {},
     timeline: {

packages/libp2p-interface-compliance-tests/src/mocks/muxer.ts

@@ -191,7 +191,15 @@ class MuxedStream {
       source: this.input,
 
       // Close for reading
-      close: () => {
+      close: async () => {
+        this.input.end()
+      },
+
+      closeRead: async () => {
+        this.input.end()
+      },
+
+      closeWrite: async () => {
         this.input.end()
       },
 
@@ -242,7 +250,11 @@ class MockMuxer implements StreamMuxer {
       onEnd: (err) => {
         this.log('closing muxed streams')
         for (const stream of this.streams) {
-          stream.abort(err)
+          if (err == null) {
+            void stream.close().catch()
+          } else {
+            stream.abort(err)
+          }
         }
       }
     })
@@ -307,7 +319,7 @@ class MockMuxer implements StreamMuxer {
       muxedStream.stream.reset()
     } else if (message.type === 'close') {
       this.log('-> closing stream %s %s', muxedStream.type, muxedStream.stream.id)
-      muxedStream.stream.close()
+      void muxedStream.stream.close()
     }
   }
 

packages/libp2p-interface-compliance-tests/src/stream-muxer/close-test.ts

@@ -9,6 +9,8 @@ import delay from 'delay'
 import type { TestSetup } from '../index.js'
 import type { StreamMuxerFactory } from '@libp2p/interfaces/stream-muxer'
 import { Components } from '@libp2p/interfaces/components'
+import pDefer from 'p-defer'
+import all from 'it-all'
 
 function randomBuffer () {
   return uint8ArrayFromString(Math.random().toString())
@@ -113,5 +115,113 @@ export default (common: TestSetup<StreamMuxerFactory>) => {
       // These should now all resolve without error
       await Promise.all(streamResults)
     })
+
+    it('can close a stream for writing', async () => {
+      const deferred = pDefer<any>()
+
+      const p = duplexPair<Uint8Array>()
+      const dialerFactory = await common.setup()
+      const dialer = dialerFactory.createStreamMuxer(new Components())
+      const data = [randomBuffer(), randomBuffer()]
+
+      const listenerFactory = await common.setup()
+      const listener = listenerFactory.createStreamMuxer(new Components(), {
+        onIncomingStream: (stream) => {
+          void Promise.resolve().then(async () => {
+            // Immediate close for write
+            await stream.closeWrite()
+
+            const results = await pipe(stream, async (source) => {
+              const data = []
+              for await (const chunk of source) {
+                data.push(chunk.slice())
+              }
+              return data
+            })
+            expect(results).to.eql(data)
+
+            try {
+              await stream.sink([randomBuffer()])
+            } catch (err) {
+              deferred.resolve(err)
+            }
+
+            deferred.reject(new Error('should not support writing to closed writer'))
+          })
+        }
+      })
+
+      void pipe(p[0], dialer, p[0])
+      void pipe(p[1], listener, p[1])
+
+      const stream = dialer.newStream()
+      await stream.sink(data)
+
+      const err = await deferred.promise
+      expect(err).to.have.property('message').that.matches(/stream closed for writing/)
+    })
+
+    it('can close a stream for reading', async () => {
+      const deferred = pDefer<any>()
+
+      const p = duplexPair<Uint8Array>()
+      const dialerFactory = await common.setup()
+      const dialer = dialerFactory.createStreamMuxer(new Components())
+      const data = [randomBuffer(), randomBuffer()]
+
+      const listenerFactory = await common.setup()
+      const listener = listenerFactory.createStreamMuxer(new Components(), {
+        onIncomingStream: (stream) => {
+          void all(stream.source).then(deferred.resolve, deferred.reject)
+        }
+      })
+
+      void pipe(p[0], dialer, p[0])
+      void pipe(p[1], listener, p[1])
+
+      const stream = dialer.newStream()
+      await stream.closeRead()
+
+      // Source should be done
+      void Promise.resolve().then(async () => {
+        // @ts-expect-error next is part of the iterable protocol
+        expect(await stream.source.next()).to.have.property('done', true)
+        await stream.sink(data)
+      })
+
+      const results = await deferred.promise
+      expect(results).to.eql(data)
+    })
+
+    it('calls onStreamEnd for closed streams not previously written', async () => {
+      const deferred = pDefer()
+
+      const onStreamEnd = () => deferred.resolve()
+      const dialerFactory = await common.setup()
+      const dialer = dialerFactory.createStreamMuxer(new Components(), {
+        onStreamEnd
+      })
+
+      const stream = await dialer.newStream()
+
+      await stream.close()
+      await deferred.promise
+    })
+
+    it('calls onStreamEnd for read and write closed streams not previously written', async () => {
+      const deferred = pDefer()
+
+      const onStreamEnd = () => deferred.resolve()
+      const dialerFactory = await common.setup()
+      const dialer = dialerFactory.createStreamMuxer(new Components(), {
+        onStreamEnd
+      })
+
+      const stream = await dialer.newStream()
+
+      await stream.closeWrite()
+      await stream.closeRead()
+      await deferred.promise
+    })
   })
 }

packages/libp2p-interfaces/src/connection/index.ts

@@ -31,10 +31,39 @@ export interface Metadata {
  * configuration of the nodes.
  */
 export interface Stream extends Duplex<Uint8Array> {
-  close: () => void
-  abort: (err?: Error) => void
+  /**
+   * Close a stream for reading and writing
+   */
+  close: () => Promise<void>
+
+  /**
+   * Close a stream for reading only
+   */
+  closeRead: () => Promise<void>
+
+  /**
+   * Close a stream for writing only
+   */
+  closeWrite: () => Promise<void>
+
+  /**
+   * Call when a local error occurs, should close the stream for reading and writing
+   */
+  abort: (err: Error) => void
+
+  /**
+   * Call when a remote error occurs, should close the stream for reading and writing
+   */
   reset: () => void
+
+  /**
+   * Records stream lifecycle event timings
+   */
   timeline: Timeline
+
+  /**
+   * Unique identifier for a stream
+   */
   id: string
 }