diff --git a/.gitignore b/.gitignore index f55c55a62..a89a9ec84 100644 --- a/.gitignore +++ b/.gitignore @@ -65,6 +65,9 @@ dkms.conf *.tgz *.zip +# Python modules +*.pyc + .**/ # NOTE: make sure to make exceptions to this pattern when needed! diff --git a/stratosphere/tma/client/Main.py b/stratosphere/tma/client/Main.py new file mode 100644 index 000000000..580e0ae52 --- /dev/null +++ b/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)) \ No newline at end of file diff --git a/stratosphere/tma/client/UsbConnection.py b/stratosphere/tma/client/UsbConnection.py new file mode 100644 index 000000000..0f7d2cf8f --- /dev/null +++ b/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 . +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 + diff --git a/stratosphere/tma/client/UsbInterface.py b/stratosphere/tma/client/UsbInterface.py new file mode 100644 index 000000000..90d91cab9 --- /dev/null +++ b/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 . +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('. + */ + +#include +#include +#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. */ +} diff --git a/stratosphere/tma/source/tma_conn_connection.hpp b/stratosphere/tma/source/tma_conn_connection.hpp new file mode 100644 index 000000000..4cf3cabe3 --- /dev/null +++ b/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 . + */ + +#pragma once +#include +#include + +#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; +}; \ No newline at end of file diff --git a/stratosphere/tma/source/tma_conn_usb_connection.cpp b/stratosphere/tma/source/tma_conn_usb_connection.cpp new file mode 100644 index 000000000..d9d40a214 --- /dev/null +++ b/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 . + */ + +#include +#include +#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(_packet); + if (packet != nullptr) { + this->FreePacket(packet); + } + } +} + +void TmaUsbConnection::SendThreadFunc(void *arg) { + TmaUsbConnection *this_ptr = reinterpret_cast(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(_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(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(i++); + this_ptr->send_queue.Send(reinterpret_cast(send_packet)); + this_ptr->FreePacket(packet); + } else { + this_ptr->FreePacket(packet); + } + } + + this_ptr->SetConnected(false); + this_ptr->send_queue.Send(reinterpret_cast(nullptr)); +} + +void TmaUsbConnection::OnUsbStateChange(void *arg, u32 state) { + TmaUsbConnection *this_ptr = reinterpret_cast(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 lk(this->lock); + + if (this->IsConnected()) { + this->send_queue.Send(reinterpret_cast(packet)); + return TmaConnResult::Success; + } else { + this->FreePacket(packet); + this->Tick(); + return TmaConnResult::Disconnected; + } +} diff --git a/stratosphere/tma/source/tma_conn_usb_connection.hpp b/stratosphere/tma/source/tma_conn_usb_connection.hpp new file mode 100644 index 000000000..15876ef3b --- /dev/null +++ b/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 . + */ + +#pragma once +#include +#include + +#include "tma_conn_connection.hpp" +#include "tma_usb_comms.hpp" + +class TmaUsbConnection : public TmaConnection { + private: + HosMessageQueue send_queue = HosMessageQueue(64); + std::atomic 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; +}; \ No newline at end of file diff --git a/stratosphere/tma/source/tma_main.cpp b/stratosphere/tma/source/tma_main.cpp index e4c23777b..8fc308444 100644 --- a/stratosphere/tma/source/tma_main.cpp +++ b/stratosphere/tma/source/tma_main.cpp @@ -22,14 +22,14 @@ #include #include -#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(); - TmaPacket *packet = new TmaPacket(); - usbDsWaitReady(U64_MAX); - packet->Write(0xCAFEBABEDEADCAFEUL); - packet->Write(0xCCCCCCCCCCCCCCCCUL); - TmaUsbComms::SendPacket(packet); - packet->ClearOffset(); + TmaUsbConnection::InitializeComms(); + auto conn = new TmaUsbConnection(); + conn->Initialize(); + while (true) { - if (TmaUsbComms::ReceivePacket(packet) == TmaConnResult::Success) { - TmaUsbComms::SendPacket(packet); - } + svcSleepThread(10000000UL); } diff --git a/stratosphere/tma/source/tma_usb_comms.cpp b/stratosphere/tma/source/tma_usb_comms.cpp index 440218bef..f2cd014c8 100644 --- a/stratosphere/tma/source/tma_usb_comms.cpp +++ b/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(); +} \ No newline at end of file diff --git a/stratosphere/tma/source/tma_usb_comms.hpp b/stratosphere/tma/source/tma_usb_comms.hpp index 9815e7634..52f119140 100644 --- a/stratosphere/tma/source/tma_usb_comms.hpp +++ b/stratosphere/tma/source/tma_usb_comms.hpp @@ -14,6 +14,7 @@ * along with this program. If not, see . */ +#pragma once #include #include