Wireguard: Refactor into peer and allowed ips functions

1 files changed, 78 insertions, 34 deletions

diff --git a/proto/wireguard/wireguard.c b/proto/wireguard/wireguard.c
index f209665e..e60a0b97 100644
--- a/proto/wireguard/wireguard.c
+++ b/proto/wireguard/wireguard.c
@@ -204,6 +204,69 @@ int decode_tunnel_encap(const eattr *e, wg_key *pubkey, u32 *as4, ip_addr *remot
   return 0;
 }
 
+static wg_peer *
+add_peer(wg_device *dev, wg_key pubkey)
+{
+  struct wg_peer *peer = malloc(sizeof(struct wg_peer));
+  memset(peer, 0, sizeof(struct wg_peer));
+
+  peer->flags = WGPEER_HAS_PUBLIC_KEY;
+  memcpy(peer->public_key, pubkey, sizeof(wg_key));
+
+  if (dev->first_peer && dev->last_peer)
+    dev->last_peer->next_peer = peer;
+  else
+    dev->first_peer = peer;
+  dev->last_peer = peer;
+  return peer;
+}
+
+static int
+set_peer_remote_ep(wg_peer *peer, ip_addr remote_ep_addr, u16 udp_dest_port)
+{
+  if (udp_dest_port != 0 && ipa_nonzero(remote_ep_addr) ) {
+    if (ipa_is_ip4(remote_ep_addr)) {
+      log(L_TRACE "WG:     found ip4 ep");
+      peer->endpoint.addr4.sin_family = AF_INET;
+      put_ip4(&peer->endpoint.addr4.sin_addr.s_addr, ipa_to_ip4(remote_ep_addr));
+      put_u16(&peer->endpoint.addr4.sin_port, udp_dest_port);
+    } else {
+      log(L_TRACE "WG:     found ip6 ep");
+      peer->endpoint.addr6.sin6_family = AF_INET6;
+      put_ip6(&peer->endpoint.addr6.sin6_addr, ipa_to_ip6(remote_ep_addr));
+      put_u16(&peer->endpoint.addr6.sin6_port, udp_dest_port);
+    }
+  }
+
+  return 0;
+}
+
+static int
+add_allowed_ips(struct wg_proto *p, struct network *n, wg_peer *peer)
+{
+  // Add allowed ip
+  struct wg_allowedip *allowedip = malloc(sizeof(struct wg_allowedip));
+  memset(allowedip, 0, sizeof(struct wg_allowedip));
+
+  if (p->p.cf->net_type == NET_IP4) {
+    allowedip->family = AF_INET;
+    allowedip->ip4.s_addr = ip4_to_u32(ip4_hton(net4_prefix(n->n.addr)));
+  } else if (p->p.cf->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);
+  if (peer->first_allowedip && peer->last_allowedip)
+    peer->last_allowedip->next_allowedip = allowedip;
+  else
+    peer->first_allowedip = allowedip;
+  peer->last_allowedip = allowedip;
+
+  return 0;
+}
+
 static void
 wg_rt_notify(struct proto *P, struct channel *ch UNUSED, struct network *n,
   struct rte *new, struct rte *old UNUSED)
@@ -242,6 +305,7 @@ wg_rt_notify(struct proto *P, struct channel *ch UNUSED, struct network *n,
 	struct wg_device *dev = NULL;
 
 	if (wg_get_device(&dev, ifname) == 0) {
+	  bool dirty = false;
 	  bool found = false;
 	  struct wg_peer *peer = NULL;
 	  wg_for_each_peer(dev, peer) {
@@ -256,44 +320,24 @@ wg_rt_notify(struct proto *P, struct channel *ch UNUSED, struct network *n,
 	    log(L_TRACE "WG:     Found");
 	    found = true;
 
-	    if (found) {
-	      // Add allowed ip
-	      struct wg_allowedip *allowedip = malloc(sizeof(struct wg_allowedip));
-	      memset(allowedip, 0, sizeof(struct wg_allowedip));
-
-	      if (p->p.cf->net_type == NET_IP4) {
-		allowedip->family = AF_INET;
-		allowedip->ip4.s_addr = ip4_to_u32(ip4_hton(net4_prefix(n->n.addr)));
-	      } else if (p->p.cf->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);
-	      }
+	    set_peer_remote_ep(peer, remote_ep_addr, udp_dest_port);
+	    add_allowed_ips(p, n, peer);
+	    dirty = true;
 
-	      if (udp_dest_port != 0 && ipa_nonzero(remote_ep_addr) ) {
-		if (ipa_is_ip4(remote_ep_addr)) {
-		  log(L_TRACE "WG:     found ip4 ep");
-		  peer->endpoint.addr4.sin_family = AF_INET;
-		  put_ip4(&peer->endpoint.addr4.sin_addr.s_addr, ipa_to_ip4(remote_ep_addr));
-		  put_u16(&peer->endpoint.addr4.sin_port, udp_dest_port);
-		} else {
-		  log(L_TRACE "WG:     found ip6 ep");
-		  peer->endpoint.addr6.sin6_family = AF_INET6;
-		  put_ip6(&peer->endpoint.addr6.sin6_addr, ipa_to_ip6(remote_ep_addr));
-		  put_u16(&peer->endpoint.addr6.sin6_port, udp_dest_port);
-		}
-	      }
-
-	      allowedip->cidr = net_pxlen(n->n.addr);
-	      peer->last_allowedip->next_allowedip = allowedip;
-	      peer->last_allowedip = allowedip;
+	    break;
+	  }
 
-	      int res = wg_set_device(dev);
-	      log(L_TRACE "WG: wg_set_device %d", res);
-	      break;
-	    }
+	  if (!found) {
+	    wg_peer *peer = add_peer(dev, pubkey);
+	    set_peer_remote_ep(peer, remote_ep_addr, udp_dest_port);
+	    add_allowed_ips(p, n, peer);
+	    dirty = true;
 	  }
 
+	  if (dirty) {
+	    int res = wg_set_device(dev);
+	    log(L_TRACE "WG: wg_set_device %d", res);
+	  }
 	  wg_free_device(dev);
 	  dev = NULL;
 	}