From 1e679cc94d48a99c4145895491bd795003641754 Mon Sep 17 00:00:00 2001
From: Adam Ierymenko <adam.ierymenko@zerotier.com>
Date: Thu, 26 Sep 2024 20:13:10 -0400
Subject: [PATCH] Bug fix in extended armor.

---
 node/Metrics.cpp | 3 +--
 node/Metrics.hpp | 1 +
 node/Packet.cpp  | 4 +++-
 node/Peer.cpp    | 1 +
 4 files changed, 6 insertions(+), 3 deletions(-)

diff --git a/node/Metrics.cpp b/node/Metrics.cpp
index 1b0b55666..45353fde2 100644
--- a/node/Metrics.cpp
+++ b/node/Metrics.cpp
@@ -10,8 +10,7 @@
  * of this software will be governed by version 2.0 of the Apache License.
  */
 
-#include <prometheus/histogram.h>
-#include <prometheus/simpleapi.h>
+#include "Metrics.hpp"
 
 namespace prometheus {
 namespace simpleapi {
diff --git a/node/Metrics.hpp b/node/Metrics.hpp
index 3a70ce752..3f33a33e0 100644
--- a/node/Metrics.hpp
+++ b/node/Metrics.hpp
@@ -12,6 +12,7 @@
 #ifndef METRICS_H_
 #define METRICS_H_
 
+#include <prometheus/gauge.h>
 #include <prometheus/histogram.h>
 #include <prometheus/simpleapi.h>
 
diff --git a/node/Packet.cpp b/node/Packet.cpp
index f606eef7e..35b29ea19 100644
--- a/node/Packet.cpp
+++ b/node/Packet.cpp
@@ -1162,6 +1162,7 @@ void Packet::armor(const void* key, bool encryptPayload, bool extendedArmor, con
         AES::CTR aesCtr(cipher);
         aesCtr.init(data, 0, data + ZT_PACKET_IDX_EXTENDED_ARMOR_START);
         aesCtr.crypt(data + ZT_PACKET_IDX_EXTENDED_ARMOR_START, size() - ZT_PACKET_IDX_EXTENDED_ARMOR_START);
+        aesCtr.finish();
 
         this->append(ephemeralKeyPair.pub.data, ZT_ECC_EPHEMERAL_PUBLIC_KEY_LEN);
     }
@@ -1184,7 +1185,8 @@ bool Packet::dearmor(const void* key, const AES aesKeys[2], const Identity& iden
         AES cipher(ephemeralSymmetric);
         AES::CTR aesCtr(cipher);
         aesCtr.init(data, 0, data + ZT_PACKET_IDX_EXTENDED_ARMOR_START);
-        aesCtr.crypt(data + ZT_PACKET_IDX_EXTENDED_ARMOR_START, size() - ZT_PACKET_IDX_EXTENDED_ARMOR_START);
+        aesCtr.crypt(data + ZT_PACKET_IDX_EXTENDED_ARMOR_START, (size() - ZT_PACKET_IDX_EXTENDED_ARMOR_START) - ZT_ECC_EPHEMERAL_PUBLIC_KEY_LEN);
+        aesCtr.finish();
 
         this->setSize(size() - ZT_ECC_EPHEMERAL_PUBLIC_KEY_LEN);
 
diff --git a/node/Peer.cpp b/node/Peer.cpp
index b3d5de7da..08448f707 100644
--- a/node/Peer.cpp
+++ b/node/Peer.cpp
@@ -25,6 +25,7 @@
 #include "Switch.hpp"
 #include "Trace.hpp"
 #include "Utils.hpp"
+#include "Switch.hpp"
 
 namespace ZeroTier {