From ff9412feef27782c31668344eaca42906c098328 Mon Sep 17 00:00:00 2001
From: Michael Scire <SciresM@gmail.com>
Date: Sun, 10 Jun 2018 01:17:00 -0600
Subject: [PATCH] fs.mitm: Add hook for postprocessing mitm messages.

---
 stratosphere/fs_mitm/source/fsmitm_service.cpp    | 11 +++++++++++
 stratosphere/fs_mitm/source/fsmitm_service.hpp    |  4 +++-
 .../fs_mitm/source/imitmserviceobject.hpp         | 12 ++++++++++++
 stratosphere/fs_mitm/source/mitm_session.hpp      | 15 +++++++++++++--
 4 files changed, 39 insertions(+), 3 deletions(-)
 create mode 100644 stratosphere/fs_mitm/source/imitmserviceobject.hpp

diff --git a/stratosphere/fs_mitm/source/fsmitm_service.cpp b/stratosphere/fs_mitm/source/fsmitm_service.cpp
index 0567aeb4d..744082a57 100644
--- a/stratosphere/fs_mitm/source/fsmitm_service.cpp
+++ b/stratosphere/fs_mitm/source/fsmitm_service.cpp
@@ -3,6 +3,17 @@
 
 Result FsMitMService::dispatch(IpcParsedCommand &r, IpcCommand &out_c, u64 cmd_id, u8 *pointer_buffer, size_t pointer_buffer_size) {
     Result rc = 0xF601;
+    while (cmd_id == 18) {
+        
+    }
+    return rc;
+}
+
+Result FsMitMService::postprocess(IpcParsedCommand &r, IpcCommand &out_c, u64 cmd_id, u8 *pointer_buffer, size_t pointer_buffer_size) {
+    Result rc = 0xF601;
+    while (cmd_id == 18) {
+        
+    }
     return rc;
 }
 
diff --git a/stratosphere/fs_mitm/source/fsmitm_service.hpp b/stratosphere/fs_mitm/source/fsmitm_service.hpp
index 7f08c7b59..3ddd8c981 100644
--- a/stratosphere/fs_mitm/source/fsmitm_service.hpp
+++ b/stratosphere/fs_mitm/source/fsmitm_service.hpp
@@ -1,9 +1,11 @@
 #pragma once
 #include <switch.h>
 #include <stratosphere/iserviceobject.hpp>
+#include "imitmserviceobject.hpp"
 
-class FsMitMService : IServiceObject {        
+class FsMitMService : public IMitMServiceObject {        
     public:
         virtual Result dispatch(IpcParsedCommand &r, IpcCommand &out_c, u64 cmd_id, u8 *pointer_buffer, size_t pointer_buffer_size);
+        virtual Result postprocess(IpcParsedCommand &r, IpcCommand &out_c, u64 cmd_id, u8 *pointer_buffer, size_t pointer_buffer_size);
         virtual Result handle_deferred();
 };
\ No newline at end of file
diff --git a/stratosphere/fs_mitm/source/imitmserviceobject.hpp b/stratosphere/fs_mitm/source/imitmserviceobject.hpp
new file mode 100644
index 000000000..bba10a4c8
--- /dev/null
+++ b/stratosphere/fs_mitm/source/imitmserviceobject.hpp
@@ -0,0 +1,12 @@
+#pragma once
+#include <switch.h>
+
+#include <stratosphere.hpp>
+
+class IMitMServiceObject : public IServiceObject {
+    protected:
+        virtual ~IMitMServiceObject() { }
+        virtual Result dispatch(IpcParsedCommand &r, IpcCommand &out_c, u64 cmd_id, u8 *pointer_buffer, size_t pointer_buffer_size) = 0;
+        virtual Result postprocess(IpcParsedCommand &r, IpcCommand &out_c, u64 cmd_id, u8 *pointer_buffer, size_t pointer_buffer_size) = 0;
+        virtual Result handle_deferred() = 0;
+};
diff --git a/stratosphere/fs_mitm/source/mitm_session.hpp b/stratosphere/fs_mitm/source/mitm_session.hpp
index a286326e2..ceb1f4926 100644
--- a/stratosphere/fs_mitm/source/mitm_session.hpp
+++ b/stratosphere/fs_mitm/source/mitm_session.hpp
@@ -1,6 +1,7 @@
 #pragma once
 #include <switch.h>
 #include <stratosphere.hpp>
+#include "imitmserviceobject.hpp"
 
 #include "mitm_server.hpp"
 
@@ -12,7 +13,7 @@ class MitMServer;
 
 template <typename T>
 class MitMSession final : public IWaitable {
-    static_assert(std::is_base_of<IServiceObject, T>::value, "Service Objects must derive from IServiceObject");
+    static_assert(std::is_base_of<IMitMServiceObject, T>::value, "MitM Service Objects must derive from IMitMServiceObject");
     
     T *service_object;
     MitMServer<T> *server;
@@ -102,13 +103,19 @@ class MitMSession final : public IWaitable {
                 
                 retval = ipcParse(&r);
                 
+                u64 cmd_id = U64_MAX;
+                
                 /* TODO: Close input copy handles that we don't need. */
                                 
                 if (R_SUCCEEDED(retval)) {
                     rawdata_start = (u32 *)r.Raw;
+                    cmd_id = rawdata_start[2];
                     retval = 0xF601;
                     if (r.CommandType == IpcCommandType_Request || r.CommandType == IpcCommandType_RequestWithContext) {
-                        retval = this->service_object->dispatch(r, c, rawdata_start[2], (u8 *)this->pointer_buffer, this->pointer_buffer_size);
+                        retval = this->service_object->dispatch(r, c, cmd_id, (u8 *)this->pointer_buffer, this->pointer_buffer_size);
+                        if (R_SUCCEEDED(retval)) {
+                            ipcParse(&out_r);
+                        }
                     }
                     
                     /* 0xF601 --> Dispatch onwards. */
@@ -137,6 +144,10 @@ class MitMSession final : public IWaitable {
                     /* Session close. */
                     rc = retval;
                 } else {
+                    if (R_SUCCEEDED(rc)) {
+                        ipcInitialize(&c);
+                        retval = this->service_object->postprocess(r, c, cmd_id, (u8 *)this->pointer_buffer, this->pointer_buffer_size);
+                    }
                     Log(armGetTls(), 0x100);
                     rc = svcReplyAndReceive(&handle_index, &this->server_handle, 0, this->server_handle, 0);
                     /* Clean up copy handles. */