add redirect listen option

1 files changed, 43 insertions, 8 deletions

diff --git a/pkg/server/socket/listen.go b/pkg/server/socket/listen.go
index caf5fcf..e640e16 100644
--- a/pkg/server/socket/listen.go
+++ b/pkg/server/socket/listen.go
@@ -1,43 +1,78 @@
 package socket
 
 import (
+	"errors"
 	"fmt"
+	"log"
 	"net"
 	"strings"
 	"tunnel/pkg/server/env"
+	"tunnel/pkg/server/opts"
 )
 
-func newListenSocket(proto, addr string) (S, error) {
+type listenSocket struct {
+	proto, addr string
+	listen      net.Listener
+	redirect    bool
+}
+
+func newListenSocket(proto, addr string, opts opts.Opts) (S, error) {
+	redirect := opts.Bool("redirect")
+
 	if proto == "tcp" {
 		if !strings.Contains(addr, ":") {
 			addr = ":" + addr
 		}
 	}
 
+	if redirect && proto != "tcp" {
+		return nil, errors.New("redirect not supported")
+	}
+
 	listen, err := net.Listen(proto, addr)
 	if err != nil {
 		return nil, err
 	}
 
 	s := &listenSocket{
-		proto:  proto,
-		addr:   addr,
-		listen: listen,
+		proto:    proto,
+		addr:     addr,
+		listen:   listen,
+		redirect: redirect,
 	}
 
 	return s, nil
 }
 
-func (s *listenSocket) Open(env.Env) (Conn, error) {
+func (s *listenSocket) Open(env env.Env) (Conn, error) {
+	var original string
+
 	conn, err := s.listen.Accept()
 	if err != nil {
 		return nil, err
 	}
 
-	addr := conn.RemoteAddr()
-	info := fmt.Sprintf("<%s/%s", addr.Network(), addr)
+	la, ra := conn.LocalAddr(), conn.RemoteAddr()
+	desc := fmt.Sprintf("%s/%s->%s", la.Network(), ra, la)
+	info := fmt.Sprintf("<%s/%s", ra.Network(), ra)
+
+	if s.redirect {
+		if err := getConnOriginalAddr(conn, &original); err != nil {
+			log.Println("accept", desc, "failed")
+			conn.Close()
+			return nil, err
+		} else {
+			env.Set("original.addr", original)
+		}
+	}
+
+	if original == "" {
+		log.Println("accept", desc)
+	} else {
+		log.Println("accept", desc, "original", original)
+	}
 
-	return exported{info, newConn(conn)}, nil
+	return newConn(conn, desc, info), nil
 }
 
 func (s *listenSocket) String() string {