Skip to content

Commit 2df67b4

Browse files
mmatczukpires
authored andcommitted
protocol: avoid double buffering
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.
1 parent 9171270 commit 2df67b4

File tree

1 file changed

+28
-3
lines changed

1 file changed

+28
-3
lines changed

protocol.go

+28-3
Original file line numberDiff line numberDiff line change
@@ -44,6 +44,7 @@ type Conn struct {
4444
readErr error
4545
conn net.Conn
4646
bufReader *bufio.Reader
47+
reader io.Reader
4748
header *Header
4849
ProxyHeaderPolicy Policy
4950
Validate Validator
@@ -138,9 +139,11 @@ func NewConn(conn net.Conn, opts ...func(*Conn)) *Conn {
138139
// For v2 the header length is at most 52 bytes plus the length of the TLVs.
139140
// We use 256 bytes to be safe.
140141
const bufSize = 256
142+
br := bufio.NewReaderSize(conn, bufSize)
141143

142144
pConn := &Conn{
143-
bufReader: bufio.NewReaderSize(conn, bufSize),
145+
bufReader: br,
146+
reader: io.MultiReader(br, conn),
144147
conn: conn,
145148
}
146149

@@ -162,7 +165,7 @@ func (p *Conn) Read(b []byte) (int, error) {
162165
return 0, p.readErr
163166
}
164167

165-
return p.bufReader.Read(b)
168+
return p.reader.Read(b)
166169
}
167170

168171
// Write wraps original conn.Write
@@ -344,5 +347,27 @@ func (p *Conn) WriteTo(w io.Writer) (int64, error) {
344347
if p.readErr != nil {
345348
return 0, p.readErr
346349
}
347-
return p.bufReader.WriteTo(w)
350+
351+
b := make([]byte, p.bufReader.Buffered())
352+
if _, err := p.bufReader.Read(b); err != nil {
353+
return 0, err // this should never as we read buffered data
354+
}
355+
356+
var n int64
357+
{
358+
nn, err := w.Write(b)
359+
n += int64(nn)
360+
if err != nil {
361+
return n, err
362+
}
363+
}
364+
{
365+
nn, err := io.Copy(w, p.conn)
366+
n += nn
367+
if err != nil {
368+
return n, err
369+
}
370+
}
371+
372+
return n, nil
348373
}

0 commit comments

Comments
 (0)