github.com/10XDev/rclone@v1.52.3-0.20200626220027-16af9ab76b2a/backend/mailru/api/helpers.go (about)

     1  package api
     2  
     3  // BIN protocol helpers
     4  
     5  import (
     6  	"bufio"
     7  	"bytes"
     8  	"encoding/binary"
     9  	"io"
    10  	"log"
    11  	"time"
    12  
    13  	"github.com/pkg/errors"
    14  	"github.com/rclone/rclone/lib/readers"
    15  )
    16  
    17  // protocol errors
    18  var (
    19  	ErrorPrematureEOF  = errors.New("Premature EOF")
    20  	ErrorInvalidLength = errors.New("Invalid length")
    21  	ErrorZeroTerminate = errors.New("String must end with zero")
    22  )
    23  
    24  // BinWriter is a binary protocol writer
    25  type BinWriter struct {
    26  	b *bytes.Buffer // growing byte buffer
    27  	a []byte        // temporary buffer for next varint
    28  }
    29  
    30  // NewBinWriter creates a binary protocol helper
    31  func NewBinWriter() *BinWriter {
    32  	return &BinWriter{
    33  		b: new(bytes.Buffer),
    34  		a: make([]byte, binary.MaxVarintLen64),
    35  	}
    36  }
    37  
    38  // Bytes returns binary data
    39  func (w *BinWriter) Bytes() []byte {
    40  	return w.b.Bytes()
    41  }
    42  
    43  // Reader returns io.Reader with binary data
    44  func (w *BinWriter) Reader() io.Reader {
    45  	return bytes.NewReader(w.b.Bytes())
    46  }
    47  
    48  // WritePu16 writes a short as unsigned varint
    49  func (w *BinWriter) WritePu16(val int) {
    50  	if val < 0 || val > 65535 {
    51  		log.Fatalf("Invalid UInt16 %v", val)
    52  	}
    53  	w.WritePu64(int64(val))
    54  }
    55  
    56  // WritePu32 writes a signed long as unsigned varint
    57  func (w *BinWriter) WritePu32(val int64) {
    58  	if val < 0 || val > 4294967295 {
    59  		log.Fatalf("Invalid UInt32 %v", val)
    60  	}
    61  	w.WritePu64(val)
    62  }
    63  
    64  // WritePu64 writes an unsigned (actually, signed) long as unsigned varint
    65  func (w *BinWriter) WritePu64(val int64) {
    66  	if val < 0 {
    67  		log.Fatalf("Invalid UInt64 %v", val)
    68  	}
    69  	w.b.Write(w.a[:binary.PutUvarint(w.a, uint64(val))])
    70  }
    71  
    72  // WriteString writes a zero-terminated string
    73  func (w *BinWriter) WriteString(str string) {
    74  	buf := []byte(str)
    75  	w.WritePu64(int64(len(buf) + 1))
    76  	w.b.Write(buf)
    77  	w.b.WriteByte(0)
    78  }
    79  
    80  // Write writes a byte buffer
    81  func (w *BinWriter) Write(buf []byte) {
    82  	w.b.Write(buf)
    83  }
    84  
    85  // WriteWithLength writes a byte buffer prepended with its length as varint
    86  func (w *BinWriter) WriteWithLength(buf []byte) {
    87  	w.WritePu64(int64(len(buf)))
    88  	w.b.Write(buf)
    89  }
    90  
    91  // BinReader is a binary protocol reader helper
    92  type BinReader struct {
    93  	b     *bufio.Reader
    94  	count *readers.CountingReader
    95  	err   error // keeps the first error encountered
    96  }
    97  
    98  // NewBinReader creates a binary protocol reader helper
    99  func NewBinReader(reader io.Reader) *BinReader {
   100  	r := &BinReader{}
   101  	r.count = readers.NewCountingReader(reader)
   102  	r.b = bufio.NewReader(r.count)
   103  	return r
   104  }
   105  
   106  // Count returns number of bytes read
   107  func (r *BinReader) Count() uint64 {
   108  	return r.count.BytesRead()
   109  }
   110  
   111  // Error returns first encountered error or nil
   112  func (r *BinReader) Error() error {
   113  	return r.err
   114  }
   115  
   116  // check() keeps the first error encountered in a stream
   117  func (r *BinReader) check(err error) bool {
   118  	if err == nil {
   119  		return true
   120  	}
   121  	if r.err == nil {
   122  		// keep the first error
   123  		r.err = err
   124  	}
   125  	if err != io.EOF {
   126  		log.Fatalf("Error parsing response: %v", err)
   127  	}
   128  	return false
   129  }
   130  
   131  // ReadByteAsInt reads a single byte as uint32, returns -1 for EOF or errors
   132  func (r *BinReader) ReadByteAsInt() int {
   133  	if octet, err := r.b.ReadByte(); r.check(err) {
   134  		return int(octet)
   135  	}
   136  	return -1
   137  }
   138  
   139  // ReadByteAsShort reads a single byte as uint16, returns -1 for EOF or errors
   140  func (r *BinReader) ReadByteAsShort() int16 {
   141  	if octet, err := r.b.ReadByte(); r.check(err) {
   142  		return int16(octet)
   143  	}
   144  	return -1
   145  }
   146  
   147  // ReadIntSpl reads two bytes as little-endian uint16, returns -1 for EOF or errors
   148  func (r *BinReader) ReadIntSpl() int {
   149  	var val uint16
   150  	if r.check(binary.Read(r.b, binary.LittleEndian, &val)) {
   151  		return int(val)
   152  	}
   153  	return -1
   154  }
   155  
   156  // ReadULong returns uint64 equivalent of -1 for EOF or errors
   157  func (r *BinReader) ReadULong() uint64 {
   158  	if val, err := binary.ReadUvarint(r.b); r.check(err) {
   159  		return val
   160  	}
   161  	return 0xffffffffffffffff
   162  }
   163  
   164  // ReadPu32 returns -1 for EOF or errors
   165  func (r *BinReader) ReadPu32() int64 {
   166  	if val, err := binary.ReadUvarint(r.b); r.check(err) {
   167  		return int64(val)
   168  	}
   169  	return -1
   170  }
   171  
   172  // ReadNBytes reads given number of bytes, returns invalid data for EOF or errors
   173  func (r *BinReader) ReadNBytes(len int) []byte {
   174  	buf := make([]byte, len)
   175  	n, err := r.b.Read(buf)
   176  	if r.check(err) {
   177  		return buf
   178  	}
   179  	if n != len {
   180  		r.check(ErrorPrematureEOF)
   181  	}
   182  	return buf
   183  }
   184  
   185  // ReadBytesByLength reads buffer length and its bytes
   186  func (r *BinReader) ReadBytesByLength() []byte {
   187  	len := r.ReadPu32()
   188  	if len < 0 {
   189  		r.check(ErrorInvalidLength)
   190  		return []byte{}
   191  	}
   192  	return r.ReadNBytes(int(len))
   193  }
   194  
   195  // ReadString reads a zero-terminated string with length
   196  func (r *BinReader) ReadString() string {
   197  	len := int(r.ReadPu32())
   198  	if len < 1 {
   199  		r.check(ErrorInvalidLength)
   200  		return ""
   201  	}
   202  	buf := make([]byte, len-1)
   203  	n, err := r.b.Read(buf)
   204  	if !r.check(err) {
   205  		return ""
   206  	}
   207  	if n != len-1 {
   208  		r.check(ErrorPrematureEOF)
   209  		return ""
   210  	}
   211  	zeroByte, err := r.b.ReadByte()
   212  	if !r.check(err) {
   213  		return ""
   214  	}
   215  	if zeroByte != 0 {
   216  		r.check(ErrorZeroTerminate)
   217  		return ""
   218  	}
   219  	return string(buf)
   220  }
   221  
   222  // ReadDate reads a Unix encoded time
   223  func (r *BinReader) ReadDate() time.Time {
   224  	return time.Unix(r.ReadPu32(), 0)
   225  }