Skip to content

Commit

Permalink
protocol: avoid double buffering
Browse files Browse the repository at this point in the history
Use buffer only to read the PROXY header.
Users may use they own buffers with they own buffer sizes and pools - connection should respect that.
  • Loading branch information
mmatczuk committed Sep 18, 2024
1 parent 2e5664f commit 45b522d
Showing 1 changed file with 31 additions and 3 deletions.
34 changes: 31 additions & 3 deletions protocol.go
Original file line number Diff line number Diff line change
Expand Up @@ -44,6 +44,7 @@ type Conn struct {
readErr error
conn net.Conn
bufReader *bufio.Reader
reader io.Reader
header *Header
ProxyHeaderPolicy Policy
Validate Validator
Expand Down Expand Up @@ -129,9 +130,11 @@ func NewConn(conn net.Conn, opts ...func(*Conn)) *Conn {
// For v2 the header length is at most 52 bytes plus the length of the TLVs.
// We use 256 bytes to be safe.
const bufSize = 256
br := bufio.NewReaderSize(conn, bufSize)

pConn := &Conn{
bufReader: bufio.NewReaderSize(conn, bufSize),
bufReader: br,
reader: io.MultiReader(br, conn),
conn: conn,
}

Expand All @@ -153,7 +156,7 @@ func (p *Conn) Read(b []byte) (int, error) {
return 0, p.readErr
}

return p.bufReader.Read(b)
return p.reader.Read(b)
}

// Write wraps original conn.Write
Expand Down Expand Up @@ -335,5 +338,30 @@ func (p *Conn) WriteTo(w io.Writer) (int64, error) {
if p.readErr != nil {
return 0, p.readErr
}
return p.bufReader.WriteTo(w)

b := make([]byte, p.bufReader.Buffered())
if _, err := p.bufReader.Read(b); err != nil {
return 0, err // this should never as we read buffered data
}
if _, err := p.bufReader.Discard(len(b)); err != nil {
return 0, err // this should never as we read buffered data
}

var n int64
{
nn, err := w.Write(b)
n += int64(nn)
if err != nil {
return n, err
}
}
{
nn, err := io.Copy(w, p.conn)
n += nn
if err != nil {
return n, err
}
}

return n, nil
}

0 comments on commit 45b522d

Please sign in to comment.