🔨 Abort firmware update on transfer error (#24472)

buildroot/share/scripts/MarlinBinaryProtocol.py

         token, data = self.await_response(1000)
         if token == 'PFT:success':
             print("File closed")
-            return
+            return True
         elif token == 'PFT:ioerror':
             print("Client storage device IO error")
+            return False
         elif token == 'PFT:invalid':
             print("No open file")
+            return False
 
     def abort(self):
         self.protocol.send(FileTransferProtocol.protocol_id, FileTransferProtocol.Packet.ABORT);
             self.write(data[start:end])
             kibs = (( (i+1) * block_size) / 1024) / (millis() + 1 - start_time) * 1000
             if (i / blocks) >= dump_pctg:
-                print("\r{0:2.2f}% {1:4.2f}KiB/s {2} Errors: {3}".format((i / blocks) * 100, kibs, "[{0:4.2f}KiB/s]".format(kibs * cratio) if compression_support else "", self.protocol.errors), end='')
+                print("\r{0:2.0f}% {1:4.2f}KiB/s {2} Errors: {3}".format((i / blocks) * 100, kibs, "[{0:4.2f}KiB/s]".format(kibs * cratio) if compression_support else "", self.protocol.errors), end='')
                 dump_pctg += 0.1
-        print("\r{0:2.2f}% {1:4.2f}KiB/s {2} Errors: {3}".format(100, kibs, "[{0:4.2f}KiB/s]".format(kibs * cratio) if compression_support else "", self.protocol.errors)) # no one likes transfers finishing at 99.8%
-
-        self.close()
+            if self.protocol.errors > 0:
+                # Dump last status (errors may not be visible)
+                print("\r{0:2.0f}% {1:4.2f}KiB/s {2} Errors: {3} - Aborting...".format((i / blocks) * 100, kibs, "[{0:4.2f}KiB/s]".format(kibs * cratio) if compression_support else "", self.protocol.errors), end='')
+                print("")   # New line to break the transfer speed line
+                self.close()
+                print("Transfer aborted due to protocol errors")
+                #raise Exception("Transfer aborted due to protocol errors")
+                return False;
+        print("\r{0:2.0f}% {1:4.2f}KiB/s {2} Errors: {3}".format(100, kibs, "[{0:4.2f}KiB/s]".format(kibs * cratio) if compression_support else "", self.protocol.errors)) # no one likes transfers finishing at 99.8%
+
+        if not self.close():
+            print("Transfer failed")
+            return False
         print("Transfer complete")
+        return True
 
 
 class EchoProtocol(object):

buildroot/share/scripts/upload.py

 
 import MarlinBinaryProtocol
 
-# Internal debug flag
-Debug = False
-
 #-----------------#
 # Upload Callback #
 #-----------------#
 def Upload(source, target, env):
 
+    #-------#
+    # Debug #
+    #-------#
+    Debug = False                # Set to True to enable script debug
+    def debugPrint(data):
+        if Debug: print(f"[Debug]: {data}")
+
     #------------------#
     # Marlin functions #
     #------------------#
     # Port functions #
     #----------------#
     def _GetUploadPort(env):
-        if Debug: print('Autodetecting upload port...')
+        debugPrint('Autodetecting upload port...')
         env.AutodetectUploadPort(env)
-        port = env.subst('$UPLOAD_PORT')
-        if not port:
+        portName = env.subst('$UPLOAD_PORT')
+        if not portName:
             raise Exception('Error detecting the upload port.')
-        if Debug: print('OK')
-        return port
+        debugPrint('OK')
+        return portName
 
     #-------------------------#
     # Simple serial functions #
     #-------------------------#
+    def _OpenPort():
+        # Open serial port
+        if port.is_open: return
+        debugPrint('Opening upload port...')
+        port.open()
+        port.reset_input_buffer()
+        debugPrint('OK')
+
+    def _ClosePort():
+        # Open serial port
+        if port is None: return
+        if not port.is_open: return
+        debugPrint('Closing upload port...')
+        port.close()
+        debugPrint('OK')
+
     def _Send(data):
-        if Debug: print(f'>> {data}')
+        debugPrint(f'>> {data}')
         strdata = bytearray(data, 'utf8') + b'\n'
         port.write(strdata)
         time.sleep(0.010)
         clean_responses = []
         responses = port.readlines()
         for Resp in responses:
-            # Test: suppress invaid chars (coming from debug info)
+            # Suppress invalid chars (coming from debug info)
             try:
                 clean_response = Resp.decode('utf8').rstrip().lstrip()
                 clean_responses.append(clean_response)
             except:
                 pass
-            if Debug: print(f'<< {clean_response}')
+            debugPrint(f'<< {clean_response}')
         return clean_responses
 
     #------------------#
     # SDCard functions #
     #------------------#
     def _CheckSDCard():
-        if Debug: print('Checking SD card...')
+        debugPrint('Checking SD card...')
         _Send('M21')
         Responses = _Recv()
         if len(Responses) < 1 or not any('SD card ok' in r for r in Responses):
             raise Exception('Error accessing SD card')
-        if Debug: print('SD Card OK')
+        debugPrint('SD Card OK')
         return True
 
     #----------------#
     # File functions #
     #----------------#
     def _GetFirmwareFiles(UseLongFilenames):
-        if Debug: print('Get firmware files...')
+        debugPrint('Get firmware files...')
         _Send(f"M20 F{'L' if UseLongFilenames else ''}")
         Responses = _Recv()
         if len(Responses) < 3 or not any('file list' in r for r in Responses):
             raise Exception('Error getting firmware files')
-        if Debug: print('OK')
+        debugPrint('OK')
         return Responses
 
     def _FilterFirmwareFiles(FirmwareList, UseLongFilenames):
             raise Exception(f"Firmware file '{FirmwareFile}' not removed")
         return Removed
 
+    def _RollbackUpload(FirmwareFile):
+        if not rollback: return
+        print(f"Rollback: trying to delete firmware '{FirmwareFile}'...")
+        _OpenPort()
+        # Wait for SD card release
+        time.sleep(1)
+        # Remount SD card
+        _CheckSDCard()
+        print(' OK' if _RemoveFirmwareFile(FirmwareFile) else ' Error!')
+        _ClosePort()
+
 
     #---------------------#
     # Callback Entrypoint #
     port = None
     protocol = None
     filetransfer = None
+    rollback = False
 
     # Get Marlin evironment vars
     MarlinEnv = env['MARLIN_FEATURES']
             if not marlin_custom_firmware_upload:
                 raise Exception(f"CUSTOM_FIRMWARE_UPLOAD must be enabled in 'Configuration_adv.h' for '{marlin_motherboard}'")
 
-            # Init serial port
+            # Init & Open serial port
             port = serial.Serial(upload_port, baudrate = upload_speed, write_timeout = 0, timeout = 0.1)
-            port.reset_input_buffer()
+            _OpenPort()
 
             # Check SD card status
             _CheckSDCard()
                     print(' OK' if _RemoveFirmwareFile(OldFirmwareFile) else ' Error!')
 
             # Close serial
-            port.close()
+            _ClosePort()
 
             # Cleanup completed
-            if Debug: print('Cleanup completed')
+            debugPrint('Cleanup completed')
 
         # WARNING! The serial port must be closed here because the serial transfer that follow needs it!
 
         # Upload firmware file
-        if Debug: print(f"Copy '{upload_firmware_source_name}' --> '{upload_firmware_target_name}'")
+        debugPrint(f"Copy '{upload_firmware_source_name}' --> '{upload_firmware_target_name}'")
         protocol = MarlinBinaryProtocol.Protocol(upload_port, upload_speed, upload_blocksize, float(upload_error_ratio), int(upload_timeout))
         #echologger = MarlinBinaryProtocol.EchoProtocol(protocol)
         protocol.connect()
+        # Mark the rollback (delete broken transfer) from this point on
+        rollback = True
         filetransfer = MarlinBinaryProtocol.FileTransferProtocol(protocol)
-        filetransfer.copy(upload_firmware_source_name, upload_firmware_target_name, upload_compression, upload_test)
+        transferOK = filetransfer.copy(upload_firmware_source_name, upload_firmware_target_name, upload_compression, upload_test)
         protocol.disconnect()
 
         # Notify upload completed
-        protocol.send_ascii('M117 Firmware uploaded')
+        protocol.send_ascii('M117 Firmware uploaded' if transferOK else 'M117 Firmware upload failed')
 
         # Remount SD card
         print('Wait for SD card release...')
         print('Remount SD card')
         protocol.send_ascii('M21')
 
-        # Trigger firmware update
-        if upload_reset:
-            print('Trigger firmware update...')
-            protocol.send_ascii('M997', True)
+        # Transfer failed?
+        if not transferOK:
+            protocol.shutdown()
+            _RollbackUpload(upload_firmware_target_name)
+        else:
+            # Trigger firmware update
+            if upload_reset:
+                print('Trigger firmware update...')
+                protocol.send_ascii('M997', True)
+            protocol.shutdown()
 
-        protocol.shutdown()
-        print('Firmware update completed')
+        print('Firmware update completed' if transferOK else 'Firmware update failed')
+        return 0 if transferOK else -1
 
     except KeyboardInterrupt:
-        if port: port.close()
+        print('Aborted by user')
         if filetransfer: filetransfer.abort()
-        if protocol: protocol.shutdown()
+        if protocol: 
+            protocol.disconnect()
+            protocol.shutdown()
+        _RollbackUpload(upload_firmware_target_name)
+        _ClosePort()
         raise
 
     except serial.SerialException as se:
-        if port: port.close()
-        print(f'Serial excepion: {se}')
+        # This exception is raised only for send_ascii data (not for binary transfer)
+        print(f'Serial excepion: {se}, transfer aborted')
+        if protocol: 
+            protocol.disconnect()
+            protocol.shutdown()
+        _RollbackUpload(upload_firmware_target_name)
+        _ClosePort()
         raise Exception(se)
 
     except MarlinBinaryProtocol.FatalError:
-        if port: port.close()
-        if protocol: protocol.shutdown()
-        print('Too many retries, Abort')
+        print('Too many retries, transfer aborted')
+        if protocol: 
+            protocol.disconnect()
+            protocol.shutdown()
+        _RollbackUpload(upload_firmware_target_name)
+        _ClosePort()
         raise
 
-    except:
-        if port: port.close()
-        if protocol: protocol.shutdown()
+    except Exception as ex:
+        print(f"\nException: {ex}, transfer aborted")
+        if protocol: 
+            protocol.disconnect()
+            protocol.shutdown()
+        _RollbackUpload(upload_firmware_target_name)
+        _ClosePort()
         print('Firmware not updated')
         raise