From ce95ce259c7b147603d7f96b9a6c068eb2eefdfc Mon Sep 17 00:00:00 2001
From: Elara Musayelyan <elara@elara.ws>
Date: Wed, 6 Oct 2021 17:08:21 -0700
Subject: [PATCH] Prevent sending of extra zero bytes during DFU

---
 dfu.go | 12 ++++++++++--
 1 file changed, 10 insertions(+), 2 deletions(-)

diff --git a/dfu.go b/dfu.go
index ba89b50..e054ad3 100644
--- a/dfu.go
+++ b/dfu.go
@@ -346,8 +346,16 @@ func (dfu *DFU) stepSeven() error {
 	// While send is not done
 	for !dfu.fwSendDone {
 		for i := 0; i < DFUPktRecvInterval; i++ {
-			// Create byte slice with segment size
-			segment := make([]byte, DFUSegmentSize)
+			amtLeft := int(dfu.fwSize) - dfu.bytesSent
+			var segment []byte
+			// If amount left is less than segment size
+			if amtLeft < DFUSegmentSize {
+				// Create byte slice with amount left
+				segment = make([]byte, amtLeft)
+			} else {
+				// Create byte slice with segment size
+				segment = make([]byte, DFUSegmentSize)
+			}
 			// Write firmware image into slice
 			n, err := dfu.fwImage.Read(segment)
 			// If EOF, send is done