tractor.dev/toolkit-go@v0.0.0-20241010005851-214d91207d07/duplex/rpc/frame.go (about)

     1  package rpc
     2  
     3  import (
     4  	"bytes"
     5  	"encoding/binary"
     6  	"io"
     7  
     8  	"tractor.dev/toolkit-go/duplex/codec"
     9  )
    10  
    11  // FrameCodec is a special codec used to actually read/write other
    12  // codecs to a transport using a length prefix.
    13  type FrameCodec struct {
    14  	codec.Codec
    15  }
    16  
    17  // Encoder returns a frame encoder that first encodes a value
    18  // to a buffer using the embedded codec, prepends the encoded value
    19  // byte length as a four byte big endian uint32, then writes to
    20  // the given Writer.
    21  func (c *FrameCodec) Encoder(w io.Writer) codec.Encoder {
    22  	return &frameEncoder{
    23  		w: w,
    24  		c: c.Codec,
    25  	}
    26  }
    27  
    28  type frameEncoder struct {
    29  	w io.Writer
    30  	c codec.Codec
    31  }
    32  
    33  func (e *frameEncoder) Encode(v interface{}) error {
    34  	var buf bytes.Buffer
    35  	enc := e.c.Encoder(&buf)
    36  	err := enc.Encode(v)
    37  	if err != nil {
    38  		return err
    39  	}
    40  	b := buf.Bytes()
    41  	prefix := make([]byte, 4)
    42  	binary.BigEndian.PutUint32(prefix, uint32(len(b)))
    43  	_, err = e.w.Write(append(prefix, b...))
    44  	if err != nil {
    45  		return err
    46  	}
    47  	return nil
    48  }
    49  
    50  // Decoder returns a frame decoder that first reads a four byte frame
    51  // length value used to read the rest of the frame, then uses the
    52  // embedded codec to decode those bytes into a value.
    53  func (c *FrameCodec) Decoder(r io.Reader) codec.Decoder {
    54  	return &frameDecoder{
    55  		r: r,
    56  		c: c.Codec,
    57  	}
    58  }
    59  
    60  type frameDecoder struct {
    61  	r io.Reader
    62  	c codec.Codec
    63  }
    64  
    65  func (d *frameDecoder) Decode(v interface{}) error {
    66  	prefix := make([]byte, 4)
    67  	_, err := io.ReadFull(d.r, prefix)
    68  	if err != nil {
    69  		return err
    70  	}
    71  	size := binary.BigEndian.Uint32(prefix)
    72  	buf := make([]byte, size)
    73  	_, err = io.ReadFull(d.r, buf)
    74  	if err != nil {
    75  		return err
    76  	}
    77  	dec := d.c.Decoder(bytes.NewBuffer(buf))
    78  	err = dec.Decode(v)
    79  	if err != nil {
    80  		return err
    81  	}
    82  	return nil
    83  }