Register restart after system reboot, #2154

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");
+            }
+        }
     }
 }

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))
                 {