Skip to content

Commit eb95b25

Browse files
committed
Add macros to Configuration schema
1 parent 25f3dc2 commit eb95b25

File tree

2 files changed

+177
-9
lines changed

2 files changed

+177
-9
lines changed

proxy/config.go

Lines changed: 67 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@ import (
44
"fmt"
55
"io"
66
"os"
7+
"regexp"
78
"runtime"
89
"sort"
910
"strconv"
@@ -67,6 +68,9 @@ type Config struct {
6768
Profiles map[string][]string `yaml:"profiles"`
6869
Groups map[string]GroupConfig `yaml:"groups"` /* key is group ID */
6970

71+
// for key/value replacements in model's cmd, cmdStop, proxy, checkEndPoint
72+
Macros map[string]string `yaml:"macros"`
73+
7074
// map aliases to actual model IDs
7175
aliases map[string]string
7276

@@ -141,6 +145,30 @@ func LoadConfigFromReader(r io.Reader) (Config, error) {
141145
}
142146
}
143147

148+
/* check macro constraint rules:
149+
150+
- name must fit the regex ^[a-zA-Z0-9_-]+$
151+
- names must be less than 64 characters (no reason, just cause)
152+
- name can not be any reserved macros: PORT
153+
- macro values must be less than 1024 characters
154+
*/
155+
macroNameRegex := regexp.MustCompile(`^[a-zA-Z0-9_-]+$`)
156+
for macroName, macroValue := range config.Macros {
157+
if len(macroName) >= 64 {
158+
return Config{}, fmt.Errorf("macro name '%s' exceeds maximum length of 63 characters", macroName)
159+
}
160+
if !macroNameRegex.MatchString(macroName) {
161+
return Config{}, fmt.Errorf("macro name '%s' contains invalid characters, must match pattern ^[a-zA-Z0-9_-]+$", macroName)
162+
}
163+
if len(macroValue) >= 1024 {
164+
return Config{}, fmt.Errorf("macro value for '%s' exceeds maximum length of 1024 characters", macroName)
165+
}
166+
switch macroName {
167+
case "PORT":
168+
return Config{}, fmt.Errorf("macro name '%s' is reserved and cannot be used", macroName)
169+
}
170+
}
171+
144172
// Get and sort all model IDs first, makes testing more consistent
145173
modelIds := make([]string, 0, len(config.Models))
146174
for modelId := range config.Models {
@@ -151,19 +179,51 @@ func LoadConfigFromReader(r io.Reader) (Config, error) {
151179
nextPort := config.StartPort
152180
for _, modelId := range modelIds {
153181
modelConfig := config.Models[modelId]
154-
// iterate over the models and replace any ${PORT} with the next available port
155-
if strings.Contains(modelConfig.Cmd, "${PORT}") {
156-
modelConfig.Cmd = strings.ReplaceAll(modelConfig.Cmd, "${PORT}", strconv.Itoa(nextPort))
182+
183+
// go through model config fields: cmd, cmdStop, proxy, checkEndPoint and replace macros with macro values
184+
for macroName, macroValue := range config.Macros {
185+
macroSlug := fmt.Sprintf("${%s}", macroName)
186+
modelConfig.Cmd = strings.ReplaceAll(modelConfig.Cmd, macroSlug, macroValue)
187+
modelConfig.CmdStop = strings.ReplaceAll(modelConfig.CmdStop, macroSlug, macroValue)
188+
modelConfig.Proxy = strings.ReplaceAll(modelConfig.Proxy, macroSlug, macroValue)
189+
modelConfig.CheckEndpoint = strings.ReplaceAll(modelConfig.CheckEndpoint, macroSlug, macroValue)
190+
}
191+
192+
// only iterate over models that use ${PORT} to keep port numbers from increasing unnecessarily
193+
if strings.Contains(modelConfig.Cmd, "${PORT}") || strings.Contains(modelConfig.Proxy, "${PORT}") || strings.Contains(modelConfig.CmdStop, "${PORT}") {
157194
if modelConfig.Proxy == "" {
158-
modelConfig.Proxy = fmt.Sprintf("http://localhost:%d", nextPort)
159-
} else {
160-
modelConfig.Proxy = strings.ReplaceAll(modelConfig.Proxy, "${PORT}", strconv.Itoa(nextPort))
195+
modelConfig.Proxy = "http://localhost:${PORT}"
161196
}
197+
198+
nextPortStr := strconv.Itoa(nextPort)
199+
modelConfig.Cmd = strings.ReplaceAll(modelConfig.Cmd, "${PORT}", nextPortStr)
200+
modelConfig.CmdStop = strings.ReplaceAll(modelConfig.CmdStop, "${PORT}", nextPortStr)
201+
modelConfig.Proxy = strings.ReplaceAll(modelConfig.Proxy, "${PORT}", nextPortStr)
162202
nextPort++
163-
config.Models[modelId] = modelConfig
164203
} else if modelConfig.Proxy == "" {
165204
return Config{}, fmt.Errorf("model %s requires a proxy value when not using automatic ${PORT}", modelId)
166205
}
206+
207+
// make sure there are no unknown macros that have not been replaced
208+
macroPattern := regexp.MustCompile(`\$\{([a-zA-Z0-9_-]+)\}`)
209+
fieldMap := map[string]string{
210+
"cmd": modelConfig.Cmd,
211+
"cmdStop": modelConfig.CmdStop,
212+
"proxy": modelConfig.Proxy,
213+
"checkEndpoint": modelConfig.CheckEndpoint,
214+
}
215+
216+
for fieldName, fieldValue := range fieldMap {
217+
matches := macroPattern.FindAllStringSubmatch(fieldValue, -1)
218+
for _, match := range matches {
219+
macroName := match[1]
220+
if _, exists := config.Macros[macroName]; !exists {
221+
return Config{}, fmt.Errorf("unknown macro '${%s}' found in %s.%s", macroName, modelId, fieldName)
222+
}
223+
}
224+
}
225+
226+
config.Models[modelId] = modelConfig
167227
}
168228

169229
config = AddDefaultGroupToConfig(config)

proxy/config_test.go

Lines changed: 110 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,8 @@ func TestConfig_Load(t *testing.T) {
1919

2020
tempFile := filepath.Join(tempDir, "config.yaml")
2121
content := `
22+
macros:
23+
svr-path: "path/to/server"
2224
models:
2325
model1:
2426
cmd: path/to/cmd --arg1 one
@@ -31,7 +33,7 @@ models:
3133
- "VAR2=value2"
3234
checkEndpoint: "/health"
3335
model2:
34-
cmd: path/to/cmd --arg1 one
36+
cmd: ${svr-path} --arg1 one
3537
proxy: "http://localhost:8081"
3638
aliases:
3739
- "m2"
@@ -76,6 +78,9 @@ groups:
7678

7779
expected := Config{
7880
StartPort: 5800,
81+
Macros: map[string]string{
82+
"svr-path": "path/to/server",
83+
},
7984
Models: map[string]ModelConfig{
8085
"model1": {
8186
Cmd: "path/to/cmd --arg1 one",
@@ -85,7 +90,7 @@ groups:
8590
CheckEndpoint: "/health",
8691
},
8792
"model2": {
88-
Cmd: "path/to/cmd --arg1 one",
93+
Cmd: "path/to/server --arg1 one",
8994
Proxy: "http://localhost:8081",
9095
Aliases: []string{"m2"},
9196
Env: nil,
@@ -331,3 +336,106 @@ models:
331336
assert.Equal(t, "model model1 requires a proxy value when not using automatic ${PORT}", err.Error())
332337
})
333338
}
339+
340+
func TestConfig_MacroReplacement(t *testing.T) {
341+
content := `
342+
startPort: 9990
343+
macros:
344+
svr-path: "path/to/server"
345+
argOne: "--arg1"
346+
argTwo: "--arg2"
347+
autoPort: "--port ${PORT}"
348+
349+
models:
350+
model1:
351+
cmd: |
352+
${svr-path} ${argTwo}
353+
# the automatic ${PORT} is replaced
354+
${autoPort}
355+
${argOne}
356+
--arg3 three
357+
cmdStop: |
358+
/path/to/stop.sh --port ${PORT} ${argTwo}
359+
`
360+
361+
config, err := LoadConfigFromReader(strings.NewReader(content))
362+
assert.NoError(t, err)
363+
sanitizedCmd, err := SanitizeCommand(config.Models["model1"].Cmd)
364+
assert.NoError(t, err)
365+
assert.Equal(t, "path/to/server --arg2 --port 9990 --arg1 --arg3 three", strings.Join(sanitizedCmd, " "))
366+
367+
sanitizedCmdStop, err := SanitizeCommand(config.Models["model1"].CmdStop)
368+
assert.NoError(t, err)
369+
assert.Equal(t, "/path/to/stop.sh --port 9990 --arg2", strings.Join(sanitizedCmdStop, " "))
370+
}
371+
372+
func TestConfig_MacroErrorOnUnknownMacros(t *testing.T) {
373+
tests := []struct {
374+
name string
375+
field string
376+
content string
377+
}{
378+
{
379+
name: "unknown macro in cmd",
380+
field: "cmd",
381+
content: `
382+
startPort: 9990
383+
macros:
384+
svr-path: "path/to/server"
385+
models:
386+
model1:
387+
cmd: |
388+
${svr-path} --port ${PORT}
389+
${unknownMacro}
390+
`,
391+
},
392+
{
393+
name: "unknown macro in cmdStop",
394+
field: "cmdStop",
395+
content: `
396+
startPort: 9990
397+
macros:
398+
svr-path: "path/to/server"
399+
models:
400+
model1:
401+
cmd: "${svr-path} --port ${PORT}"
402+
cmdStop: "kill ${unknownMacro}"
403+
`,
404+
},
405+
{
406+
name: "unknown macro in proxy",
407+
field: "proxy",
408+
content: `
409+
startPort: 9990
410+
macros:
411+
svr-path: "path/to/server"
412+
models:
413+
model1:
414+
cmd: "${svr-path} --port ${PORT}"
415+
proxy: "http://localhost:${unknownMacro}"
416+
`,
417+
},
418+
{
419+
name: "unknown macro in checkEndpoint",
420+
field: "checkEndpoint",
421+
content: `
422+
startPort: 9990
423+
macros:
424+
svr-path: "path/to/server"
425+
models:
426+
model1:
427+
cmd: "${svr-path} --port ${PORT}"
428+
checkEndpoint: "http://localhost:${unknownMacro}/health"
429+
`,
430+
},
431+
}
432+
433+
for _, tt := range tests {
434+
t.Run(tt.name, func(t *testing.T) {
435+
_, err := LoadConfigFromReader(strings.NewReader(tt.content))
436+
assert.Error(t, err)
437+
assert.Contains(t, err.Error(), "unknown macro '${unknownMacro}' found in model1."+tt.field)
438+
//t.Log(err)
439+
})
440+
}
441+
}

0 commit comments

Comments
 (0)