From a49e280e033238563475ce579bc695c5c332eb5e Mon Sep 17 00:00:00 2001
From: ChsBuffer <33744752+chsbuffer@users.noreply.github.com>
Date: Sun, 11 Oct 2020 00:51:41 +0800
Subject: [PATCH] refactor: NetTraffic record pid fix: #399

---
 Netch/Controllers/MainController.cs       |  2 ++
 Netch/Servers/Shadowsocks/SSController.cs | 11 +++++----
 Netch/Utils/Bandwidth.cs                  | 28 +++++++++++++++--------
 3 files changed, 27 insertions(+), 14 deletions(-)

diff --git a/Netch/Controllers/MainController.cs b/Netch/Controllers/MainController.cs
index 7faf0d15..42da26c5 100644
--- a/Netch/Controllers/MainController.cs
+++ b/Netch/Controllers/MainController.cs
@@ -184,6 +184,8 @@ namespace Netch.Controllers
                 Task.Run(() => ModeController?.Stop()),
             };
             await Task.WhenAll(tasks);
+            ModeController = null;
+            ServerController = null;
         }
 
 
diff --git a/Netch/Servers/Shadowsocks/SSController.cs b/Netch/Servers/Shadowsocks/SSController.cs
index 1ba9475d..5584b3fc 100644
--- a/Netch/Servers/Shadowsocks/SSController.cs
+++ b/Netch/Servers/Shadowsocks/SSController.cs
@@ -14,16 +14,16 @@ namespace Netch.Servers.Shadowsocks
         public int? Socks5LocalPort { get; set; }
         public string LocalAddress { get; set; }
 
+        private Mode _savedMode;
+        public bool DllFlag => Global.Settings.BootShadowsocksFromDLL && (_savedMode.Type == 0 || _savedMode.Type == 1 || _savedMode.Type == 2);
+
         public bool Start(Server s, Mode mode)
         {
-            bool DllFlag()
-            {
-                return Global.Settings.BootShadowsocksFromDLL && (mode.Type == 0 || mode.Type == 1 || mode.Type == 2);
-            }
+            _savedMode = mode;
 
             var server = (Shadowsocks) s;
             //从DLL启动Shaowsocks
-            if (DllFlag())
+            if (DllFlag)
             {
                 State = State.Starting;
                 var client = Encoding.UTF8.GetBytes($"{LocalAddress ?? Global.Settings.LocalAddress}:{Socks5LocalPort ?? Global.Settings.Socks5LocalPort}");
@@ -79,6 +79,7 @@ namespace Netch.Servers.Shadowsocks
                 ShadowsocksDLL.Stop();
             else
                 StopInstance();
+            _savedMode = null;
         }
 
 
diff --git a/Netch/Utils/Bandwidth.cs b/Netch/Utils/Bandwidth.cs
index 8d9a6c87..efcdae3a 100644
--- a/Netch/Utils/Bandwidth.cs
+++ b/Netch/Utils/Bandwidth.cs
@@ -7,6 +7,7 @@ using Microsoft.Diagnostics.Tracing.Parsers;
 using Microsoft.Diagnostics.Tracing.Session;
 using Netch.Controllers;
 using Netch.Models;
+using Netch.Servers.Shadowsocks;
 
 namespace Netch.Utils
 {
@@ -66,20 +67,29 @@ namespace Netch.Utils
             {
                 instances.Add(((HTTPController) MainController.ModeController).pPrivoxyController.Instance);
             }
-            else if (server.Type.Equals("SS") && Global.Settings.BootShadowsocksFromDLL &&
-                     (mode.Type == 0 || mode.Type == 1 || mode.Type == 2))
-            {
-                instances.Add(Process.GetCurrentProcess());
-            }
             else if (MainController.ServerController != null)
             {
-                if (MainController.ServerController is Guard instanceController)
-                    instances.Add(instanceController.Instance);
+                switch (MainController.ServerController)
+                {
+                    case SSController ssController when ssController.DllFlag:
+                        instances.Add(Process.GetCurrentProcess());
+                        break;
+                    case Guard instanceController:
+                        instances.Add(instanceController.Instance);
+                        break;
+                }
             }
             else if (MainController.ModeController != null)
             {
-                if (MainController.ModeController is Guard instanceController)
-                    instances.Add(instanceController.Instance);
+                switch (MainController.ModeController)
+                {
+                    case NFController _:
+                        instances.Add(Process.GetCurrentProcess());
+                        break;
+                    case Guard instanceController:
+                        instances.Add(instanceController.Instance);
+                        break;
+                }
             }
 
             var processList = instances.Select(instance => instance.Id).ToList();