From 95fbe85ba2fbfc1958d6757059e3c09a47bcbc6c Mon Sep 17 00:00:00 2001
From: Daniel <dhaavi@users.noreply.github.com>
Date: Tue, 9 Nov 2021 10:09:30 +0100
Subject: [PATCH] Add Withdraw function to database controllers

---
 database/controller.go  |  4 +++-
 database/controllers.go | 27 ++++++++++++++++++++++-----
 2 files changed, 25 insertions(+), 6 deletions(-)

diff --git a/database/controller.go b/database/controller.go
index cb89cf5..47f7444 100644
--- a/database/controller.go
+++ b/database/controller.go
@@ -14,6 +14,7 @@ import (
 
 // A Controller takes care of all the extra database logic.
 type Controller struct {
+	database     *Database
 	storage      storage.Interface
 	shadowDelete bool
 
@@ -25,8 +26,9 @@ type Controller struct {
 }
 
 // newController creates a new controller for a storage.
-func newController(storageInt storage.Interface, shadowDelete bool) *Controller {
+func newController(database *Database, storageInt storage.Interface, shadowDelete bool) *Controller {
 	return &Controller{
+		database:     database,
 		storage:      storageInt,
 		shadowDelete: shadowDelete,
 	}
diff --git a/database/controllers.go b/database/controllers.go
index 9b2ac68..13d8fde 100644
--- a/database/controllers.go
+++ b/database/controllers.go
@@ -8,6 +8,8 @@ import (
 	"github.com/safing/portbase/database/storage"
 )
 
+const StorageTypeInjected = "injected"
+
 var (
 	controllers     = make(map[string]*Controller)
 	controllersLock sync.RWMutex
@@ -39,6 +41,11 @@ func getController(name string) (*Controller, error) {
 		return nil, fmt.Errorf(`could not start database %s: %s`, name, err)
 	}
 
+	// Check if database is injected.
+	if registeredDB.StorageType == StorageTypeInjected {
+		return nil, fmt.Errorf("database storage is not injected")
+	}
+
 	// get location
 	dbLocation, err := getLocation(name, registeredDB.StorageType)
 	if err != nil {
@@ -51,7 +58,7 @@ func getController(name string) (*Controller, error) {
 		return nil, fmt.Errorf(`could not start database %s (type %s): %s`, name, registeredDB.StorageType, err)
 	}
 
-	controller = newController(storageInt, registeredDB.ShadowDelete)
+	controller = newController(registeredDB, storageInt, registeredDB.ShadowDelete)
 	controllers[name] = controller
 	return controller, nil
 }
@@ -76,13 +83,23 @@ func InjectDatabase(name string, storageInt storage.Interface) (*Controller, err
 	// check if database is registered
 	registeredDB, ok := registry[name]
 	if !ok {
-		return nil, fmt.Errorf(`database "%s" not registered`, name)
+		return nil, fmt.Errorf("database %q not registered", name)
 	}
-	if registeredDB.StorageType != "injected" {
-		return nil, fmt.Errorf(`database not of type "injected"`)
+	if registeredDB.StorageType != StorageTypeInjected {
+		return nil, fmt.Errorf("database not of type %q", StorageTypeInjected)
 	}
 
-	controller := newController(storageInt, false)
+	controller := newController(registeredDB, storageInt, false)
 	controllers[name] = controller
 	return controller, nil
 }
+
+// Withdraw withdraws an injected database, but leaves the database registered.
+func (c *Controller) Withdraw() {
+	if c != nil && c.Injected() {
+		controllersLock.Lock()
+		defer controllersLock.Unlock()
+
+		delete(controllers, c.database.Name)
+	}
+}