From 228d43440589aacd33d359291c6d31c0261f71cc Mon Sep 17 00:00:00 2001
From: Matthew Nickson <mnickson@sidingsmedia.com>
Date: Sun, 13 Mar 2022 23:21:00 +0000
Subject: [PATCH] Changed GetAllocatedIPs to use IStore

Previously the GetAllocatedIPs function used it's own database
connection. It now uses a pointer with the store.IStore interface to
access the DB.

Signed-off-by: Matthew Nickson <mnickson@sidingsmedia.com>
---
 handler/routes.go |  6 +++---
 util/util.go      | 24 ++++++++----------------
 2 files changed, 11 insertions(+), 19 deletions(-)

diff --git a/handler/routes.go b/handler/routes.go
index 9c18fc4..6f83923 100644
--- a/handler/routes.go
+++ b/handler/routes.go
@@ -148,7 +148,7 @@ func NewClient(db store.IStore) echo.HandlerFunc {
 		}
 
 		// validate the input Allocation IPs
-		allocatedIPs, err := util.GetAllocatedIPs("")
+		allocatedIPs, err := util.GetAllocatedIPs(db, "")
 		check, err := util.ValidateIPAllocation(server.Interface.Addresses, allocatedIPs, client.AllocatedIPs)
 		if !check {
 			return c.JSON(http.StatusBadRequest, jsonHTTPResponse{false, fmt.Sprintf("%s", err)})
@@ -268,7 +268,7 @@ func UpdateClient(db store.IStore) echo.HandlerFunc {
 		}
 		client := *clientData.Client
 		// validate the input Allocation IPs
-		allocatedIPs, err := util.GetAllocatedIPs(client.ID)
+		allocatedIPs, err := util.GetAllocatedIPs(db, client.ID)
 		check, err := util.ValidateIPAllocation(server.Interface.Addresses, allocatedIPs, _client.AllocatedIPs)
 		if !check {
 			return c.JSON(http.StatusBadRequest, jsonHTTPResponse{false, fmt.Sprintf("%s", err)})
@@ -624,7 +624,7 @@ func SuggestIPAllocation(db store.IStore) echo.HandlerFunc {
 		// we take the first available ip address from
 		// each server's network addresses.
 		suggestedIPs := make([]string, 0)
-		allocatedIPs, err := util.GetAllocatedIPs("")
+		allocatedIPs, err := util.GetAllocatedIPs(db, "")
 		if err != nil {
 			log.Error("Cannot suggest ip allocation. Failed to get list of allocated ip addresses: ", err)
 			return c.JSON(http.StatusInternalServerError, jsonHTTPResponse{
diff --git a/util/util.go b/util/util.go
index 7c347a9..0146f98 100644
--- a/util/util.go
+++ b/util/util.go
@@ -1,7 +1,6 @@
 package util
 
 import (
-	"encoding/json"
 	"errors"
 	"fmt"
 	"net"
@@ -15,7 +14,7 @@ import (
 	externalip "github.com/glendc/go-external-ip"
 	"github.com/labstack/gommon/log"
 	"github.com/ngoduykhanh/wireguard-ui/model"
-	"github.com/sdomino/scribble"
+	"github.com/ngoduykhanh/wireguard-ui/store"
 )
 
 // BuildClientConfig to create wireguard client config string
@@ -214,20 +213,15 @@ func GetIPFromCIDR(cidr string) (string, error) {
 }
 
 // GetAllocatedIPs to get all ip addresses allocated to clients and server
-func GetAllocatedIPs(ignoreClientID string) ([]string, error) {
+func GetAllocatedIPs(db store.IStore, ignoreClientID string) ([]string, error) {
 	allocatedIPs := make([]string, 0)
-
-	// initialize database directory
-	dir := "./db"
-	db, err := scribble.New(dir, nil)
-	if err != nil {
-		return nil, err
-	}
-
+	
 	// read server information
 	serverInterface := model.ServerInterface{}
-	if err := db.Read("server", "interfaces", &serverInterface); err != nil {
+	if server, err := db.GetServer(); err != nil {
 		return nil, err
+	} else{
+		serverInterface = *server.Interface
 	}
 
 	// append server's addresses to the result
@@ -240,7 +234,7 @@ func GetAllocatedIPs(ignoreClientID string) ([]string, error) {
 	}
 
 	// read client information
-	records, err := db.ReadAll("clients")
+	records, err := db.GetClients(false)
 	if err != nil {
 		return nil, err
 	}
@@ -248,9 +242,7 @@ func GetAllocatedIPs(ignoreClientID string) ([]string, error) {
 	// append client's addresses to the result
 	for _, f := range records {
 		client := model.Client{}
-		if err := json.Unmarshal([]byte(f), &client); err != nil {
-			return nil, err
-		}
+		client = *f.Client
 
 		if client.ID != ignoreClientID {
 			for _, cidr := range client.AllocatedIPs {