add origin host and port to env

1 files changed, 24 insertions, 12 deletions

diff --git a/pkg/server/socket/listen.go b/pkg/server/socket/listen.go
index 9d19677..6848620 100644
--- a/pkg/server/socket/listen.go
+++ b/pkg/server/socket/listen.go
@@ -52,36 +52,48 @@ func (s *listenSocket) Prepare(e env.Env) error {
 	return nil
 }
 
+func setOriginHostPort(env env.Env, origin, host, port string) {
+	env.Set("origin", origin)
+	env.Set("origin.host", host)
+	env.Set("origin.port", port)
+}
+
 func (s *listenSocket) New(env env.Env) (Conn, error) {
-	var original string
+	var origin string
 
 	conn, err := s.listen.Accept()
 	if err != nil {
 		return nil, err
 	}
 
-	la, ra := conn.LocalAddr(), conn.RemoteAddr()
-	desc := fmt.Sprintf("%s/%s->%s", la.Network(), ra, la)
+	ra := conn.RemoteAddr()
+	desc := fmt.Sprintf("%s/%s->%s", s.Proto, ra, s.Addr)
 	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", original)
+		host, port, err := getConnOriginalHostPort(conn)
+		if err != nil {
+			defer conn.Close()
+			return nil, fmt.Errorf("accept %s failed: %s", desc, err)
 		}
+		origin = net.JoinHostPort(host, port)
+		setOriginHostPort(env, origin, host, port)
 	}
 
 	if s.Tproxy {
-		env.Set("original", la.String())
+		origin = conn.LocalAddr().String()
+		host, port, err := net.SplitHostPort(origin)
+		if err != nil {
+			defer conn.Close()
+			return nil, fmt.Errorf("accept %s failed: %s", desc, err)
+		}
+		setOriginHostPort(env, origin, host, port)
 	}
 
-	if original == "" {
+	if origin == "" {
 		log.Println("accept", desc)
 	} else {
-		log.Println("accept", desc, "original", original)
+		log.Println("accept", desc, "origin", origin)
 	}
 
 	return newConn(conn, desc, info), nil