add bind for relay

pkg/connector/relay/conn.go

@@ -6,45 +6,53 @@ import (
 	"errors"
 	"fmt"
 	"io"
+	"math"
 	"net"
 	"sync"
 
-	"github.com/go-gost/gost/pkg/logger"
 	"github.com/go-gost/relay"
 )
 
-type conn struct {
+type tcpConn struct {
 	net.Conn
-	udp        bool
-	wbuf       bytes.Buffer
-	once       sync.Once
-	headerSent bool
-	logger     logger.Logger
+	wbuf bytes.Buffer
+	once sync.Once
 }
 
-func (c *conn) Read(b []byte) (n int, err error) {
+func (c *tcpConn) Read(b []byte) (n int, err error) {
 	c.once.Do(func() {
-		resp := relay.Response{}
-		_, err = resp.ReadFrom(c.Conn)
-		if err != nil {
-			return
-		}
-		if resp.Version != relay.Version1 {
-			err = relay.ErrBadVersion
-			return
-		}
-		if resp.Status != relay.StatusOK {
-			err = fmt.Errorf("status %d", resp.Status)
-			return
-		}
+		err = readResponse(c.Conn)
 	})
 
 	if err != nil {
 		return
 	}
+	return c.Conn.Read(b)
+}
 
-	if !c.udp {
-		return c.Conn.Read(b)
+func (c *tcpConn) Write(b []byte) (n int, err error) {
+	n = len(b) // force byte length consistent
+	if c.wbuf.Len() > 0 {
+		c.wbuf.Write(b) // append the data to the cached header
+		_, err = c.wbuf.WriteTo(c.Conn)
+		return
+	}
+	_, err = c.Conn.Write(b)
+	return
+}
+
+type udpConn struct {
+	net.Conn
+	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 err != nil {
+		return
 	}
 
 	var bb [2]byte
@@ -52,6 +60,7 @@ func (c *conn) Read(b []byte) (n int, err error) {
 	if err != nil {
 		return
 	}
+
 	dlen := int(binary.BigEndian.Uint16(bb[:]))
 	if len(b) >= dlen {
 		return io.ReadFull(c.Conn, b[:dlen])
@@ -59,59 +68,64 @@ func (c *conn) Read(b []byte) (n int, err error) {
 	buf := make([]byte, dlen)
 	_, err = io.ReadFull(c.Conn, buf)
 	n = copy(b, buf)
+
 	return
 }
 
-func (c *conn) ReadFrom(b []byte) (n int, addr net.Addr, err error) {
-	n, err = c.Read(b)
-	addr = c.Conn.RemoteAddr()
-	return
-}
-
-func (c *conn) Write(b []byte) (n int, err error) {
-	if len(b) > 0xFFFF {
+func (c *udpConn) Write(b []byte) (n int, err error) {
+	if len(b) > math.MaxUint16 {
 		err = errors.New("write: data maximum exceeded")
 		return
 	}
-	n = len(b) // force byte length consistent
+
+	n = len(b)
 	if c.wbuf.Len() > 0 {
-		if c.udp {
-			var bb [2]byte
-			binary.BigEndian.PutUint16(bb[:2], uint16(len(b)))
-			c.wbuf.Write(bb[:])
-			c.headerSent = true
-		}
+		var bb [2]byte
+		binary.BigEndian.PutUint16(bb[:], uint16(len(b)))
+		c.wbuf.Write(bb[:])
 		c.wbuf.Write(b) // append the data to the cached header
-		// _, err = c.Conn.Write(c.wbuf.Bytes())
-		// c.wbuf.Reset()
 		_, err = c.wbuf.WriteTo(c.Conn)
 		return
 	}
 
-	if !c.udp {
-		return c.Conn.Write(b)
-	}
-	if !c.headerSent {
-		c.headerSent = true
-		b2 := make([]byte, len(b)+2)
-		copy(b2, b)
-		_, err = c.Conn.Write(b2)
+	var bb [2]byte
+	binary.BigEndian.PutUint16(bb[:], uint16(len(b)))
+	_, err = c.Conn.Write(bb[:])
+	if err != nil {
 		return
 	}
-	nsize := 2 + len(b)
-	var buf []byte
-	if nsize <= mediumBufferSize {
-		buf = mPool.Get().([]byte)
-		defer mPool.Put(buf)
-	} else {
-		buf = make([]byte, nsize)
-	}
-	binary.BigEndian.PutUint16(buf[:2], uint16(len(b)))
-	n = copy(buf[2:], b)
-	_, err = c.Conn.Write(buf[:nsize])
-	return
+	return c.Conn.Write(b)
 }
 
-func (c *relayConn) WriteTo(b []byte, addr net.Addr) (n int, err error) {
-	return c.Write(b)
+func readResponse(r io.Reader) (err error) {
+	resp := relay.Response{}
+	_, err = resp.ReadFrom(r)
+	if err != nil {
+		return
+	}
+
+	if resp.Version != relay.Version1 {
+		err = relay.ErrBadVersion
+		return
+	}
+
+	if resp.Status != relay.StatusOK {
+		err = fmt.Errorf("status %d", resp.Status)
+		return
+	}
+	return nil
+}
+
+type bindConn struct {
+	net.Conn
+	localAddr  net.Addr
+	remoteAddr net.Addr
+}
+
+func (c *bindConn) LocalAddr() net.Addr {
+	return c.localAddr
+}
+
+func (c *bindConn) RemoteAddr() net.Addr {
+	return c.remoteAddr
 }