github.com/qri-io/qri@v0.10.1-0.20220104210721-c771715036cb/p2p/p2putil/wrapped_stream.go (about)

     1  package p2putil
     2  
     3  import (
     4  	"bufio"
     5  
     6  	net "github.com/libp2p/go-libp2p-core/network"
     7  	json "github.com/qri-io/dag/dsync/p2putil/json"
     8  	multicodec "github.com/qri-io/dag/dsync/p2putil/multicodec_old"
     9  )
    10  
    11  // HandlerFunc is the signature of a function that can handle p2p messages
    12  type HandlerFunc func(ws *WrappedStream, msg Message) (hangup bool)
    13  
    14  // WrappedStream wraps a libp2p stream. We encode/decode whenever we
    15  // write/read from a stream, so we can just carry the encoders
    16  // and bufios with us
    17  type WrappedStream struct {
    18  	stream net.Stream
    19  	Enc    multicodec.Encoder
    20  	Dec    multicodec.Decoder
    21  	W      *bufio.Writer
    22  	R      *bufio.Reader
    23  }
    24  
    25  // WrapStream takes a stream and complements it with r/w bufios and
    26  // decoder/encoder. In order to write raw data to the stream we can use
    27  // wrap.w.Write(). To encode something into it we can wrap.enc.Encode().
    28  // Finally, we should wrap.w.Flush() to actually send the data. Handling
    29  // incoming data works similarly with wrap.r.Read() for raw-reading and
    30  // wrap.dec.Decode() to decode.
    31  func WrapStream(s net.Stream) *WrappedStream {
    32  	reader := bufio.NewReader(s)
    33  	writer := bufio.NewWriter(s)
    34  	// This is where we pick our specific multicodec. In order to change the
    35  	// codec, we only need to change this place.
    36  	// See https://godoc.org/github.com/multiformats/go-multicodec/json
    37  	dec := json.Multicodec(false).Decoder(reader)
    38  	enc := json.Multicodec(false).Encoder(writer)
    39  	return &WrappedStream{
    40  		stream: s,
    41  		R:      reader,
    42  		W:      writer,
    43  		Enc:    enc,
    44  		Dec:    dec,
    45  	}
    46  }
    47  
    48  // ReceiveMessage reads and decodes a message from the stream
    49  func (ws *WrappedStream) ReceiveMessage() (msg Message, err error) {
    50  	err = ws.Dec.Decode(&msg)
    51  	msg.provider = ws.stream.Conn().RemotePeer()
    52  	// log.Debugf("%s '%s' <- %s", ws.stream.Conn().LocalPeer(), msg.Type, ws.stream.Conn().RemotePeer())
    53  	return
    54  }
    55  
    56  // SendMessage encodes and writes a message to the stream
    57  func (ws *WrappedStream) SendMessage(msg Message) error {
    58  	err := ws.Enc.Encode(&msg)
    59  	// Because output is buffered with bufio, we need to flush!
    60  	ws.W.Flush()
    61  	// log.Debugf("%s '%s' -> %s", ws.stream.Conn().LocalPeer(), msg.Type, ws.stream.Conn().RemotePeer())
    62  	return err
    63  }
    64  
    65  // Stream exposes the underlying libp2p net.Stream
    66  func (ws *WrappedStream) Stream() net.Stream {
    67  	return ws.stream
    68  }
    69  
    70  // Close closes the stream connection
    71  func (ws *WrappedStream) Close() error {
    72  	return ws.stream.Close()
    73  }