From ef737e194afd0c03186e3eeae82bcecb7cb127c4 Mon Sep 17 00:00:00 2001
From: Myzel394 <50424412+Myzel394@users.noreply.github.com>
Date: Tue, 15 Oct 2024 16:20:18 +0200
Subject: [PATCH] feat(sshd_config): Add support for token completions

---
 .../handlers/sshd_config/analyzer/options.go  |  2 +-
 server/handlers/sshd_config/ast/listener.go   |  3 +-
 .../handlers/sshd_config/ast/sshd_config.go   |  4 +-
 server/handlers/sshd_config/fields/tokens.go  | 71 +++++++++++++++++++
 .../sshd_config/handlers/completions.go       | 42 +++++++++--
 .../sshd_config/handlers/completions_match.go |  8 +--
 .../sshd_config/handlers/formatting_nodes.go  |  7 +-
 server/handlers/sshd_config/handlers/hover.go |  2 +-
 .../sshd_config/indexes/indexes_handlers.go   |  2 +-
 .../lsp/text-document-completion.go           |  2 +-
 10 files changed, 123 insertions(+), 20 deletions(-)
 create mode 100644 server/handlers/sshd_config/fields/tokens.go

diff --git a/server/handlers/sshd_config/analyzer/options.go b/server/handlers/sshd_config/analyzer/options.go
index 0c13f62..8337ade 100644
--- a/server/handlers/sshd_config/analyzer/options.go
+++ b/server/handlers/sshd_config/analyzer/options.go
@@ -40,7 +40,7 @@ func checkOption(
 	checkIsUsingDoubleQuotes(ctx, option.Key.Value, option.Key.LocationRange)
 	checkQuotesAreClosed(ctx, option.Key.Value, option.Key.LocationRange)
 
-	key := fields.CreateNormalizedName(option.Key.Key)
+	key := option.Key.Key
 	docOption, found := fields.Options[key]
 
 	if !found {
diff --git a/server/handlers/sshd_config/ast/listener.go b/server/handlers/sshd_config/ast/listener.go
index cc5a414..287c1fb 100644
--- a/server/handlers/sshd_config/ast/listener.go
+++ b/server/handlers/sshd_config/ast/listener.go
@@ -4,6 +4,7 @@ import (
 	"config-lsp/common"
 	commonparser "config-lsp/common/parser"
 	"config-lsp/handlers/sshd_config/ast/parser"
+	"config-lsp/handlers/sshd_config/fields"
 	"config-lsp/handlers/sshd_config/match-parser"
 	"strings"
 
@@ -76,7 +77,7 @@ func (s *sshdParserListener) EnterKey(ctx *parser.KeyContext) {
 	s.sshdContext.currentOption.Key = &SSHDKey{
 		LocationRange: location,
 		Value:         value,
-		Key:           key,
+		Key:           fields.CreateNormalizedName(key),
 	}
 }
 
diff --git a/server/handlers/sshd_config/ast/sshd_config.go b/server/handlers/sshd_config/ast/sshd_config.go
index 73513e2..f29285c 100644
--- a/server/handlers/sshd_config/ast/sshd_config.go
+++ b/server/handlers/sshd_config/ast/sshd_config.go
@@ -3,14 +3,16 @@ package ast
 import (
 	"config-lsp/common"
 	commonparser "config-lsp/common/parser"
+	"config-lsp/handlers/sshd_config/fields"
 	"config-lsp/handlers/sshd_config/match-parser"
+
 	"github.com/emirpasic/gods/maps/treemap"
 )
 
 type SSHDKey struct {
 	common.LocationRange
 	Value commonparser.ParsedString
-	Key   string
+	Key   fields.NormalizedOptionName
 }
 
 type SSHDValue struct {
diff --git a/server/handlers/sshd_config/fields/tokens.go b/server/handlers/sshd_config/fields/tokens.go
new file mode 100644
index 0000000..ee9d875
--- /dev/null
+++ b/server/handlers/sshd_config/fields/tokens.go
@@ -0,0 +1,71 @@
+package fields
+
+var AvailableTokens = map[string]string{
+	"%%": "A literal ‘%’.",
+	"%C": "Identifies the connection endpoints, containing four space-separated values: client address, client port number, server address, and server port number.",
+	"%D": "The routing domain in which the incoming connection was received.",
+	"%F": "The fingerprint of the CA key.",
+	"%f": "The fingerprint of the key or certificate.",
+	"%h": "The home directory of the user.",
+	"%i": "The key ID in the certificate.",
+	"%K": "The base64-encoded CA key.",
+	"%k": "The base64-encoded key or certificate for authentication.",
+	"%s": "The serial number of the certificate.",
+	"%T": "The type of the CA key.",
+	"%t": "The key or certificate type.",
+	"%U": "The numeric user ID of the target user.",
+	"%u": "The username.",
+}
+
+// A map of <option name> to <list of supported tokens>
+// This is derived from the TOKENS section of ssh_config
+var OptionsTokensMap = map[NormalizedOptionName][]string{
+	"authorizedkeyscommand": {
+		"%%",
+		"%C",
+		"%D",
+		"%f",
+		"%h",
+		"%k",
+		"%t",
+		"%U",
+		"%u",
+	},
+	"authorizedkeysfile":          {
+		"%%",
+		"%h",
+		"%U",
+		"%u",
+	},
+	"authorizedprincipalscommand": {
+		"%%",
+		"%C",
+		"%D",
+		"%F",
+		"%f",
+		"%h",
+		"%i",
+		"%K",
+		"%k",
+		"%s",
+		"%T",
+		"%t",
+		"%U",
+		"%u",
+	},
+	"authorizedprincipalsfile":    {
+		"%%",
+		"%h",
+		"%U",
+		"%u",
+	},
+	"chrootdirectory":             {
+		"%%",
+		"%h",
+		"%U",
+		"%u",
+	},
+	"routingdomain":               {
+		"%D",
+	},
+}
diff --git a/server/handlers/sshd_config/handlers/completions.go b/server/handlers/sshd_config/handlers/completions.go
index 5553aa6..9854e43 100644
--- a/server/handlers/sshd_config/handlers/completions.go
+++ b/server/handlers/sshd_config/handlers/completions.go
@@ -69,12 +69,12 @@ func GetOptionCompletions(
 	entry *ast.SSHDOption,
 	matchBlock *ast.SSHDMatchBlock,
 	cursor common.CursorPosition,
-) ([]protocol.CompletionItem, error) {
-	key := fields.CreateNormalizedName(entry.Key.Key)
+) []protocol.CompletionItem {
+	key := entry.Key.Key
 	option, found := fields.Options[key]
 
 	if !found {
-		return nil, nil
+		return nil
 	}
 
 	if entry.Key.Key == "Match" {
@@ -86,18 +86,48 @@ func GetOptionCompletions(
 	}
 
 	if entry.OptionValue == nil {
-		return option.DeprecatedFetchCompletions("", 0), nil
+		return option.DeprecatedFetchCompletions("", 0)
 	}
 
+	// token completions
+	completions := getTokenCompletions(entry, cursor)
+
 	// Hello wo|rld
 	line := entry.OptionValue.Value.Raw
 	// NEW: docvalues index
-	return option.DeprecatedFetchCompletions(
+	completions = append(completions, option.DeprecatedFetchCompletions(
 		line,
 		common.DeprecatedImprovedCursorToIndex(
 			cursor,
 			line,
 			entry.OptionValue.Start.Character,
 		),
-	), nil
+	)...)
+
+	return completions
+}
+
+func getTokenCompletions(
+	entry *ast.SSHDOption,
+	cursor common.CursorPosition,
+) []protocol.CompletionItem {
+	completions := make([]protocol.CompletionItem, 0)
+	index := common.CursorToCharacterIndex(uint32(cursor))
+
+	if entry.Value.Raw[index] == '%' {
+		if tokens, found := fields.OptionsTokensMap[entry.Key.Key]; found {
+			for _, token := range tokens {
+				description := fields.AvailableTokens[token]
+				kind := protocol.CompletionItemKindConstant
+
+				completions = append(completions, protocol.CompletionItem{
+					Label:         token,
+					Kind:          &kind,
+					Documentation: description,
+				})
+			}
+		}
+	}
+
+	return completions
 }
diff --git a/server/handlers/sshd_config/handlers/completions_match.go b/server/handlers/sshd_config/handlers/completions_match.go
index 79e7d77..e716f34 100644
--- a/server/handlers/sshd_config/handlers/completions_match.go
+++ b/server/handlers/sshd_config/handlers/completions_match.go
@@ -12,21 +12,21 @@ func getMatchCompletions(
 	d *sshdconfig.SSHDDocument,
 	cursor common.CursorPosition,
 	match *matchparser.Match,
-) ([]protocol.CompletionItem, error) {
+) []protocol.CompletionItem {
 	if match == nil || len(match.Entries) == 0 {
 		completions := getMatchCriteriaCompletions()
 		completions = append(completions, getMatchAllKeywordCompletion())
 
-		return completions, nil
+		return completions
 	}
 
 	entry := match.GetEntryAtPosition(cursor)
 
 	if entry == nil || entry.Criteria.ContainsPosition(cursor) {
-		return getMatchCriteriaCompletions(), nil
+		return getMatchCriteriaCompletions()
 	}
 
-	return getMatchValueCompletions(entry, cursor), nil
+	return getMatchValueCompletions(entry, cursor)
 }
 
 func getMatchCriteriaCompletions() []protocol.CompletionItem {
diff --git a/server/handlers/sshd_config/handlers/formatting_nodes.go b/server/handlers/sshd_config/handlers/formatting_nodes.go
index dc1f0ec..5deac6c 100644
--- a/server/handlers/sshd_config/handlers/formatting_nodes.go
+++ b/server/handlers/sshd_config/handlers/formatting_nodes.go
@@ -30,11 +30,10 @@ func formatSSHDOption(
 	var key string
 
 	if option.Key != nil {
-		key = option.Key.Key
-		normalizedKey := fields.CreateNormalizedName(key)
-
-		if formattedName, found := fields.FieldsNameFormattedMap[normalizedKey]; found {
+		if formattedName, found := fields.FieldsNameFormattedMap[option.Key.Key]; found {
 			key = formattedName
+		} else {
+			key = string(option.Key.Key)
 		}
 	} else {
 		key = ""
diff --git a/server/handlers/sshd_config/handlers/hover.go b/server/handlers/sshd_config/handlers/hover.go
index ce3529d..a3afa3b 100644
--- a/server/handlers/sshd_config/handlers/hover.go
+++ b/server/handlers/sshd_config/handlers/hover.go
@@ -17,7 +17,7 @@ func GetHoverInfoForOption(
 	index common.IndexPosition,
 ) (*protocol.Hover, error) {
 	var docValue *docvalues.DocumentationValue
-	key := fields.CreateNormalizedName(option.Key.Key)
+	key := option.Key.Key
 
 	// Either root level or in the line of a match block
 	if matchBlock == nil || matchBlock.Start.Line == line {
diff --git a/server/handlers/sshd_config/indexes/indexes_handlers.go b/server/handlers/sshd_config/indexes/indexes_handlers.go
index c3218dd..97069c1 100644
--- a/server/handlers/sshd_config/indexes/indexes_handlers.go
+++ b/server/handlers/sshd_config/indexes/indexes_handlers.go
@@ -89,7 +89,7 @@ func addOption(
 	matchBlock *ast.SSHDMatchBlock,
 ) []common.LSPError {
 	var errs []common.LSPError
-	key := fields.CreateNormalizedName(option.Key.Key)
+	key := option.Key.Key
 
 	if optionsMap, found := i.AllOptionsPerName[key]; found {
 		if options, found := optionsMap[matchBlock]; found {
diff --git a/server/handlers/sshd_config/lsp/text-document-completion.go b/server/handlers/sshd_config/lsp/text-document-completion.go
index 96c1ce9..72a9a38 100644
--- a/server/handlers/sshd_config/lsp/text-document-completion.go
+++ b/server/handlers/sshd_config/lsp/text-document-completion.go
@@ -43,7 +43,7 @@ func TextDocumentCompletion(context *glsp.Context, params *protocol.CompletionPa
 			entry,
 			matchBlock,
 			cursor,
-		)
+		), nil
 	}
 
 	return nil, nil