From 09c01175446d517a36208ce5c9911a2b8d80fb61 Mon Sep 17 00:00:00 2001
From: rsteube <rsteube@users.noreply.github.com>
Date: Fri, 28 Jul 2023 12:34:32 +0200
Subject: [PATCH] lexer: patch shlex to support pipeline separators

---
 internal/lexer/lexer.go                      |  2 +-
 internal/lexer/lexer_test.go                 | 25 ++++++++++++
 third_party/github.com/google/shlex/shlex.go | 40 +++++++++++++++++++-
 3 files changed, 65 insertions(+), 2 deletions(-)

diff --git a/internal/lexer/lexer.go b/internal/lexer/lexer.go
index 5afac048..8f03b23a 100644
--- a/internal/lexer/lexer.go
+++ b/internal/lexer/lexer.go
@@ -44,7 +44,7 @@ func Split(s string) (*Tokenset, error) {
 }
 
 func split(s string) (*Tokenset, error) {
-	splitted, err := shlex.Split(s)
+	splitted, err := shlex.SplitP(s)
 	if strings.HasSuffix(s, " ") {
 		splitted = append(splitted, "")
 	}
diff --git a/internal/lexer/lexer_test.go b/internal/lexer/lexer_test.go
index 212dea0f..c945b4fc 100644
--- a/internal/lexer/lexer_test.go
+++ b/internal/lexer/lexer_test.go
@@ -96,4 +96,29 @@ func TestSplit(t *testing.T) {
 		Tokens: []string{"example", "action", "--", ""},
 		Prefix: `example 'action' -- `,
 	})
+
+	_test(`example 'action' -- | echo `, Tokenset{
+		Tokens: []string{"echo", ""},
+		Prefix: `example 'action' -- | echo `,
+	})
+
+	_test(`example 'action' -- || echo `, Tokenset{
+		Tokens: []string{"echo", ""},
+		Prefix: `example 'action' -- || echo `,
+	})
+
+	_test(`example 'action' -- && echo `, Tokenset{
+		Tokens: []string{"echo", ""},
+		Prefix: `example 'action' -- && echo `,
+	})
+
+	_test(`example 'action' -- ; echo `, Tokenset{
+		Tokens: []string{"echo", ""},
+		Prefix: `example 'action' -- ; echo `,
+	})
+
+	_test(`example 'action' -- & echo `, Tokenset{
+		Tokens: []string{"echo", ""},
+		Prefix: `example 'action' -- & echo `,
+	})
 }
diff --git a/third_party/github.com/google/shlex/shlex.go b/third_party/github.com/google/shlex/shlex.go
index 04d2deda..8ac2609f 100644
--- a/third_party/github.com/google/shlex/shlex.go
+++ b/third_party/github.com/google/shlex/shlex.go
@@ -80,6 +80,7 @@ const (
 	nonEscapingQuoteRunes = "'"
 	escapeRunes           = `\`
 	commentRunes          = "#"
+	terminateRunes        = "|&;"
 )
 
 // Classes of rune token
@@ -90,6 +91,7 @@ const (
 	nonEscapingQuoteRuneClass
 	escapeRuneClass
 	commentRuneClass
+	pipelineRuneClass
 	eofRuneClass
 )
 
@@ -99,6 +101,7 @@ const (
 	WordToken
 	SpaceToken
 	CommentToken
+	PipelineToken
 )
 
 // Lexer state machine states
@@ -129,6 +132,7 @@ func newDefaultClassifier() tokenClassifier {
 	t.addRuneClass(nonEscapingQuoteRunes, nonEscapingQuoteRuneClass)
 	t.addRuneClass(escapeRunes, escapeRuneClass)
 	t.addRuneClass(commentRunes, commentRuneClass)
+	t.addRuneClass(terminateRunes, pipelineRuneClass)
 	return t
 }
 
@@ -146,6 +150,12 @@ func NewLexer(r io.Reader) *Lexer {
 	return (*Lexer)(NewTokenizer(r))
 }
 
+type PipelineSeparatorError struct{}
+
+func (m *PipelineSeparatorError) Error() string {
+	return "encountered a pipeline separator like `|`"
+}
+
 // Next returns the next word, or an error. If there are no more words,
 // the error will be io.EOF.
 func (l *Lexer) Next() (string, error) {
@@ -159,6 +169,9 @@ func (l *Lexer) Next() (string, error) {
 			return token.value, nil
 		case CommentToken:
 			// skip comments
+		case PipelineToken:
+			// return token but with pseudo err to mark end of pipeline
+			return token.value, &PipelineSeparatorError{}
 		default:
 			return "", fmt.Errorf("Unknown token type: %v", token.tokenType)
 		}
@@ -232,6 +245,12 @@ func (t *Tokenizer) scanStream() (*Token, error) {
 						tokenType = CommentToken
 						state = commentState
 					}
+				case pipelineRuneClass:
+					{
+						tokenType = PipelineToken
+						value = append(value, nextRune)
+						state = inWordState
+					}
 				default:
 					{
 						tokenType = WordToken
@@ -400,6 +419,19 @@ func (t *Tokenizer) Next() (*Token, error) {
 
 // Split partitions a string into a slice of strings.
 func Split(s string) ([]string, error) {
+	return split(s, false)
+}
+
+// Split is like Split but only returns the last pipeline.
+//
+//	`echo example | bat -`
+//	# [bat, -]
+func SplitP(s string) ([]string, error) {
+	return split(s, true)
+}
+
+// Split partitions a string into a slice of strings.
+func split(s string, resetOnPipe bool) ([]string, error) {
 	l := NewLexer(strings.NewReader(s))
 	subStrings := make([]string, 0)
 	for {
@@ -408,7 +440,13 @@ func Split(s string) ([]string, error) {
 			if err == io.EOF {
 				return subStrings, nil
 			}
-			return subStrings, err
+			if _, ok := err.(*PipelineSeparatorError); !ok {
+				return subStrings, err
+			}
+			if resetOnPipe {
+				subStrings = make([]string, 0)
+				continue
+			}
 		}
 		subStrings = append(subStrings, word)
 	}