From 730dc35d94734a41bf59eeed4955ff8c87da7b60 Mon Sep 17 00:00:00 2001
From: lawl <github@dumbinter.net>
Date: Sun, 31 Jan 2021 07:44:22 +0100
Subject: [PATCH] Require capabilities on startup

---
 capability.go |  6 +++++-
 main.go       |  2 ++
 ui.go         | 55 +++++++++++++++++++++++++++++++++++++++++++++++++++
 3 files changed, 62 insertions(+), 1 deletion(-)

diff --git a/capability.go b/capability.go
index f8dadce..65cb17a 100644
--- a/capability.go
+++ b/capability.go
@@ -58,10 +58,11 @@ func makeBinarySetcapped() error {
 	return nil
 }
 
-func pkexecSetcapSelf() {
+func pkexecSetcapSelf() error {
 	self, err := os.Executable()
 	if err != nil {
 		log.Fatalf("Couldn't find path to own binary\n")
+		return err
 	}
 
 	cmd := exec.Command("pkexec", self, "-setcap")
@@ -69,5 +70,8 @@ func pkexecSetcapSelf() {
 	err = cmd.Run()
 	if err != nil {
 		log.Printf("Couldn't setcap self as root: %v\n", err)
+		return err
 	}
+
+	return nil
 }
diff --git a/main.go b/main.go
index 7b3f20a..1d997c3 100644
--- a/main.go
+++ b/main.go
@@ -200,6 +200,8 @@ func main() {
 
 	go paConnectionWatchdog(&ctx)
 
+	ctx.haveCapabilities = hasCapSysResource(getCurrentCaps())
+
 	wnd := nucular.NewMasterWindowSize(0, appName, image.Point{600, 400}, func(w *nucular.Window) {
 		updatefn(&ctx, w)
 	})
diff --git a/ui.go b/ui.go
index 0343a6d..cf171bc 100644
--- a/ui.go
+++ b/ui.go
@@ -8,7 +8,9 @@ import (
 	"image/draw"
 	"image/png"
 	"log"
+	"os"
 	"os/exec"
+	"syscall"
 	"time"
 
 	"github.com/aarzilli/nucular"
@@ -31,6 +33,8 @@ type ntcontext struct {
 	masterWindow             *nucular.MasterWindow
 	update                   updateui
 	reloadRequired           bool
+	haveCapabilities         bool
+	errorMsg                 string
 }
 
 var green = color.RGBA{34, 187, 69, 255}
@@ -41,6 +45,18 @@ var patreonImg *image.RGBA
 
 func updatefn(ctx *ntcontext, w *nucular.Window) {
 
+	//TODO: this is disgusting
+
+	if ctx.errorMsg != "" {
+		errorScreen(ctx, w)
+		return
+	}
+
+	if !ctx.haveCapabilities {
+		capabilitiesScreen(ctx, w)
+		return
+	}
+
 	if !ctx.paClient.Connected() {
 		connectScreen(ctx, w)
 		return
@@ -373,6 +389,45 @@ func connectScreen(ctx *ntcontext, w *nucular.Window) {
 	w.Label("Connecting to pulseaudio...", "CB")
 }
 
+func capabilitiesScreen(ctx *ntcontext, w *nucular.Window) {
+	w.Row(15).Dynamic(1)
+	w.Label("NoiseTorch currently does not have the capabilities to function properly.", "CB")
+	w.Row(15).Dynamic(1)
+	w.Label("We require CAP_SYS_RESOURCE. If that doesn't mean anything to you, don't worry. I'll fix it for you.", "CB")
+	w.Row(40).Dynamic(1)
+	w.Row(25).Dynamic(1)
+	if w.ButtonText("Grant capability (requires root)") {
+		err := pkexecSetcapSelf()
+		if err != nil {
+			ctx.errorMsg = err.Error()
+			return
+		}
+		self, err := os.Executable()
+		if err != nil {
+			ctx.errorMsg = err.Error()
+			return
+		}
+		err = syscall.Exec(self, []string{""}, os.Environ())
+		if err != nil {
+			ctx.errorMsg = err.Error()
+			return
+		}
+	}
+}
+
+func errorScreen(ctx *ntcontext, w *nucular.Window) {
+	w.Row(15).Dynamic(1)
+	w.Label("Error", "CB")
+	w.Row(15).Dynamic(1)
+	w.Label(ctx.errorMsg, "CB")
+	w.Row(40).Dynamic(1)
+	w.Row(25).Dynamic(1)
+	if w.ButtonText("OK") {
+		ctx.errorMsg = ""
+		return
+	}
+}
+
 func resetUI(ctx *ntcontext) {
 	ctx.loadingScreen = false