From ea488067fb9b80d65c12449b702861aec933d1af Mon Sep 17 00:00:00 2001
From: Arsen Musayelyan <moussaelianarsen@gmail.com>
Date: Sat, 21 Aug 2021 00:04:29 -0700
Subject: [PATCH] Add progress to DFU

---
 dfu.go | 31 ++++++++++++++++++++++++++++++-
 1 file changed, 30 insertions(+), 1 deletion(-)

diff --git a/dfu.go b/dfu.go
index 9e990ce..ba89b50 100644
--- a/dfu.go
+++ b/dfu.go
@@ -55,6 +55,12 @@ var (
 
 var btOptsCmd = map[string]interface{}{"type": "command"}
 
+type DFUProgress struct {
+	Sent     int   `json:"sent"`
+	Received int   `json:"recvd"`
+	Total    int64 `json:"total"`
+}
+
 // DFU stores everything required for doing firmware upgrades
 type DFU struct {
 	initPacket    fs.File
@@ -64,6 +70,7 @@ type DFU struct {
 	bytesSent     int
 	bytesRecvd    int
 	fwSendDone    bool
+	progress      chan DFUProgress
 	ctrlPointChar *gatt.GattCharacteristic1
 	packetChar    *gatt.GattCharacteristic1
 }
@@ -83,7 +90,7 @@ func (dfu *DFU) LoadFiles(initPath, fwPath string) error {
 		return err
 	}
 	dfu.fwImage = fwImgFl
-	
+
 	// Get firmware file size
 	dfu.fwSize, err = getFlSize(dfu.fwImage)
 	if err != nil {
@@ -169,8 +176,28 @@ func getFlSize(fl fs.File) (int64, error) {
 	return flInfo.Size(), nil
 }
 
+func (dfu *DFU) Progress() <-chan DFUProgress {
+	if dfu.progress == nil {
+		dfu.progress = make(chan DFUProgress, 5)
+	}
+	return dfu.progress
+}
+
+func (dfu *DFU) sendProgress() {
+	dfu.progress <- DFUProgress{
+		Sent:     dfu.bytesSent,
+		Received: dfu.bytesRecvd,
+		Total:    dfu.fwSize,
+	}
+}
+
 // Start DFU process
 func (dfu *DFU) Start() error {
+	if dfu.progress == nil {
+		dfu.progress = make(chan DFUProgress, 5)
+	}
+	defer close(dfu.progress)
+
 	if dfu.fwImage == nil || dfu.initPacket == nil {
 		return ErrDFUNoFilesLoaded
 	}
@@ -325,6 +352,7 @@ func (dfu *DFU) stepSeven() error {
 			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 {
@@ -345,6 +373,7 @@ func (dfu *DFU) stepSeven() error {
 			if dfu.bytesRecvd != dfu.bytesSent {
 				return ErrDFUSizeMismatch
 			}
+			dfu.sendProgress()
 			return nil
 		})
 		if err != nil {