From 4a1039e5b170c68ce919d1fe5dea8a3867e61c0e Mon Sep 17 00:00:00 2001
From: HeshamTB <hishaminv@gmail.com>
Date: Wed, 13 Mar 2024 00:26:19 +0300
Subject: [PATCH] setup: add and remove peer on setup as runtime test

---
 cmd/hvpn-node/hvpn-node.go | 60 ++++++++++++++++++++++++++++++++++++--
 1 file changed, 57 insertions(+), 3 deletions(-)

diff --git a/cmd/hvpn-node/hvpn-node.go b/cmd/hvpn-node/hvpn-node.go
index ec2aacb..76afe4e 100644
--- a/cmd/hvpn-node/hvpn-node.go
+++ b/cmd/hvpn-node/hvpn-node.go
@@ -3,6 +3,7 @@ package main
 import (
 	"fmt"
 	"log/slog"
+	"net"
 	"net/http"
 	"net/netip"
 	"os"
@@ -244,7 +245,7 @@ func setup() error {
         slog.Error(fmt.Sprintf("main.IPPool: %s", err))
         os.Exit(1)
     }
-    slog.Info(fmt.Sprintf("Init ip pool %s", VPNIPCIDR))
+    slog.Debug(fmt.Sprintf("Init ip pool %s", VPNIPCIDR))
 
     testVip, err := ipPool.Allocate()
     if err != nil {
@@ -252,13 +253,13 @@ func setup() error {
         os.Exit(1)
     }
 
-    slog.Info(fmt.Sprintf("main.testVip: IP Pool Test IP: %s", testVip.String()))
+    slog.Debug(fmt.Sprintf("main.testVip: IP Pool Test IP: %s", testVip.String()))
     err = ipPool.Free(testVip)
     if err != nil {
         slog.Error("main.testVip: Could not free test Vip from IPPool!", err)
         os.Exit(1)
     }
-    slog.Info("main.testVip: Test IP Freed")
+    slog.Debug("main.testVip: Test IP Freed")
 
     IPPool = ipPool
 
@@ -273,6 +274,59 @@ func setup() error {
         os.Exit(0)
     }()
 
+    err = testWgPeerAdd(wgLink)
+    if err != nil {
+        return err
+    }
+
+    return nil
+}
+
+func testWgPeerAdd(wgLink *hvpnnode3.WGLink) error {
+    privateKey, err := wgtypes.GeneratePrivateKey()
+    if err != nil {
+        return err
+    }
+    publicKey := privateKey.PublicKey()
+
+    ip, err := IPPool.Allocate()
+    if err != nil {
+        return err
+    }
+
+    peerConfig := wgtypes.PeerConfig{
+        PublicKey: publicKey,
+        AllowedIPs: []net.IPNet{
+            {
+                IP: ip,
+                Mask: net.IPv4Mask(255, 255, 255, 255),
+            },
+        },
+    }
+
+    wgConfig := wgtypes.Config{
+        Peers: []wgtypes.PeerConfig{peerConfig},
+    }
+
+    err = wgLink.ConfigureDevice(wgLink.Name, wgConfig)
+
+    if err != nil {
+        return err
+    }
+    slog.Debug(fmt.Sprintf("Added test peer %v", peerConfig.PublicKey))
+
+    wgConfig.ReplacePeers = true
+    wgConfig.Peers = []wgtypes.PeerConfig{}
+
+    err = wgLink.ConfigureDevice(wgLink.Name, wgConfig)
+    if err != nil {
+        return err
+    }
+    slog.Debug("Removed test peer")
+    IPPool.Free(ip)
+    slog.Debug("Freed test peer ip")
+
+    
     return nil
 }