From ce446dac929282790dc03923b94db2bab2b99f39 Mon Sep 17 00:00:00 2001
From: mickabrig7 <mickaelgarabello7@gmail.com>
Date: Wed, 19 Aug 2020 10:40:16 +0200
Subject: [PATCH] Improve TPM2.NET code

Add support for receiving TPM2.NET packets out of order, and hopefully prevent crashes if faced with protocol errors
---
 wled00/udp.cpp | 24 +++++++++++++++++-------
 wled00/wled.h  |  3 ++-
 2 files changed, 19 insertions(+), 8 deletions(-)

diff --git a/wled00/udp.cpp b/wled00/udp.cpp
index 0a3f2d78..886b6fe3 100644
--- a/wled00/udp.cpp
+++ b/wled00/udp.cpp
@@ -219,6 +219,8 @@ void handleNotifications()
   //TPM2.NET
   if (udpIn[0] == 0x9c)
   {
+    //WARNING: this code assumes that the final TMP2.NET payload is evenly distributed if using multiple packets (ie. frame size is constant)
+    //if the number of LEDs in your installation doesn't allow that, please include padding bytes at the end of the last packet
     byte tpmType = udpIn[1];
     if (tpmType == 0xaa) { //TPM2.NET polling, expect answer
       sendTPM2Ack(); return;
@@ -229,18 +231,26 @@ void handleNotifications()
     realtimeLock(realtimeTimeoutMs, REALTIME_MODE_TPM2NET);
     if (realtimeOverride) return;
 
-    uint16_t frameSize = (udpIn[2] << 8) + udpIn[3];
+    tpmPacketCount++; //increment the packet count
+    if (tpmPacketCount == 1) tpmPayloadFrameSize = (udpIn[2] << 8) + udpIn[3]; //save frame size for the whole payload if this is the first packet
     byte packetNum = udpIn[4]; //starts with 1!
     byte numPackets = udpIn[5];
 
-    uint16_t id = ((tpmFirstFrameSize/3)*(packetNum-1)) / 3; //start LED
-    for (uint16_t i = 6; i < frameSize + 4; i += 3)
+    uint16_t id = (tpmPayloadFrameSize/3)*(packetNum-1); //start LED
+    for (uint16_t i = 6; i < tpmPayloadFrameSize + 4; i += 3)
     {
-      setRealtimePixel(id, udpIn[i], udpIn[i+1], udpIn[i+2], 0);
-      id++; if (id >= ledCount) break;
+      if (id < ledCount)
+      {
+        setRealtimePixel(id, udpIn[i], udpIn[i+1], udpIn[i+2], 0);
+        id++;
+      }
+      else break;
+    }
+    if (tpmPacketCount == numPackets) //reset packet count and show if all packets were received
+    {
+      tpmPacketCount = 0;
+      strip.show();
     }
-    if (packetNum == 1) tpmFirstFrameSize = frameSize;
-    if (packetNum == numPackets) {strip.show(); } //show if last packet 
   }
 
   //UDP realtime: 1 warls 2 drgb 3 drgbw
diff --git a/wled00/wled.h b/wled00/wled.h
index a7a1c92b..a0ead883 100644
--- a/wled00/wled.h
+++ b/wled00/wled.h
@@ -428,7 +428,8 @@ WLED_GLOBAL byte realtimeMode _INIT(REALTIME_MODE_INACTIVE);
 WLED_GLOBAL byte realtimeOverride _INIT(REALTIME_OVERRIDE_NONE);
 WLED_GLOBAL IPAddress realtimeIP _INIT((0, 0, 0, 0));
 WLED_GLOBAL unsigned long realtimeTimeout _INIT(0);
-WLED_GLOBAL uint16_t tpmFirstFrameSize _INIT(0);
+WLED_GLOBAL uint8_t tpmPacketCount _INIT(0);
+WLED_GLOBAL uint16_t tpmPayloadFrameSize _INIT(0);
 
 // mqtt
 WLED_GLOBAL long lastMqttReconnectAttempt _INIT(0);