Skip to content

Commit d8f5bc2

Browse files
achingbrainachingbrain
authored
fix: allow mss lazy select on read (#2246)
Updates the multistream lazy select to support lazy select on streams that are only read from.
1 parent 13a870c commit d8f5bc2

File tree

7 files changed

+135
-21
lines changed

7 files changed

+135
-21
lines changed

‎packages/libp2p/src/upgrader.ts

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -389,7 +389,8 @@ export class DefaultUpgrader implements Upgrader {
389389
.then(async () => {
390390
const protocols = this.components.registrar.getProtocols()
391391
const { stream, protocol } = await mss.handle(muxedStream, protocols, {
392-
log: muxedStream.log
392+
log: muxedStream.log,
393+
yieldBytes: false
393394
})
394395

395396
if (connection == null) {
@@ -458,7 +459,8 @@ export class DefaultUpgrader implements Upgrader {
458459

459460
const { stream, protocol } = await mss.select(muxedStream, protocols, {
460461
...options,
461-
log: muxedStream.log
462+
log: muxedStream.log,
463+
yieldBytes: false
462464
})
463465

464466
connection.log('negotiated protocol stream %s with id %s', protocol, muxedStream.id)

‎packages/multistream-select/package.json

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -58,6 +58,7 @@
5858
"it-length-prefixed-stream": "^1.1.1",
5959
"it-pipe": "^3.0.1",
6060
"it-stream-types": "^2.0.1",
61+
"uint8-varint": "^2.0.2",
6162
"uint8arraylist": "^2.4.3",
6263
"uint8arrays": "^4.0.6"
6364
},

‎packages/multistream-select/src/handle.ts

Lines changed: 15 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -55,24 +55,30 @@ import type { Duplex } from 'it-stream-types'
5555
*/
5656
export async function handle <Stream extends Duplex<any, any, any>> (stream: Stream, protocols: string | string[], options: MultistreamSelectInit): Promise<ProtocolStream<Stream>> {
5757
protocols = Array.isArray(protocols) ? protocols : [protocols]
58+
options.log.trace('handle: available protocols %s', protocols)
59+
5860
const lp = lpStream(stream, {
59-
maxDataLength: MAX_PROTOCOL_LENGTH
61+
...options,
62+
maxDataLength: MAX_PROTOCOL_LENGTH,
63+
maxLengthLength: 2 // 2 bytes is enough to length-prefix MAX_PROTOCOL_LENGTH
6064
})
6165

6266
while (true) {
63-
options.log.trace('handle - available protocols %s', protocols)
67+
options?.log.trace('handle: reading incoming string')
6468
const protocol = await multistream.readString(lp, options)
65-
options.log.trace('read "%s"', protocol)
69+
options.log.trace('handle: read "%s"', protocol)
6670

6771
if (protocol === PROTOCOL_ID) {
68-
options.log.trace('respond with "%s" for "%s"', PROTOCOL_ID, protocol)
72+
options.log.trace('handle: respond with "%s" for "%s"', PROTOCOL_ID, protocol)
6973
await multistream.write(lp, uint8ArrayFromString(`${PROTOCOL_ID}\n`), options)
74+
options.log.trace('handle: responded with "%s" for "%s"', PROTOCOL_ID, protocol)
7075
continue
7176
}
7277

7378
if (protocols.includes(protocol)) {
74-
options.log.trace('respond with "%s" for "%s"', protocol, protocol)
79+
options.log.trace('handle: respond with "%s" for "%s"', protocol, protocol)
7580
await multistream.write(lp, uint8ArrayFromString(`${protocol}\n`), options)
81+
options.log.trace('handle: responded with "%s" for "%s"', protocol, protocol)
7682

7783
return { stream: lp.unwrap(), protocol }
7884
}
@@ -84,12 +90,14 @@ export async function handle <Stream extends Duplex<any, any, any>> (stream: Str
8490
uint8ArrayFromString('\n')
8591
)
8692

93+
options.log.trace('handle: respond with "%s" for %s', protocols, protocol)
8794
await multistream.write(lp, protos, options)
88-
options.log.trace('respond with "%s" for %s', protocols, protocol)
95+
options.log.trace('handle: responded with "%s" for %s', protocols, protocol)
8996
continue
9097
}
9198

92-
options.log('respond with "na" for "%s"', protocol)
99+
options.log('handle: respond with "na" for "%s"', protocol)
93100
await multistream.write(lp, uint8ArrayFromString('na\n'), options)
101+
options.log('handle: responded with "na" for "%s"', protocol)
94102
}
95103
}

‎packages/multistream-select/src/index.ts

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,7 @@
2222

2323
import { PROTOCOL_ID } from './constants.js'
2424
import type { AbortOptions, LoggerOptions } from '@libp2p/interface'
25+
import type { LengthPrefixedStreamOpts } from 'it-length-prefixed-stream'
2526

2627
export { PROTOCOL_ID }
2728

@@ -30,7 +31,7 @@ export interface ProtocolStream<Stream> {
3031
protocol: string
3132
}
3233

33-
export interface MultistreamSelectInit extends AbortOptions, LoggerOptions {
34+
export interface MultistreamSelectInit extends AbortOptions, LoggerOptions, Partial<LengthPrefixedStreamOpts> {
3435

3536
}
3637

‎packages/multistream-select/src/select.ts

Lines changed: 42 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,7 @@
11
import { CodeError } from '@libp2p/interface/errors'
22
import { lpStream } from 'it-length-prefixed-stream'
3+
import * as varint from 'uint8-varint'
4+
import { Uint8ArrayList } from 'uint8arraylist'
35
import { fromString as uint8ArrayFromString } from 'uint8arrays/from-string'
46
import { MAX_PROTOCOL_LENGTH } from './constants.js'
57
import * as multistream from './multistream.js'
@@ -53,6 +55,7 @@ import type { Duplex } from 'it-stream-types'
5355
export async function select <Stream extends Duplex<any, any, any>> (stream: Stream, protocols: string | string[], options: MultistreamSelectInit): Promise<ProtocolStream<Stream>> {
5456
protocols = Array.isArray(protocols) ? [...protocols] : [protocols]
5557
const lp = lpStream(stream, {
58+
...options,
5659
maxDataLength: MAX_PROTOCOL_LENGTH
5760
})
5861
const protocol = protocols.shift()
@@ -109,29 +112,58 @@ export async function select <Stream extends Duplex<any, any, any>> (stream: Str
109112
export function lazySelect <Stream extends Duplex<any, any, any>> (stream: Stream, protocol: string, options: MultistreamSelectInit): ProtocolStream<Stream> {
110113
const originalSink = stream.sink.bind(stream)
111114
const originalSource = stream.source
115+
let selected = false
112116

113117
const lp = lpStream({
114118
sink: originalSink,
115119
source: originalSource
116120
}, {
121+
...options,
117122
maxDataLength: MAX_PROTOCOL_LENGTH
118123
})
119124

120125
stream.sink = async source => {
121-
options?.log.trace('lazy: write ["%s", "%s"]', PROTOCOL_ID, protocol)
122-
123-
await lp.writeV([
124-
uint8ArrayFromString(`${PROTOCOL_ID}\n`),
125-
uint8ArrayFromString(`${protocol}\n`)
126-
])
127-
128-
options?.log.trace('lazy: writing rest of "%s" stream', protocol)
129-
await lp.unwrap().sink(source)
126+
const { sink } = lp.unwrap()
127+
128+
await sink(async function * () {
129+
for await (const buf of source) {
130+
// if writing before selecting, send selection with first data chunk
131+
if (!selected) {
132+
selected = true
133+
options?.log.trace('lazy: write ["%s", "%s", data] in sink', PROTOCOL_ID, protocol)
134+
135+
const protocolString = `${protocol}\n`
136+
137+
// send protocols in first chunk of data written to transport
138+
yield new Uint8ArrayList(
139+
Uint8Array.from([19]), // length of PROTOCOL_ID plus newline
140+
uint8ArrayFromString(`${PROTOCOL_ID}\n`),
141+
varint.encode(protocolString.length),
142+
uint8ArrayFromString(protocolString),
143+
buf
144+
).subarray()
145+
146+
options?.log.trace('lazy: wrote ["%s", "%s", data] in sink', PROTOCOL_ID, protocol)
147+
} else {
148+
yield buf
149+
}
150+
}
151+
}())
130152
}
131153

132154
stream.source = (async function * () {
133-
options?.log.trace('lazy: reading multistream select header')
155+
// if reading before selecting, send selection before first data chunk
156+
if (!selected) {
157+
selected = true
158+
options?.log.trace('lazy: write ["%s", "%s", data] in source', PROTOCOL_ID, protocol)
159+
await lp.writeV([
160+
uint8ArrayFromString(`${PROTOCOL_ID}\n`),
161+
uint8ArrayFromString(`${protocol}\n`)
162+
])
163+
options?.log.trace('lazy: wrote ["%s", "%s", data] in source', PROTOCOL_ID, protocol)
164+
}
134165

166+
options?.log.trace('lazy: reading multistream select header')
135167
let response = await multistream.readString(lp, options)
136168
options?.log.trace('lazy: read multistream select header "%s"', response)
137169

‎packages/multistream-select/test/integration.spec.ts

Lines changed: 70 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -113,6 +113,76 @@ describe('Dialer and Listener integration', () => {
113113
expect(new Uint8ArrayList(...dialerOut).slice()).to.eql(new Uint8ArrayList(...input).slice())
114114
})
115115

116+
it('should handle and lazySelect that fails', async () => {
117+
const protocol = '/echo/1.0.0'
118+
const otherProtocol = '/echo/2.0.0'
119+
const pair = duplexPair<Uint8ArrayList | Uint8Array>()
120+
121+
const dialerSelection = mss.lazySelect(pair[0], protocol, {
122+
log: logger('mss:test')
123+
})
124+
expect(dialerSelection.protocol).to.equal(protocol)
125+
126+
// the listener handles the incoming stream
127+
void mss.handle(pair[1], otherProtocol, {
128+
log: logger('mss:test')
129+
})
130+
131+
// should fail when we interact with the stream
132+
const input = [randomBytes(10), randomBytes(64), randomBytes(3)]
133+
await expect(pipe(input, dialerSelection.stream, async source => all(source)))
134+
.to.eventually.be.rejected.with.property('code', 'ERR_UNSUPPORTED_PROTOCOL')
135+
})
136+
137+
it('should handle and lazySelect only by reading', async () => {
138+
const protocol = '/echo/1.0.0'
139+
const pair = duplexPair<Uint8ArrayList | Uint8Array>()
140+
141+
const dialerSelection = mss.lazySelect(pair[0], protocol, {
142+
log: logger('mss:dialer')
143+
})
144+
expect(dialerSelection.protocol).to.equal(protocol)
145+
146+
// ensure stream is usable after selection
147+
const input = [randomBytes(10), randomBytes(64), randomBytes(3)]
148+
149+
const [, dialerOut] = await Promise.all([
150+
// the listener handles the incoming stream
151+
mss.handle(pair[1], protocol, {
152+
log: logger('mss:listener')
153+
}).then(async result => {
154+
// the listener writes to the incoming stream
155+
await pipe(input, result.stream)
156+
}),
157+
158+
// the dialer just reads from the stream
159+
pipe(dialerSelection.stream, async source => all(source))
160+
])
161+
162+
expect(new Uint8ArrayList(...dialerOut).slice()).to.eql(new Uint8ArrayList(...input).slice())
163+
})
164+
165+
it('should handle and lazySelect only by reading that fails', async () => {
166+
const protocol = '/echo/1.0.0'
167+
const otherProtocol = '/echo/2.0.0'
168+
const pair = duplexPair<Uint8ArrayList | Uint8Array>()
169+
170+
// lazy succeeds
171+
const dialerSelection = mss.lazySelect(pair[0], protocol, {
172+
log: logger('mss:dialer')
173+
})
174+
expect(dialerSelection.protocol).to.equal(protocol)
175+
176+
// the listener handles the incoming stream
177+
void mss.handle(pair[1], otherProtocol, {
178+
log: logger('mss:listener')
179+
})
180+
181+
// should fail when we interact with the stream
182+
await expect(pipe(dialerSelection.stream, async source => all(source)))
183+
.to.eventually.be.rejected.with.property('code', 'ERR_UNSUPPORTED_PROTOCOL')
184+
})
185+
116186
it('should abort an unhandled lazySelect', async () => {
117187
const protocol = '/echo/1.0.0'
118188
const pair = duplexPair<Uint8ArrayList | Uint8Array>()

‎packages/multistream-select/test/multistream.spec.ts

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -61,7 +61,7 @@ describe('Multistream', () => {
6161
const duplexes = duplexPair<Uint8Array>()
6262
const inputStream = lpStream(duplexes[0])
6363
const outputStream = lpStream(duplexes[1], {
64-
maxDataLength: 100
64+
maxDataLength: 9999
6565
})
6666

6767
void inputStream.write(input)

0 commit comments

Comments
 (0)