Joyit 1 тиждень тому
батько
коміт
645155f3db
1 змінених файлів з 23 додано та 4 видалено
  1. 23 4
      conn/ws2/ws2.go

+ 23 - 4
conn/ws2/ws2.go

@@ -27,6 +27,7 @@ type Ws2 struct {
 	conn     *websocket.Conn
 	cipher   *util.Cipher // 记录当前的加解密类,可以保证在没有ssl的情况下数据安全
 	compress bool         // 是否启用压缩
+	remoteIp net.IP
 }
 
 var upgrader = websocket.Upgrader{
@@ -55,6 +56,22 @@ func Server(cf *config.Config, bind string, path string, hash string, fn conn.Se
 		}
 	}
 	http.HandleFunc(path, func(w http.ResponseWriter, r *http.Request) {
+		// Get real remote address from proxy headers
+		remoteIp := r.Header.Get("X-Real-IP")
+		if remoteIp == "" {
+			// X-Forwarded-For can contain a list of comma-separated IPs; the first is the original client
+			forwarded := r.Header.Get("X-Forwarded-For")
+			if forwarded != "" {
+				remoteIp = strings.Split(forwarded, ",")[0]
+				remoteIp = strings.TrimSpace(remoteIp)
+			}
+		}
+
+		// Fallback to standard RemoteAddr if no proxy headers exist
+		if remoteIp == "" {
+			remoteIp = r.RemoteAddr
+		}
+
 		conn, err := upgrader.Upgrade(w, r, nil)
 		if err != nil {
 			log.Println("[ws2 Server Upgrade ERROR]", err)
@@ -104,9 +121,10 @@ func Server(cf *config.Config, bind string, path string, hash string, fn conn.Se
 			return
 		}
 		ws := &Ws2{
-			cf:     cf,
-			conn:   conn,
-			cipher: cipher,
+			cf:       cf,
+			conn:     conn,
+			cipher:   cipher,
+			remoteIp: net.ParseIP(remoteIp),
 		}
 		fn(ws)
 	})
@@ -298,7 +316,8 @@ func (c *Ws2) ReadMessage(deadline int) (msgType conn.MsgType, id uint16, cmd st
 
 // 获取远程的地址
 func (c *Ws2) RemoteIP() net.IP {
-	return util.AddrToIP(c.conn.RemoteAddr())
+	return c.remoteIp
+	// return util.AddrToIP(c.conn.RemoteAddr())
 }
 
 // 获取本地的地址