From 3f1b435b4ec31d54a31c36c381f62f5afa9951f4 Mon Sep 17 00:00:00 2001
From: Ludvig Strigeus <strigeus@gmail.com>
Date: Sun, 12 Aug 2018 12:38:54 +0200
Subject: [PATCH] Handle packet overload on non-windows

 - Also ensure RunAllMainThreadScheduled(); gets called on non-windows
---
 installer/ChangeLog.txt |  2 ++
 network_bsd.cpp         | 76 +++++++++++++++++++++++------------------
 network_bsd_mt.cpp      |  2 ++
 3 files changed, 46 insertions(+), 34 deletions(-)

diff --git a/installer/ChangeLog.txt b/installer/ChangeLog.txt
index 22120bd..f3b5b91 100644
--- a/installer/ChangeLog.txt
+++ b/installer/ChangeLog.txt
@@ -1,4 +1,6 @@
 2018-08-11 - TunSafe v1.4-rc1
+
+Changes:
 1.Subfolders in the Config/ directory now show up as submenus.
 2.Added a way to run TunSafe as a Windows Service.
   Foreground Mode: The service will disconnect when TunSafe closes.
diff --git a/network_bsd.cpp b/network_bsd.cpp
index 3616a27..40e987c 100644
--- a/network_bsd.cpp
+++ b/network_bsd.cpp
@@ -68,7 +68,7 @@ public:
   virtual void HandleExit() override { exit_ = true; }
   
 private:
-  bool ReadFromUdp();
+  bool ReadFromUdp(bool overload);
   bool ReadFromTun();
   bool WriteToUdp();
   bool WriteToTun();
@@ -131,7 +131,7 @@ void TunsafeBackendBsdImpl::SetTunFd(int fd) {
 }
 
 
-bool TunsafeBackendBsdImpl::ReadFromUdp() {
+bool TunsafeBackendBsdImpl::ReadFromUdp(bool overload) {
   socklen_t sin_len;
   sin_len = sizeof(read_packet_->addr.sin);
   int r = recvfrom(udp_fd_, read_packet_->data, kPacketCapacity, 0,
@@ -141,7 +141,7 @@ bool TunsafeBackendBsdImpl::ReadFromUdp() {
     read_packet_->sin_size = sin_len;
     read_packet_->size = r;
     if (processor_) {
-      processor_->HandleUdpPacket(read_packet_, false);
+      processor_->HandleUdpPacket(read_packet_, overload);
       read_packet_ = AllocPacket();
     }
     return true;        
@@ -277,62 +277,70 @@ void TunsafeBackendBsdImpl::WriteUdpPacket(Packet *packet) override {
 
 void TunsafeBackendBsdImpl::RunLoopInner() {
   int free_packet_interval = 10;
+  int overload_ctr = 0;
 
   while (!exit_) {
     int n = -1;
 
-//    printf("entering sleep %d,%d,%d %d\n", udp_fd_, tun_fd_, max_fd_, FD_ISSET(tun_fd_, &readfds_));
-    // Wait for sockets to become usable
-    if (!got_sig_alarm_) {
-
-      if (tun_fd_ >= 0) {
-        FD_SET(tun_fd_, &readfds_);
-        if (tun_writable_) FD_CLR(tun_fd_, &writefds_); else FD_SET(tun_fd_, &writefds_);
-      }
-
-      if (udp_fd_ >= 0) {
-        FD_SET(udp_fd_, &readfds_);
-        if (udp_writable_) FD_CLR(udp_fd_, &writefds_); else FD_SET(udp_fd_, &writefds_);
-      }
-
-      n = select(max_fd_, &readfds_, &writefds_, NULL, NULL);
-      if (n == -1) {
-        if (errno != EINTR) {
-          fprintf(stderr, "select failed\n");
-          break;
-        }
-      }
-    }
     // This is not fully signal safe.
     if (got_sig_alarm_) {
       got_sig_alarm_ = false;
       processor_->SecondLoop();
+
       if (free_packet_interval == 0) {
         FreePackets();
         free_packet_interval = 10;
       }
       free_packet_interval--;
+
+      overload_ctr -= (overload_ctr != 0);
     }
-    if (n < 0) continue;
 
     if (tun_fd_ >= 0) {
-      tun_readable_ = (FD_ISSET(tun_fd_, &readfds_) != 0);
-      tun_writable_ |= (FD_ISSET(tun_fd_, &writefds_) != 0);
-    }
-    if (udp_fd_ >= 0) {
-      udp_readable_ = (FD_ISSET(udp_fd_, &readfds_) != 0);
-      udp_writable_ |= (FD_ISSET(udp_fd_, &writefds_) != 0);
+      FD_SET(tun_fd_, &readfds_);
+      if (tun_writable_) FD_CLR(tun_fd_, &writefds_); else FD_SET(tun_fd_, &writefds_);
     }
 
-    for(int loop = 0; loop < 256; loop++) {
+    if (udp_fd_ >= 0) {
+      FD_SET(udp_fd_, &readfds_);
+      if (udp_writable_) FD_CLR(udp_fd_, &writefds_); else FD_SET(udp_fd_, &writefds_);
+    }
+
+    n = select(max_fd_, &readfds_, &writefds_, NULL, NULL);
+    if (n == -1) {
+      if (errno != EINTR) {
+        fprintf(stderr, "select failed\n");
+        break;
+      }
+    } else {
+      if (tun_fd_ >= 0) {
+        tun_readable_ = (FD_ISSET(tun_fd_, &readfds_) != 0);
+        tun_writable_ |= (FD_ISSET(tun_fd_, &writefds_) != 0);
+      }
+      if (udp_fd_ >= 0) {
+        udp_readable_ = (FD_ISSET(udp_fd_, &readfds_) != 0);
+        udp_writable_ |= (FD_ISSET(udp_fd_, &writefds_) != 0);
+      }
+    }
+    
+    bool overload = (overload_ctr != 0);
+
+    for(int loop = 0; ; loop++) {
+      // Whenever we don't finish set overload ctr.
+      if (loop == 256) {
+        overload_ctr = 4;
+        break;
+      }
       bool more_work = false;
       if (tun_queue_ != NULL && tun_writable_) more_work |= WriteToTun();
       if (udp_queue_ != NULL && udp_writable_) more_work |= WriteToUdp();
       if (tun_readable_)                       more_work |= ReadFromTun();
-      if (udp_readable_)                       more_work |= ReadFromUdp();
+      if (udp_readable_)                       more_work |= ReadFromUdp(overload);
       if (!more_work)
         break;
     }    
+
+    processor_->RunAllMainThreadScheduled();
   }  
 }
 
diff --git a/network_bsd_mt.cpp b/network_bsd_mt.cpp
index f7134a8..af74540 100644
--- a/network_bsd_mt.cpp
+++ b/network_bsd_mt.cpp
@@ -229,6 +229,7 @@ void *WorkerLoop::ThreadMain() {
         got_sig_alarm_ = false;
         lock_.Release();
         processor_->SecondLoop();
+        processor_->RunAllMainThreadScheduled();
         lock_.Acquire();
       }
       if (shutting_down_ || queue_ != NULL)
@@ -252,6 +253,7 @@ void *WorkerLoop::ThreadMain() {
       }
       packet_queue = next;
     }
+    processor_->RunAllMainThreadScheduled();
     lock_.Acquire();
   }
   lock_.Release();