Skip to content

Commit 84edf7e

Browse files
nacxnacx
authored
mcp: allow configuring the session encryption iterations (#1598)
**Description** Allow configuring the number of iterations to be executed for the MCP session encryption. MCP sessions are encrypted using PBKDF2 for key derivation and AES-GCM for encryption. This provides some flexibility to allow users to balance between security and performance. The default implementation defaults to 100.000 iterations, which should be a reasonable tradeoff between security and speed. This PR introduces a new Helm value to allow users change this value to make the MCP proxy faster if needed. The following benchmarks show the behaviour for different values: ``` Running tool: /opt/homebrew/opt/[email protected]/bin/go test -benchmem -run=^$ -bench ^BenchmarkPBKDF2AesGcmSessionCrypto$ github.com/envoyproxy/ai-gateway/internal/mcpproxy goos: darwin goarch: arm64 pkg: github.com/envoyproxy/ai-gateway/internal/mcpproxy cpu: Apple M1 BenchmarkPBKDF2AesGcmSessionCrypto/encrypt_100-8 85375 14220 ns/op 2560 B/op 20 allocs/op BenchmarkPBKDF2AesGcmSessionCrypto/decrypt_100-8 87002 14356 ns/op 2260 B/op 16 allocs/op BenchmarkPBKDF2AesGcmSessionCrypto/encrypt_1000-8 8824 126737 ns/op 2560 B/op 20 allocs/op BenchmarkPBKDF2AesGcmSessionCrypto/decrypt_1000-8 9558 116771 ns/op 2260 B/op 16 allocs/op BenchmarkPBKDF2AesGcmSessionCrypto/encrypt_10000-8 1056 1512627 ns/op 2560 B/op 20 allocs/op BenchmarkPBKDF2AesGcmSessionCrypto/decrypt_10000-8 799 1350919 ns/op 2260 B/op 16 allocs/op BenchmarkPBKDF2AesGcmSessionCrypto/encrypt_50000-8 211 5655061 ns/op 2560 B/op 20 allocs/op BenchmarkPBKDF2AesGcmSessionCrypto/decrypt_50000-8 208 5703031 ns/op 2260 B/op 16 allocs/op BenchmarkPBKDF2AesGcmSessionCrypto/encrypt_100000-8 100 11378411 ns/op 2560 B/op 20 allocs/op BenchmarkPBKDF2AesGcmSessionCrypto/decrypt_100000-8 100 11343512 ns/op 2260 B/op 16 allocs/op BenchmarkPBKDF2AesGcmSessionCrypto/encrypt_200000-8 52 23629280 ns/op 2560 B/op 20 allocs/op BenchmarkPBKDF2AesGcmSessionCrypto/decrypt_200000-8 50 26281746 ns/op 2260 B/op 16 allocs/op PASS ok github.com/envoyproxy/ai-gateway/internal/mcpproxy 15.551s ``` **Related Issues/PRs (if applicable)** N/A **Special notes for reviewers (if applicable)** N/A --------- Signed-off-by: Ignasi Barrera <[email protected]>
1 parent 5f56187 commit 84edf7e

File tree

12 files changed

+402
-215
lines changed

12 files changed

+402
-215
lines changed

‎cmd/controller/main.go‎

Lines changed: 62 additions & 45 deletions
Original file line numberDiff line numberDiff line change
@@ -56,10 +56,13 @@ type flags struct {
5656
// extProcMaxRecvMsgSize is the maximum message size in bytes that the gRPC server can receive.
5757
extProcMaxRecvMsgSize int
5858
// maxRecvMsgSize is the maximum message size in bytes that the gRPC extension server can receive.
59-
maxRecvMsgSize int
60-
mcpSessionEncryptionSeed string
61-
watchNamespaces []string
62-
cacheSyncTimeout time.Duration
59+
maxRecvMsgSize int
60+
mcpSessionEncryptionSeed string
61+
mcpFallbackSessionEncryptionSeed string
62+
mcpSessionEncryptionIterations int
63+
mcpFallbackSessionEncryptionIterations int
64+
watchNamespaces []string
65+
cacheSyncTimeout time.Duration
6366
}
6467

6568
// parsePullPolicy parses string into a k8s PullPolicy.
@@ -189,13 +192,14 @@ func parseAndValidateFlags(args []string) (flags, error) {
189192
2*time.Minute, // This is the controller-runtime default
190193
"Maximum time to wait for k8s caches to sync",
191194
)
192-
mcpSessionEncryptionSeed := fs.String(
193-
"mcpSessionEncryptionSeed",
194-
"seed",
195-
"Arbitrary string seed used to derive the MCP session encryption key. "+
196-
"Do not include commas as they are used as separators. You can optionally pass \"fallback\" seed after the first one to allow for key rotation. "+
197-
"For example: \"new-seed,old-seed-for-fallback\". The fallback seed is only used for decryption.",
198-
)
195+
mcpSessionEncryptionSeed := fs.String("mcpSessionEncryptionSeed", "default-insecure-seed",
196+
"Seed used to derive the MCP session encryption key. This should be changed and set to a secure value.")
197+
mcpSessionEncryptionIterations := fs.Int("mcpSessionEncryptionIterations", 100_000,
198+
"Number of iterations to use for PBKDF2 key derivation for MCP session encryption.")
199+
mcpFallbackSessionEncryptionSeed := fs.String("mcpFallbackSessionEncryptionSeed", "",
200+
"Optional fallback seed used for MCP session key rotation")
201+
mcpFallbackSessionEncryptionIterations := fs.Int("mcpFallbackSessionEncryptionIterations", 100_000,
202+
"Number of iterations used in the fallback PBKDF2 key derivation for MCP session encryption.")
199203

200204
if err := fs.Parse(args); err != nil {
201205
err = fmt.Errorf("failed to parse flags: %w", err)
@@ -258,28 +262,38 @@ func parseAndValidateFlags(args []string) (flags, error) {
258262
}
259263
}
260264

265+
if *mcpSessionEncryptionIterations <= 0 {
266+
return flags{}, fmt.Errorf("mcp session encryption iterations must be positive: %d", *mcpSessionEncryptionIterations)
267+
}
268+
if *mcpFallbackSessionEncryptionSeed != "" && *mcpFallbackSessionEncryptionIterations <= 0 {
269+
return flags{}, fmt.Errorf("mcp fallback session encryption iterations must be positive: %d", *mcpFallbackSessionEncryptionIterations)
270+
}
271+
261272
return flags{
262-
extProcLogLevel: *extProcLogLevelPtr,
263-
extProcImage: *extProcImagePtr,
264-
extProcImagePullPolicy: extProcPullPolicy,
265-
enableLeaderElection: *enableLeaderElectionPtr,
266-
logLevel: zapLogLevel,
267-
extensionServerPort: *extensionServerPortPtr,
268-
tlsCertDir: *tlsCertDir,
269-
tlsCertName: *tlsCertName,
270-
tlsKeyName: *tlsKeyName,
271-
caBundleName: *caBundleName,
272-
metricsRequestHeaderAttributes: *metricsRequestHeaderAttributes,
273-
spanRequestHeaderAttributes: *spanRequestHeaderAttributes,
274-
endpointPrefixes: *endpointPrefixes,
275-
rootPrefix: *rootPrefix,
276-
extProcExtraEnvVars: *extProcExtraEnvVars,
277-
extProcImagePullSecrets: *extProcImagePullSecrets,
278-
extProcMaxRecvMsgSize: *extProcMaxRecvMsgSize,
279-
maxRecvMsgSize: *maxRecvMsgSize,
280-
watchNamespaces: parseWatchNamespaces(*watchNamespaces),
281-
cacheSyncTimeout: *cacheSyncTimeout,
282-
mcpSessionEncryptionSeed: *mcpSessionEncryptionSeed,
273+
extProcLogLevel: *extProcLogLevelPtr,
274+
extProcImage: *extProcImagePtr,
275+
extProcImagePullPolicy: extProcPullPolicy,
276+
enableLeaderElection: *enableLeaderElectionPtr,
277+
logLevel: zapLogLevel,
278+
extensionServerPort: *extensionServerPortPtr,
279+
tlsCertDir: *tlsCertDir,
280+
tlsCertName: *tlsCertName,
281+
tlsKeyName: *tlsKeyName,
282+
caBundleName: *caBundleName,
283+
metricsRequestHeaderAttributes: *metricsRequestHeaderAttributes,
284+
spanRequestHeaderAttributes: *spanRequestHeaderAttributes,
285+
endpointPrefixes: *endpointPrefixes,
286+
rootPrefix: *rootPrefix,
287+
extProcExtraEnvVars: *extProcExtraEnvVars,
288+
extProcImagePullSecrets: *extProcImagePullSecrets,
289+
extProcMaxRecvMsgSize: *extProcMaxRecvMsgSize,
290+
maxRecvMsgSize: *maxRecvMsgSize,
291+
watchNamespaces: parseWatchNamespaces(*watchNamespaces),
292+
cacheSyncTimeout: *cacheSyncTimeout,
293+
mcpSessionEncryptionSeed: *mcpSessionEncryptionSeed,
294+
mcpFallbackSessionEncryptionSeed: *mcpFallbackSessionEncryptionSeed,
295+
mcpSessionEncryptionIterations: *mcpSessionEncryptionIterations,
296+
mcpFallbackSessionEncryptionIterations: *mcpFallbackSessionEncryptionIterations,
283297
}, nil
284298
}
285299

