ms-select2: abbreviation on the server side

Author: ppopth

abbrevTree.go

@@ -41,6 +41,23 @@ func (at *abbrevTree[T]) Abbreviate(pid T) []byte {
 	return nil
 }
 
+func (at *abbrevTree[T]) GetProtocolID(prefix []byte) (T, error) {
+	if at.root == nil {
+		return "", ErrUnknownPrefix
+	}
+	current := at.root
+	for _, b := range prefix {
+		if current.children[b] == nil {
+			return "", ErrUnknownPrefix
+		}
+		current = current.children[b]
+	}
+	if current.p == nil {
+		return "", ErrUnknownPrefix
+	}
+	return current.p.protocolID, nil
+}
+
 func (at *abbrevTree[T]) AddProtocol(pid T) {
 	hash := sha256.Sum256([]byte(pid))
 

multistream.go

@@ -5,6 +5,7 @@ package multistream
 
 import (
 	"bufio"
+	"encoding/hex"
 	"errors"
 	"fmt"
 	"io"
@@ -18,6 +19,9 @@ import (
 // ErrTooLarge is an error to signal that an incoming message was too large
 var ErrTooLarge = errors.New("incoming message was too large")
 
+// ErrUnknownPrefix is an error to signal that the protocol hash prefix is unknown
+var ErrUnknownPrefix = errors.New("unknown protocol hash prefix")
+
 // ProtocolID identifies the multistream protocol itself and makes sure
 // the multistream muxers on both sides of a channel can work with each other.
 const ProtocolID = "/multistream/1.0.0"
@@ -55,6 +59,7 @@ type Handler[T StringLike] struct {
 type MultistreamMuxer[T StringLike] struct {
 	handlerlock sync.RWMutex
 	handlers    []Handler[T]
+	abbrevTree  abbrevTree[T]
 }
 
 // NewMultistreamMuxer creates a muxer.
@@ -137,6 +142,7 @@ func (msm *MultistreamMuxer[T]) AddHandlerWithFunc(protocol T, match func(T) boo
 	msm.handlerlock.Lock()
 	defer msm.handlerlock.Unlock()
 
+	msm.abbrevTree.AddProtocol(protocol)
 	msm.removeHandler(protocol)
 	msm.handlers = append(msm.handlers, Handler[T]{
 		MatchFunc: match,
@@ -150,6 +156,7 @@ func (msm *MultistreamMuxer[T]) RemoveHandler(protocol T) {
 	msm.handlerlock.Lock()
 	defer msm.handlerlock.Unlock()
 
+	msm.abbrevTree.RemoveProtocol(protocol)
 	msm.removeHandler(protocol)
 }
 
@@ -179,6 +186,21 @@ 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) {
+	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()
@@ -225,7 +247,12 @@ loop:
 			return "", nil, err
 		}
 
-		h := msm.findHandler(tok)
+		p, err := msm.decodeProtocol(tok)
+		if err != nil {
+			return "", nil, err
+		}
+
+		h := msm.findHandler(p)
 		if h == nil {
 			if err := delimWriteBuffered(rwc, []byte("na")); err != nil {
 				return "", nil, err
@@ -236,10 +263,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, []byte(p))
 
 		// hand off processing to the sub-protocol handler
-		return tok, h.Handle, nil
+		return p, h.Handle, nil
 	}
 
 }