update forward handler
This commit is contained in:
@ -17,13 +17,15 @@ import (
|
||||
|
||||
type tcpConn struct {
|
||||
net.Conn
|
||||
wbuf bytes.Buffer
|
||||
wbuf *bytes.Buffer
|
||||
once sync.Once
|
||||
}
|
||||
|
||||
func (c *tcpConn) Read(b []byte) (n int, err error) {
|
||||
c.once.Do(func() {
|
||||
err = readResponse(c.Conn)
|
||||
if c.wbuf != nil {
|
||||
err = readResponse(c.Conn)
|
||||
}
|
||||
})
|
||||
|
||||
if err != nil {
|
||||
@ -34,7 +36,7 @@ func (c *tcpConn) Read(b []byte) (n int, err error) {
|
||||
|
||||
func (c *tcpConn) Write(b []byte) (n int, err error) {
|
||||
n = len(b) // force byte length consistent
|
||||
if c.wbuf.Len() > 0 {
|
||||
if c.wbuf != nil && c.wbuf.Len() > 0 {
|
||||
c.wbuf.Write(b) // append the data to the cached header
|
||||
_, err = c.Conn.Write(c.wbuf.Bytes())
|
||||
c.wbuf.Reset()
|
||||
@ -46,13 +48,15 @@ func (c *tcpConn) Write(b []byte) (n int, err error) {
|
||||
|
||||
type udpConn struct {
|
||||
net.Conn
|
||||
wbuf bytes.Buffer
|
||||
wbuf *bytes.Buffer
|
||||
once sync.Once
|
||||
}
|
||||
|
||||
func (c *udpConn) Read(b []byte) (n int, err error) {
|
||||
c.once.Do(func() {
|
||||
err = readResponse(c.Conn)
|
||||
if c.wbuf != nil {
|
||||
err = readResponse(c.Conn)
|
||||
}
|
||||
})
|
||||
if err != nil {
|
||||
return
|
||||
@ -84,7 +88,7 @@ func (c *udpConn) Write(b []byte) (n int, err error) {
|
||||
}
|
||||
|
||||
n = len(b)
|
||||
if c.wbuf.Len() > 0 {
|
||||
if c.wbuf != nil && c.wbuf.Len() > 0 {
|
||||
var bb [2]byte
|
||||
binary.BigEndian.PutUint16(bb[:], uint16(len(b)))
|
||||
c.wbuf.Write(bb[:])
|
||||
|
@ -1,6 +1,7 @@
|
||||
package relay
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"fmt"
|
||||
"net"
|
||||
@ -121,8 +122,9 @@ func (c *relayConnector) Connect(ctx context.Context, conn net.Conn, network, ad
|
||||
if !c.md.noDelay {
|
||||
cc := &tcpConn{
|
||||
Conn: conn,
|
||||
wbuf: &bytes.Buffer{},
|
||||
}
|
||||
if _, err := req.WriteTo(&cc.wbuf); err != nil {
|
||||
if _, err := req.WriteTo(cc.wbuf); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
conn = cc
|
||||
@ -132,7 +134,8 @@ func (c *relayConnector) Connect(ctx context.Context, conn net.Conn, network, ad
|
||||
Conn: conn,
|
||||
}
|
||||
if !c.md.noDelay {
|
||||
if _, err := req.WriteTo(&cc.wbuf); err != nil {
|
||||
cc.wbuf = &bytes.Buffer{}
|
||||
if _, err := req.WriteTo(cc.wbuf); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
}
|
||||
|
Reference in New Issue
Block a user