Wireguard: Refactor remove_allowed_ip

1 files changed, 66 insertions, 66 deletions

diff --git a/proto/wireguard/wireguard.c b/proto/wireguard/wireguard.c
index 73fd3b86..f4bc56c6 100644
--- a/proto/wireguard/wireguard.c
+++ b/proto/wireguard/wireguard.c
@@ -228,11 +228,9 @@ set_peer_tunnel_ep(wg_peer *peer, ip_addr tunnel_ep_addr, u16 udp_dest_port)
   return 0;
 }
 
-static int
-add_allowed_ips(u8 net_type, struct network *n, wg_peer *peer)
+static void
+init_allowed_ip(struct wg_allowedip *allowedip, u8 net_type, struct network *n)
 {
-  // Add allowed ip
-  struct wg_allowedip *allowedip = malloc(sizeof(struct wg_allowedip));
   memset(allowedip, 0, sizeof(struct wg_allowedip));
 
   if (net_type == NET_IP4) {
@@ -245,6 +243,15 @@ add_allowed_ips(u8 net_type, struct network *n, wg_peer *peer)
   }
 
   allowedip->cidr = net_pxlen(n->n.addr);
+}
+
+static int
+add_allowed_ip(u8 net_type, struct network *n, wg_peer *peer)
+{
+  // Add allowed ip
+  struct wg_allowedip *allowedip = malloc(sizeof(struct wg_allowedip));
+  init_allowed_ip(allowedip, net_type, n);
+
   if (peer->first_allowedip && peer->last_allowedip)
     peer->last_allowedip->next_allowedip = allowedip;
   else
@@ -255,6 +262,56 @@ add_allowed_ips(u8 net_type, struct network *n, wg_peer *peer)
 }
 
 static void
+remove_allowed_ip(wg_peer *peer, struct wg_allowedip *allowedip)
+{
+  struct wg_allowedip *ip = NULL;
+  struct wg_allowedip *previp = NULL;
+
+  wg_for_each_allowedip(peer, ip) {
+    if (allowedip->family != ip->family) {
+      debug("WG: family no match\n");
+      previp = ip;
+      continue;
+    }
+
+    if (allowedip->cidr != ip->cidr) {
+      debug("WG: cidr no match\n");
+      previp = ip;
+      continue;
+    }
+
+    if (memcmp(&allowedip->ip6, &ip->ip6, sizeof(struct in6_addr))) {
+      debug("WG: ip no match\n");
+      dump(&allowedip->ip6, sizeof(struct in6_addr));
+      dump(&ip->ip6, sizeof(struct in6_addr));
+      previp = ip;
+      continue;
+    }
+
+    debug("WG: found ip\n");
+
+    if (!previp) {
+      debug("WG: remove first\n");
+      peer->first_allowedip = ip->next_allowedip;
+      if (peer->last_allowedip == ip) {
+	peer->last_allowedip = NULL;
+	peer->first_allowedip = NULL;
+      }
+    } else {
+      debug("WG: remove middle\n");
+      // Remove
+      if (peer->last_allowedip == ip)
+	peer->last_allowedip = previp;
+
+      previp->next_allowedip = ip->next_allowedip;
+    }
+
+    free(ip);
+    break;
+  }
+}
+
+static void
 wg_rt_notify(struct proto *P, struct channel *CH, struct network *n,
   struct rte *new, struct rte *old UNUSED)
 {
@@ -330,7 +387,7 @@ wg_rt_notify(struct proto *P, struct channel *CH, struct network *n,
 	    found = true;
 
 	    set_peer_tunnel_ep(peer, encap.ep.ip, encap.udp_dest_port);
-	    add_allowed_ips(ch->c.net_type, n, peer);
+	    add_allowed_ip(ch->c.net_type, n, peer);
 	    dirty = true;
 
 	    break;
@@ -339,7 +396,7 @@ wg_rt_notify(struct proto *P, struct channel *CH, struct network *n,
 	  if (!found) {
 	    wg_peer *peer = add_peer(dev, pubkey);
 	    set_peer_tunnel_ep(peer, encap.ep.ip, encap.udp_dest_port);
-	    add_allowed_ips(ch->c.net_type, n, peer);
+	    add_allowed_ip(ch->c.net_type, n, peer);
 	    dirty = true;
 	  }
 
@@ -454,67 +511,10 @@ wg_rt_notify(struct proto *P, struct channel *CH, struct network *n,
 	    found = true;
 
 	    if (found) {
-	      // Add allowed ip
-	      struct wg_allowedip *allowedip = malloc(sizeof(struct wg_allowedip));
-	      memset(allowedip, 0, sizeof(struct wg_allowedip));
-
-	      if (ch->c.net_type == NET_IP4) {
-		allowedip->family = AF_INET;
-		allowedip->ip4.s_addr = ip4_to_u32(ip4_hton(net4_prefix(n->n.addr)));
-	      } else if (ch->c.net_type == NET_IP6) {
-		allowedip->family = AF_INET6;
-		ip6_addr addr = ip6_hton(net6_prefix(n->n.addr));
-		memcpy(allowedip->ip6.s6_addr, &addr, 16);
-	      }
-
-	      allowedip->cidr = net_pxlen(n->n.addr);
-
-	      struct wg_allowedip *ip = NULL;
-	      struct wg_allowedip *previp = NULL;
-	      
-	      wg_for_each_allowedip(peer, ip) {
-		if (allowedip->family != ip->family) {
-		  debug("WG: family no match\n");
-		  previp = ip;
-		  continue;
-		}
+	      struct wg_allowedip *allowedip = alloca(sizeof(struct wg_allowedip));
+	      init_allowed_ip(allowedip, ch->c.net_type, n);
+	      remove_allowed_ip(peer, allowedip);
 
-		if (allowedip->cidr != ip->cidr) {
-		  debug("WG: cidr no match\n");
-		  previp = ip;
-		  continue;
-		}
-		
-		if (memcmp(&allowedip->ip6, &ip->ip6, sizeof(struct in6_addr))) {
-		  debug("WG: ip no match\n");
-		  dump(&allowedip->ip6, sizeof(struct in6_addr));
-		  dump(&ip->ip6, sizeof(struct in6_addr));
-		  previp = ip;
-		  continue;
-		}
-
-		debug("WG: found ip\n");
-
-		if (!previp) {
-		  debug("WG: remove first\n");
-		  peer->first_allowedip = ip->next_allowedip;
-		  if (peer->last_allowedip == ip) {
-		    peer->last_allowedip = NULL;
-		    peer->first_allowedip = NULL;
-		  }
-		} else {
-		  debug("WG: remove middle\n");
-		  // Remove
-		  if (peer->last_allowedip == ip)
-		      peer->last_allowedip = previp;
-
-		  previp->next_allowedip = ip->next_allowedip;
-		}
-
-		free(ip);
-		break;
-	      }
-	      
 	      peer->flags |= WGPEER_REPLACE_ALLOWEDIPS;
 	    }
 	  }