From 5e0ccffa9680b28d1e4bd574e852ea9c76bcce87 Mon Sep 17 00:00:00 2001
From: Student Main <scncgz@gmail.com>
Date: Thu, 17 Jan 2019 16:55:46 +0800
Subject: [PATCH] Register restart after system reboot, #2154

---
 .../Controller/System/AutoStartup.cs          | 55 ++++++++++++++++---
 shadowsocks-csharp/Program.cs                 |  1 +
 2 files changed, 49 insertions(+), 7 deletions(-)

diff --git a/shadowsocks-csharp/Controller/System/AutoStartup.cs b/shadowsocks-csharp/Controller/System/AutoStartup.cs
index 7c7b0dc4..923c16f5 100644
--- a/shadowsocks-csharp/Controller/System/AutoStartup.cs
+++ b/shadowsocks-csharp/Controller/System/AutoStartup.cs
@@ -1,5 +1,6 @@
 using System;
 using System.Reflection;
+using System.Runtime.InteropServices;
 using System.Windows.Forms;
 using Microsoft.Win32;
 using Shadowsocks.Util;
@@ -20,8 +21,9 @@ namespace Shadowsocks.Controller
             try
             {
                 runKey = Utils.OpenRegKey(@"Software\Microsoft\Windows\CurrentVersion\Run", true);
-                if ( runKey == null ) {
-                    Logging.Error( @"Cannot find HKCU\Software\Microsoft\Windows\CurrentVersion\Run" );
+                if (runKey == null)
+                {
+                    Logging.Error(@"Cannot find HKCU\Software\Microsoft\Windows\CurrentVersion\Run");
                     return false;
                 }
                 if (enabled)
@@ -32,6 +34,8 @@ namespace Shadowsocks.Controller
                 {
                     runKey.DeleteValue(Key);
                 }
+                // When autostartup setting change, change RegisterForRestart state to avoid start 2 times
+                RegisterForRestart(!enabled);
                 return true;
             }
             catch (Exception e)
@@ -43,10 +47,12 @@ namespace Shadowsocks.Controller
             {
                 if (runKey != null)
                 {
-                    try {
+                    try
+                    {
                         runKey.Close();
                         runKey.Dispose();
-                    } catch (Exception e)
+                    }
+                    catch (Exception e)
                     { Logging.LogUsefulException(e); }
                 }
             }
@@ -58,7 +64,8 @@ namespace Shadowsocks.Controller
             try
             {
                 runKey = Utils.OpenRegKey(@"Software\Microsoft\Windows\CurrentVersion\Run", true);
-                if (runKey == null) {
+                if (runKey == null)
+                {
                     Logging.Error(@"Cannot find HKCU\Software\Microsoft\Windows\CurrentVersion\Run");
                     return false;
                 }
@@ -89,13 +96,47 @@ namespace Shadowsocks.Controller
             {
                 if (runKey != null)
                 {
-                    try {
+                    try
+                    {
                         runKey.Close();
                         runKey.Dispose();
-                    } catch (Exception e)
+                    }
+                    catch (Exception e)
                     { Logging.LogUsefulException(e); }
                 }
             }
         }
+
+        [DllImport("kernel32.dll", SetLastError = true)]
+        static extern int RegisterApplicationRestart([MarshalAs(UnmanagedType.LPWStr)] string commandLineArgs, int Flags);
+
+        [DllImport("kernel32.dll", SetLastError = true)]
+        static extern int UnregisterApplicationRestart();
+
+        enum ApplicationRestartFlags
+        {
+            RESTART_NO_CRASH = 1,
+            RESTART_NO_HANG = 2,
+            RESTART_NO_PATCH = 4,
+            RESTART_NO_REBOOT = 8,
+        }
+
+        // regist restart after system reboot/update
+        public static void RegisterForRestart(bool regist)
+        {
+            // requested regist and not autostart
+            if (regist && !Check())
+            {
+                // param 1 is process command param
+                RegisterApplicationRestart(null, (int)ApplicationRestartFlags.RESTART_NO_CRASH | (int)ApplicationRestartFlags.RESTART_NO_HANG);
+                Logging.Debug("Register restart after system reboot");
+            }
+            // request unregist, unregist has no side effect
+            else if (!regist)
+            {
+                UnregisterApplicationRestart();
+                Logging.Debug("Unregister restart after system reboot");
+            }
+        }
     }
 }
diff --git a/shadowsocks-csharp/Program.cs b/shadowsocks-csharp/Program.cs
index 87da70c6..8914ccef 100755
--- a/shadowsocks-csharp/Program.cs
+++ b/shadowsocks-csharp/Program.cs
@@ -54,6 +54,7 @@ namespace Shadowsocks
                 SystemEvents.PowerModeChanged += SystemEvents_PowerModeChanged;
                 Application.EnableVisualStyles();
                 Application.SetCompatibleTextRenderingDefault(false);
+                AutoStartup.RegisterForRestart(true);
 
                 if (!mutex.WaitOne(0, false))
                 {