fix: support spaces in directive argument values

internal/injector/aspect/advice/code/dot_directive.go

@@ -6,6 +6,7 @@
 package code
 
 import (
+	"bufio"
 	"regexp"
 	"strings"
 )
@@ -22,25 +23,85 @@ var spaces = regexp.MustCompile(`\s+`)
 // DirectiveArgs returns arguments provided to the named directive. A directive is a single-line
 // comment with the directive immediately following the leading `//`, without any spacing in
 // between; followed by optional arguments formatted as `key:value`, separated by spaces.
-func (d *dot) DirectiveArgs(directive string) (args []DirectiveArgument) {
+//
+// Values might contain spaces, and in that case they need to be quoted either using single or double quotes as
+// `key:"value with spaces"` or `key:'value with spaces'`.
+func (d *dot) DirectiveArgs(directive string) []DirectiveArgument {
 	prefix := "//" + directive
 
 	for curr := d.context.Chain(); curr != nil; curr = curr.Parent() {
 		for _, dec := range curr.Node().Decorations().Start {
-			if !strings.HasPrefix(dec, prefix) {
-				continue
+			args, ok := parseDirectiveArgs(prefix, dec)
+			if ok {
+				return args
 			}
-			parts := spaces.Split(dec, -1)
-			if parts[0] != prefix {
-				// This is not the directive we're looking for -- its name only starts the same.
-				continue
+		}
+	}
+	return nil
+}
+
+func parseDirectiveArgs(prefix string, comment string) ([]DirectiveArgument, bool) {
+	if !strings.HasPrefix(comment, prefix) {
+		return nil, false
+	}
+	parts := spaces.Split(comment, -1)
+	if parts[0] != prefix {
+		// This is not the directive we're looking for -- its name only starts the same.
+		return nil, false
+	}
+
+	// Strip the prefix from the comment.
+	argsStr := strings.TrimSpace(strings.TrimPrefix(comment, prefix))
+	if argsStr == "" {
+		return nil, true
+	}
+
+	scanner := bufio.NewScanner(strings.NewReader(argsStr))
+	scanner.Split(splitArgs)
+
+	var res []DirectiveArgument
+	for scanner.Scan() {
+		part := scanner.Text()
+		if key, value, ok := strings.Cut(part, ":"); ok {
+			if (strings.HasPrefix(value, `"`) && strings.HasSuffix(value, `"`)) ||
+				(strings.HasPrefix(value, `'`) && strings.HasSuffix(value, `'`)) {
+				value = value[1 : len(value)-1]
+			}
+			res = append(res, DirectiveArgument{Key: key, Value: value})
+		} else {
+			res = append(res, DirectiveArgument{Key: part, Value: ""})
+		}
+	}
+	return res, true
+}
+
+func splitArgs(data []byte, atEOF bool) (advance int, token []byte, err error) {
+	var (
+		doubleQuote = false
+		singleQuote = false
+		start       = 0
+	)
+	for i := 0; i < len(data); i++ {
+		switch data[i] {
+		case '"':
+			if !singleQuote {
+				doubleQuote = !doubleQuote
 			}
-			for _, part := range parts[1:] {
-				key, value, _ := strings.Cut(part, ":")
-				args = append(args, DirectiveArgument{Key: key, Value: value})
+		case '\'':
+			if !doubleQuote {
+				singleQuote = !singleQuote
+			}
+		case ' ':
+			if !doubleQuote && !singleQuote {
+				if start < i {
+					return i + 1, data[start:i], nil
+				}
+				start = i + 1
 			}
-			return
 		}
 	}
-	return
+	if atEOF && start < len(data) {
+		return len(data), data[start:], nil
+	}
+	return 0, nil, nil
 }

internal/injector/aspect/advice/code/dot_directive_test.go

@@ -0,0 +1,156 @@
+// Unless explicitly stated otherwise all files in this repository are licensed
+// under the Apache License Version 2.0.
+// This product includes software developed at Datadog (https://www.datadoghq.com/).
+// Copyright 2023-present Datadog, Inc.
+
+package code
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func Test_parseDirectiveArgs(t *testing.T) {
+	testCases := []struct {
+		name    string
+		prefix  string
+		comment string
+		want    []DirectiveArgument
+		wantOk  bool
+	}{
+		{
+			name:    "valid directive with two args",
+			prefix:  "//dd:span",
+			comment: "//dd:span span.name:rootHandler resource.name:\"GET /\"",
+			want: []DirectiveArgument{
+				{Key: "span.name", Value: "rootHandler"},
+				{Key: "resource.name", Value: "GET /"},
+			},
+			wantOk: true,
+		},
+		{
+			name:    "args with spaces double quote",
+			prefix:  "//dd:span",
+			comment: "//dd:span span.name:rootHandler resource.name:\"GET /\" foo:\"bar\" ",
+			want: []DirectiveArgument{
+				{Key: "span.name", Value: "rootHandler"},
+				{Key: "resource.name", Value: "GET /"},
+				{Key: "foo", Value: "bar"},
+			},
+			wantOk: true,
+		},
+		{
+			name:    "args with spaces single quote",
+			prefix:  "//dd:span",
+			comment: "//dd:span span.name:rootHandler resource.name:'GET /' foo:'bar'",
+			want: []DirectiveArgument{
+				{Key: "span.name", Value: "rootHandler"},
+				{Key: "resource.name", Value: "GET /"},
+				{Key: "foo", Value: "bar"},
+			},
+			wantOk: true,
+		},
+		{
+			name:    "single and double quotes",
+			prefix:  "//dd:span",
+			comment: `//dd:span span.name:'root handler' resource.name:"GET /home"`,
+			want: []DirectiveArgument{
+				{Key: "span.name", Value: "root handler"},
+				{Key: "resource.name", Value: "GET /home"},
+			},
+			wantOk: true,
+		},
+		{
+			name:    "valid directive with one arg",
+			prefix:  "//dd:span",
+			comment: "//dd:span service.name:my-service",
+			want: []DirectiveArgument{
+				{Key: "service.name", Value: "my-service"},
+			},
+			wantOk: true,
+		},
+		{
+			name:    "prefix matches at start but is not full word",
+			prefix:  "//dd:span",
+			comment: "//dd:spanExtra service.name:my-service",
+			want:    nil,
+			wantOk:  false,
+		},
+		{
+			name:    "non-matching prefix",
+			prefix:  "//dd:span",
+			comment: "//other:span span.name:foo",
+			want:    nil,
+			wantOk:  false,
+		},
+		{
+			name:    "only prefix with no arguments",
+			prefix:  "//dd:span",
+			comment: "//dd:span",
+			want:    nil,
+			wantOk:  true,
+		},
+		{
+			name:    "arg with only key and no colon",
+			prefix:  "//dd:span",
+			comment: "//dd:span standalone_arg",
+			want: []DirectiveArgument{
+				{Key: "standalone_arg", Value: ""},
+			},
+			wantOk: true,
+		},
+		{
+			name:    "arg with multiple colons",
+			prefix:  "//dd:span",
+			comment: "//dd:span foo:bar:baz",
+			want: []DirectiveArgument{
+				{Key: "foo", Value: "bar:baz"},
+			},
+			wantOk: true,
+		},
+		{
+			name:    "unclosed double quote value",
+			prefix:  "//dd:span",
+			comment: `//dd:span service.name:"my-service`,
+			want: []DirectiveArgument{
+				{Key: "service.name", Value: `"my-service`},
+			},
+			wantOk: true,
+		},
+		{
+			name:    "unclosed single quote value",
+			prefix:  "//dd:span",
+			comment: `//dd:span service.name:'my-service`,
+			want: []DirectiveArgument{
+				{Key: "service.name", Value: "'my-service"},
+			},
+			wantOk: true,
+		},
+		{
+			name:    "missing space between quoted args",
+			prefix:  "//dd:span",
+			comment: `//dd:span service.name:"my-service"resource.name:"GET /"`,
+			want: []DirectiveArgument{
+				{Key: "service.name", Value: `my-service"resource.name:"GET /`},
+			},
+			wantOk: true,
+		},
+		{
+			name:    "quote starts but ends with different quote type",
+			prefix:  "//dd:span",
+			comment: `//dd:span service.name:"my-service'`,
+			want: []DirectiveArgument{
+				{Key: "service.name", Value: `"my-service'`},
+			},
+			wantOk: true,
+		},
+	}
+	for _, tt := range testCases {
+		t.Run(tt.name, func(t *testing.T) {
+			got, ok := parseDirectiveArgs(tt.prefix, tt.comment)
+			assert.Equal(t, tt.wantOk, ok)
+			assert.Equal(t, tt.want, got)
+		})
+	}
+}