mirror of
https://github.com/mjl-/mox.git
synced 2024-12-27 17:03:47 +03:00
102 lines
2.7 KiB
Go
102 lines
2.7 KiB
Go
package moxio
|
|
|
|
import (
|
|
"bufio"
|
|
"errors"
|
|
"fmt"
|
|
"io"
|
|
"log/slog"
|
|
|
|
"github.com/mjl-/mox/mlog"
|
|
)
|
|
|
|
// 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(log mlog.Log, buf []byte, n int) {
|
|
if len(buf) != b.size {
|
|
log.Error("buffer with bad size returned, ignoring", slog.Int("badsize", len(buf)), slog.Int("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(log mlog.Log, r *bufio.Reader) (line string, rerr error) {
|
|
var nread int
|
|
buf := b.get()
|
|
defer func() {
|
|
b.put(log, 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++
|
|
}
|
|
}
|