Refactor Dockerfile lookup code to handle \\wsl$ URIs

CHANGELOG.md

@@ -16,7 +16,8 @@ All notable changes to the Docker Language Server will be documented in this fil
   - textDocument/codeLens
     - refactor the URI handling code so it will accept a WSL URI with a dollar sign ([#388](https://github.com/docker/docker-language-server/issues/388))
   - textDocument/definition
-    - handle WSL URIs with a dollar sign properly to fix file looks up on those hosts ([#390](https://github.com/docker/docker-language-server/issues/390))
+    - handle WSL URIs with a dollar sign properly to fix build stage lookups on those hosts ([#390](https://github.com/docker/docker-language-server/issues/390))
+    - handle WSL URIs with a dollar sign properly to fix build ARG reference lookups on those hosts ([#393](https://github.com/docker/docker-language-server/issues/393))
   - textDocument/documentLink
     - convert links properly if a WSL URI with a dollar sign is used ([#378](https://github.com/docker/docker-language-server/issues/378))
   - textDocument/inlineCompletion

internal/bake/hcl/definition.go

@@ -3,8 +3,6 @@ package hcl
 import (
 	"context"
 	"errors"
-	"fmt"
-	"path/filepath"
 	"strings"
 
 	"github.com/docker/docker-language-server/internal/pkg/document"
@@ -176,7 +174,7 @@ func ResolveExpression(ctx context.Context, definitionLinkSupport bool, manager
 		for _, item := range objectConsExpression.Items {
 			if isInsideRange(item.KeyExpr.Range(), position) && sourceBlock != nil {
 				if attributeName == "args" && sourceBlock.Type == "target" {
-					dockerfilePath, err := doc.DockerfileForTarget(sourceBlock)
+					dockerfileURI, dockerfilePath, err := doc.DockerfileDocumentPathForTarget(sourceBlock)
 					if dockerfilePath == "" || err != nil {
 						return nil
 					}
@@ -188,7 +186,7 @@ func ResolveExpression(ctx context.Context, definitionLinkSupport bool, manager
 						end--
 					}
 					arg := string(doc.Input()[start:end])
-					bytes, nodes := document.OpenDockerfile(ctx, manager, "", dockerfilePath)
+					bytes, nodes := document.OpenDockerfile(ctx, manager, dockerfileURI, dockerfilePath)
 					lines := strings.Split(string(bytes), "\n")
 					for _, child := range nodes {
 						if strings.EqualFold(child.Value, "ARG") {
@@ -228,7 +226,7 @@ func ResolveExpression(ctx context.Context, definitionLinkSupport bool, manager
 											End:   protocol.Position{Line: uint32(node.EndLine) - 1, Character: uint32(endLineLength)},
 										},
 										&originSelectionRange,
-										protocol.URI(fmt.Sprintf("file:///%v", strings.TrimPrefix(filepath.ToSlash(dockerfilePath), "/"))),
+										dockerfileURI,
 									)
 								}
 								child = child.Next

internal/bake/hcl/definition_test.go

@@ -1045,7 +1045,7 @@ func testDefinition(t *testing.T, dockerfilePath, backendDockerfilePath, bakeFil
 			},
 		},
 		{
-			name:         "attribute object value's attribute name",
+			name:         "top-level args should not lookup anything",
 			content:      "args = { var = value }",
 			line:         0,
 			character:    11,
@@ -1054,7 +1054,7 @@ func testDefinition(t *testing.T, dockerfilePath, backendDockerfilePath, bakeFil
 			links:        nil,
 		},
 		{
-			name:         "attribute object value's attribute name",
+			name:         "top-level args variable reference should work",
 			content:      "a1 = \"\"\nargs = { var = a1 }",
 			line:         1,
 			character:    16,
@@ -1817,6 +1817,138 @@ func testDefinition_WSL(t *testing.T, dockerfilePath, backendDockerfilePath stri
 				},
 			},
 		},
+		{
+			name:         "args key references Dockerfile ARG variable (unquoted key, no default value set)",
+			content:      "target default {\n  args = {\n    var = \"value\"\n  }\n}",
+			line:         2,
+			character:    6,
+			endCharacter: -1,
+			locations: []protocol.Location{
+				{
+					URI: dockerfileURIString,
+					Range: protocol.Range{
+						Start: protocol.Position{Line: 1, Character: 0},
+						End:   protocol.Position{Line: 1, Character: 7},
+					},
+				},
+			},
+			links: []protocol.LocationLink{
+				{
+					OriginSelectionRange: &protocol.Range{
+						Start: protocol.Position{Line: 2, Character: 4},
+						End:   protocol.Position{Line: 2, Character: 7},
+					},
+					TargetURI: dockerfileURIString,
+					TargetRange: protocol.Range{
+						Start: protocol.Position{Line: 1, Character: 0},
+						End:   protocol.Position{Line: 1, Character: 7},
+					},
+					TargetSelectionRange: protocol.Range{
+						Start: protocol.Position{Line: 1, Character: 0},
+						End:   protocol.Position{Line: 1, Character: 7},
+					},
+				},
+			},
+		},
+		{
+			name:         "args key references Dockerfile ARG variable in a Dockerfile from the context attribute",
+			content:      "target default {\n  context = \"backend\"\n  args = {\n    NESTED_VAR = \"value\"\n  }\n}",
+			line:         3,
+			character:    6,
+			endCharacter: -1,
+			locations: []protocol.Location{
+				{
+					URI: backendDockerfileURIString,
+					Range: protocol.Range{
+						Start: protocol.Position{Line: 0, Character: 0},
+						End:   protocol.Position{Line: 0, Character: 14},
+					},
+				},
+			},
+			links: []protocol.LocationLink{
+				{
+					OriginSelectionRange: &protocol.Range{
+						Start: protocol.Position{Line: 3, Character: 4},
+						End:   protocol.Position{Line: 3, Character: 14},
+					},
+					TargetURI: backendDockerfileURIString,
+					TargetRange: protocol.Range{
+						Start: protocol.Position{Line: 0, Character: 0},
+						End:   protocol.Position{Line: 0, Character: 14},
+					},
+					TargetSelectionRange: protocol.Range{
+						Start: protocol.Position{Line: 0, Character: 0},
+						End:   protocol.Position{Line: 0, Character: 14},
+					},
+				},
+			},
+		},
+		{
+			name:         "args key references Dockerfile ARG variable (unquoted key, default value set)",
+			content:      "target default {\n  args = {\n    defined = \"value\"\n  }\n}",
+			line:         2,
+			character:    8,
+			endCharacter: -1,
+			locations: []protocol.Location{
+				{
+					URI: dockerfileURIString,
+					Range: protocol.Range{
+						Start: protocol.Position{Line: 2, Character: 0},
+						End:   protocol.Position{Line: 2, Character: 19},
+					},
+				},
+			},
+			links: []protocol.LocationLink{
+				{
+					OriginSelectionRange: &protocol.Range{
+						Start: protocol.Position{Line: 2, Character: 4},
+						End:   protocol.Position{Line: 2, Character: 11},
+					},
+					TargetURI: dockerfileURIString,
+					TargetRange: protocol.Range{
+						Start: protocol.Position{Line: 2, Character: 0},
+						End:   protocol.Position{Line: 2, Character: 19},
+					},
+					TargetSelectionRange: protocol.Range{
+						Start: protocol.Position{Line: 2, Character: 0},
+						End:   protocol.Position{Line: 2, Character: 19},
+					},
+				},
+			},
+		},
+		{
+			name:         "args key references Dockerfile ARG variable (quoted key, no default value set)",
+			content:      "target default {\n  args = {\n    \"var\" = \"value\"\n  }\n}",
+			line:         2,
+			character:    7,
+			endCharacter: -1,
+			locations: []protocol.Location{
+				{
+					URI: dockerfileURIString,
+					Range: protocol.Range{
+						Start: protocol.Position{Line: 1, Character: 0},
+						End:   protocol.Position{Line: 1, Character: 7},
+					},
+				},
+			},
+			links: []protocol.LocationLink{
+				{
+					OriginSelectionRange: &protocol.Range{
+						Start: protocol.Position{Line: 2, Character: 5},
+						End:   protocol.Position{Line: 2, Character: 8},
+					},
+					TargetURI: dockerfileURIString,
+					TargetRange: protocol.Range{
+						Start: protocol.Position{Line: 1, Character: 0},
+						End:   protocol.Position{Line: 1, Character: 7},
+					},
+					TargetSelectionRange: protocol.Range{
+						Start: protocol.Position{Line: 1, Character: 0},
+						End:   protocol.Position{Line: 1, Character: 7},
+					},
+				},
+			},
+		},
 	}
 
 	f, err := os.Open(dockerfilePath)