tma: impl helper services, cleanup hostside packets

.gitignore

@@ -65,6 +65,9 @@ dkms.conf
 *.tgz
 *.zip
 
+# Python modules
+*.pyc
+
 .**/
 
 # NOTE: make sure to make exceptions to this pattern when needed!

stratosphere/tma/client/Main.py

@@ -0,0 +1,26 @@
+# Copyright (c) 2018 Atmosphere-NX
+#
+# This program is free software; you can redistribute it and/or modify it
+# under the terms and conditions of the GNU General Public License,
+# version 2, as published by the Free Software Foundation.
+#
+# This program is distributed in the hope it will be useful, but WITHOUT
+# ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
+# FITNESS FOR A PARTICULAR PURPOSE.  See the GNU General Public License for
+# more details.
+#
+# You should have received a copy of the GNU General Public License
+from UsbConnection import UsbConnection
+import sys, time
+
+def main(argc, argv):
+    with UsbConnection(None) as c:
+        print 'Waiting for connection...'
+        c.wait_connected()
+        print 'Connected!'
+        while True:
+            c.send_packet('AAAAAAAA')
+    return 0
+    
+if __name__ == '__main__':
+    sys.exit(main(len(sys.argv), sys.argv))

stratosphere/tma/client/UsbConnection.py

@@ -0,0 +1,117 @@
+# Copyright (c) 2018 Atmosphere-NX
+#
+# This program is free software; you can redistribute it and/or modify it
+# under the terms and conditions of the GNU General Public License,
+# version 2, as published by the Free Software Foundation.
+#
+# This program is distributed in the hope it will be useful, but WITHOUT
+# ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
+# FITNESS FOR A PARTICULAR PURPOSE.  See the GNU General Public License for
+# more details.
+#
+# You should have received a copy of the GNU General Public License
+# along with this program.  If not, see <http://www.gnu.org/licenses/>.
+from UsbInterface import UsbInterface
+from threading import Thread, Condition
+from collections import deque
+import time
+
+class UsbConnection(UsbInterface):
+    # Auto connect thread func.
+    def auto_connect(connection):
+        while not connection.is_connected():
+            try:
+                connection.connect(UsbInterface())
+            except ValueError as e:
+                continue
+    def recv_thread(connection):
+        if connection.is_connected():
+            try:
+                # If we've previously been connected, PyUSB will read garbage...
+                connection.recv_packet()
+            except ValueError:
+                pass
+        while connection.is_connected():
+            try:
+                connection.recv_packet()
+            except Exception as e:
+                print 'An exception occurred:'
+                print 'Type: '+e.__class__.__name__
+                print 'Msg: '+str(e)
+                connection.disconnect()
+                connection.send_packet(None)
+    def send_thread(connection):
+        while connection.is_connected():
+            try:
+                next_packet = connection.get_next_send_packet()
+                if next_packet is not None:
+                    connection.intf.send_packet(next_packet)
+                else:
+                    connection.disconnect()
+            except Exception as e:
+                print 'An exception occurred:'
+                print 'Type: '+e.__class__.__name__
+                print 'Msg: '+str(e)
+                connection.disconnect()
+    def __init__(self, manager):
+        self.manager = manager
+        self.connected = False
+        self.intf = None
+        self.conn_lock, self.send_lock = Condition(), Condition()
+        self.send_queue = deque()
+    def __enter__(self):
+        self.conn_thrd = Thread(target=UsbConnection.auto_connect, args=(self,))
+        self.conn_thrd.daemon = True
+        self.conn_thrd.start()
+        return self
+    def __exit__(self, type, value, traceback):
+        time.sleep(1)
+        print 'Closing!'
+        time.sleep(1)
+    def wait_connected(self):
+        self.conn_lock.acquire()
+        if not self.is_connected():
+            self.conn_lock.wait()
+        self.conn_lock.release()
+    def is_connected(self):
+        return self.connected
+    def connect(self, intf):
+        # This indicates we have a connection.
+        self.conn_lock.acquire()
+        assert not self.connected
+        self.intf = intf
+        self.connected = True
+        self.conn_lock.notify()
+        self.conn_lock.release()
+        self.recv_thrd = Thread(target=UsbConnection.recv_thread, args=(self,))
+        self.send_thrd = Thread(target=UsbConnection.send_thread, args=(self,))
+        self.recv_thrd.daemon = True
+        self.send_thrd.daemon = True
+        self.recv_thrd.start()
+        self.send_thrd.start()
+    def disconnect(self):
+        self.conn_lock.acquire()
+        if self.connected:
+            self.connected = False
+        self.conn_lock.release()
+    def recv_packet(self):
+        hdr, body = self.intf.read_packet()
+        print('Got Packet: %s' % body.encode('hex'))
+    def send_packet(self, packet):
+        self.send_lock.acquire()
+        if len(self.send_queue) == 0x40:
+            self.send_lock.wait()
+        self.send_queue.append(packet)
+        if len(self.send_queue) == 1:
+            self.send_lock.notify()
+        self.send_lock.release()
+    def get_next_send_packet(self):
+        self.send_lock.acquire()
+        if len(self.send_queue) == 0:
+            self.send_lock.wait()
+        packet = self.send_queue.popleft()
+        if len(self.send_queue) == 0x3F:
+            self.send_lock.notify()
+        self.send_lock.release()
+        return packet
+

