1package moxio

2

3import (

4 "bufio"

5 "errors"

6 "fmt"

7 "io"

8 "log/slog"

9

10 "github.com/mjl-/mox/mlog"

11)

12

13// 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.

14

15var ErrLineTooLong = errors.New("line from remote too long") // Returned by Bufpool.Readline.

16

17// Bufpool caches byte slices for reuse during parsing of line-terminated commands.

18type Bufpool struct {

19 c chan []byte

20 size int

21}

22

23// NewBufpool makes a new pool, initially empty, but holding at most "max" buffers of "size" bytes each.

24func NewBufpool(max, size int) *Bufpool {

25 return &Bufpool{

26 c: make(chan []byte, max),

27 size: size,

28 }

29}

30

31// get returns a buffer from the pool if available, otherwise allocates a new buffer.

32// The buffer should be returned with a call to put.

33func (b *Bufpool) get() []byte {

34 var buf []byte

35

36 // Attempt to get buffer from pool. Otherwise create new buffer.

37 select {

38 case buf = <-b.c:

39 default:

40 }

41 if buf == nil {

42 buf = make([]byte, b.size)

43 }

44 return buf

45}

46

47// put puts a "buf" back in the pool. Put clears the first "n" bytes, which should

48// be all the bytes that have been read in the buffer. If the pool is full, the

49// buffer is discarded, and will be cleaned up by the garbage collector.

50// The caller should no longer reference "buf" after a call to put.

51func (b *Bufpool) put(log mlog.Log, buf []byte, n int) {

52 if len(buf) != b.size {

53 log.Error("buffer with bad size returned, ignoring", slog.Int("badsize", len(buf)), slog.Int("expsize", b.size))

54 return

55 }

56

57 for i := 0; i < n; i++ {

58 buf[i] = 0

59 }

60 select {

61 case b.c <- buf:

62 default:

63 }

64}

65

66// Readline reads a \n- or \r\n-terminated line. Line is returned without \n or \r\n.

67// If the line was too long, ErrLineTooLong is returned.

68// If an EOF is encountered before a \n, io.ErrUnexpectedEOF is returned.

69func (b *Bufpool) Readline(log mlog.Log, r *bufio.Reader) (line string, rerr error) {

70 var nread int

71 buf := b.get()

72 defer func() {

73 b.put(log, buf, nread)

74 }()

75

76 // Read until newline. If we reach the end of the buffer first, we write back an

77 // error and abort the connection because our protocols cannot be recovered. We

78 // don't want to consume data until we finally see a newline, which may be never.

79 for {

80 if nread >= len(buf) {

81 return "", fmt.Errorf("%w: no newline after all %d bytes", ErrLineTooLong, nread)

82 }

83 c, err := r.ReadByte()

84 if err == io.EOF {

85 return "", io.ErrUnexpectedEOF

86 } else if err != nil {

87 return "", fmt.Errorf("reading line from remote: %w", err)

88 }

89 if c == '\n' {

90 var s string

91 if nread > 0 && buf[nread-1] == '\r' {

92 s = string(buf[:nread-1])

93 } else {

94 s = string(buf[:nread])

95 }

96 nread++

97 return s, nil

98 }

99 buf[nread] = c

100 nread++

101 }

102}

103