w.go

   1 // netstring -- netstring format serialization library
   2 // Copyright (C) 2015-2024 Sergey Matveev <stargrave@stargrave.org>
   3 //
   4 // This program is free software: you can redistribute it and/or modify
   5 // it under the terms of the GNU General Public License as published by
   6 // the Free Software Foundation, version 3 of the License.
   7 //
   8 // This program is distributed in the hope that it will be useful,
   9 // but WITHOUT ANY WARRANTY; without even the implied warranty of
  10 // MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
  11 // GNU General Public License for more details.
  12 //
  13 // You should have received a copy of the GNU General Public License
  14 // along with this program.  If not, see <http://www.gnu.org/licenses/>.
  15
  16 package netstring
  17
  18 import (
  19         "bufio"
  20         "errors"
  21         "io"
  22         "strconv"
  23 )
  24
  25 type Writer struct {
  26         w    *bufio.Writer
  27         left uint64
  28 }
  29
  30 // Create new Writer.
  31 // Pay attention that bufio.Writer is used to write to it.
  32 func NewWriter(w io.Writer) *Writer {
  33         return &Writer{w: bufio.NewWriter(w)}
  34 }
  35
  36 // Write size of the data going to be supplied to Write method.
  37 // It returns number of length prefixed written (possibly just
  38 // buffered).
  39 func (w *Writer) WriteSize(size uint64) (n int, err error) {
  40         if w.left > 0 {
  41                 return 0, errors.New("current chunk in not written")
  42         }
  43         w.left = size
  44         return w.w.WriteString(strconv.FormatUint(size, 10) + ":")
  45 }
  46
  47 // Write the chunk data. WriteSize must preceed this call.
  48 // Write could be buffered for writing. Only when the last bytes are
  49 // written, then terminator is appended and all the data flushed.
  50 // Terminator is not taken in account of written bytes count!
  51 func (w *Writer) Write(buf []byte) (written int, err error) {
  52         if w.left == 0 && len(buf) > 0 {
  53                 return 0, errors.New("chunk in already written")
  54         }
  55         written, err = w.w.Write(buf)
  56         if err != nil {
  57                 return
  58         }
  59         w.left -= uint64(written)
  60         if w.left == 0 {
  61                 _, err = w.w.Write([]byte{','})
  62                 if err != nil {
  63                         return
  64                 }
  65                 err = w.w.Flush()
  66         }
  67         return
  68 }
  69
  70 // Write the whole chunk at once. It could be convenient to use instead
  71 // of WriteSize/Write invocations
  72 func (w *Writer) WriteChunk(buf []byte) (n int, err error) {
  73         n, err = w.WriteSize(uint64(len(buf)))
  74         if err != nil {
  75                 return
  76         }
  77         var nw int
  78         nw, err = w.Write(buf)
  79         n += nw
  80         if err != nil {
  81                 return
  82         }
  83         n++
  84         return
  85 }