Split listener into subpackage

listen/listen.go

@@ -0,0 +1,51 @@
+package listen
+
+import (
+	"errors"
+	"net"
+)
+
+var listeners = make(map[string]func(network, address string) (net.Listener, error))
+
+var ErrUnsupported = errors.New("unsupported")
+
+func Listen(kind, network, address string) (net.Listener, error) {
+	f, ok := listeners[kind]
+	if ok {
+		return f(network, address)
+	}
+	return nil, ErrUnsupported
+}
+
+type strAddr string
+
+func (a strAddr) Network() string { return "tcp" }
+func (a strAddr) String() string  { return string(a) }
+
+type targetConn struct {
+	net.Conn
+	target net.Addr
+}
+
+func (c *targetConn) LocalAddr() net.Addr { return c.target }
+
+type targetListener struct {
+	net.Listener
+	target net.Addr
+}
+
+func ListenTo(network, address, target string) (net.Listener, error) {
+	l, err := net.Listen(network, address)
+	if err != nil {
+		return nil, err
+	}
+	return &targetListener{l, strAddr(target)}, err
+}
+
+func (l *targetListener) Accept() (net.Conn, error) {
+	c, err := l.Listener.Accept()
+	if err != nil {
+		return nil, err
+	}
+	return &targetConn{c, l.target}, nil
+}

listener_linux.go → listen/nf_linux.go

@@ -1,4 +1,4 @@
-package main
+package listen
 
 import (
 	"context"
@@ -54,9 +54,5 @@ func tproxyListen(network, address string) (net.Listener, error) {
 		}
 		return err2
 	}}
-	l, err := lcfg.Listen(context.Background(), network, address)
-	if err != nil {
-		return nil, err
-	}
-	return l, nil
+	return lcfg.Listen(context.Background(), network, address)
 }

listener_darwin.go → listen/pf_darwin.go

@@ -1,4 +1,4 @@
-package main
+package listen
 
 import (
 	"net"

listen/socks.go

@@ -0,0 +1,39 @@
+package listen
+
+import (
+	"net"
+
+	"github.com/riobard/go-shadowsocks2/socks"
+)
+
+func init() {
+	listeners["socks"] = socksListen
+}
+
+type socksConn struct{ net.Conn }
+
+func (sc socksConn) LocalAddr() net.Addr {
+	addr, err := socks.Handshake(sc.Conn)
+	if err != nil {
+		return nil
+	}
+	return strAddr(addr.String())
+}
+
+type socksListener struct{ net.Listener }
+
+func (l *socksListener) Accept() (net.Conn, error) {
+	c, err := l.Listener.Accept()
+	if err != nil {
+		return nil, err
+	}
+	return socksConn{c}, nil
+}
+
+func socksListen(network, addr string) (net.Listener, error) {
+	l, err := net.Listen(network, addr)
+	if err != nil {
+		return nil, err
+	}
+	return &socksListener{l}, nil
+}

main.go

@@ -11,6 +11,7 @@ import (
 	"time"
 
 	"github.com/riobard/go-shadowsocks2/core"
+	"github.com/riobard/go-shadowsocks2/listen"
 )
 
 var config struct {
@@ -33,7 +34,6 @@ func main() {
 		UDPTun    PairList
 		Socks     string
 		RedirTCP  string
-		RedirTCP6 string
 		TproxyTCP string
 	}
 
@@ -82,7 +82,7 @@ func main() {
 
 		if len(flags.TCPTun) > 0 {
 			for _, p := range flags.TCPTun {
-				l, err := tunListen("tcp", p[0], p[1])
+				l, err := listen.ListenTo("tcp", p[0], p[1])
 				if err != nil {
 					log.Fatal(err)
 				}
@@ -92,7 +92,7 @@ func main() {
 		}
 
 		if flags.Socks != "" {
-			l, err := socksListen("tcp", flags.Socks)
+			l, err := listen.Listen("socks", "tcp", flags.Socks)
 			if err != nil {
 				log.Fatal(err)
 			}
@@ -101,7 +101,7 @@ func main() {
 		}
 
 		if flags.RedirTCP != "" {
-			l, err := listen("redir", "tcp", flags.RedirTCP)
+			l, err := listen.Listen("redir", "tcp", flags.RedirTCP)
 			if err != nil {
 				log.Fatal(err)
 			}
@@ -110,7 +110,7 @@ func main() {
 		}
 
 		if flags.TproxyTCP != "" {
-			l, err := listen("tproxy", "tcp", flags.TproxyTCP)
+			l, err := listen.Listen("tproxy", "tcp", flags.TproxyTCP)
 			if err != nil {
 				log.Fatal(err)
 			}

tcp.go

@@ -19,6 +19,10 @@ func tcpLocal(l net.Listener, d Dialer) {
 		go func() {
 			defer c.Close()
 			laddr := c.LocalAddr()
+			if laddr == nil {
+				logf("failed to determine target address")
+				return
+			}
 			rc, err := d.Dial(laddr.Network(), laddr.String())
 			if err != nil {
 				logf("failed to connect: %v", err)