From ab309b39d2e30481ae6c5056a412bad3b880e810 Mon Sep 17 00:00:00 2001
From: Frank <91616163+softhack007@users.noreply.github.com>
Date: Mon, 20 Nov 2023 21:13:35 +0100
Subject: [PATCH] Merge pull request #3536 from Aircoookie/ntp_errorchecking

NTP validation, and rejecting malformed responses (related to #3515)
---
 wled00/const.h  |  3 ++-
 wled00/ntp.cpp  | 31 ++++++++++++++++++++++++++-----
 wled00/set.cpp  |  2 +-
 wled00/wled.cpp |  2 +-
 wled00/wled.h   |  5 +++--
 5 files changed, 33 insertions(+), 10 deletions(-)

diff --git a/wled00/const.h b/wled00/const.h
index 4db4d898..2f210508 100644
--- a/wled00/const.h
+++ b/wled00/const.h
@@ -350,7 +350,8 @@
 #define NL_MODE_SUN               3            //Sunrise/sunset. Target brightness is set immediately, then Sunrise effect is started. Max 60 min.
 
 
-#define NTP_PACKET_SIZE 48
+#define NTP_PACKET_SIZE 48       // size of NTP recive buffer
+#define NTP_MIN_PACKET_SIZE 48   // min expected size - NTP v4 allows for "extended information" appended to the standard fields
 
 //maximum number of rendered LEDs - this does not have to match max. physical LEDs, e.g. if there are virtual busses
 #ifndef MAX_LEDS
diff --git a/wled00/ntp.cpp b/wled00/ntp.cpp
index 7d056d3c..ecc11ce0 100644
--- a/wled00/ntp.cpp
+++ b/wled00/ntp.cpp
@@ -199,6 +199,9 @@ void handleNetworkTime()
   {
     if (millis() - ntpPacketSentTime > 10000)
     {
+      #ifdef ARDUINO_ARCH_ESP32   // I had problems using udp.flush() on 8266
+      while (ntpUdp.parsePacket() > 0) ntpUdp.flush(); // flush any existing packets
+      #endif
       sendNTPPacket();
       ntpPacketSentTime = millis();
     }
@@ -239,23 +242,41 @@ void sendNTPPacket()
   ntpUdp.endPacket();
 }
 
+static bool isValidNtpResponse(byte * ntpPacket) {
+  // Perform a few validity checks on the packet
+  //   based on https://github.com/taranais/NTPClient/blob/master/NTPClient.cpp
+  if((ntpPacket[0] & 0b11000000) == 0b11000000) return false; //reject LI=UNSYNC
+  // if((ntpPacket[0] & 0b00111000) >> 3 < 0b100) return false; //reject Version < 4
+  if((ntpPacket[0] & 0b00000111) != 0b100)      return false; //reject Mode != Server
+  if((ntpPacket[1] < 1) || (ntpPacket[1] > 15)) return false; //reject invalid Stratum
+  if( ntpPacket[16] == 0 && ntpPacket[17] == 0 && 
+      ntpPacket[18] == 0 && ntpPacket[19] == 0 &&
+      ntpPacket[20] == 0 && ntpPacket[21] == 0 &&
+      ntpPacket[22] == 0 && ntpPacket[23] == 0)               //reject ReferenceTimestamp == 0
+    return false;
+
+  return true;
+}
+
 bool checkNTPResponse()
 {
 #ifdef ARDUINO_ARCH_ESP32
   ntpUdp.flush();
 #endif
   int cb = ntpUdp.parsePacket();
-#ifdef ARDUINO_ARCH_ESP32
-  if (!cb) {ntpUdp.flush(); return false;}    // WLEDMM flush buffer
-#else
-  if (!cb) {return false;}    // WLEDMM do not flush buffer
-#endif
+  if (cb < NTP_MIN_PACKET_SIZE) {
+    #ifdef ARDUINO_ARCH_ESP32   // I had problems using udp.flush() on 8266
+    if (cb > 0) ntpUdp.flush();  // this avoids memory leaks on esp32
+    #endif
+    return false;
+  }
 
   uint32_t ntpPacketReceivedTime = millis();
   DEBUG_PRINT(F("NTP recv, l="));
   DEBUG_PRINTLN(cb);
   byte pbuf[NTP_PACKET_SIZE];
   ntpUdp.read(pbuf, NTP_PACKET_SIZE); // read the packet into the buffer
+  if (!isValidNtpResponse(pbuf)) return false;  // verify we have a valid response to client
 
   Toki::Time arrived  = toki.fromNTP(pbuf + 32);
   Toki::Time departed = toki.fromNTP(pbuf + 40);
diff --git a/wled00/set.cpp b/wled00/set.cpp
index 580857ef..cc68bf45 100644
--- a/wled00/set.cpp
+++ b/wled00/set.cpp
@@ -426,7 +426,7 @@ void handleSettingsSet(AsyncWebServerRequest *request, byte subPage)
 
     //start ntp if not already connected
     if (ntpEnabled && WLED_CONNECTED && !ntpConnected) ntpConnected = ntpUdp.begin(ntpLocalPort);
-    ntpLastSyncTime = 0; // force new NTP query
+    ntpLastSyncTime = NTP_NEVER; // force new NTP query
 
     longitude = request->arg(F("LN")).toFloat();
     latitude = request->arg(F("LT")).toFloat();
diff --git a/wled00/wled.cpp b/wled00/wled.cpp
index dba77245..2fd88117 100644
--- a/wled00/wled.cpp
+++ b/wled00/wled.cpp
@@ -232,7 +232,7 @@ void WLED::loop()
   if (lastMqttReconnectAttempt > millis()) {
     rolloverMillis++;
     lastMqttReconnectAttempt = 0;
-    ntpLastSyncTime = 0;
+    ntpLastSyncTime = NTP_NEVER;  // force new NTP query
     strip.restartRuntime();
   }
   if (millis() - lastMqttReconnectAttempt > 30000 || lastMqttReconnectAttempt == 0) { // lastMqttReconnectAttempt==0 forces immediate broadcast
diff --git a/wled00/wled.h b/wled00/wled.h
index e8890695..6f21e453 100644
--- a/wled00/wled.h
+++ b/wled00/wled.h
@@ -670,10 +670,11 @@ WLED_GLOBAL String escapedMac;
 WLED_GLOBAL DNSServer dnsServer;
 
 // network time
+#define NTP_NEVER 999000000L
 WLED_GLOBAL bool ntpConnected _INIT(false);
 WLED_GLOBAL time_t localTime _INIT(0);
-WLED_GLOBAL unsigned long ntpLastSyncTime _INIT(999000000L);
-WLED_GLOBAL unsigned long ntpPacketSentTime _INIT(999000000L);
+WLED_GLOBAL unsigned long ntpLastSyncTime _INIT(NTP_NEVER);
+WLED_GLOBAL unsigned long ntpPacketSentTime _INIT(NTP_NEVER);
 WLED_GLOBAL IPAddress ntpServerIP;
 WLED_GLOBAL uint16_t ntpLocalPort _INIT(2390);
 WLED_GLOBAL uint16_t rolloverMillis _INIT(0);