From a4d063e260800a427a1da5ee03ae9dc1877076ad Mon Sep 17 00:00:00 2001 From: El Mehdi Abenhazou Date: Wed, 3 Jun 2026 00:06:01 +0000 Subject: [PATCH] security: reject caller writes to reserved internal session state keys The dev web API accepted arbitrary session state from callers through both initial session creation and per-run stateDelta. Internal framework keys such as _code_execution_context are stored in the same mutable session-state namespace that callers can write to. The code execution path in CodeExecution.getOrSetExecutionId() reads _code_execution_context.execution_session_id from session state and prefers any pre-existing value over deriving a new one from the ADK session id. This means a caller who seeds that key via createSession state or stateDelta can steer which backing interpreter session is used by VertexAiCodeExecutor, potentially accessing another session's files, variables, or prior execution context. Fix: - Add RESERVED_STATE_KEYS constant enumerating all internal executor control keys (_code_execution_context, _code_executor_input_files, _code_executor_error_counts, _code_execution_results). - Add sanitizeCallerState() helper that strips reserved keys from any caller-supplied state map and logs a warning for each dropped key. - Apply sanitizeCallerState() in createSession() so reserved keys in initial session state are never persisted. - In appendEvent(), skip writes for reserved keys in stateDelta with a warning log, so per-run stateDelta cannot overwrite them either. Signed-off-by: El Mehdi Abenhazou --- .../adk/sessions/InMemorySessionService.java | 57 ++++++++++++++++++- 1 file changed, 54 insertions(+), 3 deletions(-) diff --git a/core/src/main/java/com/google/adk/sessions/InMemorySessionService.java b/core/src/main/java/com/google/adk/sessions/InMemorySessionService.java index 72a14cc4d..c11e3d1db 100644 --- a/core/src/main/java/com/google/adk/sessions/InMemorySessionService.java +++ b/core/src/main/java/com/google/adk/sessions/InMemorySessionService.java @@ -20,6 +20,7 @@ import com.google.adk.events.Event; import com.google.adk.events.EventActions; +import com.google.common.collect.ImmutableSet; import com.google.common.collect.ImmutableList; import com.google.errorprone.annotations.CanIgnoreReturnValue; import io.reactivex.rxjava3.core.Completable; @@ -32,10 +33,13 @@ import java.util.Map; import java.util.Objects; import java.util.Optional; +import java.util.Set; import java.util.UUID; import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.ConcurrentMap; import org.jspecify.annotations.Nullable; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; /** * An in-memory implementation of {@link BaseSessionService} assuming {@link Session} objects are @@ -49,6 +53,23 @@ * during retrieval operations ({@code getSession}, {@code createSession}). */ public final class InMemorySessionService implements BaseSessionService { + + private static final Logger log = LoggerFactory.getLogger(InMemorySessionService.class); + + /** + * Reserved session-state keys that are managed internally by the ADK framework. Callers are not + * permitted to set or override these keys through the public API (initial session state or + * per-run stateDelta). Allowing external writes to these keys would let an untrusted caller steer + * internal framework behaviour, such as hijacking the code-execution session identifier used by + * {@code VertexAiCodeExecutor}. + */ + private static final Set RESERVED_STATE_KEYS = + ImmutableSet.of( + "_code_execution_context", + "_code_executor_input_files", + "_code_executor_error_counts", + "_code_execution_results"); + // Structure: appName -> userId -> sessionId -> Session private final ConcurrentMap>> sessions; @@ -65,6 +86,31 @@ public InMemorySessionService() { this.appState = new ConcurrentHashMap<>(); } + /** + * Removes reserved internal keys from a caller-supplied state map before it is persisted. + * Logs a warning for each key that is dropped. + * + * @param state The caller-supplied state map (may be null). + * @return A new {@link ConcurrentHashMap} containing only the non-reserved entries. + */ + private static ConcurrentMap sanitizeCallerState( + @Nullable Map state) { + if (state == null) { + return new ConcurrentHashMap<>(); + } + ConcurrentMap sanitized = new ConcurrentHashMap<>(); + state.forEach( + (key, value) -> { + if (RESERVED_STATE_KEYS.contains(key)) { + log.warn( + "Caller attempted to set reserved internal state key '{}'; ignoring.", key); + } else { + sanitized.put(key, value); + } + }); + return sanitized; + } + @Override public Single createSession( String appName, @@ -89,9 +135,8 @@ public Single createSession( .filter(s -> !s.isEmpty()) .orElseGet(() -> UUID.randomUUID().toString()); - // Ensure state map and events list are mutable for the new session - ConcurrentMap initialState = - (state == null) ? new ConcurrentHashMap<>() : new ConcurrentHashMap<>(state); + // Sanitize caller-supplied state: strip reserved internal keys before persisting. + ConcurrentMap initialState = sanitizeCallerState(state); // Assuming Session constructor or setters allow setting these mutable collections Session newSession = @@ -268,6 +313,12 @@ public Single appendEvent(Session session, Event event) { .put(userStateKey, value); } } else { + // Reject writes to reserved internal keys from any external stateDelta. + if (RESERVED_STATE_KEYS.contains(key)) { + log.warn( + "stateDelta contains reserved internal key '{}'; ignoring write.", key); + return; + } if (value == State.REMOVED) { session.state().remove(key); } else {