From 3430ee2c3a4f6ad700febc79d10ce7020c96b028 Mon Sep 17 00:00:00 2001
From: Cameron Gutman <aicommander@gmail.com>
Date: Sun, 14 Jan 2024 12:58:48 -0600
Subject: [PATCH] Add control stream encryption v2 support

---
 src/ControlStream.c | 28 +++++++++++++++++++++++-----
 src/SdpGenerator.c  |  9 ++++++++-
 2 files changed, 31 insertions(+), 6 deletions(-)

diff --git a/src/ControlStream.c b/src/ControlStream.c
index dc4ba4e..9736a6a 100644
--- a/src/ControlStream.c
+++ b/src/ControlStream.c
@@ -90,7 +90,7 @@ static int intervalTotalFrameCount;
 static uint64_t intervalStartTimeMs;
 static int lastIntervalLossPercentage;
 static int lastConnectionStatusUpdate;
-static int currentEnetSequenceNumber;
+static uint32_t currentEnetSequenceNumber;
 static uint64_t firstFrameTimeMs;
 
 static LINKED_BLOCKING_QUEUE invalidReferenceFrameTuples;
@@ -503,9 +503,18 @@ static bool encryptControlMessage(PNVCTL_ENCRYPTED_PACKET_HEADER encPacket, PNVC
     unsigned char iv[16] = { 0 };
     int encryptedSize = sizeof(*packet) + packet->payloadLength;
 
-    // This is a truncating cast, but it's what Nvidia does, so we have to mimic it.
     // NB: Setting the IV must happen while encPacket->seq is still in native byte-order!
-    iv[0] = (unsigned char)encPacket->seq;
+    if (EncryptionFeaturesEnabled & SS_ENC_CONTROL_V2) {
+        // Populate the IV in little endian byte order
+        iv[3] = (unsigned char)(encPacket->seq >> 24);
+        iv[2] = (unsigned char)(encPacket->seq >> 16);
+        iv[1] = (unsigned char)(encPacket->seq >> 8);
+        iv[0] = (unsigned char)(encPacket->seq >> 0);
+    }
+    else {
+        // This is a truncating cast, but it's what Nvidia does, so we have to mimic it.
+        iv[0] = (unsigned char)encPacket->seq;
+    }
 
     encPacket->encryptedHeaderType = LE16(encPacket->encryptedHeaderType);
     encPacket->length = LE16(encPacket->length);
@@ -545,8 +554,17 @@ static bool decryptControlMessageToV1(PNVCTL_ENCRYPTED_PACKET_HEADER encPacket,
         return false;
     }
 
-    // This is a truncating cast, but it's what Nvidia does, so we have to mimic it.
-    iv[0] = (unsigned char)encPacket->seq;
+    if (EncryptionFeaturesEnabled & SS_ENC_CONTROL_V2) {
+        // Populate the IV in little endian byte order
+        iv[3] = (unsigned char)(encPacket->seq >> 24);
+        iv[2] = (unsigned char)(encPacket->seq >> 16);
+        iv[1] = (unsigned char)(encPacket->seq >> 8);
+        iv[0] = (unsigned char)(encPacket->seq >> 0);
+    }
+    else {
+        // This is a truncating cast, but it's what Nvidia does, so we have to mimic it.
+        iv[0] = (unsigned char)encPacket->seq;
+    }
 
     int plaintextLength = encPacket->length - sizeof(encPacket->seq) - AES_GCM_TAG_LENGTH;
     *packet = malloc(plaintextLength);
diff --git a/src/SdpGenerator.c b/src/SdpGenerator.c
index c6d8d36..2a304da 100644
--- a/src/SdpGenerator.c
+++ b/src/SdpGenerator.c
@@ -266,11 +266,18 @@ static PSDP_OPTION getAttributesList(char*urlSafeAddr) {
     optionHead = NULL;
     err = 0;
 
-    // Send client feature flags to Sunshine hosts
     if (IS_SUNSHINE()) {
+        // Send client feature flags to Sunshine hosts
         uint32_t moonlightFeatureFlags = ML_FF_FEC_STATUS | ML_FF_SESSION_ID_V1;
         snprintf(payloadStr, sizeof(payloadStr), "%u", moonlightFeatureFlags);
         err |= addAttributeString(&optionHead, "x-ml-general.featureFlags", payloadStr);
+
+        // New-style control stream encryption is low overhead, so we enable it any time it is supported
+        if (EncryptionFeaturesSupported & SS_ENC_CONTROL_V2) {
+            EncryptionFeaturesEnabled |= SS_ENC_CONTROL_V2;
+        }
+        snprintf(payloadStr, sizeof(payloadStr), "%u", EncryptionFeaturesEnabled);
+        err |= addAttributeString(&optionHead, "x-ss-general.encryptionEnabled", payloadStr);
     }
 
     snprintf(payloadStr, sizeof(payloadStr), "%d", StreamConfig.width);