diff --git a/internal/injector/aspect/advice/code/dot_directive.go b/internal/injector/aspect/advice/code/dot_directive.go index cf224c3d..fb8ec102 100644 --- a/internal/injector/aspect/advice/code/dot_directive.go +++ b/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 } diff --git a/internal/injector/aspect/advice/code/dot_directive_test.go b/internal/injector/aspect/advice/code/dot_directive_test.go new file mode 100644 index 00000000..5d350c6f --- /dev/null +++ b/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) + }) + } +}