From 73a850c11abc68dd1c9d4168f3016450febf9702 Mon Sep 17 00:00:00 2001
From: Myzel394 <50424412+Myzel394@users.noreply.github.com>
Date: Sun, 22 Sep 2024 22:48:02 +0200
Subject: [PATCH] fix(ssh_config): Fix fetching options

---
 handlers/ssh_config/ast/ssh_config_fields.go | 32 ++++++++---
 handlers/ssh_config/document_fields_test.go  | 59 ++++++++++++++++++++
 handlers/ssh_config/fields/fields.go         |  2 +
 handlers/ssh_config/handlers/completions.go  |  9 ---
 4 files changed, 86 insertions(+), 16 deletions(-)
 create mode 100644 handlers/ssh_config/document_fields_test.go

diff --git a/handlers/ssh_config/ast/ssh_config_fields.go b/handlers/ssh_config/ast/ssh_config_fields.go
index 7902b21..9e7d3cb 100644
--- a/handlers/ssh_config/ast/ssh_config_fields.go
+++ b/handlers/ssh_config/ast/ssh_config_fields.go
@@ -2,6 +2,7 @@ package ast
 
 import (
 	"config-lsp/common"
+	"config-lsp/utils"
 
 	"github.com/emirpasic/gods/maps/treemap"
 )
@@ -145,10 +146,28 @@ func (c SSHConfig) FindOption(line uint32) (*SSHOption, SSHBlock) {
 }
 
 type AllOptionInfo struct {
-	Block SSHBlock
+	Block  SSHBlock
 	Option *SSHOption
 }
 
+func (c SSHConfig) GetAllOptionsForBlock(block SSHBlock) []AllOptionInfo {
+	if block == nil {
+		return c.GetAllOptions()
+	}
+
+	return utils.Map(
+		block.GetOptions().Values(),
+		func(rawOption interface{}) AllOptionInfo {
+			option := rawOption.(*SSHOption)
+
+			return AllOptionInfo{
+				Block:  block,
+				Option: option,
+			}
+		},
+	)
+}
+
 func (c SSHConfig) GetAllOptions() []AllOptionInfo {
 	options := make([]AllOptionInfo, 0, 50)
 
@@ -157,21 +176,21 @@ func (c SSHConfig) GetAllOptions() []AllOptionInfo {
 		case *SSHOption:
 			option := rawEntry.(*SSHOption)
 			options = append(options, AllOptionInfo{
-				Block: nil,
+				Block:  nil,
 				Option: option,
 			})
 		case *SSHMatchBlock:
 			block := rawEntry.(SSHBlock)
 
 			options = append(options, AllOptionInfo{
-				Block: block,
+				Block:  block,
 				Option: block.GetEntryOption(),
 			})
 
 			for _, rawOption := range block.GetOptions().Values() {
 				option := rawOption.(*SSHOption)
 				options = append(options, AllOptionInfo{
-					Block: nil,
+					Block:  block,
 					Option: option,
 				})
 			}
@@ -179,14 +198,14 @@ func (c SSHConfig) GetAllOptions() []AllOptionInfo {
 			block := rawEntry.(SSHBlock)
 
 			options = append(options, AllOptionInfo{
-				Block: block,
+				Block:  block,
 				Option: block.GetEntryOption(),
 			})
 
 			for _, rawOption := range block.GetOptions().Values() {
 				option := rawOption.(*SSHOption)
 				options = append(options, AllOptionInfo{
-					Block: nil,
+					Block:  block,
 					Option: option,
 				})
 			}
@@ -196,4 +215,3 @@ func (c SSHConfig) GetAllOptions() []AllOptionInfo {
 
 	return options
 }
-
diff --git a/handlers/ssh_config/document_fields_test.go b/handlers/ssh_config/document_fields_test.go
new file mode 100644
index 0000000..a4ce15d
--- /dev/null
+++ b/handlers/ssh_config/document_fields_test.go
@@ -0,0 +1,59 @@
+package sshconfig
+
+import (
+	"config-lsp/handlers/ssh_config/ast"
+	"config-lsp/handlers/ssh_config/indexes"
+	"config-lsp/utils"
+	"testing"
+)
+
+func TestComplexExample(
+	t *testing.T,
+) {
+	input := utils.Dedent(`
+ProxyCommand hello
+
+Host laptop
+    HostName laptop.lan
+	ProxyCommand test
+
+Match originalhost laptop exec "[[ $(/usr/bin/dig +short laptop.lan) == '' ]]"
+    HostName laptop.sdn
+`)
+	c := ast.NewSSHConfig()
+	errors := c.Parse(input)
+
+	if len(errors) != 0 {
+		t.Fatalf("Expected no errors, got %v", errors)
+	}
+
+	i, errors := indexes.CreateIndexes(*c)
+
+	if len(errors) != 0 {
+		t.Fatalf("Expected no errors, got %v", errors)
+	}
+
+	d := &SSHDocument{
+		Config:  c,
+		Indexes: i,
+	}
+
+	options := d.FindOptionsByName("ProxyCommand")
+	if !(len(options) == 2 && options[0].Option.Start.Line == 0 && options[1].Option.Start.Line == 4) {
+		t.Errorf("Expected 2 options, got %v", options)
+	}
+
+	options = d.FindOptionsByName("HostName")
+	if !(len(options) == 2 && options[0].Option.Start.Line == 3 && options[1].Option.Start.Line == 7) {
+		t.Errorf("Expected 2 options, got %v", options)
+	}
+
+	block := d.Config.FindBlock(4)
+	if !(d.FindOptionByNameAndBlock("ProxyCommand", block).Option.Start.Line == 4) {
+		t.Errorf("Expected 4, got %v", d.FindOptionByNameAndBlock("PorxyCommand", block).Option.Start.Line)
+	}
+
+	if !(d.FindOptionByNameAndBlock("ProxyCommand", nil).Option.Start.Line == 0) {
+		t.Errorf("Expected 0, got %v", d.FindOptionByNameAndBlock("ProxyCommand", nil).Option.Start.Line)
+	}
+}
diff --git a/handlers/ssh_config/fields/fields.go b/handlers/ssh_config/fields/fields.go
index d90b649..6aeed90 100644
--- a/handlers/ssh_config/fields/fields.go
+++ b/handlers/ssh_config/fields/fields.go
@@ -11,6 +11,8 @@ var MAX_PORT = 65535
 
 var AllowedDuplicateOptions = map[string]struct{}{
 	"CertificateFile": {},
+	"Match": {},
+	"Host": {},
 }
 
 var Options = map[string]docvalues.DocumentationValue{
diff --git a/handlers/ssh_config/handlers/completions.go b/handlers/ssh_config/handlers/completions.go
index c918093..e8974c7 100644
--- a/handlers/ssh_config/handlers/completions.go
+++ b/handlers/ssh_config/handlers/completions.go
@@ -28,15 +28,6 @@ func GetRootCompletions(
 		}
 	}
 
-	// Remove all fields that are already present and are not allowed to be duplicated
-	for _, info := range d.Config.GetAllOptions() {
-		if _, found := fields.AllowedDuplicateOptions[info.Option.Key.Key]; found {
-			continue
-		}
-
-		delete(availableOptions, info.Option.Key.Key)
-	}
-
 	return utils.MapMapToSlice(
 		availableOptions,
 		func(name string, doc docvalues.DocumentationValue) protocol.CompletionItem {