Merge pull request #2166 from studentmain/master

Register restart after system reboot, #2154

shadowsocks-csharp/Controller/System/AutoStartup.cs

@@ -1,5 +1,8 @@
 using System;
+using System.Collections.Generic;
+using System.Linq;
 using System.Reflection;
+using System.Runtime.InteropServices;
 using System.Windows.Forms;
 using Microsoft.Win32;
 using Shadowsocks.Util;
@@ -20,8 +23,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 +36,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 +49,12 @@ namespace Shadowsocks.Controller
             {
                 if (runKey != null)
                 {
-                    try {
+                    try
+                    {
                         runKey.Close();
                         runKey.Dispose();
-                    } catch (Exception e)
+                    }
+                    catch (Exception e)
                     { Logging.LogUsefulException(e); }
                 }
             }
@@ -58,7 +66,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 +98,54 @@ 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,
+        }
+
+        // register restart after system reboot/update
+        public static void RegisterForRestart(bool register)
+        {
+            // requested register and not autostartup
+            if (register && !Check())
+            {
+                // escape command line parameter
+                string[] args = new List<string>(Program.Args)
+                    .Select(p => p.Replace("\"", "\\\""))                   // escape " to \"
+                    .Select(p => p.IndexOf(" ") >= 0 ? "\"" + p + "\"" : p) // encapsule with "
+                    .ToArray();
+                string cmdline = string.Join(" ", args);
+                // first parameter is process command line parameter
+                // needn't include the name of the executable in the command line
+                RegisterApplicationRestart(cmdline, (int)ApplicationRestartFlags.RESTART_NO_CRASH | (int)ApplicationRestartFlags.RESTART_NO_HANG);
+                Logging.Debug("Register restart after system reboot, command line:" + cmdline);
+            }
+            // requested unregister, which has no side effect
+            else if (!register)
+            {
+                UnregisterApplicationRestart();
+                Logging.Debug("Unregister restart after system reboot");
+            }
+        }
     }
 }

shadowsocks-csharp/Program.cs

@@ -16,13 +16,16 @@ namespace Shadowsocks
     {
         public static ShadowsocksController MainController { get; private set; }
         public static MenuViewController MenuController { get; private set; }
-
+        public static string[] Args { get; private set; }
         /// <summary>
         /// 应用程序的主入口点。
         /// </summary>
+        /// </summary>
         [STAThread]
-        static void Main()
+        static void Main(string[] args)
         {
+            // store args for further use
+            Args = args;
             // Check OS since we are using dual-mode socket
             if (!Utils.IsWinVistaOrHigher())
             {
@@ -54,6 +57,7 @@ namespace Shadowsocks
                 SystemEvents.PowerModeChanged += SystemEvents_PowerModeChanged;
                 Application.EnableVisualStyles();
                 Application.SetCompatibleTextRenderingDefault(false);
+                AutoStartup.RegisterForRestart(true);
 
                 if (!mutex.WaitOne(0, false))
                 {