From 7d1a0cdbdcda9032c47541be9454eac99eb2987f Mon Sep 17 00:00:00 2001
From: HeshamTB <hishaminv@gmail.com>
Date: Thu, 28 Mar 2024 01:22:29 +0300
Subject: [PATCH] feat: grace period for new peers as a new evection policy

Signed-off-by: HeshamTB <hishaminv@gmail.com>
---
 handlers.go               |  4 +--
 link.go                   | 15 +++++++----
 monitor.go                |  7 +++--
 peer_meta.go              | 54 +++++++++++++++++++++++++++++++++++++++
 templates/client.ini.tmpl |  2 +-
 5 files changed, 72 insertions(+), 10 deletions(-)
 create mode 100644 peer_meta.go

diff --git a/handlers.go b/handlers.go
index ef65b36..a20fdea 100644
--- a/handlers.go
+++ b/handlers.go
@@ -180,6 +180,8 @@ func HandleDeletePeer(wg *WGLink) http.HandlerFunc {
 
 func HandleGetCreatePeer(wg *WGLink) http.HandlerFunc {
     return func(w http.ResponseWriter, r *http.Request) {
+        reqId := r.Context().Value(CtxReqID).(uuid.UUID)
+        info("Create peer", reqId)
         newPeer, err := wg.CreateNewPeer()
         if err != nil {
             slog.Error(err.Error())
@@ -200,8 +202,6 @@ func HandleGetCreatePeer(wg *WGLink) http.HandlerFunc {
             slog.Error(err.Error())
         }
 
-        // TODO: Add peer to device
-
     }
 }
 
diff --git a/link.go b/link.go
index 7d2385c..cb5f02a 100644
--- a/link.go
+++ b/link.go
@@ -20,6 +20,7 @@ type WGLink struct {
     *netlink.LinkAttrs
     *wgctrl.Client
     IPPool
+    meta *PeerMeta
     Endpoint string // Endpoint for clients to use when connecting to this link. Has no effect.
     StartedAT time.Time
     lock *sync.Mutex
@@ -60,6 +61,7 @@ func InitWGLink(ifName string, privateKey *wgtypes.Key, port int, endpoint strin
         return nil, err
     }
 
+    wg.meta = NewPeerMeta()
     wg.lock = &sync.Mutex{}
 
     return &wg, netlink.LinkSetUp(&wg)
@@ -146,6 +148,8 @@ func (wg *WGLink) AddPeer(publicKey string) (*wgtypes.Peer, error) {
     if err != nil {
         return nil, err
     }
+
+    wg.meta.AddPeer(pubKey)
     
     return wg.getPeer(pubKey)
 }
@@ -175,6 +179,8 @@ func (wg *WGLink) deletePeer(publickey wgtypes.Key) error {
         return err
     }
 
+    wg.meta.DeletePeer(publickey)
+
     return nil
 }
 
@@ -258,7 +264,7 @@ func (wg *WGLink) GetPublicKey() (*wgtypes.Key, error) {
 type NewPeer struct {
     Peer wgtypes.Peer // The server as a peer 
     PrivateKey wgtypes.Key
-    Address net.IP
+    Address net.IPNet
     DNS []net.IP
     MTU uint
 
@@ -266,7 +272,7 @@ type NewPeer struct {
 }
 
 // Create a new peer. Used for server-side clinet config generation
-// Peer is not added to the wglink
+// New peer is added to the wireguard device
 func (wg *WGLink) CreateNewPeer() (*NewPeer, error) {
     dev, err := wg.Device(wg.Name)
     if err != nil {
@@ -283,19 +289,18 @@ func (wg *WGLink) CreateNewPeer() (*NewPeer, error) {
         return nil, err
     }
     
-    ip, err := wg.Allocate()
+    wgPeer, err := wg.AddPeer(privKey.PublicKey().String())
     if err != nil {
         return nil, err
     }
 
     peer := NewPeer{
-        Address: ip,
+        Address: wgPeer.AllowedIPs[0],
         DNS: WG_CLIENT_DNS,
         PrivateKey: privKey,
         MTU: WG_CLIENT_MTU,
         Peer: wgtypes.Peer{
             PublicKey: *serverPubkey,
-            AllowedIPs: []net.IPNet{wg.Network()},
         },
         Endpoint: fmt.Sprintf("%s:%d", wg.Endpoint, dev.ListenPort),
     }
diff --git a/monitor.go b/monitor.go
index 1b88146..beb0b45 100644
--- a/monitor.go
+++ b/monitor.go
@@ -31,7 +31,7 @@ func monitor(wg *WGLink, log slog.Logger) {
         for _, peer := range peers {
             totalRx += peer.ReceiveBytes
             totalTx += peer.TransmitBytes
-            if !isActive(peer) {
+            if !isActive(peer, wg) {
                 log.Info(
                     fmt.Sprintf("[WGMonitor] Evecting peer %s", peer.PublicKey.String()),
                 )
@@ -48,8 +48,11 @@ func monitor(wg *WGLink, log slog.Logger) {
     }
 }
 
-func isActive(peer wgtypes.Peer) bool {
+func isActive(peer wgtypes.Peer, wg *WGLink) bool {
     if peer.LastHandshakeTime.IsZero() {
+        if time.Since(wg.meta.TimeAdded(peer.PublicKey)) < time.Duration(4 * time.Minute) {
+            return true
+        }
         return false
     } else if time.Since(peer.LastHandshakeTime.UTC()) < time.Duration(4 * time.Minute) {
         return true
diff --git a/peer_meta.go b/peer_meta.go
new file mode 100644
index 0000000..df83c20
--- /dev/null
+++ b/peer_meta.go
@@ -0,0 +1,54 @@
+package hvpnnode3
+
+import (
+	"sync"
+	"time"
+
+	"golang.zx2c4.com/wireguard/wgctrl/wgtypes"
+)
+
+// Contians data about peers
+type PeerData struct {
+    addedOn time.Time
+    deleted bool
+}
+
+// Basic implementaion to track peers based on public keys. Thraed safe.
+type PeerMeta struct {
+    data map[wgtypes.Key]PeerData
+    rwl *sync.RWMutex
+}
+
+func NewPeerMeta() *PeerMeta {
+    meta := new(PeerMeta)
+    meta.data = make(map[wgtypes.Key]PeerData)
+    meta.rwl = new(sync.RWMutex)
+    return meta
+}
+
+func (m *PeerMeta) AddPeer(publicKey wgtypes.Key) {
+    m.rwl.Lock()
+    m.data[publicKey] = PeerData{
+        addedOn: time.Now().UTC(),
+        deleted: false,
+    }
+    m.rwl.Unlock()
+}
+
+func (m *PeerMeta) DeletePeer(publicKey wgtypes.Key) {
+    m.rwl.Lock()
+    d := m.data[publicKey]
+    d.deleted = true
+    m.rwl.Unlock()
+}
+
+func (m *PeerMeta) TimeAdded(publicKey wgtypes.Key) time.Time {
+    m.rwl.RLock()
+    t := m.data[publicKey].addedOn
+    m.rwl.RUnlock()
+    return t
+}
+
+
+
+
diff --git a/templates/client.ini.tmpl b/templates/client.ini.tmpl
index ab687ed..f6f0962 100644
--- a/templates/client.ini.tmpl
+++ b/templates/client.ini.tmpl
@@ -1,5 +1,5 @@
 [Interface]
-Address = {{ .Address }}/24
+Address = {{ .Address }}
 PrivateKey = {{ .PrivateKey }} 
 DNS = {{$dnss := .DNS}}{{range $index, $element := .DNS}}{{if $index}}, {{end}}{{$element.String}}{{end}}
 MTU = {{ .MTU }}