Merge pull request #9 from safing/develop

Fix clean shutdown with portmaster-control

main.go

@@ -31,7 +31,7 @@ func init() {
 func main() {
 
 	// Set Info
-	info.Set("Portmaster", "0.3.0", "AGPLv3", true)
+	info.Set("Portmaster", "0.3.1", "AGPLv3", true)
 
 	// Start
 	err := modules.Start()

pmctl/get.go

@@ -1,46 +1,61 @@
 package main
 
 import (
+	"errors"
 	"fmt"
-	"os"
+	"time"
 
-	"github.com/safing/portbase/utils"
 	"github.com/safing/portmaster/updates"
 )
 
-func getFile(identifier string) (*updates.File, error) {
+func getFile(opts *Options) (*updates.File, error) {
 	// get newest local file
 	updates.LoadLatest()
-	file, err := updates.GetPlatformFile(identifier)
+
+	file, err := updates.GetLocalPlatformFile(opts.Identifier)
 	if err == nil {
 		return file, nil
 	}
 	if err != updates.ErrNotFound {
 		return nil, err
 	}
 
-	fmt.Printf("%s downloading %s...\n", logPrefix, identifier)
-
-	// if no matching file exists, load index
-	err = updates.LoadIndexes()
-	if err != nil {
-		if os.IsNotExist(err) {
-			// create dirs
-			err = utils.EnsureDirectory(updateStoragePath, 0755)
-			if err != nil {
-				return nil, err
-			}
-
-			// download indexes
-			err = updates.CheckForUpdates()
-			if err != nil {
-				return nil, err
-			}
-		} else {
+	// download
+	if opts.AllowDownload {
+		fmt.Printf("%s downloading %s...\n", logPrefix, opts.Identifier)
+
+		// download indexes
+		err = updates.UpdateIndexes()
+		if err != nil {
 			return nil, err
 		}
+
+		// download file
+		file, err := updates.GetPlatformFile(opts.Identifier)
+		if err != nil {
+			return nil, err
+		}
+		return file, nil
 	}
 
-	// get file
-	return updates.GetPlatformFile(identifier)
+	// wait for 30 seconds
+	fmt.Printf("%s waiting for download of %s (by Portmaster Core) to complete...\n", logPrefix, opts.Identifier)
+
+	// try every 0.5 secs
+	for tries := 0; tries < 60; tries++ {
+		time.Sleep(500 * time.Millisecond)
+
+		// reload local files
+		updates.LoadLatest()
+
+		// get file
+		file, err := updates.GetLocalPlatformFile(opts.Identifier)
+		if err == nil {
+			return file, nil
+		}
+		if err != updates.ErrNotFound {
+			return nil, err
+		}
+	}
+	return nil, errors.New("please try again later or check the Portmaster logs")
 }

pmctl/main.go

@@ -5,7 +5,10 @@ import (
 	"flag"
 	"fmt"
 	"os"
+	"os/user"
 	"path/filepath"
+	"runtime"
+	"strings"
 
 	"github.com/safing/portbase/info"
 	"github.com/safing/portbase/log"
@@ -55,7 +58,7 @@ func main() {
 	// }()
 
 	// set meta info
-	info.Set("Portmaster Control", "0.2.0", "AGPLv3", true)
+	info.Set("Portmaster Control", "0.2.1", "AGPLv3", true)
 
 	// check if meta info is ok
 	err := info.CheckVersion()
@@ -86,7 +89,23 @@ func initPmCtl(cmd *cobra.Command, args []string) error {
 		return errors.New("please supply the database directory using the --db flag")
 	}
 
-	err := removeOldBin()
+	// check if we are root/admin for self upgrade
+	userInfo, err := user.Current()
+	if err != nil {
+		return nil
+	}
+	switch runtime.GOOS {
+	case "linux":
+		if userInfo.Username != "root" {
+			return nil
+		}
+	case "windows":
+		if !strings.HasSuffix(userInfo.Username, "SYSTEM") { // is this correct?
+			return nil
+		}
+	}
+
+	err = removeOldBin()
 	if err != nil {
 		fmt.Printf("%s warning: failed to remove old upgrade: %s\n", logPrefix, err)
 	}

pmctl/run.go

@@ -5,12 +5,20 @@ import (
 	"io"
 	"os"
 	"os/exec"
+	"os/signal"
 	"runtime"
 	"strings"
+	"syscall"
 
 	"github.com/spf13/cobra"
 )
 
+// Options for starting component
+type Options struct {
+	Identifier    string
+	AllowDownload bool
+}
+
 func init() {
 	rootCmd.AddCommand(runCmd)
 	runCmd.AddCommand(runCore)
@@ -27,7 +35,10 @@ var runCore = &cobra.Command{
 	Use:   "core",
 	Short: "Run the Portmaster Core",
 	RunE: func(cmd *cobra.Command, args []string) error {
-		return run("core/portmaster-core", cmd, false)
+		return run(cmd, &Options{
+			Identifier:    "core/portmaster-core",
+			AllowDownload: true,
+		})
 	},
 	FParseErrWhitelist: cobra.FParseErrWhitelist{
 		// UnknownFlags will ignore unknown flags errors and continue parsing rest of the flags
@@ -39,7 +50,10 @@ var runApp = &cobra.Command{
 	Use:   "app",
 	Short: "Run the Portmaster App",
 	RunE: func(cmd *cobra.Command, args []string) error {
-		return run("app/portmaster-app", cmd, true)
+		return run(cmd, &Options{
+			Identifier:    "app/portmaster-app",
+			AllowDownload: false,
+		})
 	},
 	FParseErrWhitelist: cobra.FParseErrWhitelist{
 		// UnknownFlags will ignore unknown flags errors and continue parsing rest of the flags
@@ -51,94 +65,71 @@ var runNotifier = &cobra.Command{
 	Use:   "notifier",
 	Short: "Run the Portmaster Notifier",
 	RunE: func(cmd *cobra.Command, args []string) error {
-		return run("notifier/portmaster-notifier", cmd, true)
+		return run(cmd, &Options{
+			Identifier:    "notifier/portmaster-notifier",
+			AllowDownload: false,
+		})
 	},
 	FParseErrWhitelist: cobra.FParseErrWhitelist{
 		// UnknownFlags will ignore unknown flags errors and continue parsing rest of the flags
 		UnknownFlags: true,
 	},
 }
 
-func run(identifier string, cmd *cobra.Command, filterDatabaseFlag bool) error {
+func run(cmd *cobra.Command, opts *Options) error {
 
 	// get original arguments
-	if len(os.Args) <= 3 {
-		return cmd.Help()
-	}
 	var args []string
-
-	// filter out database flag
-	if filterDatabaseFlag {
-		skip := false
-		for _, arg := range os.Args[3:] {
-			if skip {
-				skip = false
-				continue
-			}
-
-			if arg == "--db" {
-				// flag is seperated, skip two arguments
-				skip = true
-				continue
-			}
-
-			if strings.HasPrefix(arg, "--db=") {
-				// flag is one string, skip one argument
-				continue
-			}
-
-			args = append(args, arg)
-		}
-	} else {
-		args = os.Args[3:]
+	if len(os.Args) < 4 {
+		return cmd.Help()
 	}
+	args = os.Args[3:]
 
 	// adapt identifier
 	if windows() {
-		identifier += ".exe"
+		opts.Identifier += ".exe"
 	}
 
 	// run
 	for {
-		file, err := getFile(identifier)
+		file, err := getFile(opts)
 		if err != nil {
-			return fmt.Errorf("%s could not get component: %s", logPrefix, err)
+			return fmt.Errorf("could not get component: %s", err)
 		}
 
 		// check permission
 		if !windows() {
 			info, err := os.Stat(file.Path())
 			if err != nil {
-				return fmt.Errorf("%s failed to get file info on %s: %s", logPrefix, file.Path(), err)
+				return fmt.Errorf("failed to get file info on %s: %s", file.Path(), err)
 			}
 			if info.Mode() != 0755 {
 				err := os.Chmod(file.Path(), 0755)
 				if err != nil {
-					return fmt.Errorf("%s failed to set exec permissions on %s: %s", logPrefix, file.Path(), err)
+					return fmt.Errorf("failed to set exec permissions on %s: %s", file.Path(), err)
 				}
 			}
 		}
 
 		fmt.Printf("%s starting %s %s\n", logPrefix, file.Path(), strings.Join(args, " "))
-		// os.Exit(0)
 
 		// create command
 		exc := exec.Command(file.Path(), args...)
 
 		// consume stdout/stderr
 		stdout, err := exc.StdoutPipe()
 		if err != nil {
-			return fmt.Errorf("%s failed to connect stdout: %s", logPrefix, err)
+			return fmt.Errorf("failed to connect stdout: %s", err)
 		}
 		stderr, err := exc.StderrPipe()
 		if err != nil {
-			return fmt.Errorf("%s failed to connect stderr: %s", logPrefix, err)
+			return fmt.Errorf("failed to connect stderr: %s", err)
 		}
 
 		// start
 		err = exc.Start()
 		if err != nil {
-			return fmt.Errorf("%s failed to start %s: %s", logPrefix, identifier, err)
+			return fmt.Errorf("failed to start %s: %s", opts.Identifier, err)
 		}
 
 		// start output writers
@@ -149,6 +140,24 @@ func run(identifier string, cmd *cobra.Command, filterDatabaseFlag bool) error {
 			io.Copy(os.Stderr, stderr)
 		}()
 
+		// catch interrupt for clean shutdown
+		signalCh := make(chan os.Signal)
+		signal.Notify(
+			signalCh,
+			os.Interrupt,
+			os.Kill,
+			syscall.SIGHUP,
+			syscall.SIGINT,
+			syscall.SIGTERM,
+			syscall.SIGQUIT,
+		)
+		go func() {
+			for {
+				sig := <-signalCh
+				fmt.Printf("%s got %s signal (ignoring), waiting for %s to exit...\n", logPrefix, sig, opts.Identifier)
+			}
+		}()
+
 		// wait for completion
 		err = exc.Wait()
 		if err != nil {
@@ -157,31 +166,30 @@ func run(identifier string, cmd *cobra.Command, filterDatabaseFlag bool) error {
 				switch exErr.ProcessState.ExitCode() {
 				case 0:
 					// clean exit
-					fmt.Printf("%s clean exit of %s, but with error: %s\n", logPrefix, identifier, err)
+					fmt.Printf("%s clean exit of %s, but with error: %s\n", logPrefix, opts.Identifier, err)
 					os.Exit(1)
 				case 1:
 					// error exit
-					fmt.Printf("%s error during execution of %s: %s\n", logPrefix, identifier, err)
+					fmt.Printf("%s error during execution of %s: %s\n", logPrefix, opts.Identifier, err)
 					os.Exit(1)
 				case 2357427: // Leet Speak for "restart"
 					// restart request
-					fmt.Printf("%s restarting %s\n", logPrefix, identifier)
+					fmt.Printf("%s restarting %s\n", logPrefix, opts.Identifier)
 					continue
 				default:
-					fmt.Printf("%s unexpected error during execution of %s: %s\n", logPrefix, identifier, err)
+					fmt.Printf("%s unexpected error during execution of %s: %s\n", logPrefix, opts.Identifier, err)
 					os.Exit(exErr.ProcessState.ExitCode())
 				}
 			} else {
-				fmt.Printf("%s unexpected error type during execution of %s: %s\n", logPrefix, identifier, err)
+				fmt.Printf("%s unexpected error type during execution of %s: %s\n", logPrefix, opts.Identifier, err)
 				os.Exit(1)
 			}
 		}
-
 		// clean exit
 		break
 	}
 
-	fmt.Printf("%s %s completed successfully\n", logPrefix, identifier)
+	fmt.Printf("%s %s completed successfully\n", logPrefix, opts.Identifier)
 	return nil
 }
 

pmctl/upgrade.go

@@ -11,6 +11,10 @@ import (
 	"github.com/safing/portmaster/updates"
 )
 
+var (
+	oldBinSuffix = "-old"
+)
+
 func checkForUpgrade() (update *updates.File) {
 	info := info.GetInfo()
 	file, err := updates.GetLocalPlatformFile("control/portmaster-control")
@@ -25,6 +29,8 @@ func checkForUpgrade() (update *updates.File) {
 
 func doSelfUpgrade(file *updates.File) error {
 
+	// FIXME: fix permissions if needed
+
 	// get destination
 	dst, err := os.Executable()
 	if err != nil {
@@ -36,7 +42,7 @@ func doSelfUpgrade(file *updates.File) error {
 	}
 
 	// mv destination
-	err = os.Rename(dst, dst+"_old")
+	err = os.Rename(dst, dst+oldBinSuffix)
 	if err != nil {
 		return err
 	}
@@ -105,7 +111,7 @@ func removeOldBin() error {
 	}
 
 	// delete old
-	err = os.Remove(dst + "_old")
+	err = os.Remove(dst + oldBinSuffix)
 	if err != nil {
 		if !os.IsNotExist(err) {
 			return err