From d2d69139b91b60e35bd2d43eb8a1870d758eb7bb Mon Sep 17 00:00:00 2001
From: Daniel <dhaavi@users.noreply.github.com>
Date: Fri, 24 Apr 2020 10:02:07 +0200
Subject: [PATCH] Add control db interface for triggering hooks, add
 shutdown/restart hooks

---
 core/control.go   | 100 ++++++++++++++++++++++++++++++++++++++++++++++
 core/core.go      |   9 +++++
 core/databases.go |   2 +-
 core/events.go    |  46 +++++++++++++++++++++
 4 files changed, 156 insertions(+), 1 deletion(-)
 create mode 100644 core/control.go
 create mode 100644 core/events.go

diff --git a/core/control.go b/core/control.go
new file mode 100644
index 00000000..b94163d9
--- /dev/null
+++ b/core/control.go
@@ -0,0 +1,100 @@
+package core
+
+import (
+	"fmt"
+	"strings"
+	"sync"
+
+	"github.com/safing/portbase/database"
+	"github.com/safing/portbase/database/record"
+	"github.com/safing/portbase/database/storage"
+)
+
+// StorageInterface provices a storage.Interface to the storage manager.
+type StorageInterface struct {
+	storage.InjectBase
+}
+
+// Get returns a database record.
+func (s *StorageInterface) Get(key string) (record.Record, error) {
+	msg := newMessage(key)
+	splittedKey := strings.Split(key, "/")
+
+	switch splittedKey[0] {
+	case "module":
+		return controlModule(msg, splittedKey)
+	default:
+		return nil, storage.ErrNotFound
+	}
+}
+
+func controlModule(msg *MessageRecord, splittedKey []string) (record.Record, error) {
+	// format: module/moduleName/action/param
+	var moduleName string
+	var action string
+	var param string
+	var err error
+
+	// parse elements
+	switch len(splittedKey) {
+	case 4:
+		param = splittedKey[3]
+		fallthrough
+	case 3:
+		moduleName = splittedKey[1]
+		action = splittedKey[2]
+	default:
+		return nil, storage.ErrNotFound
+	}
+
+	// execute
+	switch action {
+	case "trigger":
+		err = module.InjectEvent(fmt.Sprintf("user triggered the '%s/%s' event", moduleName, param), moduleName, param, nil)
+	default:
+		return nil, storage.ErrNotFound
+	}
+
+	if err != nil {
+		msg.Message = err.Error()
+	} else {
+		msg.Success = true
+	}
+
+	return msg, nil
+}
+
+func registerControlDatabase() error {
+	_, err := database.Register(&database.Database{
+		Name:        "control",
+		Description: "Control Interface for the Portmaster",
+		StorageType: "injected",
+		PrimaryAPI:  "",
+	})
+	if err != nil {
+		return err
+	}
+
+	_, err = database.InjectDatabase("control", &StorageInterface{})
+	if err != nil {
+		return err
+	}
+
+	return nil
+}
+
+// MessageRecord is a simple record used for control database feedback
+type MessageRecord struct {
+	record.Base
+	sync.Mutex
+
+	Success bool
+	Message string
+}
+
+func newMessage(key string) *MessageRecord {
+	m := &MessageRecord{}
+	m.SetKey("control:" + key)
+	m.UpdateMeta()
+	return m
+}
diff --git a/core/core.go b/core/core.go
index d70569f0..039758e0 100644
--- a/core/core.go
+++ b/core/core.go
@@ -31,10 +31,19 @@ func init() {
 	)
 }
 
+func prep() error {
+	registerEvents()
+	return nil
+}
+
 func start() error {
 	if err := startPlatformSpecific(); err != nil {
 		return fmt.Errorf("failed to start plattform-specific components: %s", err)
 	}
 
+	if err := registerEventHooks(); err != nil {
+		return err
+	}
+
 	return nil
 }
diff --git a/core/databases.go b/core/databases.go
index 941aa97a..df8f37ae 100644
--- a/core/databases.go
+++ b/core/databases.go
@@ -46,5 +46,5 @@ func registerDatabases() error {
 	//   return err
 	// }
 
-	return nil
+	return registerControlDatabase()
 }
diff --git a/core/events.go b/core/events.go
new file mode 100644
index 00000000..908347fd
--- /dev/null
+++ b/core/events.go
@@ -0,0 +1,46 @@
+package core
+
+import (
+	"context"
+
+	"github.com/safing/portbase/log"
+	"github.com/safing/portbase/modules"
+)
+
+const (
+	eventShutdown = "shutdown"
+	eventRestart  = "restart"
+	restartCode   = 23
+)
+
+func registerEvents() {
+	module.RegisterEvent(eventShutdown)
+	module.RegisterEvent(eventRestart)
+}
+
+func registerEventHooks() error {
+	err := module.RegisterEventHook("core", eventShutdown, "execute shutdown", shutdown)
+	if err != nil {
+		return err
+	}
+
+	err = module.RegisterEventHook("core", eventRestart, "execute shutdown", restart)
+	if err != nil {
+		return err
+	}
+
+	return nil
+}
+
+func shutdown(ctx context.Context, _ interface{}) error {
+	log.Warning("core: user requested shutdown")
+	go modules.Shutdown() //nolint:errcheck
+	return nil
+}
+
+func restart(ctx context.Context, data interface{}) error {
+	log.Info("core: user requested restart")
+	modules.SetExitStatusCode(restartCode)
+	go modules.Shutdown() //nolint:errcheck
+	return nil
+}