From 7edfda55d361196f4b1f291cee336293153c609e Mon Sep 17 00:00:00 2001
From: Xueming Feng <kuro@kuroa.me>
Date: Wed, 17 Apr 2024 16:48:53 +0800
Subject: [PATCH] Add InfluencedByPermissions field to Command struct

---
 command.go              |  6 ++++++
 zsh_completions.go      | 32 ++++++++++++++++++++++++++------
 zsh_completions_test.go | 12 ++++++++++++
 3 files changed, 44 insertions(+), 6 deletions(-)

diff --git a/command.go b/command.go
index b6f8f4b14..cd52fb67e 100644
--- a/command.go
+++ b/command.go
@@ -254,6 +254,12 @@ type Command struct {
 	// SuggestionsMinimumDistance defines minimum levenshtein distance to display suggestions.
 	// Must be > 0.
 	SuggestionsMinimumDistance int
+
+	// InfluencedByPermissions indicates that the output is influenced by the
+	// permission it is run with. Thus, when a command such as sudo appears on the
+	// command-line, it will use commands like sudo or doas to gain extra privileges
+	// when retrieving information for completion. Available only to Zsh.
+	InfluencedByPermissions bool
 }
 
 // Context returns underlying command context. If command was executed
diff --git a/zsh_completions.go b/zsh_completions.go
index 1856e4c7f..e0b80aa13 100644
--- a/zsh_completions.go
+++ b/zsh_completions.go
@@ -21,6 +21,15 @@ import (
 	"os"
 )
 
+const (
+	// ZstyleGainPrivileges enables the use of commands like sudo or doas to
+	// gain extra privileges when retrieving information for completion.
+	ZstyleGainPrivileges = `zstyle ':completion:*:%s\*' gain-privileges yes`
+	// ZshCompRunCmdPermFlag indicates the command output is influenced by the
+	// permissions it is run with.
+	ZshCompRunCmdPermFlag = ` -p`
+)
+
 // GenZshCompletionFile generates zsh completion file including descriptions.
 func (c *Command) GenZshCompletionFile(filename string) error {
 	return c.genZshCompletionFile(filename, true)
@@ -79,16 +88,23 @@ func (c *Command) genZshCompletionFile(filename string, includeDesc bool) error
 
 func (c *Command) genZshCompletion(w io.Writer, includeDesc bool) error {
 	buf := new(bytes.Buffer)
-	genZshComp(buf, c.Name(), includeDesc)
+	genZshComp(buf, c.Name(), includeDesc, c.InfluencedByPermissions)
 	_, err := buf.WriteTo(w)
 	return err
 }
 
-func genZshComp(buf io.StringWriter, name string, includeDesc bool) {
+func genZshComp(buf io.StringWriter, name string, includeDesc bool, permission bool) {
 	compCmd := ShellCompRequestCmd
 	if !includeDesc {
 		compCmd = ShellCompNoDescRequestCmd
 	}
+	zstyleGain := ""
+	permFlag := ""
+	if permission {
+		zstyleGain = fmt.Sprintf(ZstyleGainPrivileges, name)
+		permFlag = ZshCompRunCmdPermFlag
+	}
+
 	WriteStringAndCheck(buf, fmt.Sprintf(`#compdef %[1]s
 compdef _%[1]s %[1]s
 
@@ -145,10 +161,14 @@ _%[1]s()
         requestComp="${requestComp} \"\""
     fi
 
-    __%[1]s_debug "About to call: eval ${requestComp}"
+    # Set zstyle if gain-privileges is requested
+    __%[1]s_debug "Setting zstyle: %[10]s"
+    %[10]s
+
+    __%[1]s_debug "About to call: _call_program%[11]s %[1]s-tag ${requestComp}"
 
-    # Use eval to handle any environment variables and such
-    out=$(eval ${requestComp} 2>/dev/null)
+    # Use _call_program to call the completion code
+    out=$(_call_program%[11]s %[1]s-tag ${requestComp})
     __%[1]s_debug "completion output: ${out}"
 
     # Extract the directive integer following a : from the last line
@@ -304,5 +324,5 @@ fi
 `, name, compCmd,
 		ShellCompDirectiveError, ShellCompDirectiveNoSpace, ShellCompDirectiveNoFileComp,
 		ShellCompDirectiveFilterFileExt, ShellCompDirectiveFilterDirs, ShellCompDirectiveKeepOrder,
-		activeHelpMarker))
+		activeHelpMarker, zstyleGain, permFlag))
 }
diff --git a/zsh_completions_test.go b/zsh_completions_test.go
index fe898b3da..bef88fc28 100644
--- a/zsh_completions_test.go
+++ b/zsh_completions_test.go
@@ -31,3 +31,15 @@ func TestZshCompletionWithActiveHelp(t *testing.T) {
 	activeHelpVar := activeHelpEnvVar(c.Name())
 	checkOmit(t, output, fmt.Sprintf("%s=0", activeHelpVar))
 }
+
+func TestZshCompletionWithInfluencedPermission(t *testing.T) {
+	c := &Command{Use: "c", Run: emptyRun, InfluencedByPermissions: true}
+
+	buf := new(bytes.Buffer)
+	assertNoErr(t, c.GenZshCompletion(buf))
+	output := buf.String()
+
+	// check that related commands are being generated
+	check(t, output, fmt.Sprintf(ZstyleGainPrivileges, c.Name()))
+	check(t, output, fmt.Sprintf("_call_program -p %s-tag", c.Name()))
+}