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 authored and pires committed Oct 8, 2024
1 parent 9171270 commit 2df67b4
Showing 1 changed file with 28 additions and 3 deletions.
31 changes: 28 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 @@ -138,9 +139,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 @@ -162,7 +165,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 @@ -344,5 +347,27 @@ 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
}

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 2df67b4

Please sign in to comment.