From 1446217d1e052d834f7842476ee4b569c7b6b66a Mon Sep 17 00:00:00 2001
From: Huntereb <Huntereb@users.noreply.github.com>
Date: Fri, 25 Sep 2020 12:07:45 -0400
Subject: [PATCH] Some changes to socket handling

---
 include/netInstall.hpp       |  2 ++
 source/netInstall.cpp        | 19 +++++++++++++------
 source/ui/netInstPage.cpp    |  3 +++
 source/util/network_util.cpp |  4 ++--
 4 files changed, 20 insertions(+), 8 deletions(-)

diff --git a/include/netInstall.hpp b/include/netInstall.hpp
index 685234a..b3b20ce 100755
--- a/include/netInstall.hpp
+++ b/include/netInstall.hpp
@@ -22,6 +22,8 @@ SOFTWARE.
 #include <vector>
 
 namespace netInstStuff {
+    void OnUnwound();
+    void sendExitCommands();
     void installTitleNet(std::vector<std::string> ourUrlList, int ourStorage, std::vector<std::string> urlListAltNames, std::string ourSource);
     std::vector<std::string> OnSelected();
 }
\ No newline at end of file
diff --git a/source/netInstall.cpp b/source/netInstall.cpp
index a19fcf6..a03d76e 100755
--- a/source/netInstall.cpp
+++ b/source/netInstall.cpp
@@ -108,6 +108,17 @@ namespace netInstStuff{
         curl_global_cleanup();
     }
 
+    void sendExitCommands()
+    {
+        LOG_DEBUG("Telling the server we're done installing\n");
+        // Send 1 byte ack to close the server, OG tinfoil compatibility
+        u8 ack = 0;
+        tin::network::WaitSendNetworkData(m_clientSocket, &ack, sizeof(u8));
+        // Send 'DEAD\r\n' so ns-usbloader knows we're done
+        //u8 nsUsbAck [6] = {0x44,0x45,0x41,0x44,0x0D,0x0A};
+        //tin::network::WaitSendNetworkData(m_clientSocket, &nsUsbAck, sizeof(u8) * 6);
+    }
+
     void installTitleNet(std::vector<std::string> ourUrlList, int ourStorage, std::vector<std::string> urlListAltNames, std::string ourSource)
     {
         inst::util::initInstallServices();
@@ -178,10 +189,8 @@ namespace netInstStuff{
             inst::util::setClockSpeed(2, previousClockValues[2]);
         }
 
-        LOG_DEBUG("Telling the server we're done installing\n");
-        // Send 1 byte ack to close the server
-        u8 ack = 0;
-        tin::network::WaitSendNetworkData(m_clientSocket, &ack, sizeof(u8));
+        sendExitCommands();
+        OnUnwound();
 
         if(nspInstalled) {
             inst::ui::instPage::setInstInfoText("inst.info_page.complete"_lang);
@@ -206,8 +215,6 @@ namespace netInstStuff{
         u64 freq = armGetSystemTickFreq();
         u64 startTime = armGetSystemTick();
 
-        OnUnwound();
-
         try
         {
             ASSERT_OK(curl_global_init(CURL_GLOBAL_ALL), "Curl failed to initialized");
diff --git a/source/ui/netInstPage.cpp b/source/ui/netInstPage.cpp
index 63cb23b..c52f177 100755
--- a/source/ui/netInstPage.cpp
+++ b/source/ui/netInstPage.cpp
@@ -156,6 +156,9 @@ namespace inst::ui {
 
     void netInstPage::onInput(u64 Down, u64 Up, u64 Held, pu::ui::Touch Pos) {
         if (Down & KEY_B) {
+            if (this->menu->GetItems().size() > 0) {}
+                netInstStuff::sendExitCommands();
+            netInstStuff::OnUnwound();
             mainApp->LoadLayout(mainApp->mainPage);
         }
         if ((Down & KEY_A) || (Up & KEY_TOUCH)) {
diff --git a/source/util/network_util.cpp b/source/util/network_util.cpp
index be64e6d..e6c5944 100755
--- a/source/util/network_util.cpp
+++ b/source/util/network_util.cpp
@@ -235,7 +235,7 @@ namespace tin::network
         int ret = 0;
         size_t read = 0;
 
-        while ((((ret = recv(sockfd, (u8*)buf + read, len - read, 0)) > 0 && (read += ret) < len) || errno == EAGAIN) && !(hidKeysDown(CONTROLLER_P1_AUTO) & KEY_B)) 
+        while ((((ret = recv(sockfd, (u8*)buf + read, len - read, 0)) > 0 && (read += ret) < len) || errno == EAGAIN)) 
         {
             errno = 0;
         }
@@ -249,7 +249,7 @@ namespace tin::network
         int ret = 0;
         size_t written = 0;
 
-        while ((((ret = send(sockfd, (u8*)buf + written, len - written, 0)) > 0 && (written += ret) < len) || errno == EAGAIN) && !(hidKeysDown(CONTROLLER_P1_AUTO) & KEY_B)) 
+        while ((((ret = send(sockfd, (u8*)buf + written, len - written, 0)) > 0 && (written += ret) < len) || errno == EAGAIN)) 
         {
             errno = 0;
         }