@@ -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