github.com/decred/dcrlnd@v0.7.6/tlv/varint.go (about)

     1  package tlv
     2  
     3  import (
     4  	"encoding/binary"
     5  	"errors"
     6  	"io"
     7  
     8  	"github.com/decred/dcrd/wire"
     9  )
    10  
    11  // ErrVarIntNotCanonical signals that the decoded varint was not minimally encoded.
    12  var ErrVarIntNotCanonical = errors.New("decoded varint is not canonical")
    13  
    14  // ReadVarInt reads a variable length integer from r and returns it as a uint64.
    15  func ReadVarInt(r io.Reader, buf *[8]byte) (uint64, error) {
    16  	_, err := io.ReadFull(r, buf[:1])
    17  	if err != nil {
    18  		return 0, err
    19  	}
    20  	discriminant := buf[0]
    21  
    22  	var rv uint64
    23  	switch {
    24  	case discriminant < 0xfd:
    25  		rv = uint64(discriminant)
    26  
    27  	case discriminant == 0xfd:
    28  		_, err := io.ReadFull(r, buf[:2])
    29  		switch {
    30  		case err == io.EOF:
    31  			return 0, io.ErrUnexpectedEOF
    32  		case err != nil:
    33  			return 0, err
    34  		}
    35  		rv = uint64(binary.BigEndian.Uint16(buf[:2]))
    36  
    37  		// The encoding is not canonical if the value could have been
    38  		// encoded using fewer bytes.
    39  		if rv < 0xfd {
    40  			return 0, ErrVarIntNotCanonical
    41  		}
    42  
    43  	case discriminant == 0xfe:
    44  		_, err := io.ReadFull(r, buf[:4])
    45  		switch {
    46  		case err == io.EOF:
    47  			return 0, io.ErrUnexpectedEOF
    48  		case err != nil:
    49  			return 0, err
    50  		}
    51  		rv = uint64(binary.BigEndian.Uint32(buf[:4]))
    52  
    53  		// The encoding is not canonical if the value could have been
    54  		// encoded using fewer bytes.
    55  		if rv <= 0xffff {
    56  			return 0, ErrVarIntNotCanonical
    57  		}
    58  
    59  	default:
    60  		_, err := io.ReadFull(r, buf[:])
    61  		switch {
    62  		case err == io.EOF:
    63  			return 0, io.ErrUnexpectedEOF
    64  		case err != nil:
    65  			return 0, err
    66  		}
    67  		rv = binary.BigEndian.Uint64(buf[:])
    68  
    69  		// The encoding is not canonical if the value could have been
    70  		// encoded using fewer bytes.
    71  		if rv <= 0xffffffff {
    72  			return 0, ErrVarIntNotCanonical
    73  		}
    74  	}
    75  
    76  	return rv, nil
    77  }
    78  
    79  // WriteVarInt serializes val to w using a variable number of bytes depending
    80  // on its value.
    81  func WriteVarInt(w io.Writer, val uint64, buf *[8]byte) error {
    82  	var length int
    83  	switch {
    84  	case val < 0xfd:
    85  		buf[0] = uint8(val)
    86  		length = 1
    87  
    88  	case val <= 0xffff:
    89  		buf[0] = uint8(0xfd)
    90  		binary.BigEndian.PutUint16(buf[1:3], uint16(val))
    91  		length = 3
    92  
    93  	case val <= 0xffffffff:
    94  		buf[0] = uint8(0xfe)
    95  		binary.BigEndian.PutUint32(buf[1:5], uint32(val))
    96  		length = 5
    97  
    98  	default:
    99  		buf[0] = uint8(0xff)
   100  		_, err := w.Write(buf[:1])
   101  		if err != nil {
   102  			return err
   103  		}
   104  
   105  		binary.BigEndian.PutUint64(buf[:], val)
   106  		length = 8
   107  	}
   108  
   109  	_, err := w.Write(buf[:length])
   110  	return err
   111  }
   112  
   113  // VarIntSize returns the required number of bytes to encode a var int.
   114  func VarIntSize(val uint64) uint64 {
   115  	return uint64(wire.VarIntSerializeSize(val))
   116  }