Merge pull request #117 from xia0pin9/feat/improve-register-command

feat(client): improve register command error handling

Author: xia0pin9

client/register.go

@@ -23,14 +23,17 @@ Register will cache the key in the file system and keep it up to date using the
 -k specifies a specific key identifier to register
 -f specifies a file containing a new line separated list of key identifiers
 -t specifies a timeout for getting the key from the daemon (e.g. '5s', '500ms')
--g gets the key as well
+-g gets the key as well (in no-cache mode, only -g with -k is supported)
 
 For a machine to access a certain key, it needs permissions on that key.
 
 Note that knox register will only update the register file and will return successful
 even if the machine does not have access to the key. The daemon will actually retrieve
 the key.
 
+In no-cache mode (when using UncachedHTTPClient), only the -g flag with -k is supported
+to fetch a key directly from the server without local registration.
+
 For more about knox, see https://github.com/pinterest/knox.
 
 See also: knox unregister, knox daemon
@@ -58,51 +61,61 @@ func parseTimeout(val string) (time.Duration, error) {
 }
 
 func runRegister(cmd *Command, args []string) *ErrorStatus {
-	if _, ok := cli.(*knox.UncachedHTTPClient); ok {
-		fmt.Println("Cannot Register in No Cache mode")
-		return nil
+	// Validate -g flag usage early: -g requires -k (cannot be used with -f alone)
+	if *registerAndGet && *registerKey == "" {
+		return &ErrorStatus{fmt.Errorf("the -g flag requires -k to specify a single key to retrieve"), false}
 	}
-	timeout, err := parseTimeout(*registerTimeout)
-	if err != nil {
-		return &ErrorStatus{fmt.Errorf("Invalid value for timeout flag: %s", err.Error()), false}
+
+	_, isUncachedMode := cli.(*knox.UncachedHTTPClient)
+
+	// In uncached mode, only support -g with -k to fetch a key directly from the server
+	if isUncachedMode {
+		if !*registerAndGet {
+			return &ErrorStatus{fmt.Errorf("cannot register keys in no-cache mode; use -g with -k to fetch a key directly"), false}
+		}
+		// -k is already validated above
+		// Skip registration, go directly to fetching the key
+		return fetchAndPrintKey(*registerKey, *registerTimeout)
 	}
 
 	k := NewKeysFile(path.Join(daemonFolder, daemonToRegister))
+	// Handle `knox register -r` (without -k or -f) to remove all registered keys
 	if *registerRemove && *registerKey == "" && *registerKeyFile == "" {
-		// Short circuit & handle `knox register -r`, which is expected to remove all keys
 		err := k.Lock()
 		if err != nil {
-			return &ErrorStatus{fmt.Errorf("There was an error obtaining file lock: %s", err.Error()), false}
+			return &ErrorStatus{fmt.Errorf("error obtaining file lock: %w", err), false}
 		}
 		err = k.Overwrite([]string{})
 		if err != nil {
 			k.Unlock()
-			return &ErrorStatus{fmt.Errorf("Failed to unregister all keys: %s", err.Error()), false}
+			return &ErrorStatus{fmt.Errorf("failed to unregister all keys: %w", err), false}
 		}
 		err = k.Unlock()
 		if err != nil {
-			return &ErrorStatus{fmt.Errorf("There was an error unlocking register file: %s", err.Error()), false}
+			return &ErrorStatus{fmt.Errorf("error unlocking register file: %w", err), false}
 		}
 		logf("Successfully unregistered all keys.")
 		return nil
 	} else if *registerKey == "" && *registerKeyFile == "" {
-		return &ErrorStatus{fmt.Errorf("You must include a key or key file to register. see 'knox help register'"), false}
+		return &ErrorStatus{fmt.Errorf("you must include a key or key file to register; see 'knox help register'"), false}
 	}
 	// Get the list of keys to add
 	var ks []string
+	var err error
 	if *registerKey == "" {
 		f := NewKeysFile(*registerKeyFile)
 		ks, err = f.Get()
 		if err != nil {
-			return &ErrorStatus{fmt.Errorf("There was an error reading input key file %s", err.Error()), false}
+			return &ErrorStatus{fmt.Errorf("error reading input key file: %w", err), false}
 		}
 	} else {
 		ks = []string{*registerKey}
 	}
 	// Handle adding new keys to the registered file
+	// When -r is specified with -k or -f, this replaces all registered keys with the specified ones
 	err = k.Lock()
 	if err != nil {
-		return &ErrorStatus{fmt.Errorf("There was an error obtaining file lock: %s", err.Error()), false}
+		return &ErrorStatus{fmt.Errorf("error obtaining file lock: %w", err), false}
 	}
 	if *registerRemove {
 		logf("Attempting to overwrite existing keys with %v.", ks)
@@ -112,34 +125,74 @@ func runRegister(cmd *Command, args []string) *ErrorStatus {
 	}
 	if err != nil {
 		k.Unlock()
-		return &ErrorStatus{fmt.Errorf("There was an error registering keys %v: %s", ks, err.Error()), false}
+		return &ErrorStatus{fmt.Errorf("error registering keys %v: %w", ks, err), false}
 	}
 	err = k.Unlock()
 	if err != nil {
-		return &ErrorStatus{fmt.Errorf("There was an error unlocking register file: %s", err.Error()), false}
+		return &ErrorStatus{fmt.Errorf("error unlocking register file: %w", err), false}
 	}
-	// If specified, force retrieval of keys
+	// If specified, force retrieval of keys (already validated that -k is set when -g is used)
 	if *registerAndGet {
-		key, err := cli.CacheGetKey(*registerKey)
-		c := time.After(timeout)
-		for err != nil {
-			select {
-			case <-c:
+		return fetchAndPrintKey(*registerKey, *registerTimeout)
+	}
+	logf("Successfully registered keys %v. Keys are updated by the daemon process every %.0f minutes. Check the log for the most recent run.", ks, daemonRefreshTime.Minutes())
+	return nil
+}
+
+// fetchAndPrintKey fetches a key from the server and prints it as JSON.
+// This is used by both cached and uncached modes when -g flag is specified.
+//
+// Note: The timeout bounds the total retry time, but individual CacheGetKey calls
+// may block beyond the deadline since CacheGetKey doesn't support context cancellation.
+// The deadline is checked before each retry attempt to minimize overage.
+func fetchAndPrintKey(keyID string, timeoutStr string) *ErrorStatus {
+	timeout, err := parseTimeout(timeoutStr)
+	if err != nil {
+		return &ErrorStatus{fmt.Errorf("invalid value for timeout flag: %w", err), false}
+	}
+
+	// Start deadline timer before first call to bound total time
+	deadline := time.After(timeout)
+	var key *knox.Key
+	var fetchErr error // Track fetch errors separately for clarity
+
+	for {
+		// Check timeout before each attempt
+		select {
+		case <-deadline:
+			if fetchErr != nil {
 				return &ErrorStatus{fmt.Errorf(
-					"Error getting key from daemon (hit timeout after %s seconds); check knox logs for details (most recent error: %v)",
-					timeout.String(), err), false}
-			case <-time.After(registerRecheckTime):
-				key, err = cli.CacheGetKey(*registerKey)
+					"error getting key from server (hit timeout after %s): %w",
+					timeout.String(), fetchErr), false}
 			}
+			// Timeout on first attempt before any fetch was made
+			return &ErrorStatus{fmt.Errorf(
+				"error getting key from server (hit timeout after %s before fetch attempt)",
+				timeout.String()), false}
+		default:
+			// Continue with fetch attempt
 		}
-		// TODO: add json vs data option?
-		data, err := json.Marshal(key)
-		if err != nil {
-			return &ErrorStatus{err, true}
+
+		key, fetchErr = cli.CacheGetKey(keyID)
+		if fetchErr == nil {
+			break
+		}
+
+		// Wait before retry, but also check deadline
+		select {
+		case <-deadline:
+			return &ErrorStatus{fmt.Errorf(
+				"error getting key from server (hit timeout after %s): %w",
+				timeout.String(), fetchErr), false}
+		case <-time.After(registerRecheckTime):
+			// Continue to next attempt
 		}
-		fmt.Printf("%s", string(data))
-		return nil
 	}
-	logf("Successfully registered keys %v. Keys are updated by the daemon process every %.0f minutes. Check the log for the most recent run.", ks, daemonRefreshTime.Minutes())
+
+	data, marshalErr := json.Marshal(key)
+	if marshalErr != nil {
+		return &ErrorStatus{marshalErr, true}
+	}
+	fmt.Printf("%s", string(data))
 	return nil
 }