@@ -44,6 +44,7 @@ type Conn struct {
44
44
readErr error
45
45
conn net.Conn
46
46
bufReader * bufio.Reader
47
+ reader io.Reader
47
48
header * Header
48
49
ProxyHeaderPolicy Policy
49
50
Validate Validator
@@ -138,9 +139,11 @@ func NewConn(conn net.Conn, opts ...func(*Conn)) *Conn {
138
139
// For v2 the header length is at most 52 bytes plus the length of the TLVs.
139
140
// We use 256 bytes to be safe.
140
141
const bufSize = 256
142
+ br := bufio .NewReaderSize (conn , bufSize )
141
143
142
144
pConn := & Conn {
143
- bufReader : bufio .NewReaderSize (conn , bufSize ),
145
+ bufReader : br ,
146
+ reader : io .MultiReader (br , conn ),
144
147
conn : conn ,
145
148
}
146
149
@@ -162,7 +165,7 @@ func (p *Conn) Read(b []byte) (int, error) {
162
165
return 0 , p .readErr
163
166
}
164
167
165
- return p .bufReader .Read (b )
168
+ return p .reader .Read (b )
166
169
}
167
170
168
171
// Write wraps original conn.Write
@@ -344,5 +347,27 @@ func (p *Conn) WriteTo(w io.Writer) (int64, error) {
344
347
if p .readErr != nil {
345
348
return 0 , p .readErr
346
349
}
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
348
373
}
0 commit comments