ms-select2: compress multiselect proto id

Author: ppopth

lazyClient.go

@@ -1,7 +1,7 @@
 package multistream
 
 import (
-	"encoding/hex"
+	"bytes"
 	"fmt"
 	"io"
 )
@@ -10,7 +10,7 @@ import (
 // protocol selection with a MultistreamMuxer.
 func NewMSSelect[T StringLike](c io.ReadWriteCloser, proto T) LazyConn {
 	return &lazyClientConn[T]{
-		protos: []T{ProtocolID, proto},
+		protos: []protoInfo[T]{{ID: ProtocolID}, {ID: proto}},
 		con:    c,
 
 		rhandshakeOnce: newOnce(),
@@ -24,11 +24,13 @@ func NewMSSelect2[T StringLike](c io.ReadWriteCloser, proto T, peerProtos []T) L
 		t.AddProtocol(p)
 	}
 
-	// TODO: use a proper varint instead of a hex string later
-	abbrv := T(hex.EncodeToString(t.Abbreviate(proto)))
+	abbrv := t.Abbreviate(proto)
 	return &lazyClientConn[T]{
-		protos: []T{ProtocolID, abbrv},
-		con:    c,
+		protos: []protoInfo[T]{
+			{ID: ProtocolID, Abbrev: ProtocolAbbrev},
+			{ID: proto, Abbrev: abbrv},
+		},
+		con: c,
 
 		rhandshakeOnce: newOnce(),
 		whandshakeOnce: newOnce(),
@@ -40,7 +42,7 @@ func NewMSSelect2[T StringLike](c io.ReadWriteCloser, proto T, peerProtos []T) L
 // NewMSSelect.
 func NewMultistream[T StringLike](c io.ReadWriteCloser, proto T) LazyConn {
 	return &lazyClientConn[T]{
-		protos: []T{proto},
+		protos: []protoInfo[T]{{ID: proto}},
 		con:    c,
 
 		rhandshakeOnce: newOnce(),
@@ -76,6 +78,11 @@ func (o *once) Do(f func()) {
 	f()
 }
 
+type protoInfo[T StringLike] struct {
+	ID     T
+	Abbrev []byte
+}
+
 // lazyClientConn is a ReadWriteCloser adapter that lazily negotiates a protocol
 // using multistream-select on first use.
 //
@@ -92,7 +99,7 @@ type lazyClientConn[T StringLike] struct {
 	werr           error
 
 	// The sequence of protocols to negotiate.
-	protos []T
+	protos []protoInfo[T]
 
 	// The inner connection.
 	con io.ReadWriteCloser
@@ -122,18 +129,22 @@ func (l *lazyClientConn[T]) Read(b []byte) (int, error) {
 func (l *lazyClientConn[T]) doReadHandshake() {
 	for _, proto := range l.protos {
 		// read protocol
-		tok, err := ReadNextToken[T](l.con)
+		tok, err := ReadNextTokenBytes(l.con)
 		if err != nil {
 			l.rerr = err
 			return
 		}
 
-		if tok == "na" {
-			l.rerr = ErrNotSupported[T]{[]T{proto}}
+		if bytes.Equal(tok, []byte("na")) {
+			l.rerr = ErrNotSupported[T]{[]T{proto.ID}}
 			return
 		}
-		if tok != proto {
-			l.rerr = fmt.Errorf("protocol mismatch in lazy handshake ( %s != %s )", tok, proto)
+		if proto.Abbrev != nil && !bytes.Equal(tok, proto.Abbrev) {
+			l.rerr = fmt.Errorf("protocol mismatch in lazy handshake ( %x != %x )", tok, proto.Abbrev)
+			return
+		}
+		if proto.Abbrev == nil && T(tok) != proto.ID {
+			l.rerr = fmt.Errorf("protocol mismatch in lazy handshake ( %s != %s )", T(tok), proto.ID)
 			return
 		}
 	}
@@ -149,7 +160,11 @@ func (l *lazyClientConn[T]) doWriteHandshakeWithData(extra []byte) int {
 	defer putWriter(buf)
 
 	for _, proto := range l.protos {
-		l.werr = delimWrite(buf, []byte(proto))
+		if proto.Abbrev != nil {
+			l.werr = delimWrite(buf, proto.Abbrev)
+		} else {
+			l.werr = delimWrite(buf, []byte(proto.ID))
+		}
 		if l.werr != nil {
 			return 0
 		}

multistream.go

@@ -5,7 +5,7 @@ package multistream
 
 import (
 	"bufio"
-	"encoding/hex"
+	"bytes"
 	"errors"
 	"fmt"
 	"io"
@@ -26,6 +26,9 @@ var ErrUnknownPrefix = errors.New("unknown protocol hash prefix")
 // the multistream muxers on both sides of a channel can work with each other.
 const ProtocolID = "/multistream/1.0.0"
 
+// ProtocolID identifies the multistream protocol abbreviation support
+var ProtocolAbbrev = []byte{0xff, 0x11}
+
 // Multistream-select version that protocol abbreviation is supported
 const AbbrevSupportedMSSVersion = 2
 
@@ -186,24 +189,6 @@ func (msm *MultistreamMuxer[T]) Protocols() []T {
 // fails because of a ProtocolID mismatch.
 var ErrIncorrectVersion = errors.New("client connected with incorrect version")
 
-func (msm *MultistreamMuxer[T]) decodeProtocol(s T) (T, error) {
-	msm.handlerlock.RLock()
-	defer msm.handlerlock.RUnlock()
-
-	bytes, err := hex.DecodeString(string(s))
-	// TODO: decide whether to compare strings or use abbrevTree by looking at
-	// multistream version instead.
-	if err != nil {
-		return s, nil
-	}
-
-	proto, err := msm.abbrevTree.GetProtocolID(bytes)
-	if err != nil {
-		return "", err
-	}
-	return proto, nil
-}
-
 func (msm *MultistreamMuxer[T]) findHandler(proto T) *Handler[T] {
 	msm.handlerlock.RLock()
 	defer msm.handlerlock.RUnlock()
@@ -227,35 +212,49 @@ func (msm *MultistreamMuxer[T]) Negotiate(rwc io.ReadWriteCloser) (proto T, hand
 		}
 	}()
 
-	// Send the multistream protocol ID
-	// Ignore the error here.  We want the handshake to finish, even if the
-	// other side has closed this rwc for writing. They may have sent us a
-	// message and closed. Future writers will get an error anyways.
-	_ = delimWriteBuffered(rwc, []byte(ProtocolID))
-	line, err := ReadNextToken[T](rwc)
+	token, err := ReadNextTokenBytes(rwc)
 	if err != nil {
 		return "", nil, err
 	}
-
-	if line != ProtocolID {
+	supportAbbrev := false
+	// Send the multistream protocol ID or the mulstream protocol abbreviation
+	// Ignore the error here.  We want the handshake to finish, even if the
+	// other side has closed this rwc for writing. They may have sent us a
+	// message and closed. Future writers will get an error anyways.
+	if bytes.Equal(token, ProtocolAbbrev) {
+		supportAbbrev = true
+		_ = delimWriteBuffered(rwc, ProtocolAbbrev)
+	} else if T(token) == ProtocolID {
+		_ = delimWriteBuffered(rwc, []byte(ProtocolID))
+	} else {
 		rwc.Close()
 		return "", nil, ErrIncorrectVersion
 	}
 
 loop:
 	for {
 		// Now read and respond to commands until they send a valid protocol id
-		tok, err := ReadNextToken[T](rwc)
+		var proto T
+
+		tok, err := ReadNextTokenBytes(rwc)
 		if err != nil {
 			return "", nil, err
 		}
 
-		p, err := msm.decodeProtocol(tok)
-		if err != nil {
-			return "", nil, err
+		if supportAbbrev {
+			// decode the protocol abbreviation using the abbreviation tree
+			msm.handlerlock.RLock()
+			proto, err = msm.abbrevTree.GetProtocolID(tok)
+			msm.handlerlock.RUnlock()
+
+			if err != nil {
+				return "", nil, err
+			}
+		} else {
+			proto = T(tok)
 		}
 
-		h := msm.findHandler(p)
+		h := msm.findHandler(proto)
 		if h == nil {
 			if err := delimWriteBuffered(rwc, []byte("na")); err != nil {
 				return "", nil, err
@@ -266,10 +265,10 @@ loop:
 		// Ignore the error here.  We want the handshake to finish, even if the
 		// other side has closed this rwc for writing. They may have sent us a
 		// message and closed. Future writers will get an error anyways.
-		_ = delimWriteBuffered(rwc, []byte(tok))
+		_ = delimWriteBuffered(rwc, tok)
 
 		// hand off processing to the sub-protocol handler
-		return p, h.Handle, nil
+		return proto, h.Handle, nil
 	}
 
 }