rdctl: Use new function when terminating RD for factory-reset

Since we needed a generalized function to deal with terminating extensions,
we might as well re-use it for terminating the process for the rest of
factory reset.

Signed-off-by: Mark Yen <mark.yen@suse.com>

src/go/rdctl/pkg/directories/directories.go

@@ -0,0 +1,53 @@
+/*
+Copyright © 2024 SUSE LLC
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+		http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package directories
+
+import (
+	"os"
+	"path/filepath"
+	"runtime"
+)
+
+// GetApplicationDirectory returns the installation directory of the application.
+func GetApplicationDirectory() (string, error) {
+	exePathWithSymlinks, err := os.Executable()
+	if err != nil {
+		return "", err
+	}
+
+	exePath, err := filepath.EvalSymlinks(exePathWithSymlinks)
+	if err != nil {
+		return "", err
+	}
+
+	platform := runtime.GOOS
+	if runtime.GOOS == "windows" {
+		// On Windows, we use "win32" instead of "windows".
+		platform = "win32"
+	}
+
+	// Given the path to the exe, find its directory, and drop the
+	// "resources\win32\bin" suffix (possibly with another "resources" in front).
+	// On mac, we need to drop "Contents/Resources/resources/darwin/bin".
+	resultDir := filepath.Dir(exePath)
+	for _, part := range []string{"bin", platform, "resources", "Resources", "Contents"} {
+		for filepath.Base(resultDir) == part {
+			resultDir = filepath.Dir(resultDir)
+		}
+	}
+	return resultDir, nil
+}

src/go/rdctl/pkg/directories/directories_test.go

@@ -0,0 +1,31 @@
+/*
+Copyright © 2024 SUSE LLC
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+		http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package directories_test
+
+import (
+	"testing"
+
+	"github.com/rancher-sandbox/rancher-desktop/src/go/rdctl/pkg/directories"
+	"github.com/stretchr/testify/assert"
+)
+
+func TestGetApplicationDirectory(t *testing.T) {
+	_, err := directories.GetApplicationDirectory()
+	assert.NoError(t, err)
+	// `go test` makes a temporary directory, so we can't sensibly test the
+	// return value.
+}

src/go/rdctl/pkg/directories/directories_windows.go

@@ -19,7 +19,6 @@ package directories
 import (
 	"errors"
 	"fmt"
-	"path/filepath"
 	"unsafe"
 
 	"golang.org/x/sys/windows"
@@ -44,43 +43,6 @@ func InvokeWin32WithBuffer(cb func(size int) error) error {
 	}
 }
 
-// GetApplicationDirectory returns the installation directory of the application.
-func GetApplicationDirectory() (string, error) {
-	var exePath string
-	err := InvokeWin32WithBuffer(func(bufSize int) error {
-		buf := make([]uint16, bufSize)
-		n, err := windows.GetModuleFileName(windows.Handle(0), &buf[0], uint32(bufSize))
-		if err != nil {
-			return err
-		}
-		if n == uint32(bufSize) {
-			// If the buffer is too small, GetModuleFileName returns the buffer size,
-			// and the result includes the null character. If the buffer is large
-			// enough, GetModuleFileName returns the string length, _excluding_ the
-			// null character.
-			if buf[bufSize-1] == 0 {
-				// The buffer contains a null character
-				return windows.ERROR_INSUFFICIENT_BUFFER
-			}
-		}
-		exePath = windows.UTF16ToString(buf[:n])
-		return nil
-	})
-	if err != nil {
-		return "", err
-	}
-
-	// Given the path to the exe, find its directory, and drop the
-	// "resources\win32\bin" suffix (possibly with another "resources" in front).
-	resultDir := filepath.Dir(exePath)
-	for _, part := range []string{"bin", "win32", "resources"} {
-		for filepath.Base(resultDir) == part {
-			resultDir = filepath.Dir(resultDir)
-		}
-	}
-	return resultDir, nil
-}
-
 func GetLocalAppDataDirectory() (string, error) {
 	dir, err := getKnownFolder(windows.FOLDERID_LocalAppData)
 	if err != nil {

src/go/rdctl/pkg/directories/directories_windows_test.go

@@ -23,13 +23,6 @@ import (
 	"golang.org/x/sys/windows"
 )
 
-func TestGetApplicationDirectory(t *testing.T) {
-	_, err := GetApplicationDirectory()
-	assert.NoError(t, err)
-	// `go test` makes a temporary directory, so we can't sensibly test the
-	// return value.
-}
-
 func TestGetKnownFolder(t *testing.T) {
 	t.Run("AppData", func(t *testing.T) {
 		expected := os.Getenv("APPDATA")

src/go/rdctl/pkg/factoryreset/delete_data_darwin.go

@@ -16,7 +16,7 @@ func DeleteData(ctx context.Context, appPaths paths.Paths, removeKubernetesCache
 		logrus.Errorf("Failed to remove autostart configuration: %s", err)
 	}
 
-	if err := process.TerminateProcessInDirectory(ctx, appPaths.ExtensionRoot); err != nil {
+	if err := process.TerminateProcessInDirectory(appPaths.ExtensionRoot, false); err != nil {
 		logrus.Errorf("Failed to stop extension processes, ignoring: %s", err)
 	}
 

src/go/rdctl/pkg/factoryreset/delete_data_linux.go

@@ -21,7 +21,7 @@ func DeleteData(ctx context.Context, appPaths paths.Paths, removeKubernetesCache
 		logrus.Errorf("Error getting home directory: %s", err)
 	}
 
-	if err := process.TerminateProcessInDirectory(ctx, appPaths.ExtensionRoot); err != nil {
+	if err := process.TerminateProcessInDirectory(appPaths.ExtensionRoot, false); err != nil {
 		logrus.Errorf("Failed to stop extension processes, ignoring: %s", err)
 	}
 

src/go/rdctl/pkg/factoryreset/delete_data_windows.go

@@ -19,7 +19,7 @@ func DeleteData(ctx context.Context, appPaths paths.Paths, removeKubernetesCache
 		logrus.Errorf("could not unregister WSL: %s", err)
 		return err
 	}
-	if err := process.TerminateProcessInDirectory(ctx, appPaths.ExtensionRoot); err != nil {
+	if err := process.TerminateProcessInDirectory(appPaths.ExtensionRoot, false); err != nil {
 		logrus.Errorf("Failed to stop extension processes, ignoring: %s", err)
 	}
 	if err := deleteWindowsData(!removeKubernetesCache, "rancher-desktop"); err != nil {

src/go/rdctl/pkg/factoryreset/factory_reset_windows.go

@@ -25,20 +25,13 @@ import (
 	"os"
 	"os/exec"
 	"path/filepath"
-	"sort"
-	"strings"
-	"unsafe"
 
 	"github.com/rancher-sandbox/rancher-desktop/src/go/rdctl/pkg/directories"
+	"github.com/rancher-sandbox/rancher-desktop/src/go/rdctl/pkg/process"
 	"github.com/sirupsen/logrus"
 	"golang.org/x/sys/windows"
 )
 
-var (
-	pKernel32      = windows.NewLazySystemDLL("kernel32.dll")
-	pEnumProcesses = pKernel32.NewProc("K32EnumProcesses")
-)
-
 // CheckProcessWindows - returns true if Rancher Desktop is still running, false if it isn't
 // along with an error condition if there's a problem detecting that.
 //
@@ -75,106 +68,9 @@ func KillRancherDesktop() error {
 		return fmt.Errorf("could not find application directory: %w", err)
 	}
 
-	var processes []uint32
-	err = directories.InvokeWin32WithBuffer(func(size int) error {
-		processes = make([]uint32, size)
-		var bytesReturned uint32
-		// We can't use `windows.EnumProcesses`, because it passes in an incorrect
-		// value for the second argument (`cb`).
-		elementSize := unsafe.Sizeof(uint32(0))
-		bufferSize := uintptr(len(processes)) * elementSize
-		n, _, err := pEnumProcesses.Call(
-			uintptr(unsafe.Pointer(&processes[0])),
-			bufferSize,
-			uintptr(unsafe.Pointer(&bytesReturned)),
-		)
-		if n == 0 {
-			return err
-		}
-		if uintptr(bytesReturned) >= bufferSize {
-			return windows.ERROR_INSUFFICIENT_BUFFER
-		}
-		processesFound := uintptr(bytesReturned) / elementSize
-		logrus.Tracef("got %d processes", processesFound)
-		processes = processes[:processesFound]
-		return nil
-	})
+	err = process.TerminateProcessInDirectory(appDir, true)
 	if err != nil {
-		return fmt.Errorf("could not get process list: %w", err)
-	}
-
-	sort.Slice(processes, func(i, j int) bool {
-		return processes[i] < processes[j]
-	})
-	var processesToKill []uint32
-	for _, pid := range processes {
-		// Add a scope to help with defer
-		(func(pid uint32) {
-			if pid == uint32(os.Getpid()) {
-				// Skip the current process.
-				return
-			}
-
-			hProc, err := windows.OpenProcess(windows.PROCESS_QUERY_LIMITED_INFORMATION, false, pid)
-			if err != nil {
-				// We can't open privileged processes, processes that have exited since,
-				// idle process, etc.; so we log this at trace level instead.
-				logrus.Tracef("failed to open pid %d: %s (skipping)", pid, err)
-				return
-			}
-			defer func() { _ = windows.CloseHandle(hProc) }()
-
-			var imageName string
-			err = directories.InvokeWin32WithBuffer(func(size int) error {
-				nameBuf := make([]uint16, size)
-				charsWritten := uint32(size)
-				err := windows.QueryFullProcessImageName(hProc, 0, &nameBuf[0], &charsWritten)
-				if err != nil {
-					logrus.Tracef("failed to get image name for pid %d: %s", pid, err)
-					return err
-				}
-				if charsWritten >= uint32(size)-1 {
-					logrus.Tracef("buffer too small for pid %d image name", pid)
-					return windows.ERROR_INSUFFICIENT_BUFFER
-				}
-				imageName = windows.UTF16ToString(nameBuf)
-				return nil
-			})
-			if err != nil {
-				logrus.Debugf("failed to get process name of pid %d: %s (skipping)", pid, err)
-				return
-			}
-
-			relPath, err := filepath.Rel(appDir, imageName)
-			if err != nil {
-				// This may be because they're on different drives, network shares, etc.
-				logrus.Tracef("failed to make pid %d image %s relative to %s: %s", pid, imageName, appDir, err)
-				return
-			}
-			if strings.HasPrefix(relPath, "..") {
-				// Relative path includes "../" prefix, not a child of appDir
-				logrus.Tracef("skipping pid %d (%s), not in app %s", pid, imageName, appDir)
-				return
-			}
-
-			logrus.Tracef("will terminate pid %d image %s", pid, imageName)
-			processesToKill = append(processesToKill, pid)
-		})(pid)
-	}
-
-	for _, pid := range processesToKill {
-		(func() {
-			hProc, err := windows.OpenProcess(windows.PROCESS_TERMINATE, false, pid)
-			if err != nil {
-				logrus.Infof("failed to open process %d for termination, skipping", pid)
-				return
-			}
-			defer func() { _ = windows.CloseHandle(hProc) }()
-
-			if err = windows.TerminateProcess(hProc, 0); err != nil {
-				logrus.Infof("failed to terminate process %d: %s", pid, err)
-			}
-		})()
+		return err
 	}
 
 	return nil

src/go/rdctl/pkg/process/process_darwin.go

@@ -1,7 +1,6 @@
 package process
 
 import (
-	"context"
 	"errors"
 	"fmt"
 	"os"
@@ -19,14 +18,19 @@ const (
 )
 
 // TerminateProcessInDirectory terminates all processes where the executable
-// resides within the given directory, as gracefully as possible.
-func TerminateProcessInDirectory(ctx context.Context, directory string) error {
+// resides within the given directory, as gracefully as possible.  If `force` is
+// set, SIGKILL is used instead.
+func TerminateProcessInDirectory(directory string, force bool) error {
 	procs, err := unix.SysctlKinfoProcSlice("kern.proc.all")
 	if err != nil {
 		return fmt.Errorf("failed to list processes: %w", err)
 	}
 	for _, proc := range procs {
 		pid := int(proc.Proc.P_pid)
+		// Don't kill the current process
+		if pid == os.Getpid() {
+			continue
+		}
 		buf, err := unix.SysctlRaw(CTL_KERN, KERN_PROCARGS, pid)
 		if err != nil {
 			if !errors.Is(err, unix.EINVAL) {
@@ -50,7 +54,11 @@ func TerminateProcessInDirectory(ctx context.Context, directory string) error {
 		if err != nil {
 			continue
 		}
-		err = process.Signal(unix.SIGTERM)
+		if force {
+			err = process.Signal(unix.SIGKILL)
+		} else {
+			err = process.Signal(unix.SIGTERM)
+		}
 		if err == nil {
 			logrus.Infof("Terminated process %d (%s)", pid, procPath)
 		} else if !errors.Is(err, unix.EINVAL) {

src/go/rdctl/pkg/process/process_linux.go

@@ -1,7 +1,6 @@
 package process
 
 import (
-	"context"
 	"errors"
 	"fmt"
 	"os"
@@ -14,8 +13,9 @@ import (
 )
 
 // TerminateProcessInDirectory terminates all processes where the executable
-// resides within the given directory, as gracefully as possible.
-func TerminateProcessInDirectory(ctx context.Context, directory string) error {
+// resides within the given directory, as gracefully as possible.  If `force` is
+// set, SIGKILL is used instead.
+func TerminateProcessInDirectory(directory string, force bool) error {
 	// Check /proc/<pid>/exe to see if they're the correct file.
 	pidfds, err := os.ReadDir("/proc")
 	if err != nil {
@@ -29,6 +29,10 @@ func TerminateProcessInDirectory(ctx context.Context, directory string) error {
 		if err != nil {
 			continue
 		}
+		// Don't kill the current process
+		if pid == os.Getpid() {
+			continue
+		}
 		procPath, err := os.Readlink(filepath.Join("/proc", pidfd.Name(), "exe"))
 		if err != nil {
 			continue
@@ -41,7 +45,11 @@ func TerminateProcessInDirectory(ctx context.Context, directory string) error {
 		if err != nil {
 			continue
 		}
-		err = proc.Signal(unix.SIGTERM)
+		if force {
+			err = proc.Signal(unix.SIGKILL)
+		} else {
+			err = proc.Signal(unix.SIGTERM)
+		}
 		if err == nil {
 			logrus.Infof("Terminated process %d", pid)
 		} else if !errors.Is(err, unix.EINVAL) {