From d340eeb50735080811707569aee7c27194b269ca Mon Sep 17 00:00:00 2001
From: Derek Nola <derek.nola@suse.com>
Date: Mon, 8 Nov 2021 10:54:48 -0800
Subject: [PATCH] Match to last After keyword for parser (#4383)

* Made parser able to skip over subcommands
* Edge case coverage, reworked regex with groups
Signed-off-by: Derek Nola <derek.nola@suse.com>
---
 pkg/configfilearg/defaultparser.go |  2 +-
 pkg/configfilearg/parser.go        | 26 ++++++++++++++++++++-
 pkg/configfilearg/parser_test.go   | 37 +++++++++++++++++++++++++++++-
 3 files changed, 62 insertions(+), 3 deletions(-)

diff --git a/pkg/configfilearg/defaultparser.go b/pkg/configfilearg/defaultparser.go
index 3b17b95115ea..d81b9ea52d73 100644
--- a/pkg/configfilearg/defaultparser.go
+++ b/pkg/configfilearg/defaultparser.go
@@ -7,7 +7,7 @@ import (
 
 func MustParse(args []string) []string {
 	parser := &Parser{
-		After:         []string{"server", "agent", "etcd-snapshot"},
+		After:         []string{"server", "agent", "etcd-snapshot:1"},
 		FlagNames:     []string{"--config", "-c"},
 		EnvName:       version.ProgramUpper + "_CONFIG_FILE",
 		DefaultConfig: "/etc/rancher/" + version.Program + "/config.yaml",
diff --git a/pkg/configfilearg/parser.go b/pkg/configfilearg/parser.go
index 7994ca37ac5b..2440e8276b43 100644
--- a/pkg/configfilearg/parser.go
+++ b/pkg/configfilearg/parser.go
@@ -7,6 +7,8 @@ import (
 	"net/url"
 	"os"
 	"path/filepath"
+	"regexp"
+	"strconv"
 	"strings"
 
 	"github.com/rancher/k3s/pkg/agent/util"
@@ -101,14 +103,36 @@ func (p *Parser) findStart(args []string) ([]string, []string, bool) {
 		return []string{}, args, true
 	}
 
+	afterIndex := make(map[string]int)
+	re, err := regexp.Compile(`(.+):(\d+)`)
+	if err != nil {
+		return args, nil, false
+	}
+	// After keywords ending with ":<NUM>" can set + NUM of arguments as the split point.
+	// used for matching on subcommmands
+	for i, arg := range p.After {
+		if match := re.FindAllStringSubmatch(arg, -1); match != nil {
+			p.After[i] = match[0][1]
+			afterIndex[match[0][1]], err = strconv.Atoi(match[0][2])
+			if err != nil {
+				return args, nil, false
+			}
+		}
+	}
+
 	for i, val := range args {
 		for _, test := range p.After {
 			if val == test {
+				if skip := afterIndex[test]; skip != 0 {
+					if len(args) <= i+skip || strings.HasPrefix(args[i+skip], "-") {
+						return args[0 : i+1], args[i+1:], true
+					}
+					return args[0 : i+skip+1], args[i+skip+1:], true
+				}
 				return args[0 : i+1], args[i+1:], true
 			}
 		}
 	}
-
 	return args, nil, false
 }
 
diff --git a/pkg/configfilearg/parser_test.go b/pkg/configfilearg/parser_test.go
index 18128be350cd..7582a9db11f2 100644
--- a/pkg/configfilearg/parser_test.go
+++ b/pkg/configfilearg/parser_test.go
@@ -48,11 +48,46 @@ func Test_UnitParser_findStart(t *testing.T) {
 			prefix: []string{"not-server", "foo", "bar"},
 			found:  false,
 		},
+		{
+			name:   "command (with optional subcommands) but no flags",
+			args:   []string{"etcd-snapshot"},
+			prefix: []string{"etcd-snapshot"},
+			suffix: []string{},
+			found:  true,
+		},
+		{
+			name:   "command (with optional subcommands) and flags",
+			args:   []string{"etcd-snapshot", "-f"},
+			prefix: []string{"etcd-snapshot"},
+			suffix: []string{"-f"},
+			found:  true,
+		},
+		{
+			name:   "command and subcommand with no flags",
+			args:   []string{"etcd-snapshot", "list"},
+			prefix: []string{"etcd-snapshot", "list"},
+			suffix: []string{},
+			found:  true,
+		},
+		{
+			name:   "command and subcommand with flags",
+			args:   []string{"etcd-snapshot", "list", "-f"},
+			prefix: []string{"etcd-snapshot", "list"},
+			suffix: []string{"-f"},
+			found:  true,
+		},
+		{
+			name:   "command and too many subcommands",
+			args:   []string{"etcd-snapshot", "list", "delete", "foo", "bar"},
+			prefix: []string{"etcd-snapshot", "list"},
+			suffix: []string{"delete", "foo", "bar"},
+			found:  true,
+		},
 	}
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
 			p := Parser{
-				After: []string{"server", "agent"},
+				After: []string{"server", "agent", "etcd-snapshot:1"},
 			}
 			prefix, suffix, found := p.findStart(tt.args)
 			if !reflect.DeepEqual(prefix, tt.prefix) {