From e6458aba0bc3dfa93e1d93baf4256c334e8cbf8e Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Krist=C3=B3f=20T=C3=B3th?= <mrtoth@strongds.hu>
Date: Thu, 20 Aug 2020 15:57:17 +0200
Subject: [PATCH] Ensure only a single instance of after-lock runs at a time

---
 cmd/after-lock/after-lock.go | 31 ++++++++++++++++
 lockfile/lockfile.go         | 49 +++++++++++++++++++++++++
 lockfile/lockfile_test.go    | 71 ++++++++++++++++++++++++++++++++++++
 3 files changed, 151 insertions(+)
 create mode 100644 lockfile/lockfile.go
 create mode 100644 lockfile/lockfile_test.go

diff --git a/cmd/after-lock/after-lock.go b/cmd/after-lock/after-lock.go
index c7af146..317df44 100644
--- a/cmd/after-lock/after-lock.go
+++ b/cmd/after-lock/after-lock.go
@@ -2,19 +2,50 @@ package main
 
 import (
 	"after-lock/afterlock"
+	"after-lock/lockfile"
+	"fmt"
+	"runtime/debug"
+	"os"
 )
 
 const (
 	initialDelay = 3
 	loopDelay = 6
 	keyboardDeviceNode = "/dev/input/by-id/usb-Dell_Dell_USB_Entry_Keyboard-event-kbd"
+	lockfilePath = "/tmp/after-lock.lock"
+	printStackTraces = false
 )
 
 
 func main() {
+	defer handleErrors()
+
+	lock := grabExclusiveProcessLock()
+	defer lock.Unlock()
+
 	al := afterlock.New(keyboardDeviceNode)
 	al.InitialDelay = initialDelay
 	al.LoopDelay = loopDelay
 
 	al.Start()
+}
+
+func handleErrors() {
+	if err := recover(); err != nil {
+		fmt.Printf("Crashed with unexpected error: %v!\n", err)
+		if printStackTraces {
+			fmt.Printf("Stack trace:\n\n")
+			debug.PrintStack()
+		}
+		os.Exit(1)
+	}
+}
+
+func grabExclusiveProcessLock() *lockfile.Lockfile {
+	lf := lockfile.New(lockfilePath)
+	err := lf.Lock()
+	if err != nil {
+		panic(err)
+	}
+	return lf
 }
\ No newline at end of file
diff --git a/lockfile/lockfile.go b/lockfile/lockfile.go
new file mode 100644
index 0000000..4e1852e
--- /dev/null
+++ b/lockfile/lockfile.go
@@ -0,0 +1,49 @@
+package lockfile
+
+import (
+	"errors"
+	"os"
+	"syscall"
+)
+
+
+type Lockfile struct {
+	path string
+	fd int
+}
+
+func New(path string) *Lockfile {
+	return &Lockfile{path: path}
+}
+
+func (lf *Lockfile) Lock() error {
+	fd, err := syscall.Open(lf.path, syscall.O_CREAT | syscall.O_RDWR, 0600)
+	if err != nil {
+		return err
+	}
+	lf.fd = fd
+	err = syscall.Flock(lf.fd, syscall.LOCK_EX | syscall.LOCK_NB)
+	if err != nil {
+		return err
+	}
+	return nil
+}
+
+func (lf *Lockfile) Unlock() error {
+	if lf.fd == 0 {
+		panic(errors.New("Lockfile was not locked"))
+	}
+	err := syscall.Flock(lf.fd, syscall.LOCK_UN)
+	if err != nil {
+		return err
+	}
+	err = syscall.Close(lf.fd)
+	if err != nil {
+		return err
+	}
+	os.Remove(lf.path)
+	if err != nil {
+		return err
+	}
+	return nil
+}
diff --git a/lockfile/lockfile_test.go b/lockfile/lockfile_test.go
new file mode 100644
index 0000000..5a73274
--- /dev/null
+++ b/lockfile/lockfile_test.go
@@ -0,0 +1,71 @@
+package lockfile_test
+
+import (
+	"io/ioutil"
+	"os"
+	"testing"
+
+	"after-lock/lockfile"
+)
+
+
+func tmpFile() string {
+	tf, err := ioutil.TempFile("", "lockfile_test*")
+	if err != nil {
+		panic(err)
+	}
+	return tf.Name()
+}
+
+func TestNotLockedUnlockPanics(t *testing.T) {
+	defer func() {
+        if r := recover(); r == nil {
+            t.Errorf("Unlock() did not panic")
+        }
+    }()
+
+	tf := tmpFile()
+	defer os.Remove(tf)
+
+
+	err := lockfile.New(tf).Unlock()
+	if err != nil {
+		panic(err)
+	}
+}
+
+func TestSecondLockFails(t *testing.T) {
+	tf := tmpFile()
+	defer os.Remove(tf)
+
+	fl1 := lockfile.New(tf)
+	err := fl1.Lock()
+	if err != nil {
+		panic(err)
+	}
+	defer fl1.Unlock()
+
+	fl2 := lockfile.New(tf)
+	err = fl2.Lock()
+	if err == nil {
+		t.Errorf("Second Lock() did not fail")
+	}
+}
+
+func TestLockUnlock(t *testing.T) {
+	tf := tmpFile()
+	defer os.Remove(tf)
+
+	fl := lockfile.New(tf)
+	err := fl.Lock()
+	if err != nil {
+		panic(err)
+	}
+	err = fl.Unlock()
+	if err != nil {
+		panic(err)
+	}
+	if _, err := os.Stat(tf); !os.IsNotExist(err) {
+		t.Errorf("Lockfile %s still exists", tf)
+	}
+}