Skip to content

Commit b4098fd

Browse files
committed
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 7fe8e98 commit b4098fd

File tree

1 file changed

+27
-3
lines changed

1 file changed

+27
-3
lines changed

protocol.go

Lines changed: 27 additions & 3 deletions
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
@@ -129,9 +130,11 @@ func NewConn(conn net.Conn, opts ...func(*Conn)) *Conn {
129130
// For v2 the header length is at most 52 bytes plus the length of the TLVs.
130131
// We use 256 bytes to be safe.
131132
const bufSize = 256
133+
br := bufio.NewReaderSize(conn, bufSize)
132134

133135
pConn := &Conn{
134-
bufReader: bufio.NewReaderSize(conn, bufSize),
136+
bufReader: br,
137+
reader: io.MultiReader(br, conn),
135138
conn: conn,
136139
}
137140

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

156-
return p.bufReader.Read(b)
159+
return p.reader.Read(b)
157160
}
158161

159162
// Write wraps original conn.Write
@@ -335,5 +338,26 @@ func (p *Conn) WriteTo(w io.Writer) (int64, error) {
335338
if p.readErr != nil {
336339
return 0, p.readErr
337340
}
338-
return p.bufReader.WriteTo(w)
341+
342+
b := make([]byte, p.bufReader.Buffered())
343+
p.bufReader.Read(b)
344+
p.bufReader.Discard(len(b))
345+
346+
var n int64
347+
{
348+
nn, err := w.Write(b)
349+
n += int64(nn)
350+
if err != nil {
351+
return n, err
352+
}
353+
}
354+
{
355+
nn, err := io.Copy(w, p.conn)
356+
n += nn
357+
if err != nil {
358+
return n, err
359+
}
360+
}
361+
362+
return n, nil
339363
}

0 commit comments

Comments
 (0)