stratosphere/tma/client/UsbInterface.py

@@ -0,0 +1,69 @@
+# Copyright (c) 2018 Atmosphere-NX
+#
+# This program is free software; you can redistribute it and/or modify it
+# under the terms and conditions of the GNU General Public License,
+# version 2, as published by the Free Software Foundation.
+#
+# This program is distributed in the hope it will be useful, but WITHOUT
+# ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
+# FITNESS FOR A PARTICULAR PURPOSE.  See the GNU General Public License for
+# more details.
+#
+# You should have received a copy of the GNU General Public License
+# along with this program.  If not, see <http://www.gnu.org/licenses/>.
+import usb, zlib
+from struct import unpack as up, pack as pk
+
+def crc32(s):
+    return zlib.crc32(s) & 0xFFFFFFFF
+
+class UsbInterface():
+    def __init__(self):
+        self.dev = usb.core.find(idVendor=0x057e, idProduct=0x3000)
+        if self.dev is None:
+            raise ValueError('Device not found')
+
+        self.dev.set_configuration()
+        self.cfg = self.dev.get_active_configuration()
+        self.intf = usb.util.find_descriptor(self.cfg, bInterfaceClass=0xff, bInterfaceSubClass=0xff, bInterfaceProtocol=0xfc)
+        assert self.intf is not None
+        
+        self.ep_in = usb.util.find_descriptor(
+            self.intf,
+            custom_match = \
+            lambda e: \
+                usb.util.endpoint_direction(e.bEndpointAddress) == \
+                usb.util.ENDPOINT_IN)
+        assert self.ep_in is not None
+
+        self.ep_out = usb.util.find_descriptor(
+            self.intf,
+            custom_match = \
+            lambda e: \
+                usb.util.endpoint_direction(e.bEndpointAddress) == \
+                usb.util.ENDPOINT_OUT)
+        assert self.ep_out is not None
+    def close(self):
+        usb.util.dispose_resources(self.dev)
+    def blocking_read(self, size):
+        return ''.join(chr(x) for x in self.ep_in.read(size, 0xFFFFFFFFFFFFFFFF))
+    def blocking_write(self, data):
+        self.ep_out.write(data, 0xFFFFFFFFFFFFFFFF)
+    def read_packet(self):
+        hdr = self.blocking_read(0x28)
+        _, _, _, body_size, _, _, _, _, body_chk, hdr_chk = up('<IIIIIIIIII', hdr)
+        if crc32(hdr[:-4]) != hdr_chk:
+            raise ValueError('Invalid header checksum in received packet!')
+        body = self.blocking_read(body_size)
+        if len(body) != body_size:
+            raise ValueError('Failed to receive packet body!')
+        elif crc32(body) != body_chk:
+            raise ValueError('Invalid body checksum in received packet!')
+        return (hdr, body)
+    def send_packet(self, body):
+        hdr = pk('<IIIIIIIII', 0, 0, 0, len(body), 0, 0, 0, 0, crc32(body))
+        hdr += pk('<I', crc32(hdr))
+        self.blocking_write(hdr)
+        self.blocking_write(body)
+
+

