diff --git a/net/peer.go b/net/peer.go index c5dd763..f3127e3 100644 --- a/net/peer.go +++ b/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 } diff --git a/net/ws.go b/net/ws.go index 5db95ce..acf8112 100644 --- a/net/ws.go +++ b/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)