From 3bbc442955ff04e275b2cf96490ead7a8e01c9fa Mon Sep 17 00:00:00 2001
From: programmer10110 <programmer10110@yandex.ru>
Date: Wed, 28 Oct 2020 16:25:06 +0300
Subject: [PATCH] fix read_qr command lock

---
 airgapped/airgapped.go |  4 +++
 cmd/airgapped/main.go  | 22 ++++++++----
 qr/qr.go               | 77 +++++++++++++++++++++++++-----------------
 3 files changed, 66 insertions(+), 37 deletions(-)

diff --git a/airgapped/airgapped.go b/airgapped/airgapped.go
index 2f8ee9d..b72c888 100644
--- a/airgapped/airgapped.go
+++ b/airgapped/airgapped.go
@@ -78,6 +78,10 @@ func (am *Machine) SetQRProcessorFramesDelay(delay int) {
 	am.qrProcessor.SetDelay(delay)
 }
 
+func (am *Machine) CloseCameraReader() {
+	am.qrProcessor.CloseCameraReader()
+}
+
 func (am *Machine) SetQRProcessorChunkSize(chunkSize int) {
 	am.qrProcessor.SetChunkSize(chunkSize)
 }
diff --git a/cmd/airgapped/main.go b/cmd/airgapped/main.go
index 52d5910..33903e3 100644
--- a/cmd/airgapped/main.go
+++ b/cmd/airgapped/main.go
@@ -33,10 +33,12 @@ type terminal struct {
 	reader    *bufio.Reader
 	airgapped *airgapped.Machine
 	commands  map[string]*terminalCommand
+
+	currentCommand string
 }
 
 func NewTerminal(machine *airgapped.Machine) *terminal {
-	t := terminal{bufio.NewReader(os.Stdin), machine, make(map[string]*terminalCommand)}
+	t := terminal{bufio.NewReader(os.Stdin), machine, make(map[string]*terminalCommand), ""}
 	t.addCommand("read_qr", &terminalCommand{
 		commandHandler: t.readQRCommand,
 		description:    "Reads QR chunks from camera, handle a decoded operation and returns paths to qr chunks of operation's result",
@@ -229,7 +231,9 @@ func (t *terminal) run() error {
 		if err != nil {
 			return fmt.Errorf("failed to read command: %w", err)
 		}
-		handler, ok := t.commands[strings.Trim(command, "\n")]
+
+		clearCommand := strings.Trim(command, "\n")
+		handler, ok := t.commands[clearCommand]
 		if !ok {
 			fmt.Printf("unknown command: %s\n", command)
 			continue
@@ -238,11 +242,12 @@ func (t *terminal) run() error {
 			return err
 		}
 		t.airgapped.Lock()
+
+		t.currentCommand = clearCommand
 		if err := handler.commandHandler(); err != nil {
 			fmt.Printf("failled to execute command %s: %v \n", command, err)
-			t.airgapped.Unlock()
-			continue
 		}
+		t.currentCommand = ""
 		t.airgapped.Unlock()
 	}
 }
@@ -288,14 +293,19 @@ func main() {
 
 	c := make(chan os.Signal, 1)
 	signal.Notify(c, os.Interrupt)
+
+	t := NewTerminal(air)
+	go t.dropSensitiveData(passwordLifeDuration)
 	go func() {
 		for range c {
+			if t.currentCommand == "read_qr" {
+				t.airgapped.CloseCameraReader()
+				continue
+			}
 			fmt.Printf("Intercepting SIGINT, please type `exit` to stop the machine\n>>> ")
 		}
 	}()
 
-	t := NewTerminal(air)
-	go t.dropSensitiveData(passwordLifeDuration)
 	if err = t.run(); err != nil {
 		log.Fatalf(err.Error())
 	}
diff --git a/qr/qr.go b/qr/qr.go
index 3da357f..5fcd114 100644
--- a/qr/qr.go
+++ b/qr/qr.go
@@ -30,15 +30,24 @@ type Processor interface {
 	WriteQR(path string, data []byte) error
 	SetDelay(delay int)
 	SetChunkSize(chunkSize int)
+	CloseCameraReader()
 }
 
 type CameraProcessor struct {
 	gifFramesDelay int
 	chunkSize      int
+
+	closeCameraReader chan bool
 }
 
-func NewCameraProcessor() *CameraProcessor {
-	return &CameraProcessor{}
+func NewCameraProcessor() Processor {
+	return &CameraProcessor{
+		closeCameraReader: make(chan bool),
+	}
+}
+
+func (p *CameraProcessor) CloseCameraReader() {
+	p.closeCameraReader <- true
 }
 
 func (p *CameraProcessor) SetDelay(delay int) {
@@ -73,37 +82,43 @@ func (p *CameraProcessor) ReadQR() ([]byte, error) {
 	chunks := make([]*chunk, 0)
 	decodedChunksCount := uint(0)
 	// detects and scans QR-cods from camera until we scan successfully
+READER:
 	for {
-		webcam.Read(&img)
-		window.IMShow(img)
-		window.WaitKey(1)
+		select {
+		case <-p.closeCameraReader:
+			return nil, fmt.Errorf("camera reader was closed")
+		default:
+			webcam.Read(&img)
+			window.IMShow(img)
+			window.WaitKey(1)
 
-		imgObject, err := img.ToImage()
-		if err != nil {
-			return nil, fmt.Errorf("failed to get image object: %w", err)
-		}
-		data, err := ReadDataFromQR(imgObject)
-		if err != nil {
-			continue
-		}
-		decodedChunk, err := decodeChunk(data)
-		if err != nil {
-			return nil, err
-		}
-		if cap(chunks) == 0 {
-			chunks = make([]*chunk, decodedChunk.Total)
-		}
-		if decodedChunk.Index > decodedChunk.Total {
-			return nil, fmt.Errorf("invalid QR-code chunk")
-		}
-		if chunks[decodedChunk.Index] != nil {
-			continue
-		}
-		chunks[decodedChunk.Index] = decodedChunk
-		decodedChunksCount++
-		window.SetWindowTitle(fmt.Sprintf("Read %d/%d chunks", decodedChunksCount, decodedChunk.Total))
-		if decodedChunksCount == decodedChunk.Total {
-			break
+			imgObject, err := img.ToImage()
+			if err != nil {
+				return nil, fmt.Errorf("failed to get image object: %w", err)
+			}
+			data, err := ReadDataFromQR(imgObject)
+			if err != nil {
+				continue
+			}
+			decodedChunk, err := decodeChunk(data)
+			if err != nil {
+				return nil, err
+			}
+			if cap(chunks) == 0 {
+				chunks = make([]*chunk, decodedChunk.Total)
+			}
+			if decodedChunk.Index > decodedChunk.Total {
+				return nil, fmt.Errorf("invalid QR-code chunk")
+			}
+			if chunks[decodedChunk.Index] != nil {
+				continue
+			}
+			chunks[decodedChunk.Index] = decodedChunk
+			decodedChunksCount++
+			window.SetWindowTitle(fmt.Sprintf("Read %d/%d chunks", decodedChunksCount, decodedChunk.Total))
+			if decodedChunksCount == decodedChunk.Total {
+				break READER
+			}
 		}
 	}
 	window.SetWindowTitle("QR-code chunks successfully read!")