stratosphere/tma/source/tma_conn_connection.cpp

@@ -0,0 +1,64 @@
+/*
+ * Copyright (c) 2018 Atmosphère-NX
+ *
+ * This program is free software; you can redistribute it and/or modify it
+ * under the terms and conditions of the GNU General Public License,
+ * version 2, as published by the Free Software Foundation.
+ *
+ * This program is distributed in the hope it will be useful, but WITHOUT
+ * ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
+ * FITNESS FOR A PARTICULAR PURPOSE.  See the GNU General Public License for
+ * more details.
+ *
+ * You should have received a copy of the GNU General Public License
+ * along with this program.  If not, see <http://www.gnu.org/licenses/>.
+ */
+
+#include <switch.h>
+#include <stratosphere.hpp>
+#include "tma_conn_connection.hpp"
+
+/* Packet management. */
+TmaPacket *TmaConnection::AllocateSendPacket() {
+    /* TODO: Service Manager. */
+    return new TmaPacket();
+}
+
+TmaPacket *TmaConnection::AllocateRecvPacket() {
+    /* TODO: Service Manager. */
+    return new TmaPacket();
+}
+
+void TmaConnection::FreePacket(TmaPacket *packet) {
+    /* TODO: Service Manager. */
+    delete packet;
+}
+
+void TmaConnection::OnReceivePacket(TmaPacket *packet) {
+    /* TODO: Service Manager. */
+}
+
+void TmaConnection::OnDisconnected() {
+    if (!this->is_initialized) {
+        std::abort();
+    }
+    
+    /* TODO: Service Manager. */
+    
+    this->has_woken_up = false;
+    this->OnConnectionEvent(ConnectionEvent::Disconnected);
+}
+
+void TmaConnection::OnConnectionEvent(ConnectionEvent event) {
+    if (this->connection_event_callback != nullptr) {
+        this->connection_event_callback(this->connection_event_arg, event);
+    }
+}
+
+void TmaConnection::CancelTasks() {
+    /* TODO: Service Manager. */
+}
+
+void TmaConnection::Tick() {
+    /* TODO: Service Manager. */
+}

stratosphere/tma/source/tma_conn_connection.hpp

@@ -0,0 +1,80 @@
+/*
+ * Copyright (c) 2018 Atmosphère-NX
+ *
+ * This program is free software; you can redistribute it and/or modify it
+ * under the terms and conditions of the GNU General Public License,
+ * version 2, as published by the Free Software Foundation.
+ *
+ * This program is distributed in the hope it will be useful, but WITHOUT
+ * ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
+ * FITNESS FOR A PARTICULAR PURPOSE.  See the GNU General Public License for
+ * more details.
+ *
+ * You should have received a copy of the GNU General Public License
+ * along with this program.  If not, see <http://www.gnu.org/licenses/>.
+ */
+ 
+#pragma once
+#include <switch.h>
+#include <stratosphere.hpp>
+
+#include "tma_conn_result.hpp"
+#include "tma_conn_packet.hpp"
+
+enum class ConnectionEvent : u32 {
+    Connected,
+    Disconnected
+};
+
+class TmaConnection {
+    protected:
+        HosMutex lock;
+        void (*connection_event_callback)(void *, ConnectionEvent) = nullptr;
+        void *connection_event_arg = nullptr;
+        bool has_woken_up = false;
+        bool is_initialized = false;
+    protected:
+        void OnReceivePacket(TmaPacket *packet);
+        void OnDisconnected();
+        void OnConnectionEvent(ConnectionEvent event);
+        void CancelTasks();
+        void Tick();
+    public:
+        /* Setup */
+        TmaConnection() { }
+        virtual ~TmaConnection() { }
+        
+        void Initialize() { 
+            if (this->is_initialized) { 
+                std::abort();
+            }
+            this->is_initialized = true;
+        }
+        
+        void SetConnectionEventCallback(void (*callback)(void *, ConnectionEvent), void *arg) {
+            this->connection_event_callback = callback;
+            this->connection_event_arg = arg;
+        }
+        
+        void Finalize() {
+            if (this->is_initialized) {
+                this->StopListening();
+                if (this->IsConnected()) {
+                    this->Disconnect();
+                }
+                this->is_initialized = false;
+            }
+        }
+        
+        /* Packet management. */
+        TmaPacket *AllocateSendPacket();
+        TmaPacket *AllocateRecvPacket();
+        void FreePacket(TmaPacket *packet);
+
+        /* For sub-interfaces to implement, connection management. */
+        virtual void StartListening() { }
+        virtual void StopListening() { }
+        virtual bool IsConnected() = 0;
+        virtual TmaConnResult Disconnect() = 0;
+        virtual TmaConnResult SendPacket(TmaPacket *packet) = 0;
+};

