From d3f48d633922ec95759d0b05afc28fa2d3b96caf Mon Sep 17 00:00:00 2001
From: Elara Musayelyan <elara@elara.ws>
Date: Thu, 21 Oct 2021 20:17:44 -0700
Subject: [PATCH] Fix issue where DFU responses are missed causing DFU to time
 out intermittently

---
 dfu.go | 85 +++++++++++++++++++++++++++++++++++++++++-----------------
 1 file changed, 61 insertions(+), 24 deletions(-)

diff --git a/dfu.go b/dfu.go
index e054ad3..5a7365f 100644
--- a/dfu.go
+++ b/dfu.go
@@ -6,7 +6,6 @@ import (
 	"encoding/binary"
 	"encoding/json"
 	"errors"
-	"io"
 	"io/fs"
 	"io/ioutil"
 	"os"
@@ -191,6 +190,23 @@ func (dfu *DFU) sendProgress() {
 	}
 }
 
+// bufferPropsCh resends all messages on in to a new channel that is buffered with 5 elements
+func bufferPropsCh(in chan *bluez.PropertyChanged) chan *bluez.PropertyChanged {
+	// Create new buffered channel
+	out := make(chan *bluez.PropertyChanged, 5)
+	go func() {
+		// Close channel when underlying channel closed
+		defer close(out)
+		// For every property change on in
+		for prop := range in {
+			// Write to out
+			out <- prop
+		}
+	}()
+	// Return new buffered channel
+	return out
+}
+
 // Start DFU process
 func (dfu *DFU) Start() error {
 	if dfu.progress == nil {
@@ -209,10 +225,12 @@ func (dfu *DFU) Start() error {
 	}
 
 	// Watch for property changes on control point
-	dfu.ctrlRespCh, err = dfu.ctrlPointChar.WatchProperties()
+	unbufferedCh, err := dfu.ctrlPointChar.WatchProperties()
 	if err != nil {
 		return err
 	}
+	// Buffer properties channel so no changes are missed
+	dfu.ctrlRespCh = bufferPropsCh(unbufferedCh)
 
 	// Run step one
 	err = dfu.stepOne()
@@ -279,24 +297,42 @@ func (dfu *DFU) Start() error {
 	return nil
 }
 
+// Reset reverts all values back to default to prepare for
+// the next DFU.
+func (dfu *DFU) Reset() {
+	dfu.bytesRecvd = 0
+	dfu.bytesSent = 0
+	dfu.fwSize = 0
+	dfu.fwSendDone = false
+	dfu.fwImage = nil
+	dfu.initPacket = nil
+	dfu.ctrlRespCh = nil
+	dfu.progress = nil
+}
+
 // on waits for the given command to be received on
 // the control point characteristic, then runs the callback.
 func (dfu *DFU) on(cmd []byte, onCmdCb func(data []byte) error) error {
-	select {
-	case propChanged := <-dfu.ctrlRespCh:
-		if propChanged.Name != "Value" {
+	// Use for loop in case of invalid property
+	for {
+		select {
+		case propChanged := <-dfu.ctrlRespCh:
+			// If property was invalid
+			if propChanged.Name != "Value" {
+				// Keep waiting
+				continue
+			}
+			// Assert propery value as byte slice
+			data := propChanged.Value.([]byte)
+			// If command has prefix of given command
+			if bytes.HasPrefix(data, cmd) {
+				// Return callback with data after command
+				return onCmdCb(data[len(cmd):])
+			}
 			return ErrDFUInvalidResponse
+		case <-time.After(50 * time.Second):
+			return ErrDFUTimeout
 		}
-		// Assert propery value as byte slice
-		data := propChanged.Value.([]byte)
-		// If command has prefix of given command
-		if bytes.HasPrefix(data, cmd) {
-			// Return callback with data after command
-			return onCmdCb(data[len(cmd):])
-		}
-		return ErrDFUInvalidResponse
-	case <-time.After(50 * time.Second):
-		return ErrDFUTimeout
 	}
 }
 
@@ -346,7 +382,13 @@ func (dfu *DFU) stepSeven() error {
 	// While send is not done
 	for !dfu.fwSendDone {
 		for i := 0; i < DFUPktRecvInterval; i++ {
-			amtLeft := int(dfu.fwSize) - dfu.bytesSent
+			amtLeft := dfu.fwSize - int64(dfu.bytesSent)
+			// If no bytes left to send, end transfer
+			if amtLeft == 0 {
+				dfu.sendProgress()
+				dfu.fwSendDone = true
+				return nil
+			}
 			var segment []byte
 			// If amount left is less than segment size
 			if amtLeft < DFUSegmentSize {
@@ -357,13 +399,8 @@ func (dfu *DFU) stepSeven() error {
 				segment = make([]byte, DFUSegmentSize)
 			}
 			// Write firmware image into slice
-			n, err := dfu.fwImage.Read(segment)
-			// If EOF, send is done
-			if err == io.EOF {
-				dfu.sendProgress()
-				dfu.fwSendDone = true
-				return nil
-			} else if err != nil {
+			_, err := dfu.fwImage.Read(segment)
+			if err != nil {
 				return err
 			}
 			// Write segment to packet characteristic
@@ -372,7 +409,7 @@ func (dfu *DFU) stepSeven() error {
 				return err
 			}
 			// Increment bytes sent by amount read
-			dfu.bytesSent += n
+			dfu.bytesSent += len(segment)
 		}
 		// On 0x11, verify packet receipt size
 		err := dfu.on(DFUNotifPktRecvd, func(data []byte) error {