w.go

   1 /*
   2 netstring -- netstring format serialization library
   3 Copyright (C) 2015-2020 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         "errors"
  23         "io"
  24         "strconv"
  25 )
  26
  27 type Writer struct {
  28         w    *bufio.Writer
  29         left uint64
  30 }
  31
  32 // Create new Writer.
  33 // Pay attention that bufio.Writer is used to write to it.
  34 func NewWriter(w io.Writer) *Writer {
  35         return &Writer{w: bufio.NewWriter(w)}
  36 }
  37
  38 // Write size of the data going to be supplied to Write method.
  39 // It returns number of length prefixed written (possibly just
  40 // buffered).
  41 func (w *Writer) WriteSize(size uint64) (n int, err error) {
  42         if w.left > 0 {
  43                 return 0, errors.New("current chunk in not written")
  44         }
  45         w.left = size
  46         return w.w.WriteString(strconv.FormatUint(size, 10) + ":")
  47 }
  48
  49 // Write the chunk data. WriteSize must preceed this call.
  50 // Write could be buffered for writing. Only when the last bytes are
  51 // written, then terminator is appended and all the data flushed.
  52 // Terminator is not taken in account of written bytes count!
  53 func (w *Writer) Write(buf []byte) (written int, err error) {
  54         if w.left == 0 && len(buf) > 0 {
  55                 return 0, errors.New("chunk in already written")
  56         }
  57         written, err = w.w.Write(buf)
  58         if err != nil {
  59                 return
  60         }
  61         w.left -= uint64(written)
  62         if w.left == 0 {
  63                 _, err = w.w.Write([]byte{','})
  64                 if err != nil {
  65                         return
  66                 }
  67                 err = w.w.Flush()
  68         }
  69         return
  70 }
  71
  72 // Write the whole chunk at once. It could be convenient to use instead
  73 // of WriteSize/Write invocations
  74 func (w *Writer) WriteChunk(buf []byte) (n int, err error) {
  75         n, err = w.WriteSize(uint64(len(buf)))
  76         if err != nil {
  77                 return
  78         }
  79         var nw int
  80         nw, err = w.Write(buf)
  81         n += nw
  82         if err != nil {
  83                 return
  84         }
  85         n++
  86         return
  87 }