Add simple RLE compression for ArduinoOTA

earlephilhower · earlephilhower · commit 75dad16b3cdc · 2019-10-05T17:09:28.000-07:00
When uploading large, sparse filesystems, there are often many repeated, empty sectors. These take a long time to transmit, even over WiFi. Implement a basic RLE compression in espota.py (and only use it if it actually saves upload size), and a simple decompressor in ArduinoOTA. The actual flash update image is decompressed as received to the staging location, so the bootloader/etc. do not need to be aware of it (i.e. a 1MB filesystem still takes 1MB of flash in the update area, even if it only took 100KB to upload it). Fixes esp8266#4277
diff --git a/libraries/ArduinoOTA/ArduinoOTA.cpp b/libraries/ArduinoOTA/ArduinoOTA.cpp
@@ -30,10 +30,114 @@ extern "C" {
 #endif
 #endif
 
+
+/* This class is only used in OTA and implements a dumb, low-memory decompression engine */
+class RLEDecompressor : public Stream {
+public:
+  RLEDecompressor(WiFiClient client) {
+    _client = client;
+    _blockLen = 0;
+    _blockIdx = 0;
+    _block = new uint8_t[128];
+  }
+
+  virtual ~RLEDecompressor() {
+    delete[] _block;
+  }
+
+  virtual int read() {
+    int ret = -1;  // Default to EOF
+    if (_blockLen == _blockIdx) {
+        _refill();
+    }
+    if (_blockIdx < _blockLen) {
+        ret = _block[_blockIdx++];
+    }
+    return ret;
+  }
+
+  virtual int peek() {
+    return -1; // Not implemented, not needed for Updater
+  }
+
+  size_t read(uint8_t*a, size_t&b) { return readBytes((char*)a, b); }
+
+  virtual size_t readBytes(char *buffer, size_t length) {
+    if (_blockLen == _blockIdx) {
+      _refill();
+    }
+    int toRead = std::min((int)(_blockLen - _blockIdx), (int)length);
+    memcpy(buffer, _block + _blockIdx, toRead);
+    _blockIdx += toRead;
+    return toRead;
+  }
+    
+  virtual size_t write(uint8_t b) { return _client.write(b); }
+
+  virtual int available() {
+    if (_blockLen == _blockIdx) {
+      _refill();
+    }
+    return _blockLen - _blockIdx;
+  }
+
+
+private:
+  bool _refill() {
+    int c = -1;
+    while (_client.connected() && (c < 0)) {
+      c = _client.read();
+      yield();
+    }
+    if (c < 0) {
+       return false;
+    }
+    if (c < 128) {
+      int l = -1;
+      while (_client.connected() && (l < 0)) {
+        l = _client.read();
+        yield();
+      }
+      if (l < 0) {
+         return false;
+      }
+      memset(_block, l, c);
+      _blockLen = c;
+      _blockIdx = 0;
+      return true;
+    } else {
+      c = c - 128;
+      _blockLen = c;
+      _blockIdx = 0;
+      while (_client.connected() && c) {
+        int ret = _client.readBytes(_block + _blockIdx, c);
+        if (ret > 0) {
+          c -= ret;
+          _blockIdx += ret;
+        } else {
+          _blockIdx = _blockLen = 0;
+          return false;
+        }
+      }
+      _blockIdx = 0;
+      return true;
+    }
+  }
+
+private:
+  WiFiClient _client;
+  int        _blockLeft;
+  int        _blockIdx;
+  int        _blockLen;
+  uint8_t    *_block;
+};
+
+
 ArduinoOTAClass::ArduinoOTAClass()
 : _port(0)
 , _udp_ota(0)
 , _initialized(false)
+, _useCompression(false)
 , _rebootOnSuccess(true)
 , _useMDNS(true)
 , _state(OTA_IDLE)
@@ -199,6 +303,15 @@ void ArduinoOTAClass::_onRx(){
     if(_md5.length() != 32)
       return;
 
+    String compress = readStringUntil('\n');
+    compress.trim();
+    if (compress == "COMPRESSRLE") {
+        _useCompression = true;
+#ifdef OTA_DEBUG
+        OTA_DEBUG.println("Compressed upload requested");
+#endif
+    }
+
     ota_ip = _ota_ip;
 
     if (_password.length()){
@@ -273,7 +386,11 @@ void ArduinoOTAClass::_runUpdate() {
     _state = OTA_IDLE;
     return;
   }
-  _udp_ota->append("OK", 2);
+  if (_useCompression) {
+    _udp_ota->append("COMPOK", 6);
+  } else {
+    _udp_ota->append("OK", 2);
+  }
   _udp_ota->send(ota_ip, _ota_udp_port);
   delay(100);
 
@@ -317,7 +434,12 @@ void ArduinoOTAClass::_runUpdate() {
       }
       _state = OTA_IDLE;
     }
-    written = Update.write(client);
+    if (_useCompression) {
+      RLEDecompressor decomp(client);
+      written = Update.write(decomp);
+    } else {
+      written = Update.write(client);
+    }
     if (written > 0) {
       client.print(written, DEC);
       total += written;
diff --git a/libraries/ArduinoOTA/ArduinoOTA.h b/libraries/ArduinoOTA/ArduinoOTA.h
@@ -75,6 +75,7 @@ class ArduinoOTAClass
     String _nonce;
     UdpContext *_udp_ota;
     bool _initialized;
+    bool _useCompression;
     bool _rebootOnSuccess;
     bool _useMDNS;
     ota_state_t _state;
diff --git a/libraries/SoftwareSerial b/libraries/SoftwareSerial
@@ -1 +1 @@
-Subproject commit 67836ae6eecabafd760d32168144f82d46f3b201
+Subproject commit 6d9f1159555f3025078b3510605b5a2721708535
diff --git a/tools/espota.py b/tools/espota.py
@@ -69,6 +69,43 @@ def update_progress(progress):
     sys.stderr.write('.')
     sys.stderr.flush()
 
+def rle(data):
+    "RLE compress the input bitstream and return it as a list"
+    buf = [0] * 128
+    ret = []
+    src = 0
+    runlen = 0
+    repeat = False
+    while src < len(data):
+        buf[runlen] = data[src]
+        runlen = runlen + 1
+        src = src + 1
+        if runlen < 2:
+            continue
+        if repeat:
+            if buf[runlen - 1] != buf[runlen - 2]:
+                repeat = False
+            if (not repeat) or (runlen == 128):
+                ret += [runlen - 1] + [buf[0]]
+                buf[0] = buf[runlen - 1]
+                runlen = 1
+        else:
+            if buf[runlen - 1] == buf[runlen - 2]:
+                repeat = True
+                if runlen > 2:
+                    ret += [128 + runlen - 2] + buf[0:runlen - 2]
+                    buf[0] = buf[runlen - 1]
+                    buf[1] = buf[runlen - 1]
+                    runlen = 2
+                continue
+            if runlen == 128:
+                ret += [128 + runlen - 1] + buf[0:runlen]
+                runlen = 0
+    if runlen:
+        ret += [128 + runlen] + buf[0:runlen]
+    return ret
+
+
 def serve(remoteAddr, localAddr, remotePort, localPort, password, filename, command = FLASH):
   # Create a TCP/IP socket
   sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
@@ -89,12 +126,18 @@ def serve(remoteAddr, localAddr, remotePort, localPort, password, filename, comm
     sys.stderr.flush()
     logging.info(file_check_msg)
   
-  content_size = os.path.getsize(filename)
-  f = open(filename,'rb')
-  file_md5 = hashlib.md5(f.read()).hexdigest()
-  f.close()
+  with open(filename, "rb") as f:
+      content = f.read()
+  content_size = len(content)
+  content_rle = rle(content)
+  request_rle = len(content_rle) < content_size
+#  request_rle = True
+  file_md5 = hashlib.md5(content).hexdigest()
   logging.info('Upload size: %d', content_size)
   message = '%d %d %d %s\n' % (command, localPort, content_size, file_md5)
+  if request_rle:
+    # Add a request for compression, ignored on earlier ArduinoOTAs
+    message += 'COMPRESSRLE\n'
 
   # Wait for a connection
   logging.info('Sending invitation to: %s', remoteAddr)
@@ -103,42 +146,50 @@ def serve(remoteAddr, localAddr, remotePort, localPort, password, filename, comm
   sent = sock2.sendto(message.encode(), remote_address)
   sock2.settimeout(10)
   try:
-    data = sock2.recv(128).decode()
+    data = sock2.recv(256).decode()
   except:
     logging.error('No Answer')
     sock2.close()
     return 1
-  if (data != "OK"):
-    if(data.startswith('AUTH')):
-      nonce = data.split()[1]
-      cnonce_text = '%s%u%s%s' % (filename, content_size, file_md5, remoteAddr)
-      cnonce = hashlib.md5(cnonce_text.encode()).hexdigest()
-      passmd5 = hashlib.md5(password.encode()).hexdigest()
-      result_text = '%s:%s:%s' % (passmd5 ,nonce, cnonce)
-      result = hashlib.md5(result_text.encode()).hexdigest()
-      sys.stderr.write('Authenticating...')
-      sys.stderr.flush()
-      message = '%d %s %s\n' % (AUTH, cnonce, result)
-      sock2.sendto(message.encode(), remote_address)
-      sock2.settimeout(10)
-      try:
-        data = sock2.recv(32).decode()
-      except:
-        sys.stderr.write('FAIL\n')
-        logging.error('No Answer to our Authentication')
-        sock2.close()
-        return 1
-      if (data != "OK"):
-        sys.stderr.write('FAIL\n')
-        logging.error('%s', data)
-        sock2.close()
-        sys.exit(1);
-        return 1
-      sys.stderr.write('OK\n')
+
+  if data == "COMPOK":
+    compress = True
+  elif data == "OK":
+    compress = False
+  elif data.startswith('AUTH'):
+    nonce = data.split()[1]
+    cnonce_text = '%s%u%s%s' % (filename, content_size, file_md5, remoteAddr)
+    cnonce = hashlib.md5(cnonce_text.encode()).hexdigest()
+    passmd5 = hashlib.md5(password.encode()).hexdigest()
+    result_text = '%s:%s:%s' % (passmd5 ,nonce, cnonce)
+    result = hashlib.md5(result_text.encode()).hexdigest()
+    sys.stderr.write('Authenticating...')
+    sys.stderr.flush()
+    message = '%d %s %s\n' % (AUTH, cnonce, result)
+    sock2.sendto(message.encode(), remote_address)
+    sock2.settimeout(10)
+    try:
+      data = sock2.recv(32).decode()
+    except:
+      sys.stderr.write('FAIL\n')
+      logging.error('No Answer to our Authentication')
+      sock2.close()
+      return 1
+    if data == "OK":
+      compress = False
+    elif data == "COMPOK":
+      compress = True
     else:
-      logging.error('Bad Answer: %s', data)
+      sys.stderr.write('FAIL\n')
+      logging.error('%s', data)
       sock2.close()
+      sys.exit(1);
       return 1
+    sys.stderr.write('OK\n')
+  else:
+    logging.error('Bad Answer: %s', data)
+    sock2.close()
+    return 1
   sock2.close()
 
   logging.info('Waiting for device...')
@@ -155,16 +206,17 @@ def serve(remoteAddr, localAddr, remotePort, localPort, password, filename, comm
   received_ok = False
 
   try:
-    f = open(filename, "rb")
     if (PROGRESS):
       update_progress(0)
     else:
       sys.stderr.write('Uploading')
       sys.stderr.flush()
     offset = 0
+    chunk_size = 1460 # MTU-safe
     while True:
-      chunk = f.read(1460)
-      if not chunk: break
+      if offset >= content_size:
+        break
+      chunk = content[offset:min(content_size, offset + chunk_size)]
       offset += len(chunk)
       update_progress(offset/float(content_size))
       connection.settimeout(10)