stratosphere/tma/source/tma_conn_usb_connection.cpp

@@ -0,0 +1,155 @@
+/*
+ * Copyright (c) 2018 Atmosphère-NX
+ *
+ * This program is free software; you can redistribute it and/or modify it
+ * under the terms and conditions of the GNU General Public License,
+ * version 2, as published by the Free Software Foundation.
+ *
+ * This program is distributed in the hope it will be useful, but WITHOUT
+ * ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
+ * FITNESS FOR A PARTICULAR PURPOSE.  See the GNU General Public License for
+ * more details.
+ *
+ * You should have received a copy of the GNU General Public License
+ * along with this program.  If not, see <http://www.gnu.org/licenses/>.
+ */
+
+#include <switch.h>
+#include <stratosphere.hpp>
+#include "tma_conn_usb_connection.hpp"
+#include "tma_usb_comms.hpp"
+
+static HosThread g_SendThread, g_RecvThread;
+
+TmaConnResult TmaUsbConnection::InitializeComms() {
+    return TmaUsbComms::Initialize();
+}
+
+TmaConnResult TmaUsbConnection::FinalizeComms() {
+    return TmaUsbComms::Finalize();
+}
+
+void TmaUsbConnection::ClearSendQueue() {
+    uintptr_t _packet;
+    while (this->send_queue.TryReceive(&_packet)) {
+        TmaPacket *packet = reinterpret_cast<TmaPacket *>(_packet);
+        if (packet != nullptr) {
+            this->FreePacket(packet);
+        }
+    }
+}
+
+void TmaUsbConnection::SendThreadFunc(void *arg) {
+    TmaUsbConnection *this_ptr = reinterpret_cast<TmaUsbConnection *>(arg);
+    TmaConnResult res = TmaConnResult::Success;
+    TmaPacket *packet = nullptr;
+    
+    while (res == TmaConnResult::Success) {
+        /* Receive a packet from the send queue. */
+        {
+            uintptr_t _packet;
+            this_ptr->send_queue.Receive(&_packet);
+            packet = reinterpret_cast<TmaPacket *>(_packet);
+        }
+        
+        if (packet != nullptr) {
+            /* Send the packet if we're connected. */
+            if (this_ptr->IsConnected()) {
+                res = TmaUsbComms::SendPacket(packet);
+            }
+            
+            this_ptr->FreePacket(packet);
+            this_ptr->Tick();
+        } else {
+            res = TmaConnResult::Disconnected;
+        }
+    }
+    
+    this_ptr->SetConnected(false);
+    this_ptr->OnDisconnected();
+}
+
+void TmaUsbConnection::RecvThreadFunc(void *arg) {
+    TmaUsbConnection *this_ptr = reinterpret_cast<TmaUsbConnection *>(arg);
+    TmaConnResult res = TmaConnResult::Success;
+    u64 i = 0;
+    this_ptr->SetConnected(true);
+    
+    while (res == TmaConnResult::Success) {
+        if (!this_ptr->IsConnected()) {
+            break;
+        }
+        TmaPacket *packet = this_ptr->AllocateRecvPacket();
+        if (packet == nullptr) { std::abort(); }
+        
+        res = TmaUsbComms::ReceivePacket(packet);
+        
+        if (res == TmaConnResult::Success) {
+            TmaPacket *send_packet = this_ptr->AllocateSendPacket();
+            send_packet->Write<u64>(i++);
+            this_ptr->send_queue.Send(reinterpret_cast<uintptr_t>(send_packet));
+            this_ptr->FreePacket(packet);
+        } else {
+            this_ptr->FreePacket(packet);
+        }
+    }
+    
+    this_ptr->SetConnected(false);
+    this_ptr->send_queue.Send(reinterpret_cast<uintptr_t>(nullptr));
+}
+
+void TmaUsbConnection::OnUsbStateChange(void *arg, u32 state) {
+    TmaUsbConnection *this_ptr = reinterpret_cast<TmaUsbConnection *>(arg);
+    switch (state) {
+        case 0:
+        case 6:
+            this_ptr->StopThreads();
+            break;
+        case 5:
+            this_ptr->StartThreads();
+            break;
+    }
+}
+
+void TmaUsbConnection::StartThreads() {
+    g_SendThread.Join();
+    g_RecvThread.Join();
+    
+    g_SendThread.Initialize(&TmaUsbConnection::SendThreadFunc, this, 0x4000, 38);
+    g_RecvThread.Initialize(&TmaUsbConnection::RecvThreadFunc, this, 0x4000, 38);
+    
+    this->ClearSendQueue();
+    g_SendThread.Start();
+    g_RecvThread.Start();
+}
+
+void TmaUsbConnection::StopThreads() {
+    TmaUsbComms::CancelComms();
+    g_SendThread.Join();
+    g_RecvThread.Join();
+}
+
+bool TmaUsbConnection::IsConnected() {
+    return this->is_connected;
+}
+
+TmaConnResult TmaUsbConnection::Disconnect() {
+    TmaUsbComms::SetStateChangeCallback(nullptr, nullptr);
+    
+    this->StopThreads();
+    
+    return TmaConnResult::Success;
+}
+
+TmaConnResult TmaUsbConnection::SendPacket(TmaPacket *packet) {
+    std::scoped_lock<HosMutex> lk(this->lock);
+    
+    if (this->IsConnected()) {
+        this->send_queue.Send(reinterpret_cast<uintptr_t>(packet));
+        return TmaConnResult::Success;
+    } else {
+        this->FreePacket(packet);
+        this->Tick();
+        return TmaConnResult::Disconnected;
+    }
+}

