fix: ensure that transformed table does not leak to session liveness …

…or permanent liveness scope
deephaven · Dec 10, 2024 · f44eb01 · f44eb01
1 parent f5adda6
commit f44eb01
Showing 1 changed file with 25 additions and 16 deletions.
diff --git a/server/src/main/java/io/deephaven/server/console/ScopeTicketResolver.java b/server/src/main/java/io/deephaven/server/console/ScopeTicketResolver.java
@@ -8,6 +8,7 @@
 import io.deephaven.base.string.EncodingInfo;
 import io.deephaven.engine.context.ExecutionContext;
 import io.deephaven.engine.context.QueryScope;
+import io.deephaven.engine.liveness.LivenessScopeStack;
 import io.deephaven.engine.table.Table;
 import io.deephaven.proto.backplane.grpc.Ticket;
 import io.deephaven.proto.flight.util.TicketRouterHelper;
@@ -18,6 +19,7 @@
 import io.deephaven.server.session.SessionState;
 import io.deephaven.server.session.TicketResolverBase;
 import io.deephaven.server.session.TicketRouter;
+import io.deephaven.util.SafeCloseable;
 import io.grpc.StatusRuntimeException;
 import org.apache.arrow.flight.impl.Flight;
 import org.jetbrains.annotations.NotNull;
@@ -61,14 +63,14 @@ public SessionState.ExportObject<Flight.FlightInfo> flightInfoFor(
         if (!(scopeVar instanceof Table)) {
             throw newNotFoundSRE(logId, scopeName);
         }
-
-        final Table transformed = authorization.transform((Table) scopeVar);
-        if (transformed == null) {
-            throw newNotFoundSRE(logId, scopeName);
+        final Flight.FlightInfo flightInfo;
+        try (final SafeCloseable ignored = LivenessScopeStack.open()) {
+            final Table transformed = authorization.transform((Table) scopeVar);
+            if (transformed == null) {
+                throw newNotFoundSRE(logId, scopeName);
+            }
+            flightInfo = TicketRouter.getFlightInfo(transformed, descriptor, flightTicketForName(scopeName));
         }
-        final Flight.FlightInfo flightInfo =
-                TicketRouter.getFlightInfo(transformed, descriptor, flightTicketForName(scopeName));
-
         return SessionState.wrapAsExport(flightInfo);
     }
 
@@ -79,11 +81,16 @@ public void forAllFlightInfo(@Nullable final SessionState session, final Consume
         }
         final QueryScope queryScope = ExecutionContext.getContext().getQueryScope();
         queryScope.toMap(queryScope::unwrapObject, (n, t) -> t instanceof Table).forEach((name, table) -> {
-            final Table transformedTable = authorization.transform((Table) table);
-            if (transformedTable != null) {
-                visitor.accept(TicketRouter.getFlightInfo(
-                        transformedTable, descriptorForName(name), flightTicketForName(name)));
+            final Flight.FlightInfo flightInfo;
+            try (final SafeCloseable ignored = LivenessScopeStack.open()) {
+                final Table transformedTable = authorization.transform((Table) table);
+                if (transformedTable == null) {
+                    return;
+                }
+                flightInfo = TicketRouter.getFlightInfo(transformedTable, descriptorForName(name),
+                        flightTicketForName(name));
             }
+            visitor.accept(flightInfo);
         });
     }
 
@@ -108,14 +115,16 @@ private <T> SessionState.ExportObject<T> resolve(final String scopeName, final S
             export = (T) queryScope.unwrapObject(queryScope.readParamValue(scopeName));
         } catch (QueryScope.MissingVariableException ignored) {
         }
-
-        export = authorization.transform(export);
-
         if (export == null) {
             return SessionState.wrapAsFailedExport(newNotFoundSRE(logId, scopeName));
         }
-
-        return SessionState.wrapAsExport(export);
+        try (final SafeCloseable ignored = LivenessScopeStack.open()) {
+            export = authorization.transform(export);
+            if (export == null) {
+                return SessionState.wrapAsFailedExport(newNotFoundSRE(logId, scopeName));
+            }
+            return SessionState.wrapAsExport(export);
+        }
     }
 
     @Override