@libp2p/utils 4.0.7-9ad8f8686 → 4.0.7-adea7bbbf

package/src/stream-to-ma-conn.ts

@@ -1,4 +1,4 @@
-import type { AbortOptions, ComponentLogger } from '@libp2p/interface'
+import type { ComponentLogger } from '@libp2p/interface'
 import type { MultiaddrConnection, Stream } from '@libp2p/interface/connection'
 import type { Multiaddr } from '@multiformats/multiaddr'
 
@@ -15,49 +15,74 @@ export interface StreamProperties {
  */
 export function streamToMaConnection (props: StreamProperties): MultiaddrConnection {
   const { stream, remoteAddr, logger } = props
-  const { sink, source } = stream
   const log = logger.forComponent('libp2p:stream:converter')
 
-  const mapSource = (async function * () {
-    for await (const list of source) {
-      if (list instanceof Uint8Array) {
-        yield list
-      } else {
-        yield * list
+  let closedRead = false
+  let closedWrite = false
+
+  // piggyback on `stream.close` invocations to close maconn
+  const streamClose = stream.close.bind(stream)
+  stream.close = async (options) => {
+    await streamClose(options)
+    close(true)
+  }
+
+  // piggyback on `stream.abort` invocations to close maconn
+  const streamAbort = stream.abort.bind(stream)
+  stream.abort = (err) => {
+    streamAbort(err)
+    close(true)
+  }
+
+  // piggyback on `stream.sink` invocations to close maconn
+  const streamSink = stream.sink.bind(stream)
+  stream.sink = async (source) => {
+    try {
+      await streamSink(source)
+    } catch (err: any) {
+      // If aborted we can safely ignore
+      if (err.type !== 'aborted') {
+        // If the source errored the socket will already have been destroyed by
+        // toIterable.duplex(). If the socket errored it will already be
+        // destroyed. There's nothing to do here except log the error & return.
+        log(err)
       }
+    } finally {
+      closedWrite = true
+      close()
     }
-  }())
+  }
 
   const maConn: MultiaddrConnection = {
-    async sink (source) {
+    log,
+    sink: stream.sink,
+    source: (async function * () {
       try {
-        await sink(source)
-        close()
-      } catch (err: any) {
-        // If aborted we can safely ignore
-        if (err.type !== 'aborted') {
-          // If the source errored the socket will already have been destroyed by
-          // toIterable.duplex(). If the socket errored it will already be
-          // destroyed. There's nothing to do here except log the error & return.
-          log(err)
+        for await (const list of stream.source) {
+          if (list instanceof Uint8Array) {
+            yield list
+          } else {
+            yield * list
+          }
         }
+      } finally {
+        closedRead = true
+        close()
       }
-    },
-    source: mapSource,
+    }()),
     remoteAddr,
     timeline: { open: Date.now(), close: undefined },
-    async close (options?: AbortOptions) {
-      close()
-      await stream.close(options)
-    },
-    abort (err: Error): void {
-      close()
-      stream.abort(err)
-    }
+    close: stream.close,
+    abort: stream.abort
   }
 
-  function close (): void {
-    if (maConn.timeline.close == null) {
+  function close (force?: boolean): void {
+    if (force === true) {
+      closedRead = true
+      closedWrite = true
+    }
+
+    if (closedRead && closedWrite && maConn.timeline.close == null) {
       maConn.timeline.close = Date.now()
     }
   }