Wireguard: Remove wg peer on withdraw

2 files changed, 85 insertions, 6 deletions

diff --git a/proto/wireguard/wireguard.c b/proto/wireguard/wireguard.c
index 815db91a..73fd3b86 100644
--- a/proto/wireguard/wireguard.c
+++ b/proto/wireguard/wireguard.c
@@ -115,6 +115,8 @@ wg_init_entry(void *e_)
 {
   struct wg_entry *e UNUSED = e_;
 //  debug("wg_init_entry\n");
+
+  memset(e, 0, sizeof(struct wg_entry) - sizeof(struct fib_node));
 }
 
 static void
@@ -172,6 +174,40 @@ add_peer(wg_device *dev, const wg_key pubkey)
   return peer;
 }
 
+static void
+remove_marked_peer(wg_device *dev)
+{
+  struct wg_peer *peer = NULL;
+  struct wg_peer *prevpeer = NULL;
+
+  log(L_TRACE "WG: remove_marked_peer");
+  wg_for_each_peer(dev, peer) {
+    if (peer->flags & WGPEER_REMOVE_ME) {
+      if (!prevpeer) {
+	debug("WG: remove first peer\n");
+	dev->first_peer = peer->next_peer;
+	if (dev->last_peer == peer) {
+	  dev->last_peer = NULL;
+	  dev->first_peer = NULL;
+	}
+      } else {
+	debug("WG: remove middle peer\n");
+	// Remove
+	if (dev->last_peer == peer)
+	  dev->last_peer = prevpeer;
+
+	prevpeer->next_peer = peer->next_peer;
+      }
+
+      free(peer);
+      return;
+    }
+
+    prevpeer = peer;
+  }
+  log(L_WARN "WG: marked peer not found");
+}
+
 static int
 set_peer_tunnel_ep(wg_peer *peer, ip_addr tunnel_ep_addr, u16 udp_dest_port)
 {
@@ -225,7 +261,6 @@ wg_rt_notify(struct proto *P, struct channel *CH, struct network *n,
   struct wg_proto *p = (struct wg_proto *) P;
   struct wg_config *c = (struct wg_config *) P->cf;
   struct wg_channel *ch = (struct wg_channel *) CH;
-  struct wg_entry *en;
   struct iface *iface = NULL;
   const char *ifname = NULL;
 
@@ -249,22 +284,31 @@ wg_rt_notify(struct proto *P, struct channel *CH, struct network *n,
 	debug("WG: he %p %I %I %p %p %I\n", t, he->addr, he->link, he->next, he->src->hostentry, he->src->nh.gw);
       }
 
-      en = fib_get(&ch->rtable, n->n.addr);
-
       debug("WG: notify new %d %N\n",
         new->attrs->dest, n->n.addr);
 
+      bool is_tunnel_ep = false;
       struct tunnel_encap encap;
       memset(&encap, 0, sizeof(encap));
       encap.ep.ip = IPA_NONE;
 
       t = ea_find(new->attrs->eattrs, EA_CODE(PROTOCOL_BGP, BA_TUNNEL_ENCAP));
+      if (t) {
+	log(L_TRACE "WG: Set is tunnel");
+	is_tunnel_ep = true;
+      }
       if (!t && he && he->src) {
 	t = ea_find(he->src->eattrs, EA_CODE(PROTOCOL_BGP, BA_TUNNEL_ENCAP));
       }
       if (t && t->u.ptr && decode_tunnel_encap(t, &encap, P->pool) == 0 && encap.type == c->tunnel_type && encap.encap_len == sizeof(wg_key)) {
 	const wg_key *pubkey = encap.encap;
 
+	if (!fib_find(&ch->rtable, n->n.addr)) {
+	  struct wg_entry *en = fib_get(&ch->rtable, n->n.addr);
+	  en->is_tunnel_ep = is_tunnel_ep;
+	  memcpy(en->public_key, pubkey, sizeof(wg_key));
+	}
+
 	log(L_TRACE "WG:     Attr %x %x %d %04x", t->flags, t->type, t->u.ptr->length, encap.flags);
 
 	struct wg_device *dev = p->dev;
@@ -366,7 +410,7 @@ wg_rt_notify(struct proto *P, struct channel *CH, struct network *n,
 	  n->n.addr);
 
     /* Withdraw */
-    en = fib_find(&ch->rtable, n->n.addr);
+    struct wg_entry *en = fib_find(&ch->rtable, n->n.addr);
 
     if (!en) { // || en->valid != RIP_ENTRY_VALID)
       debug("WG: fib not found\n");
@@ -376,9 +420,36 @@ wg_rt_notify(struct proto *P, struct channel *CH, struct network *n,
 	struct wg_device *dev = p->dev;
 
 	if (dev != NULL) {
+	  bool marked_peer = false;
 	  bool found = false;
 	  struct wg_peer *peer = NULL;
 	  wg_for_each_peer(dev, peer) {
+	    if (en->is_tunnel_ep && !marked_peer) {
+	      log(L_TRACE "WG: Is tunnel");
+	      if (memcmp(peer->public_key, en->public_key, sizeof(wg_key)) == 0) {
+		struct peer_config *pc = NULL;
+		bool remove_me = true;
+		WALK_LIST(pc,c->peers)
+		{
+		  wg_key pc_key;
+		  if (pc->public_key) {
+		    wg_key_from_base64(pc_key, pc->public_key);
+		    if (memcmp(pc_key, peer->public_key, sizeof(wg_key)) == 0) {
+		      /* Don't remove preconfigured peer */
+		      remove_me = false;
+		      break;
+		    }
+		  }
+		}
+
+		if (remove_me) {
+		  log(L_TRACE "WG: Remove peer");
+		  peer->flags |= WGPEER_REMOVE_ME;
+		  marked_peer = true;
+		  continue;
+		}
+	      }
+	    }
 	    // Remove from all peers
 	    found = true;
 
@@ -445,11 +516,17 @@ wg_rt_notify(struct proto *P, struct channel *CH, struct network *n,
 	      }
 	      
 	      peer->flags |= WGPEER_REPLACE_ALLOWEDIPS;
-	      int res = set_device(p);
-	      log(L_TRACE "WG: wg_set_device %d", res);
 	    }
 	  }
 
+	  if (marked_peer) {
+	    remove_marked_peer(p->dev);
+	  }
+	  int res = set_device(p);
+	  log(L_TRACE "WG: wg_set_device %d", res);
+
+	  fib_delete(&ch->rtable, en);
+	  en = NULL;
 	  /*
     old_metric = en->metric;
 
diff --git a/proto/wireguard/wireguard.h b/proto/wireguard/wireguard.h
index c2f780ca..e1cce73f 100644
--- a/proto/wireguard/wireguard.h
+++ b/proto/wireguard/wireguard.h
@@ -41,6 +41,8 @@ struct wg_channel {
 };
 
 struct wg_entry {
+  bool is_tunnel_ep;
+  wg_key public_key;
   struct fib_node n;
 };