1 files changed, 30 insertions, 36 deletions

diff --git a/pkg/server/socket/listen.go b/pkg/server/socket/listen.go
index 910e5de..2c2f184 100644
--- a/pkg/server/socket/listen.go
+++ b/pkg/server/socket/listen.go
@@ -5,64 +5,54 @@ import (
 	"fmt"
 	"log"
 	"net"
-	"strings"
+
 	"tunnel/pkg/server/env"
-	"tunnel/pkg/server/opts"
 )
 
 type listenSocket struct {
-	proto, addr string
-	listen      net.Listener
-	redirect    bool
-	tproxy      bool
-}
+	Proto string `opts:"default:tcp"`
+	Addr  string `opts:"required"`
 
-func newListenSocket(proto, addr string, opts opts.Opts) (S, error) {
-	redirect := opts.Bool("redirect")
-	tproxy := opts.Bool("tproxy")
+	Redirect bool
+	Tproxy   bool
 
-	if proto == "tcp" {
-		if !strings.Contains(addr, ":") {
-			addr = ":" + addr
-		}
-	}
+	listen net.Listener
+}
 
-	if redirect && proto != "tcp" {
-		return nil, errors.New("redirect not supported")
+func (s *listenSocket) Prepare(e env.Env) error {
+	if s.Redirect && s.Proto != "tcp" {
+		return errors.New("redirect not supported")
 	}
 
-	if tproxy && proto != "tcp" {
-		return nil, errors.New("tproxy not supported")
+	if s.Tproxy && s.Proto != "tcp" {
+		return errors.New("tproxy not supported")
 	}
 
-	if redirect && tproxy {
-		return nil, errors.New("redirect and tproxy cannot be used together")
+	if s.Redirect && s.Tproxy {
+		return errors.New("redirect and tproxy cannot be used together")
 	}
 
+	proto, addr := parseProtoAddr(s.Proto, s.Addr)
 	listen, err := net.Listen(proto, addr)
 	if err != nil {
-		return nil, err
+		return err
 	}
 
-	if tproxy {
+	e.Set("listen", listen.Addr().String())
+
+	if s.Tproxy {
 		if err := setConnTransparent(listen); err != nil {
 			listen.Close()
-			return nil, err
+			return err
 		}
 	}
 
-	s := &listenSocket{
-		proto:    proto,
-		addr:     addr,
-		listen:   listen,
-		redirect: redirect,
-		tproxy:   tproxy,
-	}
+	s.listen = listen
 
-	return s, nil
+	return nil
 }
 
-func (s *listenSocket) Open(env env.Env) (Conn, error) {
+func (s *listenSocket) New(env env.Env) (Conn, error) {
 	var original string
 
 	conn, err := s.listen.Accept()
@@ -74,7 +64,7 @@ func (s *listenSocket) Open(env env.Env) (Conn, error) {
 	desc := fmt.Sprintf("%s/%s->%s", la.Network(), ra, la)
 	info := fmt.Sprintf("<%s/%s", ra.Network(), ra)
 
-	if s.redirect {
+	if s.Redirect {
 		if err := getConnOriginalAddr(conn, &original); err != nil {
 			log.Println("accept", desc, "failed")
 			conn.Close()
@@ -84,7 +74,7 @@ func (s *listenSocket) Open(env env.Env) (Conn, error) {
 		}
 	}
 
-	if s.tproxy {
+	if s.Tproxy {
 		env.Set("original", la.String())
 	}
 
@@ -98,9 +88,13 @@ func (s *listenSocket) Open(env env.Env) (Conn, error) {
 }
 
 func (s *listenSocket) String() string {
-	return fmt.Sprintf("%s/%s,listen", s.proto, s.addr)
+	return fmt.Sprintf("%s/%s,listen", s.Proto, s.Addr)
 }
 
 func (s *listenSocket) Close() {
 	s.listen.Close()
 }
+
+func init() {
+	register("listen", listenSocket{})
+}