From bae3b95b6c7981bc60241c9bf6fdaf7096693764 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Krist=C3=B3f=20T=C3=B3th?= <mrtoth@strongds.hu>
Date: Fri, 15 Oct 2021 22:22:00 +0200
Subject: [PATCH] Improve cleanup logic (handle singals gracefully) and use
 config magic

---
 cmd/after-lock/after-lock.go | 101 +++++++++++++++++++++++++----------
 cmd/after-lock/logging.go    |  18 +++----
 2 files changed, 81 insertions(+), 38 deletions(-)

diff --git a/cmd/after-lock/after-lock.go b/cmd/after-lock/after-lock.go
index a27164b..94f8e62 100644
--- a/cmd/after-lock/after-lock.go
+++ b/cmd/after-lock/after-lock.go
@@ -2,49 +2,70 @@ package main
 
 import (
 	"kdelsd/afterlock"
+	"kdelsd/config"
+	"kdelsd/lockfile"
 	"kdelsd/xdg/display"
 	"kdelsd/xdg/keypressdetector"
 	"kdelsd/xdg/lockscreen"
-	"kdelsd/lockfile"
-	"os"
-	"runtime/debug"
 	"log"
+	"os"
+	"os/signal"
+	"runtime/debug"
+	"syscall"
 )
 
-const (
-	initialDelay = 3
-	loopDelay = 7
-	keyboardDeviceNode = "/dev/input/by-id/usb-Dell_Dell_USB_Entry_Keyboard-event-kbd"
-	lockfilePath = "/tmp/after-lock.lock"
-	printStackTraces = false
-	logFilePath = "/home/superuser/Desktop/after-lock.log"
-)
+
+type afterLockConfig struct {
+	INITIAL_DELAY int
+	LOOP_DELAY int
+	DEV_NODE string
+	LOCK_PATH string
+	LOG_PATH string
+	PRINT_STACKTRACES bool
+}
+
+var configuration = afterLockConfig{
+	INITIAL_DELAY: 3,
+	LOOP_DELAY: 7,
+	LOCK_PATH: "/tmp/after-lock.lock",
+	PRINT_STACKTRACES: true,
+}
+
+var mainDone chan bool
+var cleanupDone chan bool
+var sig chan os.Signal
+
+
+func init() {
+	mainDone = make(chan bool)
+	cleanupDone = make(chan bool)
+	sig = make(chan os.Signal)
+}
 
 
 func main() {
 	defer handleErrors()
-	f := initLogging()
-	if f != nil {
-		defer teardownLogging(f)
+	err := config.Build("AFL", &configuration)
+	if err != nil {
+		panic(err)
 	}
+	f := initLogging(configuration.LOG_PATH)
+	if f == nil {
+		panic("Failed to open logfile")
+	}
+	initCleanup(grabExclusiveProcessLock(configuration.LOCK_PATH), f)
 
-	lock := grabExclusiveProcessLock()
-	defer lock.Unlock()
-
-	al := afterlock.New(keyboardDeviceNode)
-	al.InitialDelay = initialDelay
-	al.LoopDelay = loopDelay
-	al.Display = display.Xset{}
-	al.LockScreen = lockscreen.XDG{}
-	al.KeypressDetector = keypressdetector.Evdev{}
-
+	al := buildAfterlock()
 	al.Start()
+
+	mainDone<-true
+	<-cleanupDone
 }
 
 func handleErrors() {
 	if err := recover(); err != nil {
 		log.Printf("Crashed with unexpected error: %v!\n", err)
-		if printStackTraces {
+		if configuration.PRINT_STACKTRACES {
 			log.Printf("Stack trace:\n\n")
 			debug.PrintStack()
 		}
@@ -52,11 +73,37 @@ func handleErrors() {
 	}
 }
 
-func grabExclusiveProcessLock() *lockfile.Lockfile {
-	lf := lockfile.New(lockfilePath)
+func grabExclusiveProcessLock(lockfile_path string) *lockfile.Lockfile {
+	lf := lockfile.New(lockfile_path)
 	err := lf.Lock()
 	if err != nil {
 		panic(err)
 	}
 	return lf
 }
+
+func initCleanup(lock *lockfile.Lockfile, logfile *os.File) {
+	signal.Notify(sig, syscall.SIGINT, syscall.SIGTERM)
+
+	go func() {
+		select{
+		case <-mainDone:
+		case <-sig:
+		}
+
+		lock.Unlock()
+		teardownLogging(logfile)
+
+		cleanupDone<-true
+	}()
+}
+
+func buildAfterlock() *afterlock.AfterLock {
+	al := afterlock.New(configuration.DEV_NODE)
+	al.InitialDelay = uint(configuration.INITIAL_DELAY)
+	al.LoopDelay = uint(configuration.LOOP_DELAY)
+	al.Display = display.Xset{}
+	al.LockScreen = lockscreen.XDG{}
+	al.KeypressDetector = keypressdetector.Evdev{}
+	return al
+}
\ No newline at end of file
diff --git a/cmd/after-lock/logging.go b/cmd/after-lock/logging.go
index 4eb48f4..88b2419 100644
--- a/cmd/after-lock/logging.go
+++ b/cmd/after-lock/logging.go
@@ -5,11 +5,7 @@ import (
 	"log"
 )
 
-func initLogging() *os.File {
-	if len(logFilePath) == 0 {
-		return nil
-	}
-
+func initLogging(logFilePath string) *os.File {
 	f, err := os.OpenFile(logFilePath, os.O_RDWR | os.O_CREATE | os.O_APPEND, 0600)
 	if err != nil {
 		panic(err)
@@ -22,16 +18,16 @@ func initLogging() *os.File {
 func teardownLogging(f *os.File) {
 	fi, err := f.Stat()
 	if err != nil {
-			panic(err)
+		panic(err)
 	}
 	if fi.Size() == 0 {
-			err = os.Remove(f.Name())
-			if err != nil {
-					panic(err)
-			}
+		err = os.Remove(f.Name())
+		if err != nil {
+			panic(err)
+		}
 	}
 	err = f.Close()
 	if err != nil {
-			panic(err)
+		panic(err)
 	}
 }
\ No newline at end of file