r.go

   1 /*
   2 netstring -- netstring format serialization library
   3 Copyright (C) 2015-2021 Sergey Matveev <stargrave@stargrave.org>
   4
   5 This program is free software: you can redistribute it and/or modify
   6 it under the terms of the GNU General Public License as published by
   7 the Free Software Foundation, version 3 of the License.
   8
   9 This program is distributed in the hope that it will be useful,
  10 but WITHOUT ANY WARRANTY; without even the implied warranty of
  11 MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
  12 GNU General Public License for more details.
  13
  14 You should have received a copy of the GNU General Public License
  15 along with this program.  If not, see <http://www.gnu.org/licenses/>.
  16 */
  17
  18 package netstring
  19
  20 import (
  21         "bufio"
  22         "bytes"
  23         "errors"
  24         "io"
  25         "strconv"
  26 )
  27
  28 type Reader struct {
  29         r    *bufio.Reader
  30         left uint64
  31         eof  bool
  32 }
  33
  34 // Create new Reader.
  35 // Pay attention that bufio.Reader is used to read from it.
  36 func NewReader(r io.Reader) *Reader {
  37         return &Reader{r: bufio.NewReader(r), eof: true}
  38 }
  39
  40 // Parse netstring prefix. It returns data length.
  41 // After this method you have to call either Read() or Discard().
  42 // io.EOF is returned when underlying reader has no data anymore.
  43 func (r *Reader) Next() (uint64, error) {
  44         if !r.eof {
  45                 return 0, errors.New("current chunk is unread")
  46         }
  47         p, _ := r.r.Peek(21)
  48         if len(p) == 0 {
  49                 return 0, io.EOF
  50         }
  51         idx := bytes.Index(p, []byte{':'})
  52         if idx == -1 {
  53                 return 0, errors.New("no length separator found")
  54         }
  55         size, err := strconv.ParseUint(string(p[:idx]), 10, 64)
  56         if err != nil {
  57                 return 0, err
  58         }
  59         if _, err = r.r.Discard(idx + 1); err != nil {
  60                 return 0, err
  61         }
  62         r.left = size
  63         r.eof = false
  64         return size, nil
  65 }
  66
  67 func (r *Reader) checkTerminator() error {
  68         b, err := r.r.ReadByte()
  69         if err != nil {
  70                 return err
  71         }
  72         if b != ',' {
  73                 return errors.New("no terminator found")
  74         }
  75         return nil
  76 }
  77
  78 // Read current netstring chunk.
  79 // This method must be called after Next().
  80 // Terminator check and error raising are performed only at the end.
  81 func (r *Reader) Read(buf []byte) (n int, err error) {
  82         if r.eof {
  83                 return 0, io.EOF
  84         }
  85         if uint64(len(buf)) > r.left {
  86                 buf = buf[:r.left]
  87         }
  88         n, err = r.r.Read(buf)
  89         r.left -= uint64(n)
  90         if r.left == 0 {
  91                 err = r.checkTerminator()
  92                 if err == nil {
  93                         r.eof = true
  94                 }
  95         }
  96         return
  97 }
  98
  99 // Discard remaining bytes in the chunk, possibly fully skipping it.
 100 // This method must be called after Next().
 101 func (r *Reader) Discard() (err error) {
 102         _, err = r.r.Discard(int(r.left))
 103         if err != nil {
 104                 return
 105         }
 106         err = r.checkTerminator()
 107         if err != nil {
 108                 return
 109         }
 110         r.eof = true
 111         return
 112 }