mox!

moxio/bufpool.go

@@ -0,0 +1,103 @@
+package moxio
+
+import (
+	"bufio"
+	"errors"
+	"fmt"
+	"io"
+
+	"github.com/mjl-/mox/mlog"
+)
+
+var xlog = mlog.New("moxio")
+
+// todo: instead of a bufpool, should maybe just make an alternative to bufio.Reader with a big enough buffer that we can fully use to read a line.
+
+var ErrLineTooLong = errors.New("line from remote too long") // Returned by Bufpool.Readline.
+
+// Bufpool caches byte slices for reuse during parsing of line-terminated commands.
+type Bufpool struct {
+	c    chan []byte
+	size int
+}
+
+// NewBufpool makes a new pool, initially empty, but holding at most "max" buffers of "size" bytes each.
+func NewBufpool(max, size int) *Bufpool {
+	return &Bufpool{
+		c:    make(chan []byte, max),
+		size: size,
+	}
+}
+
+// get returns a buffer from the pool if available, otherwise allocates a new buffer.
+// The buffer should be returned with a call to put.
+func (b *Bufpool) get() []byte {
+	var buf []byte
+
+	// Attempt to get buffer from pool. Otherwise create new buffer.
+	select {
+	case buf = <-b.c:
+	default:
+	}
+	if buf == nil {
+		buf = make([]byte, b.size)
+	}
+	return buf
+}
+
+// put puts a "buf" back in the pool. Put clears the first "n" bytes, which should
+// be all the bytes that have been read in the buffer. If the pool is full, the
+// buffer is discarded, and will be cleaned up by the garbage collector.
+// The caller should no longer reference "buf" after a call to put.
+func (b *Bufpool) put(buf []byte, n int) {
+	if len(buf) != b.size {
+		xlog.Error("buffer with bad size returned, ignoring", mlog.Field("badsize", len(buf)), mlog.Field("expsize", b.size))
+		return
+	}
+
+	for i := 0; i < n; i++ {
+		buf[i] = 0
+	}
+	select {
+	case b.c <- buf:
+	default:
+	}
+}
+
+// Readline reads a \n- or \r\n-terminated line. Line is returned without \n or \r\n.
+// If the line was too long, ErrLineTooLong is returned.
+// If an EOF is encountered before a \n, io.ErrUnexpectedEOF is returned.
+func (b *Bufpool) Readline(r *bufio.Reader) (line string, rerr error) {
+	var nread int
+	buf := b.get()
+	defer func() {
+		b.put(buf, nread)
+	}()
+
+	// Read until newline. If we reach the end of the buffer first, we write back an
+	// error and abort the connection because our protocols cannot be recovered. We
+	// don't want to consume data until we finally see a newline, which may be never.
+	for {
+		if nread >= len(buf) {
+			return "", fmt.Errorf("%w: no newline after all %d bytes", ErrLineTooLong, nread)
+		}
+		c, err := r.ReadByte()
+		if err == io.EOF {
+			return "", io.ErrUnexpectedEOF
+		} else if err != nil {
+			return "", fmt.Errorf("reading line from remote: %w", err)
+		}
+		if c == '\n' {
+			var s string
+			if nread > 0 && buf[nread-1] == '\r' {
+				s = string(buf[:nread-1])
+			} else {
+				s = string(buf[:nread])
+			}
+			nread++
+			return s, nil
+		}
+		buf[nread] = c
+		nread++
+	}
+}