fix: fixed race conditions

Author: MarvinJWendt

keyboard.go

@@ -3,6 +3,8 @@ package keyboard
 import (
 	"fmt"
 	"os"
+	"sync"
+	"sync/atomic"
 
 	"github.com/containerd/console"
 
@@ -64,27 +66,39 @@ func stopListener() error {
 //		return false, nil // Return false to continue listening
 //	})
 func Listen(onKeyPress func(key keys.Key) (stop bool, err error)) error {
-	cancel := make(chan bool)
-	stopRoutine := false
-
-	go func() {
-		for {
-			select {
-			case c := <-cancel:
-				if c {
-					return
+	var (
+		callbackMu  sync.Mutex
+		closeOnce   sync.Once
+		stopOnce    sync.Once
+		stopRoutine atomic.Bool
+	)
+
+	stopped := make(chan struct{})
+	stop := func() {
+		stopOnce.Do(func() {
+			stopRoutine.Store(true)
+			closeOnce.Do(func() {
+				closeInput()
+
+				if inputTTY != nil {
+					_ = inputTTY.Close()
 				}
+			})
+			close(stopped)
+		})
+	}
 
-			case keyInfo := <-mockChannel:
-				stopRoutine, _ = onKeyPress(keyInfo)
-				if stopRoutine {
-					closeInput()
+	handleKeyPress := func(key keys.Key) (bool, error) {
+		callbackMu.Lock()
+		defer callbackMu.Unlock()
 
-					_ = inputTTY.Close()
-				}
-			}
+		shouldStop, err := onKeyPress(key)
+		if shouldStop {
+			stop()
 		}
-	}()
+
+		return shouldStop, err
+	}
 
 	err := startListener()
 	if err != nil {
@@ -93,28 +107,53 @@ func Listen(onKeyPress func(key keys.Key) (stop bool, err error)) error {
 		}
 	}
 
-	for !stopRoutine {
-		key, keyErr := getKeyPress()
-		if keyErr != nil {
-			return keyErr
-		}
+	cancel := make(chan struct{})
+	var mockWG sync.WaitGroup
+	mockWG.Add(1)
+
+	go func() {
+		defer mockWG.Done()
 
-		// check if returned key is empty
-		// if reflect.DeepEqual(key, keys.Key{}) {
-		// 	return nil
-		// }
+		for {
+			select {
+			case <-cancel:
+				return
 
-		stop, stopErr := onKeyPress(key)
-		if stopErr != nil {
-			return stopErr
+			case keyInfo := <-mockChannel:
+				shouldStop, _ := handleKeyPress(keyInfo)
+				if shouldStop {
+					return
+				}
+			}
 		}
+	}()
 
-		if stop {
-			closeInput()
+	defer func() {
+		close(cancel)
+		mockWG.Wait()
+	}()
+
+	if inputTTY == nil {
+		<-stopped
+	} else {
+		for !stopRoutine.Load() {
+			key, keyErr := getKeyPress()
+			if keyErr != nil {
+				return keyErr
+			}
+
+			if stopRoutine.Load() {
+				break
+			}
 
-			_ = inputTTY.Close()
+			shouldStop, stopErr := handleKeyPress(key)
+			if stopErr != nil {
+				return stopErr
+			}
 
-			break
+			if shouldStop {
+				break
+			}
 		}
 	}
 
@@ -123,8 +162,6 @@ func Listen(onKeyPress func(key keys.Key) (stop bool, err error)) error {
 		return err
 	}
 
-	cancel <- true
-
 	return nil
 }