fix: allow mss lazy select on read

achingbrain · achingbrain · commit db5925b982b2 · 2023-11-20T12:33:30.000Z
Updates the multistream lazy select to support lazy select on streams
that are only read from.
diff --git a/packages/multistream-select/src/handle.ts b/packages/multistream-select/src/handle.ts
@@ -55,24 +55,29 @@ import type { Duplex } from 'it-stream-types'
  */
 export async function handle <Stream extends Duplex<any, any, any>> (stream: Stream, protocols: string | string[], options: MultistreamSelectInit): Promise<ProtocolStream<Stream>> {
   protocols = Array.isArray(protocols) ? protocols : [protocols]
+  options.log.trace('available protocols %s', protocols)
+
   const lp = lpStream(stream, {
-    maxDataLength: MAX_PROTOCOL_LENGTH
+    maxDataLength: MAX_PROTOCOL_LENGTH,
+    maxLengthLength: 2 // 2 bytes is enough to length-prefix MAX_PROTOCOL_LENGTH
   })
 
   while (true) {
-    options.log.trace('handle - available protocols %s', protocols)
+    options?.log.trace('reading incoming string')
     const protocol = await multistream.readString(lp, options)
     options.log.trace('read "%s"', protocol)
 
     if (protocol === PROTOCOL_ID) {
       options.log.trace('respond with "%s" for "%s"', PROTOCOL_ID, protocol)
       await multistream.write(lp, uint8ArrayFromString(`${PROTOCOL_ID}\n`), options)
+      options.log.trace('responded with "%s" for "%s"', PROTOCOL_ID, protocol)
       continue
     }
 
     if (protocols.includes(protocol)) {
       options.log.trace('respond with "%s" for "%s"', protocol, protocol)
       await multistream.write(lp, uint8ArrayFromString(`${protocol}\n`), options)
+      options.log.trace('responded with "%s" for "%s"', protocol, protocol)
 
       return { stream: lp.unwrap(), protocol }
     }
@@ -84,8 +89,9 @@ export async function handle <Stream extends Duplex<any, any, any>> (stream: Str
         uint8ArrayFromString('\n')
       )
 
-      await multistream.write(lp, protos, options)
       options.log.trace('respond with "%s" for %s', protocols, protocol)
+      await multistream.write(lp, protos, options)
+      options.log.trace('responded with "%s" for %s', protocols, protocol)
       continue
     }
 
diff --git a/packages/multistream-select/src/select.ts b/packages/multistream-select/src/select.ts
@@ -1,5 +1,7 @@
 import { CodeError } from '@libp2p/interface/errors'
 import { lpStream } from 'it-length-prefixed-stream'
+import * as varint from 'uint8-varint'
+import { Uint8ArrayList } from 'uint8arraylist'
 import { fromString as uint8ArrayFromString } from 'uint8arrays/from-string'
 import { MAX_PROTOCOL_LENGTH } from './constants.js'
 import * as multistream from './multistream.js'
@@ -109,6 +111,7 @@ export async function select <Stream extends Duplex<any, any, any>> (stream: Str
 export function lazySelect <Stream extends Duplex<any, any, any>> (stream: Stream, protocol: string, options: MultistreamSelectInit): ProtocolStream<Stream> {
   const originalSink = stream.sink.bind(stream)
   const originalSource = stream.source
+  let selected = false
 
   const lp = lpStream({
     sink: originalSink,
@@ -118,20 +121,47 @@ export function lazySelect <Stream extends Duplex<any, any, any>> (stream: Strea
   })
 
   stream.sink = async source => {
-    options?.log.trace('lazy: write ["%s", "%s"]', PROTOCOL_ID, protocol)
-
-    await lp.writeV([
-      uint8ArrayFromString(`${PROTOCOL_ID}\n`),
-      uint8ArrayFromString(`${protocol}\n`)
-    ])
-
-    options?.log.trace('lazy: writing rest of "%s" stream', protocol)
-    await lp.unwrap().sink(source)
+    const { sink } = lp.unwrap()
+
+    await sink(async function * () {
+      for await (const buf of source) {
+        // if writing before selecting, send selection with first data chunk
+        if (!selected) {
+          selected = true
+          options?.log.trace('lazy: write ["%s", "%s", data] in sink', PROTOCOL_ID, protocol)
+
+          const protocolString = `${protocol}\n`
+
+          // send protocols in first chunk of data written to transport
+          yield new Uint8ArrayList(
+            Uint8Array.from([19]), // length of PROTOCOL_ID plus newline
+            uint8ArrayFromString(`${PROTOCOL_ID}\n`),
+            varint.encode(protocolString.length),
+            uint8ArrayFromString(protocolString),
+            buf
+          ).subarray()
+
+          options?.log.trace('lazy: wrote ["%s", "%s", data] in sink', PROTOCOL_ID, protocol)
+        } else {
+          yield buf
+        }
+      }
+    }())
   }
 
   stream.source = (async function * () {
-    options?.log.trace('lazy: reading multistream select header')
+    // if reading before selecting, send selection before first data chunk
+    if (!selected) {
+      selected = true
+      options?.log.trace('lazy: write ["%s", "%s", data] in source', PROTOCOL_ID, protocol)
+      await lp.writeV([
+        uint8ArrayFromString(`${PROTOCOL_ID}\n`),
+        uint8ArrayFromString(`${protocol}\n`)
+      ])
+      options?.log.trace('lazy: wrote ["%s", "%s", data] in source', PROTOCOL_ID, protocol)
+    }
 
+    options?.log.trace('lazy: reading multistream select header')
     let response = await multistream.readString(lp, options)
     options?.log.trace('lazy: read multistream select header "%s"', response)
 
diff --git a/packages/multistream-select/test/integration.spec.ts b/packages/multistream-select/test/integration.spec.ts
@@ -113,6 +113,76 @@ describe('Dialer and Listener integration', () => {
     expect(new Uint8ArrayList(...dialerOut).slice()).to.eql(new Uint8ArrayList(...input).slice())
   })
 
+  it('should handle and lazySelect that fails', async () => {
+    const protocol = '/echo/1.0.0'
+    const otherProtocol = '/echo/2.0.0'
+    const pair = duplexPair<Uint8ArrayList | Uint8Array>()
+
+    const dialerSelection = mss.lazySelect(pair[0], protocol, {
+      log: logger('mss:test')
+    })
+    expect(dialerSelection.protocol).to.equal(protocol)
+
+    // the listener handles the incoming stream
+    void mss.handle(pair[1], otherProtocol, {
+      log: logger('mss:test')
+    })
+
+    // should fail when we interact with the stream
+    const input = [randomBytes(10), randomBytes(64), randomBytes(3)]
+    await expect(pipe(input, dialerSelection.stream, async source => all(source)))
+      .to.eventually.be.rejected.with.property('code', 'ERR_UNSUPPORTED_PROTOCOL')
+  })
+
+  it('should handle and lazySelect only by reading', async () => {
+    const protocol = '/echo/1.0.0'
+    const pair = duplexPair<Uint8ArrayList | Uint8Array>()
+
+    const dialerSelection = mss.lazySelect(pair[0], protocol, {
+      log: logger('mss:dialer')
+    })
+    expect(dialerSelection.protocol).to.equal(protocol)
+
+    // ensure stream is usable after selection
+    const input = [randomBytes(10), randomBytes(64), randomBytes(3)]
+
+    const [, dialerOut] = await Promise.all([
+      // the listener handles the incoming stream
+      mss.handle(pair[1], protocol, {
+        log: logger('mss:listener')
+      }).then(async result => {
+        // the listener writes to the incoming stream
+        await pipe(input, result.stream)
+      }),
+
+      // the dialer just reads from the stream
+      pipe(dialerSelection.stream, async source => all(source))
+    ])
+
+    expect(new Uint8ArrayList(...dialerOut).slice()).to.eql(new Uint8ArrayList(...input).slice())
+  })
+
+  it('should handle and lazySelect only by reading that fails', async () => {
+    const protocol = '/echo/1.0.0'
+    const otherProtocol = '/echo/2.0.0'
+    const pair = duplexPair<Uint8ArrayList | Uint8Array>()
+
+    // lazy succeeds
+    const dialerSelection = mss.lazySelect(pair[0], protocol, {
+      log: logger('mss:dialer')
+    })
+    expect(dialerSelection.protocol).to.equal(protocol)
+
+    // the listener handles the incoming stream
+    void mss.handle(pair[1], otherProtocol, {
+      log: logger('mss:listener')
+    })
+
+    // should fail when we interact with the stream
+    await expect(pipe(dialerSelection.stream, async source => all(source)))
+      .to.eventually.be.rejected.with.property('code', 'ERR_UNSUPPORTED_PROTOCOL')
+  })
+
   it('should abort an unhandled lazySelect', async () => {
     const protocol = '/echo/1.0.0'
     const pair = duplexPair<Uint8ArrayList | Uint8Array>()
