refactor: add SetIP and optionally return it if it exists in the ws peer

net/peer.go

@@ -47,6 +47,7 @@ type Peer interface {
 	SetHandshakeDone(status bool)
 	GetIPString() string
 	GetIP() net.Addr
+	SetIP(ip net.Addr)
 	Abuser() bool
 }
 
@@ -93,6 +94,10 @@ func (b *BasePeer) GetIP() net.Addr {
 	panic("must implement in child class")
 }
 
+func (b *BasePeer) SetIP(ip net.Addr) {
+	panic("must implement in child class")
+}
+
 func (b *BasePeer) Challenge() []byte {
 	return b.challenge
 }

net/ws.go

@@ -19,6 +19,7 @@ type WebSocketPeer struct {
 	BasePeer
 	socket *websocket.Conn
 	abuser bool
+	ip     net.Addr
 }
 
 func (p *WebSocketPeer) Connect(uri *url.URL) (interface{}, error) {
@@ -139,6 +140,10 @@ func (p *WebSocketPeer) GetChallenge() []byte {
 }
 
 func (p *WebSocketPeer) GetIP() net.Addr {
+	if p.ip != nil {
+		return p.ip
+	}
+
 	ctx, cancel := context.WithCancel(context.Background())
 	netConn := websocket.NetConn(ctx, p.socket, websocket.MessageBinary)