@@ -352,19 +366,22 @@ func main() {
352366

353367
// Start the controller.
354368
if err := controller.StartControllers(ctx, mgr, k8sConfig, ctrl.Log.WithName("controller"), controller.Options{
355-
ExtProcImage: parsedFlags.extProcImage,
356-
ExtProcImagePullPolicy: parsedFlags.extProcImagePullPolicy,
357-
ExtProcLogLevel: parsedFlags.extProcLogLevel,
358-
EnableLeaderElection: parsedFlags.enableLeaderElection,
359-
UDSPath: extProcUDSPath,
360-
MetricsRequestHeaderAttributes: parsedFlags.metricsRequestHeaderAttributes,
361-
TracingRequestHeaderAttributes: parsedFlags.spanRequestHeaderAttributes,
362-
EndpointPrefixes: parsedFlags.endpointPrefixes,
363-
RootPrefix: parsedFlags.rootPrefix,
364-
ExtProcExtraEnvVars: parsedFlags.extProcExtraEnvVars,
365-
ExtProcImagePullSecrets: parsedFlags.extProcImagePullSecrets,
366-
ExtProcMaxRecvMsgSize: parsedFlags.extProcMaxRecvMsgSize,
367-
MCPSessionEncryptionSeed: parsedFlags.mcpSessionEncryptionSeed,
369+
ExtProcImage: parsedFlags.extProcImage,
370+
ExtProcImagePullPolicy: parsedFlags.extProcImagePullPolicy,
371+
ExtProcLogLevel: parsedFlags.extProcLogLevel,
372+
EnableLeaderElection: parsedFlags.enableLeaderElection,
373+
UDSPath: extProcUDSPath,
374+
MetricsRequestHeaderAttributes: parsedFlags.metricsRequestHeaderAttributes,
375+
TracingRequestHeaderAttributes: parsedFlags.spanRequestHeaderAttributes,
376+
EndpointPrefixes: parsedFlags.endpointPrefixes,
377+
RootPrefix: parsedFlags.rootPrefix,
378+
ExtProcExtraEnvVars: parsedFlags.extProcExtraEnvVars,
379+
ExtProcImagePullSecrets: parsedFlags.extProcImagePullSecrets,
380+
ExtProcMaxRecvMsgSize: parsedFlags.extProcMaxRecvMsgSize,
381+
MCPSessionEncryptionSeed: parsedFlags.mcpSessionEncryptionSeed,
382+
MCPSessionEncryptionIterations: parsedFlags.mcpSessionEncryptionIterations,
383+
MCPFallbackSessionEncryptionSeed: parsedFlags.mcpFallbackSessionEncryptionSeed,
384+
MCPFallbackSessionEncryptionIterations: parsedFlags.mcpFallbackSessionEncryptionIterations,
368385
}); err != nil {
369386
setupLog.Error(err, "failed to start controller")
370387
}

‎cmd/controller/main_test.go‎

Lines changed: 28 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -56,6 +56,10 @@ func Test_parseAndValidateFlags(t *testing.T) {
5656
tc.dash + "maxRecvMsgSize=33554432",
5757
tc.dash + "watchNamespaces=default,envoy-ai-gateway-system",
5858
tc.dash + "cacheSyncTimeout=5m",
59+
tc.dash + "mcpSessionEncryptionSeed=my-seed",
60+
tc.dash + "mcpSessionEncryptionIterations=100",
61+
tc.dash + "mcpFallbackSessionEncryptionSeed=my-fallback-seed",
62+
tc.dash + "mcpFallbackSessionEncryptionIterations=200",
5963
}
6064
f, err := parseAndValidateFlags(args)
6165
require.Equal(t, "debug", f.extProcLogLevel)
@@ -70,6 +74,10 @@ func Test_parseAndValidateFlags(t *testing.T) {
7074
require.Equal(t, 32*1024*1024, f.maxRecvMsgSize)
7175
require.Equal(t, []string{"default", "envoy-ai-gateway-system"}, f.watchNamespaces)
7276
require.Equal(t, 5*time.Minute, f.cacheSyncTimeout)
77+
require.Equal(t, "my-seed", f.mcpSessionEncryptionSeed)
78+
require.Equal(t, 100, f.mcpSessionEncryptionIterations)
79+
require.Equal(t, "my-fallback-seed", f.mcpFallbackSessionEncryptionSeed)
80+
require.Equal(t, 200, f.mcpFallbackSessionEncryptionIterations)
7381
require.NoError(t, err)
7482
})
7583
}
@@ -126,6 +134,26 @@ func Test_parseAndValidateFlags(t *testing.T) {
126134
flags: []string{"--endpointPrefixes=openai"},
127135
expErr: "invalid endpoint prefixes",
128136
},
137+
{
138+
name: "invalid mcp session encryption iterations",
139+
flags: []string{"--mcpSessionEncryptionIterations=invalid"},
140+
expErr: `invalid value "invalid" for flag -mcpSessionEncryptionIterations: parse error`,
141+
},
142+
{
143+
name: "negative mcp session encryption iterations",
144+
flags: []string{"--mcpSessionEncryptionIterations=-1"},
145+
expErr: "mcp session encryption iterations must be positive: -1",
146+
},
147+
{
148+
name: "invalid mcp fallback session encryption iterations",
149+
flags: []string{"--mcpFallbackSessionEncryptionSeed=fallback", "--mcpFallbackSessionEncryptionIterations=invalid"},
150+
expErr: `invalid value "invalid" for flag -mcpFallbackSessionEncryptionIterations: parse error`,
151+
},
152+
{
153+
name: "negative mcp fallback session encryption iterations",
154+
flags: []string{"--mcpFallbackSessionEncryptionSeed=fallback", "--mcpFallbackSessionEncryptionIterations=-1"},
155+
expErr: "mcp fallback session encryption iterations must be positive: -1",
156+
},
129157
} {
130158
t.Run(tc.name, func(t *testing.T) {
131159
_, err := parseAndValidateFlags(tc.flags)

‎cmd/extproc/mainlib/main.go‎

Lines changed: 31 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -36,15 +36,18 @@ import (
3636

3737
// extProcFlags is the struct that holds the flags passed to the external processor.
3838
type extProcFlags struct {
39-
configPath string // path to the configuration file.
40-
extProcAddr string // gRPC address for the external processor.
41-
logLevel slog.Level // log level for the external processor.
42-
adminPort int // HTTP port for the admin server (metrics and health).
43-
metricsRequestHeaderAttributes string // comma-separated key-value pairs for mapping HTTP request headers to otel metric attributes.
44-
spanRequestHeaderAttributes string // comma-separated key-value pairs for mapping HTTP request headers to otel span attributes.
45-
mcpAddr string // address for the MCP proxy server which can be either tcp or unix domain socket.
46-
mcpSessionEncryptionSeed string // Seed for deriving the key for encrypting MCP sessions.
47-
mcpWriteTimeout time.Duration // the maximum duration before timing out writes of the MCP response.
39+
configPath string // path to the configuration file.
40+
extProcAddr string // gRPC address for the external processor.
41+
logLevel slog.Level // log level for the external processor.
42+
adminPort int // HTTP port for the admin server (metrics and health).
43+
metricsRequestHeaderAttributes string // comma-separated key-value pairs for mapping HTTP request headers to otel metric attributes.
44+
spanRequestHeaderAttributes string // comma-separated key-value pairs for mapping HTTP request headers to otel span attributes.
45+
mcpAddr string // address for the MCP proxy server which can be either tcp or unix domain socket.
46+
mcpSessionEncryptionSeed string // Seed for deriving the key for encrypting MCP sessions.
47+
mcpSessionEncryptionIterations int // Number of iterations to use for PBKDF2 key derivation for MCP session encryption.
48+
mcpFallbackSessionEncryptionSeed string // Fallback seed for deriving the key for encrypting MCP sessions.
49+
mcpFallbackSessionEncryptionIterations int // Number of iterations to use for PBKDF2 key derivation for fallback MCP session encryption.
50+
mcpWriteTimeout time.Duration // the maximum duration before timing out writes of the MCP response.
4851
// rootPrefix is the root prefix for all the processors.
4952
rootPrefix string
5053
// maxRecvMsgSize is the maximum message size in bytes that the gRPC server can receive.
@@ -104,13 +107,14 @@ func parseAndValidateFlags(args []string) (extProcFlags, error) {
104107
"Maximum message size in bytes that the gRPC server can receive. Default is 4MB.",
105108
)
106109
fs.StringVar(&flags.mcpAddr, "mcpAddr", "", "the address (TCP or UDS) for the MCP proxy server, such as :1063 or unix:///tmp/ext_proc.sock. Optional.")
107-
fs.StringVar(&flags.mcpSessionEncryptionSeed,
108-
"mcpSessionEncryptionSeed",
109-
"default-insecure-seed",
110-
"Arbitrary string seed used to derive the MCP session encryption key. "+
111-
"Do not include commas as they are used as separators. You can optionally pass \"fallback\" seed after the first one to allow for key rotation. "+
112-
"For example: \"new-seed,old-seed-for-fallback\". The fallback seed is only used for decryption.",
113-
)
110+
fs.StringVar(&flags.mcpSessionEncryptionSeed, "mcpSessionEncryptionSeed", "default-insecure-seed",
111+
"Seed used to derive the MCP session encryption key. This should be changed and set to a secure value.")
112+
fs.IntVar(&flags.mcpSessionEncryptionIterations, "mcpSessionEncryptionIterations", 100_000,
113+
"Number of iterations to use for PBKDF2 key derivation for MCP session encryption.")
114+
fs.StringVar(&flags.mcpFallbackSessionEncryptionSeed, "mcpFallbackSessionEncryptionSeed", "",
115+
"Optional fallback seed used for MCP session key rotation.")
116+
fs.IntVar(&flags.mcpFallbackSessionEncryptionIterations, "mcpFallbackSessionEncryptionIterations", 100_000,
117+
"Number of iterations used in the fallback PBKDF2 key derivation for MCP session encryption.")
114118
fs.DurationVar(&flags.mcpWriteTimeout, "mcpWriteTimeout", 120*time.Second,
115119
"The maximum duration before timing out writes of the MCP response")
116120

@@ -269,8 +273,17 @@ func Main(ctx context.Context, args []string, stderr io.Writer) (err error) {
269273

270274
var mcpServer *http.Server
271275
if mcpLis != nil {
272-
seed, fallbackSeed, _ := strings.Cut(flags.mcpSessionEncryptionSeed, ",")
273-
mcpSessionCrypto := mcpproxy.DefaultSessionCrypto(seed, fallbackSeed)
276+
mcpSessionCrypto := mcpproxy.NewPBKDF2AesGcmSessionCrypto(flags.mcpSessionEncryptionSeed, flags.mcpSessionEncryptionIterations)
277+
if flags.mcpFallbackSessionEncryptionSeed != "" {
278+
mcpSessionCrypto = &mcpproxy.FallbackEnabledSessionCrypto{
279+
Primary: mcpSessionCrypto,
280+
Fallback: mcpproxy.NewPBKDF2AesGcmSessionCrypto(
281+
flags.mcpFallbackSessionEncryptionSeed,
282+
flags.mcpFallbackSessionEncryptionIterations,
283+
),
284+
}
285+
}
286+
274287
var mcpProxyMux *http.ServeMux
275288
var mcpProxyConfig *mcpproxy.ProxyConfig
276289
mcpProxyConfig, mcpProxyMux, err = mcpproxy.NewMCPProxy(l.With("component", "mcp-proxy"), mcpMetrics,

‎internal/controller/controller.go‎

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -87,6 +87,12 @@ type Options struct {
8787
ExtProcMaxRecvMsgSize int
8888
// MCPSessionEncryptionSeed is the seed used to derive the encryption key for MCP session encryption.
8989
MCPSessionEncryptionSeed string
90+
// MCPSessionEncryptionIterations is the number of iterations to use for PBKDF2 key derivation for MCP session encryption.
91+
MCPSessionEncryptionIterations int
92+
// MCPFallbackSessionEncryptionSeed is the optional fallback seed used for MCP session key rotation.
93+
MCPFallbackSessionEncryptionSeed string
94+
// MCPFallbackSessionEncryptionIterations is the number of iterations used in the fallback PBKDF2 key derivation for MCP session encryption.
95+
MCPFallbackSessionEncryptionIterations int
9096
// EndpointPrefixes is the comma-separated key-value pairs for endpoint prefixes.
9197
EndpointPrefixes string
9298
}
@@ -228,6 +234,9 @@ func StartControllers(ctx context.Context, mgr manager.Manager, config *rest.Con
228234
options.ExtProcMaxRecvMsgSize,
229235
isKubernetes133OrLater(versionInfo, logger),
230236
options.MCPSessionEncryptionSeed,
237+
options.MCPSessionEncryptionIterations,
238+
options.MCPFallbackSessionEncryptionSeed,
239+
options.MCPFallbackSessionEncryptionIterations,
231240
))
232241
mgr.GetWebhookServer().Register("/mutate", &webhook.Admission{Handler: h})
233242
}

0 commit comments

Comments
 (0)