]> Cypherpunks.ru repositories - netstring.git/blobdiff - w.go
Refactoring, io.Reader/Writer friendliness, performance optimization
[netstring.git] / w.go
diff --git a/w.go b/w.go
new file mode 100644 (file)
index 0000000..aaf33d5
--- /dev/null
+++ b/w.go
@@ -0,0 +1,87 @@
+/*
+netstring -- netstring format serialization library
+Copyright (C) 2015-2020 Sergey Matveev <stargrave@stargrave.org>
+
+This program is free software: you can redistribute it and/or modify
+it under the terms of the GNU General Public License as published by
+the Free Software Foundation, version 3 of the License.
+
+This program is distributed in the hope that it will be useful,
+but WITHOUT ANY WARRANTY; without even the implied warranty of
+MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+GNU General Public License for more details.
+
+You should have received a copy of the GNU General Public License
+along with this program.  If not, see <http://www.gnu.org/licenses/>.
+*/
+
+package netstring
+
+import (
+       "bufio"
+       "errors"
+       "io"
+       "strconv"
+)
+
+type Writer struct {
+       w    *bufio.Writer
+       left uint64
+}
+
+// Create new Writer.
+// Pay attention that bufio.Writer is used to write to it.
+func NewWriter(w io.Writer) *Writer {
+       return &Writer{w: bufio.NewWriter(w)}
+}
+
+// Write size of the data going to be supplied to Write method.
+// It returns number of length prefixed written (possibly just
+// buffered).
+func (w *Writer) WriteSize(size uint64) (n int, err error) {
+       if w.left > 0 {
+               return 0, errors.New("current chunk in not written")
+       }
+       w.left = size
+       return w.w.WriteString(strconv.FormatUint(size, 10) + ":")
+}
+
+// Write the chunk data. WriteSize must preceed this call.
+// Write could be buffered for writing. Only when the last bytes are
+// written, then terminator is appended and all the data flushed.
+// Terminator is not taken in account of written bytes count!
+func (w *Writer) Write(buf []byte) (written int, err error) {
+       if w.left == 0 && len(buf) > 0 {
+               return 0, errors.New("chunk in already written")
+       }
+       written, err = w.w.Write(buf)
+       if err != nil {
+               return
+       }
+       w.left -= uint64(written)
+       if w.left == 0 {
+               _, err = w.w.Write([]byte{','})
+               if err != nil {
+                       return
+               }
+               err = w.w.Flush()
+       }
+       return
+}
+
+// Write the whole chunk at once. It could be convenient to use instead
+// of WriteSize/Write invocations
+func (w *Writer) WriteChunk(buf []byte) (n int, err error) {
+       n, err = w.WriteSize(uint64(len(buf)))
+       if err != nil {
+               return
+       }
+       var nw int
+       nw, err = w.Write(buf)
+       n += nw
+       if err != nil {
+               return
+       }
+       n++
+       return
+}