From 8ff56b5bb83c5c09d592ca22fb831135cf74efac Mon Sep 17 00:00:00 2001
From: xishang0128 <xishang02@gmail.com>
Date: Thu, 25 Apr 2024 11:48:53 +0800
Subject: [PATCH] chore: Add InUser for http/socks/mixed

---
 listener/http/proxy.go     | 13 +++++++------
 listener/socks/tcp.go      |  6 ++++--
 transport/socks4/socks4.go |  5 +++--
 transport/socks5/socks5.go |  4 ++--
 4 files changed, 16 insertions(+), 12 deletions(-)

diff --git a/listener/http/proxy.go b/listener/http/proxy.go
index 4822eabce..68e987084 100644
--- a/listener/http/proxy.go
+++ b/listener/http/proxy.go
@@ -51,8 +51,8 @@ func HandleConn(c net.Conn, tunnel C.Tunnel, cache *lru.LruCache[string, bool],
 		var resp *http.Response
 
 		if !trusted {
-			resp = authenticate(request, cache)
-
+			resp, user := authenticate(request, cache)
+			additions = append(additions, inbound.WithInUser(user))
 			trusted = resp == nil
 		}
 
@@ -130,7 +130,7 @@ func HandleConn(c net.Conn, tunnel C.Tunnel, cache *lru.LruCache[string, bool],
 	_ = conn.Close()
 }
 
-func authenticate(request *http.Request, cache *lru.LruCache[string, bool]) *http.Response {
+func authenticate(request *http.Request, cache *lru.LruCache[string, bool]) (resp *http.Response, u string) {
 	authenticator := authStore.Authenticator()
 	if inbound.SkipAuthRemoteAddress(request.RemoteAddr) {
 		authenticator = nil
@@ -140,23 +140,24 @@ func authenticate(request *http.Request, cache *lru.LruCache[string, bool]) *htt
 		if credential == "" {
 			resp := responseWith(request, http.StatusProxyAuthRequired)
 			resp.Header.Set("Proxy-Authenticate", "Basic")
-			return resp
+			return resp, ""
 		}
 
 		authed, exist := cache.Get(credential)
 		if !exist {
 			user, pass, err := decodeBasicProxyAuthorization(credential)
 			authed = err == nil && authenticator.Verify(user, pass)
+			u = user
 			cache.Set(credential, authed)
 		}
 		if !authed {
 			log.Infoln("Auth failed from %s", request.RemoteAddr)
 
-			return responseWith(request, http.StatusForbidden)
+			return responseWith(request, http.StatusForbidden), u
 		}
 	}
 
-	return nil
+	return nil, u
 }
 
 func responseWith(request *http.Request, statusCode int) *http.Response {
diff --git a/listener/socks/tcp.go b/listener/socks/tcp.go
index b6ea023a8..f2696e3fc 100644
--- a/listener/socks/tcp.go
+++ b/listener/socks/tcp.go
@@ -98,11 +98,12 @@ func HandleSocks4(conn net.Conn, tunnel C.Tunnel, additions ...inbound.Addition)
 	if inbound.SkipAuthRemoteAddr(conn.RemoteAddr()) {
 		authenticator = nil
 	}
-	addr, _, err := socks4.ServerHandshake(conn, authenticator)
+	addr, _, user, err := socks4.ServerHandshake(conn, authenticator)
 	if err != nil {
 		conn.Close()
 		return
 	}
+	additions = append(additions, inbound.WithInUser(user))
 	tunnel.HandleTCPConn(inbound.NewSocket(socks5.ParseAddr(addr), conn, C.SOCKS4, additions...))
 }
 
@@ -111,7 +112,7 @@ func HandleSocks5(conn net.Conn, tunnel C.Tunnel, additions ...inbound.Addition)
 	if inbound.SkipAuthRemoteAddr(conn.RemoteAddr()) {
 		authenticator = nil
 	}
-	target, command, err := socks5.ServerHandshake(conn, authenticator)
+	target, command, user, err := socks5.ServerHandshake(conn, authenticator)
 	if err != nil {
 		conn.Close()
 		return
@@ -121,5 +122,6 @@ func HandleSocks5(conn net.Conn, tunnel C.Tunnel, additions ...inbound.Addition)
 		io.Copy(io.Discard, conn)
 		return
 	}
+	additions = append(additions, inbound.WithInUser(user))
 	tunnel.HandleTCPConn(inbound.NewSocket(target, conn, C.SOCKS5, additions...))
 }
diff --git a/transport/socks4/socks4.go b/transport/socks4/socks4.go
index 9533a1c0a..50708eda2 100644
--- a/transport/socks4/socks4.go
+++ b/transport/socks4/socks4.go
@@ -43,7 +43,7 @@ var (
 
 var subnet = netip.PrefixFrom(netip.IPv4Unspecified(), 24)
 
-func ServerHandshake(rw io.ReadWriter, authenticator auth.Authenticator) (addr string, command Command, err error) {
+func ServerHandshake(rw io.ReadWriter, authenticator auth.Authenticator) (addr string, command Command, user string, err error) {
 	var req [8]byte
 	if _, err = io.ReadFull(rw, req[:]); err != nil {
 		return
@@ -73,6 +73,7 @@ func ServerHandshake(rw io.ReadWriter, authenticator auth.Authenticator) (addr s
 	if userID, err = readUntilNull(rw); err != nil {
 		return
 	}
+	user = string(userID)
 
 	if isReservedIP(dstIP) {
 		var target []byte
@@ -90,7 +91,7 @@ func ServerHandshake(rw io.ReadWriter, authenticator auth.Authenticator) (addr s
 	}
 
 	// SOCKS4 only support USERID auth.
-	if authenticator == nil || authenticator.Verify(string(userID), "") {
+	if authenticator == nil || authenticator.Verify(user, "") {
 		code = RequestGranted
 	} else {
 		code = RequestIdentdMismatched
diff --git a/transport/socks5/socks5.go b/transport/socks5/socks5.go
index c97c370ca..61b555f4f 100644
--- a/transport/socks5/socks5.go
+++ b/transport/socks5/socks5.go
@@ -106,7 +106,7 @@ type User struct {
 }
 
 // ServerHandshake fast-tracks SOCKS initialization to get target address to connect on server side.
-func ServerHandshake(rw net.Conn, authenticator auth.Authenticator) (addr Addr, command Command, err error) {
+func ServerHandshake(rw net.Conn, authenticator auth.Authenticator) (addr Addr, command Command, user string, err error) {
 	// Read RFC 1928 for request and reply structure and sizes.
 	buf := make([]byte, MaxAddrLen)
 	// read VER, NMETHODS, METHODS
@@ -141,7 +141,7 @@ func ServerHandshake(rw net.Conn, authenticator auth.Authenticator) (addr Addr,
 		if _, err = io.ReadFull(rw, authBuf[:userLen]); err != nil {
 			return
 		}
-		user := string(authBuf[:userLen])
+		user = string(authBuf[:userLen])
 
 		// Get password
 		if _, err = rw.Read(header[:1]); err != nil {