From 4bab068c1075f3ddda5b7f0f184f09bd4ba7dfce Mon Sep 17 00:00:00 2001
From: HeshamTB <hishaminv@gmail.com>
Date: Tue, 19 Mar 2024 02:55:31 +0300
Subject: [PATCH] fix: Assign IP to wg device

Signed-off-by: HeshamTB <hishaminv@gmail.com>
---
 cmd/hvpn-node/hvpn-node.go |  8 +++++++-
 ip_pool.go                 |  5 +++++
 link.go                    | 22 +++++++++++++++++++---
 3 files changed, 31 insertions(+), 4 deletions(-)

diff --git a/cmd/hvpn-node/hvpn-node.go b/cmd/hvpn-node/hvpn-node.go
index 2183de7..fbc1654 100644
--- a/cmd/hvpn-node/hvpn-node.go
+++ b/cmd/hvpn-node/hvpn-node.go
@@ -165,7 +165,7 @@ func createCliApp() *cli.App {
     app.Flags = append(app.Flags, &wgPort)
 
     httpListenAddr := cli.StringFlag{
-        Name: "host",
+        Name: "http-host",
         Usage: "IP address to listen on for HTTP API requests",
         Value: "0.0.0.0",
         Action: func(ctx *cli.Context, s string) error {
@@ -362,6 +362,12 @@ func setup(ctx *cli.Context) error {
     IPPool = ipPool
     wgLink.IPPool = ipPool
 
+    err = wgLink.SetIP()
+    if err != nil {
+        return err
+    }
+    slog.Debug("Assigned IP to Wiregaurd interface")
+
     //defer wgLink.Close()
     cInput := make(chan struct{})
     go handleStdin(cInput)
diff --git a/ip_pool.go b/ip_pool.go
index 3b5b27a..091a135 100644
--- a/ip_pool.go
+++ b/ip_pool.go
@@ -12,6 +12,7 @@ type IPPool interface {
 	Allocate() (net.IP, error)
 	Free(net.IP) error
 	Remove(...net.IP) error
+    Network() net.IPNet
 }
 
 // Pool is a pool of available IP numbers for allocation.
@@ -94,6 +95,10 @@ func (p *Pool) Free(ip net.IP) error {
 	return nil
 }
 
+func (p *Pool) Network() net.IPNet {
+    return *p.network
+}
+
 // ip4To6 will prefix IPv4 with the IPv6 network to create an IPv6 address.
 func ip4To6(ip4 net.IP, ip6prefix *net.IPNet) (ip6 net.IP) {
 	b6 := ip6prefix.IP.To16()
diff --git a/link.go b/link.go
index 5608813..d23ff7f 100644
--- a/link.go
+++ b/link.go
@@ -2,8 +2,6 @@ package hvpnnode3
 
 import (
 	"errors"
-	"fmt"
-	"log/slog"
 	"net"
 	"sync"
 	"time"
@@ -86,9 +84,27 @@ func (wg *WGLink) initClient() error {
     return err
 }
 
+func (wg *WGLink) SetIP() error {
+    ip, err := wg.Allocate()
+    if err != nil {
+        return err
+    }
+    ipnet := net.IPNet{
+        IP: ip,
+        Mask: wg.Network().Mask,
+    }
+    netlinkIP, err := netlink.ParseAddr(ipnet.String())
+    
+    err = netlink.AddrAdd(wg, netlinkIP)
+    if err != nil {
+        return err
+    }
+
+    return nil
+}
+
 // Adds a peer to the wireguard netlink.
 func (wg *WGLink) AddPeer(publicKey string) (*wgtypes.Peer, error) {
-    slog.Debug(fmt.Sprintf("Trying to add peer %s", publicKey))
     pubKey, err := wgtypes.ParseKey(publicKey)
     if err != nil {
         return nil, err