From 7359fac2e08aa62d573a6be484992085616ae058 Mon Sep 17 00:00:00 2001
From: Myzel394 <50424412+Myzel394@users.noreply.github.com>
Date: Wed, 2 Oct 2024 16:34:23 +0200
Subject: [PATCH] feat(ssh_config): Add tag support

---
 handlers/ssh_config/indexes/indexes.go        |  2 +
 .../ssh_config/indexes/indexes_handlers.go    | 36 +++++++++++++--
 handlers/ssh_config/indexes/indexes_test.go   | 44 +++++++++++++++++++
 3 files changed, 79 insertions(+), 3 deletions(-)

diff --git a/handlers/ssh_config/indexes/indexes.go b/handlers/ssh_config/indexes/indexes.go
index 2702d5a..e75d0e2 100644
--- a/handlers/ssh_config/indexes/indexes.go
+++ b/handlers/ssh_config/indexes/indexes.go
@@ -50,4 +50,6 @@ type SSHIndexes struct {
 	// a code action to add them to a "IgnoreUnknown" option
 	// This is a map of <line> to <option>
 	UnknownOptions map[uint32]ast.AllOptionInfo
+
+	Tags map[string]*ast.SSHMatchBlock
 }
diff --git a/handlers/ssh_config/indexes/indexes_handlers.go b/handlers/ssh_config/indexes/indexes_handlers.go
index d8860af..59d4050 100644
--- a/handlers/ssh_config/indexes/indexes_handlers.go
+++ b/handlers/ssh_config/indexes/indexes_handlers.go
@@ -4,6 +4,7 @@ import (
 	"config-lsp/common"
 	"config-lsp/handlers/ssh_config/ast"
 	"config-lsp/handlers/ssh_config/fields"
+	matchparser "config-lsp/handlers/ssh_config/match-parser"
 	"errors"
 	"fmt"
 	"regexp"
@@ -13,14 +14,16 @@ var whitespacePattern = regexp.MustCompile(`\S+`)
 
 func NewSSHIndexes() *SSHIndexes {
 	return &SSHIndexes{
-		AllOptionsPerName: make(map[fields.NormalizedOptionName](map[ast.SSHBlock]([]*ast.SSHOption)), 0),
+		AllOptionsPerName: make(map[fields.NormalizedOptionName](map[ast.SSHBlock]([]*ast.SSHOption))),
 		Includes:          make([]*SSHIndexIncludeLine, 0),
 		IgnoredOptions:    make(map[ast.SSHBlock]SSHIndexIgnoredUnknowns),
-		UnknownOptions:    make(map[uint32]ast.AllOptionInfo, 0),
+		UnknownOptions:    make(map[uint32]ast.AllOptionInfo),
+		Tags:              make(map[string]*ast.SSHMatchBlock),
 	}
 }
 
-var includeOption = fields.NormalizedOptionName("Include")
+var includeOption = fields.CreateNormalizedName("Include")
+var matchOption = fields.CreateNormalizedName("Match")
 
 func CreateIndexes(config ast.SSHConfig) (*SSHIndexes, []common.LSPError) {
 	errs := make([]common.LSPError, 0)
@@ -88,6 +91,33 @@ func CreateIndexes(config ast.SSHConfig) (*SSHIndexes, []common.LSPError) {
 		}
 	}
 
+	// Add tags
+	for rawBlock := range indexes.AllOptionsPerName[matchOption] {
+		block := rawBlock.(*ast.SSHMatchBlock)
+		for _, entry := range block.MatchValue.Entries {
+			if entry.Criteria.Type != matchparser.MatchCriteriaTypeTagged || entry.Value.Value == "" {
+				continue
+			}
+
+			for _, value := range entry.Values.Values {
+				name := value.Value.Value
+
+				if existingBlock, found := indexes.Tags[name]; found {
+					// Tag already exists
+					errs = append(errs, common.LSPError{
+						Range: entry.LocationRange,
+						Err:   fmt.Errorf("Tag %s has already been defined on line %d", name, existingBlock.Start.Line+1),
+					})
+
+					continue
+				}
+
+				// Add tag
+				indexes.Tags[name] = block
+			}
+		}
+	}
+
 	return indexes, errs
 }
 
diff --git a/handlers/ssh_config/indexes/indexes_test.go b/handlers/ssh_config/indexes/indexes_test.go
index 5b4bebf..3c81f18 100644
--- a/handlers/ssh_config/indexes/indexes_test.go
+++ b/handlers/ssh_config/indexes/indexes_test.go
@@ -3,6 +3,7 @@ package indexes
 import (
 	"config-lsp/handlers/ssh_config/ast"
 	"config-lsp/utils"
+	"fmt"
 	"testing"
 )
 
@@ -117,3 +118,46 @@ UseKeychain yes
 		t.Errorf("Expected IgnoreOptions to contain 'UseKeychain' on line 0 and from position 14-24, but got: %v", indexes.IgnoredOptions[nil].IgnoredOptions)
 	}
 }
+
+func TestTagsExample(
+	t *testing.T,
+) {
+	input := utils.Dedent(`
+Match tagged good_ip
+	AddressFamily inet
+
+Match tagged myuser
+	User root
+`)
+
+	config := ast.NewSSHConfig()
+
+	errors := config.Parse(input)
+
+	if len(errors) > 0 {
+		t.Fatalf("Expected no errors, but got %v", len(errors))
+	}
+
+	indexes, errors := CreateIndexes(*config)
+
+	if len(errors) > 0 {
+		t.Fatalf("Expected 1 error, but got %v", errors)
+	}
+
+	if !(len(indexes.Tags) == 2) {
+		t.Errorf("Expected 2 tags, but got %v", indexes.Tags)
+	}
+
+	rawFirstMatch, _ := config.Options.Get(uint32(0))
+	firstMatch := rawFirstMatch.(*ast.SSHMatchBlock)
+	println(fmt.Sprintf("%v", indexes.Tags["good_ip"]))
+	if !(indexes.Tags["good_ip"].Start.Line == firstMatch.Start.Line) {
+		t.Errorf("Expected first tag to be 'good_ip', but got %v", indexes.Tags)
+	}
+
+	rawSecondMatch, _ := config.Options.Get(uint32(3))
+	secondMatch := rawSecondMatch.(*ast.SSHMatchBlock)
+	if !(indexes.Tags["myuser"].Start.Line == secondMatch.Start.Line) {
+		t.Errorf("Expected second tag to be 'myuser', but got %v", indexes.Tags)
+	}
+}