Merge pull request #68 from lidofinance/fix/airgapped-camera-lock

fix read_qr command lock

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)
 }

cmd/airgapped/main.go

@@ -35,6 +35,7 @@ type terminal struct {
 	airgapped *airgapped.Machine
 	commands  map[string]*terminalCommand
 
+	currentCommand string
 	stopDroppingSensitiveData chan bool
 }
 
@@ -43,6 +44,7 @@ func NewTerminal(machine *airgapped.Machine) *terminal {
 		bufio.NewReader(os.Stdin),
 		machine,
 		make(map[string]*terminalCommand),
+		"",
 		make(chan bool),
 	}
 	t.addCommand("read_qr", &terminalCommand{
@@ -297,7 +299,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
@@ -306,11 +310,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()
 	}
 }
@@ -362,13 +367,17 @@ func main() {
 
 	c := make(chan os.Signal, 1)
 	signal.Notify(c, os.Interrupt)
+
+	t := NewTerminal(air)
 	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.dropSensitiveDataByTicker(passwordLifeDuration)
 	if err = t.run(); err != nil {
 		log.Fatalf(err.Error())

mocks/qrMocks/qr_mock.go

@@ -84,3 +84,15 @@ func (mr *MockProcessorMockRecorder) SetChunkSize(chunkSize interface{}) *gomock
 	mr.mock.ctrl.T.Helper()
 	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "SetChunkSize", reflect.TypeOf((*MockProcessor)(nil).SetChunkSize), chunkSize)
 }
+
+// CloseCameraReader mocks base method
+func (m *MockProcessor) CloseCameraReader() {
+	m.ctrl.T.Helper()
+	m.ctrl.Call(m, "CloseCameraReader")
+}
+
+// CloseCameraReader indicates an expected call of CloseCameraReader
+func (mr *MockProcessorMockRecorder) CloseCameraReader() *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "CloseCameraReader", reflect.TypeOf((*MockProcessor)(nil).CloseCameraReader))
+}

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!")