Support getting a reader for a message in METADAPT-B

metadaptb.go

@@ -5,6 +5,8 @@ import "net"
 import "context"
 import "git.tebibyte.media/sashakoshka/hopp/tape"
 
+// TODO: change size limit to be int64
+
 // B implements METADAPT-B over a multiplexed stream-oriented transport such as
 // QUIC.
 type b struct {
@@ -12,8 +14,8 @@ type b struct {
 	underlying MultiConn
 }
 
-// AdaptB returns a connection implementing METADAPT-B over a singular stream-
-// oriented transport such as TCP or UNIX domain stream sockets.
+// AdaptB returns a connection implementing METADAPT-B over a multiplexed
+// stream-oriented transport such as QUIC.
 func AdaptB(underlying MultiConn) Conn {
 	return &b {
 		sizeLimit:  defaultSizeLimit,
@@ -49,8 +51,8 @@ func (this *b) SetSizeLimit(limit int) {
 	this.sizeLimit = limit
 }
 
-func (this *b) newTrans(underlying Stream) transB {
-	return transB {
+func (this *b) newTrans(underlying Stream) *transB {
+	return &transB {
 		sizeLimit: this.sizeLimit,
 		underlying: underlying,
 	}
@@ -59,22 +61,49 @@ func (this *b) newTrans(underlying Stream) transB {
 type transB struct {
 	sizeLimit int
 	underlying Stream
+	currentData io.Reader
 }
 
-func (trans transB) Close() error {
-	return trans.underlying.Close()
+func (this *transB) Close() error {
+	return this.underlying.Close()
 }
 
-func (trans transB) ID() int64 {
-	return trans.underlying.ID()
+func (this *transB) ID() int64 {
+	return this.underlying.ID()
 }
 
-func (trans transB) Send(method uint16, data []byte) error {
-	return encodeMessageB(trans.underlying, trans.sizeLimit, method, data)
+func (this *transB) Send(method uint16, data []byte) error {
+	return encodeMessageB(this.underlying, this.sizeLimit, method, data)
 }
 
-func (trans transB) Receive() (uint16, []byte, error) {
-	return decodeMessageB(trans.underlying, trans.sizeLimit)
+func (this *transB) Receive() (uint16, []byte, error) {
+	// get a reader for the next message
+	method, size, data, err := this.receiveReader()
+	if err != nil { return 0, nil, err }
+	// read the entire thing
+	payloadBuffer := make([]byte, int(size))
+	_, err = io.ReadFull(data, payloadBuffer)
+	if err != nil { return 0, nil, err }
+	// we have used up the reader by now so we can forget it exists
+	this.currentData = nil
+	return method, payloadBuffer, nil
+}
+
+func (this *transB) ReceiveReader() (uint16, io.Reader, error) {
+	method, _, data, err := this.receiveReader()
+	return method, data, err
+}
+
+func (this *transB) receiveReader() (uint16, int64, io.Reader, error) {
+	// decode the message
+	method, size, data, err := decodeMessageB(this.underlying, this.sizeLimit)
+	if err != nil { return 0, 0, nil, err }
+	// discard current reader if there is one
+	if this.currentData == nil {
+		io.Copy(io.Discard, this.currentData)
+	}
+	this.currentData = data
+	return method, size, data, nil
 }
 
 // MultiConn represens a multiplexed stream-oriented transport for use in
@@ -110,19 +139,27 @@ func encodeMessageB(writer io.Writer, sizeLimit int, method uint16, data []byte)
 	return err
 }
 
-func decodeMessageB(reader io.Reader, sizeLimit int) (uint16, []byte, error) {
+func decodeMessageB(
+	reader io.Reader,
+	sizeLimit int,
+) (
+	method uint16,
+	size   int64,
+	data   io.Reader,
+	err    error,
+) {
 	headerBuffer := [10]byte { }
-	_, err := io.ReadFull(reader, headerBuffer[:])
-	if err != nil { return 0, nil, err }
-	method, err := tape.DecodeI16[uint16](headerBuffer[:2])
-	if err != nil { return 0, nil, err }
+	_, err = io.ReadFull(reader, headerBuffer[:])
+	if err != nil { return 0, 0, nil, err }
+	method, err = tape.DecodeI16[uint16](headerBuffer[:2])
+	if err != nil { return 0, 0, nil, err }
 	length, err := tape.DecodeI64[uint64](headerBuffer[2:10])
-	if err != nil { return 0, nil, err }
+	if err != nil { return 0, 0, nil, err }
 	if length > uint64(sizeLimit) {
-		return 0, nil, ErrPayloadTooLarge
+		return 0, 0, nil, ErrPayloadTooLarge
 	}
-	payloadBuffer := make([]byte, int(length))
-	_, err = io.ReadFull(reader, payloadBuffer)
-	if err != nil { return 0, nil, err }
-	return method, payloadBuffer, nil
+	return method, int64(length), &io.LimitedReader {
+		R: reader,
+		N: int64(length),
+	}, nil
 }