From dabbf2037b60e3eff03aa6a5ec8c0a80b442c17c Mon Sep 17 00:00:00 2001
From: Cuong Manh Le <cuong@windscribe.com>
Date: Mon, 19 Feb 2024 15:26:10 +0700
Subject: [PATCH] cmd/cli: do not allow running start command if pin code set

While at it, also emitting a better error message when pin code was set
but users do not provide --pin flag.
---
 cmd/cli/cli.go            | 33 +++++++++++++++++++++++++++------
 cmd/cli/control_server.go |  9 +++++++--
 2 files changed, 34 insertions(+), 8 deletions(-)

diff --git a/cmd/cli/cli.go b/cmd/cli/cli.go
index e97b53c..e7253c2 100644
--- a/cmd/cli/cli.go
+++ b/cmd/cli/cli.go
@@ -259,6 +259,13 @@ func initCLI() {
 				return
 			}
 
+			// If pin code was set, do not allow running start command.
+			if status, _ := s.Status(); status == service.StatusRunning {
+				if err := checkDeactivationPin(s); isCheckDeactivationPinErr(err) {
+					os.Exit(deactivationPinInvalidExitCode)
+				}
+			}
+
 			if router.Name() != "" && iface != "" {
 				mainLog.Load().Debug().Msg("cleaning up router before installing")
 				_ = p.router.Cleanup()
@@ -393,7 +400,7 @@ func initCLI() {
 				return
 			}
 			initLogging()
-			if err := checkDeactivationPin(s); errors.Is(err, errInvalidDeactivationPin) {
+			if err := checkDeactivationPin(s); isCheckDeactivationPinErr(err) {
 				os.Exit(deactivationPinInvalidExitCode)
 			}
 			if doTasks([]task{{s.Stop, true}}) {
@@ -552,7 +559,7 @@ NOTE: Uninstalling will set DNS to values provided by DHCP.`,
 			if iface == "" {
 				iface = "auto"
 			}
-			if err := checkDeactivationPin(s); errors.Is(err, errInvalidDeactivationPin) {
+			if err := checkDeactivationPin(s); isCheckDeactivationPinErr(err) {
 				os.Exit(deactivationPinInvalidExitCode)
 			}
 			uninstall(p, s)
@@ -834,7 +841,7 @@ func RunMobile(appConfig *AppConfig, appCallback *AppCallback, stopCh chan struc
 // CheckDeactivationPin checks if deactivation pin is valid
 func CheckDeactivationPin(pin int64) int {
 	deactivationPin = pin
-	if err := checkDeactivationPin(nil); errors.Is(err, errInvalidDeactivationPin) {
+	if err := checkDeactivationPin(nil); isCheckDeactivationPinErr(err) {
 		return deactivationPinInvalidExitCode
 	}
 	return 0
@@ -2146,6 +2153,9 @@ const deactivationPinInvalidExitCode = 126
 // errInvalidDeactivationPin indicates that the deactivation pin is invalid.
 var errInvalidDeactivationPin = errors.New("deactivation pin is invalid")
 
+// errRequiredDeactivationPin indicates that the deactivation pin is required but not provided by users.
+var errRequiredDeactivationPin = errors.New("deactivation pin is required to stop or uninstall the service")
+
 // checkDeactivationPin validates if the deactivation pin matches one in ControlD config.
 func checkDeactivationPin(s service.Service) error {
 	dir, err := socketDir()
@@ -2164,13 +2174,24 @@ func checkDeactivationPin(s service.Service) error {
 	}
 	data, _ := json.Marshal(&deactivationRequest{Pin: deactivationPin})
 	resp, _ := cc.post(deactivationPath, bytes.NewReader(data))
-	if resp != nil && resp.StatusCode == http.StatusOK {
-		return nil // valid pin
+	if resp != nil {
+		switch resp.StatusCode {
+		case http.StatusBadRequest:
+			mainLog.Load().Error().Msg(errRequiredDeactivationPin.Error())
+			return errRequiredDeactivationPin // pin is required
+		case http.StatusOK:
+			return nil // valid pin
+		}
 	}
-	mainLog.Load().Error().Msg("deactivation pin is invalid")
+	mainLog.Load().Error().Msg(errInvalidDeactivationPin.Error())
 	return errInvalidDeactivationPin
 }
 
+// isCheckDeactivationPinErr reports whether there is an error during check deactivation pin process.
+func isCheckDeactivationPinErr(err error) bool {
+	return errors.Is(err, errInvalidDeactivationPin) || errors.Is(err, errRequiredDeactivationPin)
+}
+
 // ensureUninstall ensures that s.Uninstall will remove ctrld service from system completely.
 func ensureUninstall(s service.Service) error {
 	maxAttempts := 10
diff --git a/cmd/cli/control_server.go b/cmd/cli/control_server.go
index 117174d..28c20a6 100644
--- a/cmd/cli/control_server.go
+++ b/cmd/cli/control_server.go
@@ -160,9 +160,14 @@ func (p *prog) registerControlServerHandler() {
 			mainLog.Load().Err(err).Msg("invalid deactivation request")
 			return
 		}
-		code := http.StatusBadRequest
-		if req.Pin == cdDeactivationPin {
+
+		code := http.StatusForbidden
+		switch req.Pin {
+		case cdDeactivationPin:
 			code = http.StatusOK
+		case defaultDeactivationPin:
+			// If the pin code was set, but users do not provide --pin, return proper code to client.
+			code = http.StatusBadRequest
 		}
 		w.WriteHeader(code)
 	}))