diff --git a/peer/app.go b/peer/app.go index 3c69a14..f23860b 100644 --- a/peer/app.go +++ b/peer/app.go @@ -58,6 +58,9 @@ type App struct { dev WGDevice controlConn ControlConn + // Relay-side forwarder; non-nil only when isRelay. + forwarder *Forwarder + // Peer state relay *Peer peersByKey map[wgtypes.Key]*Peer diff --git a/peer/new.go b/peer/new.go index ffe0154..556f120 100644 --- a/peer/new.go +++ b/peer/new.go @@ -31,10 +31,15 @@ func New( return nil, fmt.Errorf("configure WG device: %w", err) } + var fwd *Forwarder if state.IsRelay { if err := dev.EnableForwarding(); err != nil { return nil, fmt.Errorf("enable forwarding: %w", err) } + fwd, err = NewForwarder(state.VPNIP) + if err != nil { + return nil, fmt.Errorf("forwarder: %w", err) + } } cc, err := newUDPControlConn(state.VPNIP, ControlPort) @@ -73,6 +78,7 @@ func New( dev: dev, controlConn: cc, + forwarder: fwd, peersByKey: make(map[wgtypes.Key]*Peer), peersByIP: make(map[netip.Addr]*Peer), diff --git a/peer/on_hub.go b/peer/on_hub.go index be66b89..eb04beb 100644 --- a/peer/on_hub.go +++ b/peer/on_hub.go @@ -38,6 +38,9 @@ func (a *App) onAddPeer(p HubPeer) { if !peer.IsPublic { a.devAddPeer(peer) + if a.forwarder != nil && endpoint.IsValid() { + a.forwarder.SetEndpoint(peer.VPNIP.As4()[3], endpoint) + } return } @@ -50,6 +53,9 @@ func (a *App) onRemovePeer(key wgtypes.Key) { return } a.devRemove(peer) + if a.forwarder != nil { + a.forwarder.RemoveEndpoint(peer.VPNIP.As4()[3]) + } delete(a.peersByKey, key) delete(a.peersByIP, peer.VPNIP)