Fix net.Conn interface.

nice-journey · Aug 6, 2015 · 9f6e712 · 9f6e712
1 parent fef07e2
commit 9f6e712
Showing 1 changed file with 52 additions and 2 deletions.
diff --git a/shadowsocks/proxy.go b/shadowsocks/proxy.go
@@ -4,6 +4,8 @@ import (
 	"errors"
 	"strings"
 	"fmt"
+	"net"
+	"time"
 )
 
 type Dialer struct {
@@ -12,6 +14,16 @@ type Dialer struct {
 	support_udp bool
 }
 
+type ProxyConn struct {
+	*Conn
+	raddr *ProxyAddr
+}
+
+type ProxyAddr struct {
+	network string
+	address string
+}
+
 var ErrNilCipher = errors.New("cipher can't be nil.")
 
 func NewDialer(server string, cipher *Cipher) (dialer *Dialer, err error) {
@@ -26,9 +38,47 @@ func NewDialer(server string, cipher *Cipher) (dialer *Dialer, err error) {
 	}, nil
 }
 
-func (d *Dialer) Dial(network, addr string) (c *Conn, err error) {
+func (d *Dialer) Dial(network, addr string) (c *ProxyConn, err error) {
 	if strings.HasPrefix(network, "tcp") {
-		return Dial(addr, d.server, d.cipher)
+		conn, err := Dial(addr, d.server, d.cipher)
+		if err != nil {
+			return nil, err
+		}
+		return &ProxyConn {
+			Conn: conn,
+			raddr: &ProxyAddr {
+				network: network,
+				address: addr,
+			},
+		}, nil
 	}
 	return nil, fmt.Errorf("unsupported connection type: %s", network)
 }
+
+func (c *ProxyConn) LocalAddr() net.Addr {
+	return c.Conn.LocalAddr()
+}
+
+func (c *ProxyConn) RemoteAddr() net.Addr {
+	return c.raddr
+}
+
+func (c *ProxyConn) SetDeadline(t time.Time) error {
+	return c.Conn.SetDeadline(t)
+}
+
+func (c *ProxyConn) SetReadDeadline(t time.Time) error {
+	return c.Conn.SetReadDeadline(t)
+}
+
+func (c *ProxyConn) SetWriteDeadline(t time.Time) error {
+	return c.Conn.SetWriteDeadline(t)
+}
+
+func (a *ProxyAddr) Network() string {
+	return a.network
+}
+
+func (a *ProxyAddr) String() string {
+	return a.address
+}