From 9fc5d5b5ebfb51f9ec172a8907b99b40271846bf Mon Sep 17 00:00:00 2001
From: Benson Wong <mostlygeek@gmail.com>
Date: Sun, 1 Dec 2024 09:02:58 -0800
Subject: [PATCH] improve cmd parsing (#22)

Switch from using a naive strings.Fields() to shlex.Split() for parsing the model startup command into a string[]. This makes parsing much more reliable around newlines, quotes, etc.
---
 go.mod               |  1 +
 go.sum               |  2 ++
 proxy/config.go      |  6 +++++-
 proxy/config_test.go | 25 +++++++++++++++++--------
 4 files changed, 25 insertions(+), 9 deletions(-)

diff --git a/go.mod b/go.mod
index 7145680..7c96913 100644
--- a/go.mod
+++ b/go.mod
@@ -20,6 +20,7 @@ require (
 	github.com/go-playground/universal-translator v0.18.1 // indirect
 	github.com/go-playground/validator/v10 v10.20.0 // indirect
 	github.com/goccy/go-json v0.10.2 // indirect
+	github.com/google/shlex v0.0.0-20191202100458-e7afc7fbc510 // indirect
 	github.com/json-iterator/go v1.1.12 // indirect
 	github.com/klauspost/cpuid/v2 v2.2.7 // indirect
 	github.com/leodido/go-urn v1.4.0 // indirect
diff --git a/go.sum b/go.sum
index e8333c3..309c319 100644
--- a/go.sum
+++ b/go.sum
@@ -24,6 +24,8 @@ github.com/go-playground/validator/v10 v10.20.0/go.mod h1:dbuPbCMFw/DrkbEynArYaC
 github.com/goccy/go-json v0.10.2 h1:CrxCmQqYDkv1z7lO7Wbh2HN93uovUHgrECaO5ZrCXAU=
 github.com/goccy/go-json v0.10.2/go.mod h1:6MelG93GURQebXPDq3khkgXZkazVtN9CRI+MGFi0w8I=
 github.com/google/gofuzz v1.0.0/go.mod h1:dBl0BpW6vV/+mYPU4Po3pmUjxk6FQPldtuIdl/M65Eg=
+github.com/google/shlex v0.0.0-20191202100458-e7afc7fbc510 h1:El6M4kTTCOh6aBiKaUGG7oYTSPP8MxqL4YI3kZKwcP4=
+github.com/google/shlex v0.0.0-20191202100458-e7afc7fbc510/go.mod h1:pupxD2MaaD3pAXIBCelhxNneeOaAeabZDe5s4K6zSpQ=
 github.com/json-iterator/go v1.1.12 h1:PV8peI4a0ysnczrg+LtxykD8LfKY9ML6u2jnxaEnrnM=
 github.com/json-iterator/go v1.1.12/go.mod h1:e30LSqwooZae/UwlEbR2852Gd8hjQvJoHmT4TnhNGBo=
 github.com/klauspost/cpuid/v2 v2.0.9/go.mod h1:FInQzS24/EEf25PyTYn52gqo7WaD8xa0213Md/qVLRg=
diff --git a/proxy/config.go b/proxy/config.go
index ad0794f..f17e019 100644
--- a/proxy/config.go
+++ b/proxy/config.go
@@ -5,6 +5,7 @@ import (
 	"os"
 	"strings"
 
+	"github.com/google/shlex"
 	"gopkg.in/yaml.v3"
 )
 
@@ -81,7 +82,10 @@ func SanitizeCommand(cmdStr string) ([]string, error) {
 	cmdStr = strings.ReplaceAll(cmdStr, "\\\n", " ")
 
 	// Split the command into arguments
-	args := strings.Fields(cmdStr)
+	args, err := shlex.Split(cmdStr)
+	if err != nil {
+		return nil, err
+	}
 
 	// Ensure the command is not empty
 	if len(args) == 0 {
diff --git a/proxy/config_test.go b/proxy/config_test.go
index 871e7e0..da2eb39 100644
--- a/proxy/config_test.go
+++ b/proxy/config_test.go
@@ -148,17 +148,26 @@ func TestConfig_FindConfig(t *testing.T) {
 }
 
 func TestConfig_SanitizeCommand(t *testing.T) {
-	// Test a simple command
-	args, err := SanitizeCommand("python model1.py")
-	assert.NoError(t, err)
-	assert.Equal(t, []string{"python", "model1.py"}, args)
 
 	// Test a command with spaces and newlines
-	args, err = SanitizeCommand(`python model1.py \
-    --arg1 value1 \
-    --arg2 value2`)
+	args, err := SanitizeCommand(`python model1.py \
+    -a "double quotes" \
+    --arg2 'single quotes'
+	-s
+	--arg3 123 \
+	--arg4 '"string in string"'
+	-c "'single quoted'"
+	`)
 	assert.NoError(t, err)
-	assert.Equal(t, []string{"python", "model1.py", "--arg1", "value1", "--arg2", "value2"}, args)
+	assert.Equal(t, []string{
+		"python", "model1.py",
+		"-a", "double quotes",
+		"--arg2", "single quotes",
+		"-s",
+		"--arg3", "123",
+		"--arg4", `"string in string"`,
+		"-c", `'single quoted'`,
+	}, args)
 
 	// Test an empty command
 	args, err = SanitizeCommand("")