5 files changed, 149 insertions, 6 deletions

diff --git a/src/core/hle/service/apt/apt.cpp b/src/core/hle/service/apt/apt.cpp
index 615fe31ea..e57b19c2d 100644
--- a/src/core/hle/service/apt/apt.cpp
+++ b/src/core/hle/service/apt/apt.cpp
@@ -18,6 +18,8 @@
 #include "core/hle/service/fs/archive.h"
 #include "core/hle/service/ptm/ptm.h"
 #include "core/hle/service/service.h"
+#include "core/hw/aes/ccm.h"
+#include "core/hw/aes/key.h"
 
 namespace Service {
 namespace APT {
@@ -470,6 +472,107 @@ void GetStartupArgument(Service::Interface* self) {
     cmd_buff[2] = 0;
 }
 
+void Wrap(Service::Interface* self) {
+    IPC::RequestParser rp(Kernel::GetCommandBuffer(), 0x46, 4, 4);
+    const u32 output_size = rp.Pop<u32>();
+    const u32 input_size = rp.Pop<u32>();
+    const u32 nonce_offset = rp.Pop<u32>();
+    u32 nonce_size = rp.Pop<u32>();
+    size_t desc_size;
+    IPC::MappedBufferPermissions desc_permission;
+    const VAddr input = rp.PopMappedBuffer(&desc_size, &desc_permission);
+    ASSERT(desc_size == input_size && desc_permission == IPC::MappedBufferPermissions::R);
+    const VAddr output = rp.PopMappedBuffer(&desc_size, &desc_permission);
+    ASSERT(desc_size == output_size && desc_permission == IPC::MappedBufferPermissions::W);
+
+    // Note: real 3DS still returns SUCCESS when the sizes don't match. It seems that it doesn't
+    // check the buffer size and writes data with potential overflow.
+    ASSERT_MSG(output_size == input_size + HW::AES::CCM_MAC_SIZE,
+               "input_size (%d) doesn't match to output_size (%d)", input_size, output_size);
+
+    LOG_DEBUG(Service_APT, "called, output_size=%u, input_size=%u, nonce_offset=%u, nonce_size=%u",
+              output_size, input_size, nonce_offset, nonce_size);
+
+    // Note: This weird nonce size modification is verified against real 3DS
+    nonce_size = std::min<u32>(nonce_size & ~3, HW::AES::CCM_NONCE_SIZE);
+
+    // Reads nonce and concatenates the rest of the input as plaintext
+    HW::AES::CCMNonce nonce{};
+    Memory::ReadBlock(input + nonce_offset, nonce.data(), nonce_size);
+    u32 pdata_size = input_size - nonce_size;
+    std::vector<u8> pdata(pdata_size);
+    Memory::ReadBlock(input, pdata.data(), nonce_offset);
+    Memory::ReadBlock(input + nonce_offset + nonce_size, pdata.data() + nonce_offset,
+                      pdata_size - nonce_offset);
+
+    // Encrypts the plaintext using AES-CCM
+    auto cipher = HW::AES::EncryptSignCCM(pdata, nonce, HW::AES::KeySlotID::APTWrap);
+
+    // Puts the nonce to the beginning of the output, with ciphertext followed
+    Memory::WriteBlock(output, nonce.data(), nonce_size);
+    Memory::WriteBlock(output + nonce_size, cipher.data(), cipher.size());
+
+    IPC::RequestBuilder rb = rp.MakeBuilder(1, 4);
+    rb.Push(RESULT_SUCCESS);
+
+    // Unmap buffer
+    rb.PushMappedBuffer(input, input_size, IPC::MappedBufferPermissions::R);
+    rb.PushMappedBuffer(output, output_size, IPC::MappedBufferPermissions::W);
+}
+
+void Unwrap(Service::Interface* self) {
+    IPC::RequestParser rp(Kernel::GetCommandBuffer(), 0x47, 4, 4);
+    const u32 output_size = rp.Pop<u32>();
+    const u32 input_size = rp.Pop<u32>();
+    const u32 nonce_offset = rp.Pop<u32>();
+    u32 nonce_size = rp.Pop<u32>();
+    size_t desc_size;
+    IPC::MappedBufferPermissions desc_permission;
+    const VAddr input = rp.PopMappedBuffer(&desc_size, &desc_permission);
+    ASSERT(desc_size == input_size && desc_permission == IPC::MappedBufferPermissions::R);
+    const VAddr output = rp.PopMappedBuffer(&desc_size, &desc_permission);
+    ASSERT(desc_size == output_size && desc_permission == IPC::MappedBufferPermissions::W);
+
+    // Note: real 3DS still returns SUCCESS when the sizes don't match. It seems that it doesn't
+    // check the buffer size and writes data with potential overflow.
+    ASSERT_MSG(output_size == input_size - HW::AES::CCM_MAC_SIZE,
+               "input_size (%d) doesn't match to output_size (%d)", input_size, output_size);
+
+    LOG_DEBUG(Service_APT, "called, output_size=%u, input_size=%u, nonce_offset=%u, nonce_size=%u",
+              output_size, input_size, nonce_offset, nonce_size);
+
+    // Note: This weird nonce size modification is verified against real 3DS
+    nonce_size = std::min<u32>(nonce_size & ~3, HW::AES::CCM_NONCE_SIZE);
+
+    // Reads nonce and cipher text
+    HW::AES::CCMNonce nonce{};
+    Memory::ReadBlock(input, nonce.data(), nonce_size);
+    u32 cipher_size = input_size - nonce_size;
+    std::vector<u8> cipher(cipher_size);
+    Memory::ReadBlock(input + nonce_size, cipher.data(), cipher_size);
+
+    // Decrypts the ciphertext using AES-CCM
+    auto pdata = HW::AES::DecryptVerifyCCM(cipher, nonce, HW::AES::KeySlotID::APTWrap);
+
+    IPC::RequestBuilder rb = rp.MakeBuilder(1, 0);
+    if (!pdata.empty()) {
+        // Splits the plaintext and put the nonce in between
+        Memory::WriteBlock(output, pdata.data(), nonce_offset);
+        Memory::WriteBlock(output + nonce_offset, nonce.data(), nonce_size);
+        Memory::WriteBlock(output + nonce_offset + nonce_size, pdata.data() + nonce_offset,
+                           pdata.size() - nonce_offset);
+        rb.Push(RESULT_SUCCESS);
+    } else {
+        LOG_ERROR(Service_APT, "Failed to decrypt data");
+        rb.Push(ResultCode(static_cast<ErrorDescription>(1), ErrorModule::PS,
+                           ErrorSummary::WrongArgument, ErrorLevel::Status));
+    }
+
+    // Unmap buffer
+    rb.PushMappedBuffer(input, input_size, IPC::MappedBufferPermissions::R);
+    rb.PushMappedBuffer(output, output_size, IPC::MappedBufferPermissions::W);
+}
+
 void CheckNew3DSApp(Service::Interface* self) {
     u32* cmd_buff = Kernel::GetCommandBuffer();
 
diff --git a/src/core/hle/service/apt/apt.h b/src/core/hle/service/apt/apt.h
index 80325361f..e63b61450 100644
--- a/src/core/hle/service/apt/apt.h
+++ b/src/core/hle/service/apt/apt.h
@@ -137,6 +137,46 @@ void Initialize(Service::Interface* self);
 void GetSharedFont(Service::Interface* self);
 
 /**
+ * APT::Wrap service function
+ *  Inputs:
+ *      1 : Output buffer size
+ *      2 : Input buffer size
+ *      3 : Nonce offset to the input buffer
+ *      4 : Nonce size
+ *      5 : Buffer mapping descriptor ((input_buffer_size << 4) | 0xA)
+ *      6 : Input buffer address
+ *      7 : Buffer mapping descriptor ((input_buffer_size << 4) | 0xC)
+ *      8 : Output buffer address
+ *  Outputs:
+ *      1 : Result of function, 0 on success, otherwise error code
+ *      2 : Buffer unmapping descriptor ((input_buffer_size << 4) | 0xA)
+ *      3 : Input buffer address
+ *      4 : Buffer unmapping descriptor ((input_buffer_size << 4) | 0xC)
+ *      5 : Output buffer address
+ */
+void Wrap(Service::Interface* self);
+
+/**
+ * APT::Unwrap service function
+ *  Inputs:
+ *      1 : Output buffer size
+ *      2 : Input buffer size
+ *      3 : Nonce offset to the output buffer
+ *      4 : Nonce size
+ *      5 : Buffer mapping descriptor ((input_buffer_size << 4) | 0xA)
+ *      6 : Input buffer address
+ *      7 : Buffer mapping descriptor ((input_buffer_size << 4) | 0xC)
+ *      8 : Output buffer address
+ *  Outputs:
+ *      1 : Result of function, 0 on success, otherwise error code
+ *      2 : Buffer unmapping descriptor ((input_buffer_size << 4) | 0xA)
+ *      3 : Input buffer address
+ *      4 : Buffer unmapping descriptor ((input_buffer_size << 4) | 0xC)
+ *      5 : Output buffer address
+ */
+void Unwrap(Service::Interface* self);
+
+/**
  * APT::NotifyToWait service function
  *  Inputs:
  *      1 : AppID
diff --git a/src/core/hle/service/apt/apt_a.cpp b/src/core/hle/service/apt/apt_a.cpp
index 62dc2d61d..c496cba8d 100644
--- a/src/core/hle/service/apt/apt_a.cpp
+++ b/src/core/hle/service/apt/apt_a.cpp
@@ -78,8 +78,8 @@ const Interface::FunctionInfo FunctionTable[] = {
     {0x00430040, NotifyToWait, "NotifyToWait"},
     {0x00440000, GetSharedFont, "GetSharedFont"},
     {0x00450040, nullptr, "GetWirelessRebootInfo"},
-    {0x00460104, nullptr, "Wrap"},
-    {0x00470104, nullptr, "Unwrap"},
+    {0x00460104, Wrap, "Wrap"},
+    {0x00470104, Unwrap, "Unwrap"},
     {0x00480100, nullptr, "GetProgramInfo"},
     {0x00490180, nullptr, "Reboot"},
     {0x004A0040, nullptr, "GetCaptureInfo"},
diff --git a/src/core/hle/service/apt/apt_s.cpp b/src/core/hle/service/apt/apt_s.cpp
index effd23dce..ec5668d05 100644
--- a/src/core/hle/service/apt/apt_s.cpp
+++ b/src/core/hle/service/apt/apt_s.cpp
@@ -78,8 +78,8 @@ const Interface::FunctionInfo FunctionTable[] = {
     {0x00430040, NotifyToWait, "NotifyToWait"},
     {0x00440000, GetSharedFont, "GetSharedFont"},
     {0x00450040, nullptr, "GetWirelessRebootInfo"},
-    {0x00460104, nullptr, "Wrap"},
-    {0x00470104, nullptr, "Unwrap"},
+    {0x00460104, Wrap, "Wrap"},
+    {0x00470104, Unwrap, "Unwrap"},
     {0x00480100, nullptr, "GetProgramInfo"},
     {0x00490180, nullptr, "Reboot"},
     {0x004A0040, nullptr, "GetCaptureInfo"},
diff --git a/src/core/hle/service/apt/apt_u.cpp b/src/core/hle/service/apt/apt_u.cpp
index e06084a1e..9dd002590 100644
--- a/src/core/hle/service/apt/apt_u.cpp
+++ b/src/core/hle/service/apt/apt_u.cpp
@@ -78,8 +78,8 @@ const Interface::FunctionInfo FunctionTable[] = {
     {0x00430040, NotifyToWait, "NotifyToWait"},
     {0x00440000, GetSharedFont, "GetSharedFont"},
     {0x00450040, nullptr, "GetWirelessRebootInfo"},
-    {0x00460104, nullptr, "Wrap"},
-    {0x00470104, nullptr, "Unwrap"},
+    {0x00460104, Wrap, "Wrap"},
+    {0x00470104, Unwrap, "Unwrap"},
     {0x00480100, nullptr, "GetProgramInfo"},
     {0x00490180, nullptr, "Reboot"},
     {0x004A0040, nullptr, "GetCaptureInfo"},