github.com/andy2046/gopie@v0.7.0/pkg/tlv/tlv.go (about)

     1  // Package tlv implements Type-Length-Value encoding.
     2  package tlv
     3  
     4  import (
     5  	"bytes"
     6  	"encoding/binary"
     7  	"io"
     8  )
     9  
    10  // ByteSize const.
    11  const (
    12  	Bytes1 ByteSize = 1 << iota
    13  	Bytes2
    14  	Bytes4
    15  	Bytes8
    16  )
    17  
    18  type (
    19  	// ByteSize is the size of a field in bytes.
    20  	// Used to define the size of the type and length field in a message.
    21  	ByteSize int
    22  
    23  	// Record represents a record of data encoded in the TLV message.
    24  	Record struct {
    25  		Payload []byte
    26  		Type    uint
    27  	}
    28  
    29  	// Codec is the configuration for a TLV encoding/decoding task.
    30  	Codec struct {
    31  		// TypeBytes defines the size in bytes of the message type field.
    32  		TypeBytes ByteSize
    33  
    34  		// LenBytes defines the size in bytes of the message length field.
    35  		LenBytes ByteSize
    36  	}
    37  
    38  	// Writer encodes records into TLV format using a Codec and writes into the provided io.Writer.
    39  	Writer struct {
    40  		writer io.Writer
    41  		codec  *Codec
    42  	}
    43  
    44  	// Reader decodes records from TLV format using a Codec from the provided io.Reader.
    45  	Reader struct {
    46  		codec  *Codec
    47  		reader io.Reader
    48  	}
    49  )
    50  
    51  // NewWriter creates a new Writer.
    52  func NewWriter(w io.Writer, codec *Codec) *Writer {
    53  	return &Writer{
    54  		codec:  codec,
    55  		writer: w,
    56  	}
    57  }
    58  
    59  // Write encodes records into TLV format using a Codec and writes into the provided io.Writer.
    60  func (w *Writer) Write(rec *Record) error {
    61  	err := writeUint(w.writer, w.codec.TypeBytes, rec.Type)
    62  	if err != nil {
    63  		return err
    64  	}
    65  
    66  	ulen := uint(len(rec.Payload))
    67  	err = writeUint(w.writer, w.codec.LenBytes, ulen)
    68  	if err != nil {
    69  		return err
    70  	}
    71  
    72  	_, err = w.writer.Write(rec.Payload)
    73  	return err
    74  }
    75  
    76  func writeUint(w io.Writer, b ByteSize, i uint) error {
    77  	var num interface{}
    78  	switch b {
    79  	case Bytes1:
    80  		num = uint8(i)
    81  	case Bytes2:
    82  		num = uint16(i)
    83  	case Bytes4:
    84  		num = uint32(i)
    85  	case Bytes8:
    86  		num = uint64(i)
    87  	}
    88  	return binary.Write(w, binary.LittleEndian, num)
    89  }
    90  
    91  // NewReader creates a new Reader.
    92  func NewReader(reader io.Reader, codec *Codec) *Reader {
    93  	return &Reader{codec: codec, reader: reader}
    94  }
    95  
    96  // Next reads a single Record from the io.Reader.
    97  func (r *Reader) Next() (*Record, error) {
    98  	// get type
    99  	typeBytes := make([]byte, r.codec.TypeBytes)
   100  	_, err := r.reader.Read(typeBytes)
   101  	if err != nil {
   102  		return nil, err
   103  	}
   104  	typ := readUint(typeBytes, r.codec.TypeBytes)
   105  
   106  	// get len
   107  	payloadLenBytes := make([]byte, r.codec.LenBytes)
   108  	_, err = r.reader.Read(payloadLenBytes)
   109  	if err != nil && err != io.EOF {
   110  		return nil, err
   111  	}
   112  	payloadLen := readUint(payloadLenBytes, r.codec.LenBytes)
   113  
   114  	if err == io.EOF && payloadLen != 0 {
   115  		return nil, err
   116  	}
   117  
   118  	// get value
   119  	v := make([]byte, payloadLen)
   120  	_, err = r.reader.Read(v)
   121  	if err != nil && err != io.EOF {
   122  		return nil, err
   123  	}
   124  
   125  	return &Record{
   126  		Type:    typ,
   127  		Payload: v,
   128  	}, nil
   129  
   130  }
   131  
   132  func readUint(b []byte, sz ByteSize) uint {
   133  	reader := bytes.NewReader(b)
   134  	switch sz {
   135  	case Bytes1:
   136  		var i uint8
   137  		binary.Read(reader, binary.LittleEndian, &i)
   138  		return uint(i)
   139  	case Bytes2:
   140  		var i uint16
   141  		binary.Read(reader, binary.LittleEndian, &i)
   142  		return uint(i)
   143  	case Bytes4:
   144  		var i uint32
   145  		binary.Read(reader, binary.LittleEndian, &i)
   146  		return uint(i)
   147  	case Bytes8:
   148  		var i uint64
   149  		binary.Read(reader, binary.LittleEndian, &i)
   150  		return uint(i)
   151  	default:
   152  		return 0
   153  	}
   154  }