From 5bdbb430e7f2f055af5e4e433ff881dee1dddf4e Mon Sep 17 00:00:00 2001
From: Olaf Rempel <razzor@kopf-tisch.de>
Date: Sat, 1 Feb 2020 23:02:19 +0100
Subject: [PATCH] Fix NACK handling for flash write

---
 README.md |  8 ++++----
 main.c    | 21 +++++++++++++++++----
 2 files changed, 21 insertions(+), 8 deletions(-)

diff --git a/README.md b/README.md
index ec12bbb..147af8b 100644
--- a/README.md
+++ b/README.md
@@ -10,10 +10,10 @@ Currently the following AVR MCUs are supported:
 
 AVR MCU | Flash bytes used (.text + .data) | Bootloader region size
 --- | --- | ---
-atmega8 | 706 (0x2C2) | 512 words
-atmega88 | 726 (0x2D6) | 512 words
-atmega168 | 726 (0x2D6) | 512 words
-atmega328p | 726 (0x2D6) | 512 words
+atmega8 | 718 (0x2CE) | 512 words
+atmega88 | 740 (0x2E4) | 512 words
+atmega168 | 740 (0x2E4) | 512 words
+atmega328p | 740 (0x2E4) | 512 words
 
 (Compiled on Ubuntu 18.04 LTS (gcc 5.4.0 / avr-libc 2.0.0) with EEPROM and LED support)
 
diff --git a/main.c b/main.c
index 0859d5f..427a982 100644
--- a/main.c
+++ b/main.c
@@ -280,13 +280,21 @@ static uint8_t TWI_data_write(uint8_t bcnt, uint8_t data)
             switch (cmd)
             {
                 case CMD_ACCESS_FLASH:
-                    buf[bcnt -4] = data;
-                    if (bcnt >= sizeof(buf) +3)
+                {
+                    uint8_t pos = bcnt -4;
+
+                    buf[pos] = data;
+                    if (pos >= (sizeof(buf) -2))
                     {
-                        write_flash_page();
                         ack = 0x00;
                     }
+
+                    if (pos >= (sizeof(buf) -1))
+                    {
+                        write_flash_page();
+                    }
                     break;
+                }
 
 #if (EEPROM_SUPPORT)
                 case CMD_ACCESS_EEPROM:
@@ -364,7 +372,6 @@ static void TWI_vect(void)
             if (TWI_data_write(bcnt++, TWDR) == 0x00)
             {
                 control &= ~(1<<TWEA);
-                bcnt = 0;
             }
             break;
 
@@ -380,8 +387,14 @@ static void TWI_vect(void)
 
         /* prev. SLA+W, data received, NACK returned -> IDLE */
         case 0x88:
+            TWI_data_write(bcnt++, TWDR);
+            /* fall through */
+
         /* STOP or repeated START -> IDLE */
         case 0xA0:
+            bcnt = 0;
+            /* fall through */
+
         /* prev. SLA+R, data sent, NACK returned -> IDLE */
         case 0xC0:
             LED_RT_OFF();