stratosphere/tma/source/tma_conn_usb_connection.hpp

@@ -0,0 +1,51 @@
+/*
+ * Copyright (c) 2018 Atmosphère-NX
+ *
+ * This program is free software; you can redistribute it and/or modify it
+ * under the terms and conditions of the GNU General Public License,
+ * version 2, as published by the Free Software Foundation.
+ *
+ * This program is distributed in the hope it will be useful, but WITHOUT
+ * ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
+ * FITNESS FOR A PARTICULAR PURPOSE.  See the GNU General Public License for
+ * more details.
+ *
+ * You should have received a copy of the GNU General Public License
+ * along with this program.  If not, see <http://www.gnu.org/licenses/>.
+ */
+ 
+#pragma once
+#include <switch.h>
+#include <stratosphere.hpp>
+
+#include "tma_conn_connection.hpp"
+#include "tma_usb_comms.hpp"
+
+class TmaUsbConnection : public TmaConnection {
+    private:
+        HosMessageQueue send_queue = HosMessageQueue(64);
+        std::atomic<bool> is_connected = false;
+    private:
+        static void SendThreadFunc(void *arg);
+        static void RecvThreadFunc(void *arg);
+        static void OnUsbStateChange(void *this_ptr, u32 state);
+        void ClearSendQueue();
+        void StartThreads();
+        void StopThreads();
+        void SetConnected(bool c) { this->is_connected = c; }
+    public:
+        static TmaConnResult InitializeComms();
+        static TmaConnResult FinalizeComms();
+        
+        TmaUsbConnection() { 
+            TmaUsbComms::SetStateChangeCallback(&TmaUsbConnection::OnUsbStateChange, this);
+        }
+        
+        virtual ~TmaUsbConnection() {
+            this->Disconnect(); 
+        }
+        
+        virtual bool IsConnected() override;
+        virtual TmaConnResult Disconnect() override;
+        virtual TmaConnResult SendPacket(TmaPacket *packet) override;
+};

stratosphere/tma/source/tma_main.cpp

@@ -22,14 +22,14 @@
 #include <switch.h>
 #include <stratosphere.hpp>
 
