Skip to content

Commit 9acf1ed

Browse files
ilopezlunailopezluna
authored
Filter runtime keys (#601)
* feat(runtime): enhance runtime flag validation with backend allowlist and path safety checks * feat(runtime): add KV cache flags to runtime flags allowlist and update tests * feat(runtime): expand allowed flags for llama.cpp and vLLM backends with additional categories and safety checks * test(runtime): replace custom string contains function with strings.Contains for error message validation * feat(runtime): add flag injection validation to prevent disallowed flags in values
1 parent d2ac536 commit 9acf1ed

5 files changed

Lines changed: 785 additions & 101 deletions

File tree

‎pkg/inference/runtime_flags.go‎

Lines changed: 52 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,33 @@ import (
55
"strings"
66
)
77

8-
// ValidateRuntimeFlags ensures runtime flags don't contain paths (forward slash "/" or backslash "\")
8+
// ValidateRuntimeFlags validates runtime flags against the backend's allowlist
9+
// and checks for path characters as defense-in-depth.
10+
func ValidateRuntimeFlags(backendName string, flags []string) error {
11+
// Get allowlist for this backend
12+
allowedFlags := GetAllowedFlags(backendName)
13+
14+
// Check each flag against allowlist
15+
for _, flag := range flags {
16+
flagKey := ParseFlagKey(flag)
17+
if flagKey == "" {
18+
continue // Skip values, only validate flag keys
19+
}
20+
if !allowedFlags[flagKey] {
21+
return fmt.Errorf("runtime flag %q is not allowed for backend %q", flagKey, backendName)
22+
}
23+
}
24+
25+
// Check for flags in values (e.g., --seed=--log-file=foo or --seed=-l)
26+
if err := validateNoFlagInjection(flags); err != nil {
27+
return err
28+
}
29+
30+
// still check for path characters in values
31+
return validatePathSafety(flags)
32+
}
33+
34+
// validatePathSafety ensures runtime flags don't contain paths (forward slash "/" or backslash "\")
935
// to prevent malicious users from overwriting host files via arguments like
1036
// --log-file /some/path, --output-file /etc/passwd, or --log-file C:\Windows\file.
1137
//
@@ -17,11 +43,35 @@ import (
1743
// - UNC paths: \\network\share\file
1844
//
1945
// Returns an error if any flag contains a forward slash or backslash.
20-
func ValidateRuntimeFlags(flags []string) error {
46+
func validatePathSafety(flags []string) error {
2147
for _, flag := range flags {
2248
if strings.Contains(flag, "/") || strings.Contains(flag, "\\") {
2349
return fmt.Errorf("invalid runtime flag %q: paths are not allowed (contains '/' or '\\\\')", flag)
2450
}
2551
}
2652
return nil
2753
}
54+
55+
// validateNoFlagInjection checks for flags in values when using the = format.
56+
// This prevents attacks like --seed=--log-file=foo or --seed=-l where disallowed flags
57+
// are embedded as values.
58+
// Values starting with - are only allowed if followed by a digit (negative numbers like -1, -0.5).
59+
func validateNoFlagInjection(flags []string) error {
60+
for _, flag := range flags {
61+
if idx := strings.Index(flag, "="); idx != -1 {
62+
value := flag[idx+1:]
63+
if strings.HasPrefix(value, "-") {
64+
// Allow negative numbers (-1, -0.5) but reject flags (-l, --flag)
65+
if len(value) < 2 || !isDigit(value[1]) {
66+
return fmt.Errorf("invalid flag %q: value cannot start with '-' unless followed by a digit", flag)
67+
}
68+
}
69+
}
70+
}
71+
return nil
72+
}
73+
74+
// isDigit returns true if the byte is an ASCII digit (0-9)
75+
func isDigit(b byte) bool {
76+
return b >= '0' && b <= '9'
77+
}

‎pkg/inference/runtime_flags_allowlist.go‎

Lines changed: 236 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,236 @@
1+
package inference
2+
3+
import "strings"
4+
5+
// LlamaCppAllowedFlags contains safe flags for llama.cpp server.
6+
// This list is based on llama.cpp server documentation.
7+
// Flags involving file paths are intentionally excluded for security.
8+
var LlamaCppAllowedFlags = map[string]bool{
9+
// Threading and CPU control
10+
"-t": true, "--threads": true,
11+
"-tb": true, "--threads-batch": true,
12+
"-C": true, "--cpu-mask": true,
13+
"-Cr": true, "--cpu-range": true,
14+
"--cpu-strict": true,
15+
"--prio": true,
16+
"--poll": true,
17+
"-Cb": true, "--cpu-mask-batch": true,
18+
"-Crb": true, "--cpu-range-batch": true,
19+
"--cpu-strict-batch": true,
20+
"--prio-batch": true,
21+
"--poll-batch": true,
22+
23+
// Context and prediction
24+
"-c": true, "--ctx-size": true,
25+
"-n": true, "--predict": true, "--n-predict": true,
26+
"--keep": true,
27+
28+
// Batching and performance
29+
"-b": true, "--batch-size": true,
30+
"-ub": true, "--ubatch-size": true,
31+
"--swa-full": true,
32+
"-fa": true, "--flash-attn": true,
33+
"--perf": true, "--no-perf": true,
34+
35+
// Sampling parameters
36+
"--samplers": true,
37+
"-s": true, "--seed": true,
38+
"--temp": true, "--temperature": true,
39+
"--top-k": true,
40+
"--top-p": true,
41+
"--min-p": true,
42+
"--top-nsigma": true,
43+
"--xtc-probability": true,
44+
"--xtc-threshold": true,
45+
"--typical": true,
46+
"--repeat-last-n": true,
47+
"--repeat-penalty": true,
48+
"--presence-penalty": true,
49+
"--frequency-penalty": true,
50+
"--dry-multiplier": true,
51+
"--dry-base": true,
52+
"--dry-allowed-length": true,
53+
"--dry-penalty-last-n": true,
54+
"--mirostat": true,
55+
"--mirostat-lr": true,
56+
"--mirostat-ent": true,
57+
"--ignore-eos": true,
58+
"--dynatemp-range": true,
59+
"--dynatemp-exp": true,
60+
61+
// GPU and device management
62+
"-dev": true, "--device": true,
63+
"-ngl": true, "--gpu-layers": true, "--n-gpu-layers": true,
64+
"-sm": true, "--split-mode": true,
65+
"-ts": true, "--tensor-split": true,
66+
"-mg": true, "--main-gpu": true,
67+
"-fit": true, "--fit": true,
68+
"-fitt": true, "--fit-target": true,
69+
"-fitc": true, "--fit-ctx": true,
70+
71+
// Memory and caching
72+
"-kvo": true, "--kv-offload": true,
73+
"-nkvo": true, "--no-kv-offload": true,
74+
"--repack": true, "-nr": true, "--no-repack": true,
75+
"--no-host": true,
76+
"-ctk": true, "--cache-type-k": true,
77+
"-ctv": true, "--cache-type-v": true,
78+
"--mlock": true,
79+
"--mmap": true, "--no-mmap": true,
80+
"-dio": true, "--direct-io": true,
81+
"-ndio": true, "--no-direct-io": true,
82+
"-cram": true, "--cache-ram": true,
83+
"-kvu": true, "--kv-unified": true,
84+
"--context-shift": true, "--no-context-shift": true,
85+
86+
// RoPE scaling
87+
"--rope-scaling": true,
88+
"--rope-scale": true,
89+
"--rope-freq-base": true,
90+
"--rope-freq-scale": true,
91+
"--yarn-orig-ctx": true,
92+
"--yarn-ext-factor": true,
93+
"--yarn-attn-factor": true,
94+
"--yarn-beta-slow": true,
95+
"--yarn-beta-fast": true,
96+
97+
// Server configuration
98+
"-np": true, "--parallel": true,
99+
"-cb": true, "--cont-batching": true,
100+
"-nocb": true, "--no-cont-batching": true,
101+
"--warmup": true, "--no-warmup": true,
102+
"-to": true, "--timeout": true,
103+
"--threads-http": true,
104+
"--cache-prompt": true,
105+
"--no-cache-prompt": true,
106+
"--cache-reuse": true,
107+
"--sleep-idle-seconds": true,
108+
109+
// Multimodal (safe flags only - no file paths)
110+
"--mmproj-auto": true, "--no-mmproj": true, "--no-mmproj-auto": true,
111+
"--mmproj-offload": true, "--no-mmproj-offload": true,
112+
"--image-min-tokens": true,
113+
"--image-max-tokens": true,
114+
"--spm-infill": true,
115+
116+
// Speculative decoding (safe flags only - no file paths)
117+
"--draft": true, "--draft-n": true, "--draft-max": true,
118+
"--draft-min": true, "--draft-n-min": true,
119+
"--draft-p-min": true,
120+
"-cd": true, "--ctx-size-draft": true,
121+
"-devd": true, "--device-draft": true,
122+
"-ngld": true, "--gpu-layers-draft": true, "--n-gpu-layers-draft": true,
123+
"-td": true, "--threads-draft": true,
124+
"-tbd": true, "--threads-batch-draft": true,
125+
126+
// LoRA (safe flags only - no file paths)
127+
"--lora-init-without-apply": true,
128+
129+
// Control vectors (safe flags only - no file paths)
130+
"--control-vector-layer-range": true,
131+
132+
// Grammar and constraints (safe flags only - no file paths)
133+
"--grammar": true,
134+
"-j": true, "--json-schema": true,
135+
"-bs": true, "--backend-sampling": true,
136+
137+
// Template and format control (safe flags only - no file paths)
138+
"--chat-template": true,
139+
"--chat-template-kwargs": true,
140+
"--jinja": true, "--no-jinja": true,
141+
"--pooling": true,
142+
"--reasoning-format": true,
143+
"--reasoning-budget": true,
144+
"--prefill-assistant": true,
145+
"--no-prefill-assistant": true,
146+
147+
// Web interface and API (safe flags only - no file paths)
148+
"--api-prefix": true,
149+
"--webui": true, "--no-webui": true,
150+
"--webui-config": true,
151+
"--api-key": true,
152+
"--metrics": true,
153+
"--no-metrics": true,
154+
"--props": true,
155+
"--slots": true, "--no-slots": true,
156+
157+
// Embedding and specialized
158+
"--embedding": true, "--embeddings": true,
159+
"--rerank": true, "--reranking": true,
160+
"-sps": true, "--slot-prompt-similarity": true,
161+
162+
// Tensor and computation (safe flags only)
163+
"-cmoe": true, "--cpu-moe": true,
164+
"-ncmoe": true, "--n-cpu-moe": true,
165+
"--check-tensors": true,
166+
"--op-offload": true, "--no-op-offload": true,
167+
168+
// Verbose/debug
169+
"-v": true, "--verbose": true,
170+
}
171+
172+
// VLLMAllowedFlags contains safe flags for vLLM engine.
173+
// Flags involving file paths are intentionally excluded for security.
174+
var VLLMAllowedFlags = map[string]bool{
175+
// Parallelism
176+
"--tensor-parallel-size": true, "-tp": true,
177+
"--pipeline-parallel-size": true, "-pp": true,
178+
179+
// Model configuration
180+
"--max-model-len": true,
181+
"--max-num-batched-tokens": true,
182+
"--max-num-seqs": true,
183+
"--block-size": true,
184+
"--swap-space": true,
185+
"--seed": true,
186+
187+
// Data types and quantization
188+
"--dtype": true,
189+
"--quantization": true,
190+
"-q": true,
191+
"--kv-cache-dtype": true,
192+
193+
// Performance flags
194+
"--enforce-eager": true,
195+
"--enable-prefix-caching": true,
196+
"--enable-chunked-prefill": true,
197+
"--disable-custom-all-reduce": true,
198+
"--use-v2-block-manager": true,
199+
200+
// Tokenizer
201+
"--tokenizer-mode": true,
202+
"--trust-remote-code": true,
203+
"--max-logprobs": true,
204+
205+
// Misc
206+
"--revision": true,
207+
"--load-format": true,
208+
"--disable-log-stats": true,
209+
"--served-model-name": true,
210+
211+
// GPU memory
212+
"--gpu-memory-utilization": true,
213+
}
214+
215+
// AllowedFlags maps backend names to their allowed flag keys
216+
var AllowedFlags = map[string]map[string]bool{
217+
"llama.cpp": LlamaCppAllowedFlags,
218+
"vllm": VLLMAllowedFlags,
219+
}
220+
221+
// ParseFlagKey extracts the flag key from a flag string.
222+
// "--threads=4" -> "--threads", "-t" -> "-t", "4" -> ""
223+
func ParseFlagKey(flag string) string {
224+
if !strings.HasPrefix(flag, "-") {
225+
return "" // Not a flag, it's a value
226+
}
227+
if idx := strings.Index(flag, "="); idx != -1 {
228+
return flag[:idx]
229+
}
230+
return flag
231+
}
232+
233+
// GetAllowedFlags returns the allowlist for a backend, or nil if unknown
234+
func GetAllowedFlags(backendName string) map[string]bool {
235+
return AllowedFlags[backendName]
236+
}

0 commit comments

Comments
 (0)