File tree Expand file tree Collapse file tree 1 file changed +27
-3
lines changed Expand file tree Collapse file tree 1 file changed +27
-3
lines changed Original file line number Diff line number Diff 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}
You can’t perform that action at this time.
0 commit comments