-#include "tma_usb_comms.hpp"
+#include "tma_conn_usb_connection.hpp"
 
 extern "C" {
     extern u32 __start__;
 
     u32 __nx_applet_type = AppletType_None;
 
-    #define INNER_HEAP_SIZE 0x100000
+    #define INNER_HEAP_SIZE 0x400000
     size_t nx_inner_heap_size = INNER_HEAP_SIZE;
     char   nx_inner_heap[INNER_HEAP_SIZE];
     
@@ -111,17 +111,12 @@ int main(int argc, char **argv)
         /* TODO: Panic. */
     }
     
-    TmaUsbComms::Initialize();
+    TmaUsbConnection::InitializeComms();
-    TmaPacket *packet = new TmaPacket();
+    auto conn = new TmaUsbConnection();
-    usbDsWaitReady(U64_MAX);
+    conn->Initialize();
-    packet->Write<u64>(0xCAFEBABEDEADCAFEUL);
+    
-    packet->Write<u64>(0xCCCCCCCCCCCCCCCCUL);
-    TmaUsbComms::SendPacket(packet);
-    packet->ClearOffset();
     while (true) {
-        if (TmaUsbComms::ReceivePacket(packet) == TmaConnResult::Success) {
+        svcSleepThread(10000000UL);
-            TmaUsbComms::SendPacket(packet);
-        }
     }
 
     

stratosphere/tma/source/tma_usb_comms.cpp

@@ -258,12 +258,12 @@ TmaConnResult TmaUsbComms::Initialize() {
     }
     
     /* Start the state change thread. */
-    /*if (R_SUCCEEDED(rc)) {
+    if (R_SUCCEEDED(rc)) {
         rc = g_state_change_thread.Initialize(&TmaUsbComms::UsbStateChangeThreadFunc, nullptr, 0x4000, 38);
         if (R_SUCCEEDED(rc)) {
             rc = g_state_change_thread.Start();
         }
-    }*/
+    }
     
     /* Enable USB communication. */
     if (R_SUCCEEDED(rc)) {
@@ -277,10 +277,6 @@ TmaConnResult TmaUsbComms::Initialize() {
     if (R_FAILED(rc)) {
         /* TODO: Should I not abort here? */
         std::abort();
-        
-        // /* Cleanup, just in case. */
-        // TmaUsbComms::Finalize();
-        // res = TmaConnResult::Failure;
     }
     
     g_initialized = true;
@@ -306,6 +302,10 @@ TmaConnResult TmaUsbComms::Finalize() {
         usbDsExit();
     }
     
+    g_state_change_callback = nullptr;
+    g_interface = nullptr;
+    g_endpoint_in = nullptr;
+    g_endpoint_out = nullptr;
     g_initialized = false;
     
     return R_SUCCEEDED(rc) ? TmaConnResult::Success : TmaConnResult::ConnectionFailure;
@@ -458,3 +458,27 @@ TmaConnResult TmaUsbComms::SendPacket(TmaPacket *packet) {
 
     return res;
 }
+
+void TmaUsbComms::UsbStateChangeThreadFunc(void *arg) {
+    u32 state;
+    g_state_change_manager = new WaitableManager(1);
+    
+    auto state_change_event = LoadReadOnlySystemEvent(usbDsGetStateChangeEvent()->revent, [&](u64 timeout) {
+        if (R_SUCCEEDED(usbDsGetState(&state))) {
+            if (g_state_change_callback != nullptr) {
+                g_state_change_callback(g_state_change_arg, state);
+            }
+        }
+        return 0;
+    }, true);
+    
+    g_state_change_manager->AddWaitable(state_change_event);
+    g_state_change_manager->Process();
+    
+    /* If we get here, we're exiting. */
+    state_change_event->r_h = 0;
+    delete g_state_change_manager;
+    g_state_change_manager = nullptr;
+    
+    svcExitThread();
+}

stratosphere/tma/source/tma_usb_comms.hpp

@@ -14,6 +14,7 @@
  * along with this program.  If not, see <http://www.gnu.org/licenses/>.
  */
  
+#pragma once
 #include <switch.h>
 #include <stratosphere.hpp>