Merge pull request #1283 from chrisdennis/safe-callbacks

Provide protection from untrusted callbacks when calling user code

Author: Myron Scott

tc-client/src/main/java/com/tc/object/BinaryInvocationCallback.java

@@ -19,15 +19,14 @@
 package com.tc.object;
 
 import org.terracotta.entity.EntityResponse;
-import org.terracotta.entity.InvocationCallback;
 import org.terracotta.entity.MessageCodec;
 import org.terracotta.entity.MessageCodecException;
 
-public class BinaryInvocationCallback<R extends EntityResponse> implements InvocationCallback<byte[]> {
+public class BinaryInvocationCallback<R extends EntityResponse> implements SafeInvocationCallback<byte[]> {
   private final MessageCodec<?, R> codec;
-  private final InvocationCallback<R> callback;
+  private final SafeInvocationCallback<R> callback;
 
-  public BinaryInvocationCallback(MessageCodec<?, R> codec, InvocationCallback<R> callback) {
+  public BinaryInvocationCallback(MessageCodec<?, R> codec, SafeInvocationCallback<R> callback) {
     this.codec = codec;
     this.callback = callback;
   }

tc-client/src/main/java/com/tc/object/ClientEntityManagerImpl.java

@@ -79,6 +79,7 @@
 import org.terracotta.exception.EntityServerUncaughtException;
 
 import static com.tc.object.EntityDescriptor.createDescriptorForLifecycle;
+import static com.tc.object.SafeInvocationCallback.safe;
 import static java.util.stream.Collectors.toCollection;
 import static org.terracotta.entity.Invocation.uninterruptiblyGet;
 
@@ -196,7 +197,7 @@ private byte[] lifecycleAndRetire(EntityID entityId, long version, VoltronEntity
   }
 
   private Invocation<byte[]> lifecycle(EntityID entityID, EntityDescriptor entityDescriptor, VoltronEntityMessage.Type type, byte[] message) {
-    return (callback, callbacks) -> ClientEntityManagerImpl.this.invoke(entityID, entityDescriptor, callbacks, callback, true, type, message);
+    return (callback, callbacks) -> ClientEntityManagerImpl.this.invoke(entityID, entityDescriptor, callbacks, safe(callback), true, type, message);
   }
 
   @Override
@@ -228,12 +229,12 @@ private Set<VoltronEntityMessage.Acks> makeServerAcks(Set<InvocationCallback.Typ
 
   @Override
   public Invocation.Task invokeAction(EntityID eid, EntityDescriptor entityDescriptor, Set<InvocationCallback.Types> requestedCallbacks,
-                                      InvocationCallback<byte[]> callback, boolean requiresReplication, byte[] payload) {
+                                      SafeInvocationCallback<byte[]> callback, boolean requiresReplication, byte[] payload) {
     return invoke(eid, entityDescriptor, requestedCallbacks, callback, requiresReplication, VoltronEntityMessage.Type.INVOKE_ACTION, payload);
   }
 
   private Invocation.Task invoke(EntityID eid, EntityDescriptor entityDescriptor, Set<InvocationCallback.Types> requestedCallbacks,
-                                             InvocationCallback<byte[]> callback, boolean requiresReplication, VoltronEntityMessage.Type type, byte[] payload) {
+                                 SafeInvocationCallback<byte[]> callback, boolean requiresReplication, VoltronEntityMessage.Type type, byte[] payload) {
     Set<VoltronEntityMessage.Acks> requestedAcks = makeServerAcks(requestedCallbacks);
     return queueInFlightMessage(eid, () -> createMessageWithDescriptor(eid, entityDescriptor, requiresReplication, payload, type, requestedAcks), callback);
   }
@@ -476,7 +477,7 @@ private byte[] internalRetrieve(EntityDescriptor entityDescriptor) throws Entity
     return lifecycleAndComplete(entityDescriptor.getEntityID(), entityDescriptor, VoltronEntityMessage.Type.FETCH_ENTITY);
   }
 
-  private Invocation.Task queueInFlightMessage(EntityID eid, Supplier<NetworkVoltronEntityMessage> message, InvocationCallback<byte[]> callback) {
+  private Invocation.Task queueInFlightMessage(EntityID eid, Supplier<NetworkVoltronEntityMessage> message, SafeInvocationCallback<byte[]> callback) {
     boolean queued;
     try {
       InFlightMessage inFlight = new InFlightMessage(eid, message, callback);

tc-client/src/main/java/com/tc/object/EntityClientEndpointImpl.java

@@ -39,6 +39,8 @@
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
+import static com.tc.object.SafeInvocationCallback.safe;
+
 
 public class EntityClientEndpointImpl<M extends EntityMessage, R extends EntityResponse> implements EntityClientEndpoint<M, R> {
 
@@ -140,13 +142,13 @@ private InvocationImpl(M request) {
     public Task invoke(InvocationCallback<R> callback, Set<InvocationCallback.Types> callbacks) {
       checkInvoked();
       invoked = true;
-      InvocationCallback<byte[]> binaryCallback = new BinaryInvocationCallback(codec, callback);
+      SafeInvocationCallback<byte[]> binaryCallback = new BinaryInvocationCallback<>(codec, safe(callback));
       try {
         return invocationHandler.invokeAction(entityID, invokeDescriptor, callbacks, binaryCallback, true, codec.encodeMessage(request));
       } catch (MessageCodecException e) {
-        callback.failure(e);
-        callback.complete();
-        callback.retired();
+        binaryCallback.failure(e);
+        binaryCallback.complete();
+        binaryCallback.retired();
         return () -> false;
       }
     }

tc-client/src/main/java/com/tc/object/InFlightMessage.java

@@ -50,7 +50,6 @@
 import java.util.concurrent.atomic.AtomicReference;
 import java.util.function.Supplier;
 import com.tc.net.protocol.tcm.TCAction;
-import org.terracotta.entity.InvocationCallback;
 
 
 /**
@@ -62,7 +61,7 @@
 public class InFlightMessage implements PrettyPrintable {
   private final VoltronEntityMessage message;
   private final EntityID eid;
-  private final InvocationCallback<byte[]> callback;
+  private final SafeInvocationCallback<byte[]> callback;
 
   private final EnumSet<VoltronEntityMessage.Acks> outstandingAcks = EnumSet.allOf(VoltronEntityMessage.Acks.class);
 
@@ -86,7 +85,7 @@ public class InFlightMessage implements PrettyPrintable {
 
   private TCNetworkMessage networkMessage;
 
-  public InFlightMessage(EntityID eid, Supplier<? extends VoltronEntityMessage> message, InvocationCallback<byte[]> callback) {
+  public InFlightMessage(EntityID eid, Supplier<? extends VoltronEntityMessage> message, SafeInvocationCallback<byte[]> callback) {
     this.eid = requireNonNull(eid);
     this.message = requireNonNull(message.get());
     this.callback = callback;

tc-client/src/main/java/com/tc/object/InvocationHandler.java

@@ -29,5 +29,5 @@
  * The minimal interface, provided to the EntityClientEndpoint, to handle invocations to send to the server.
  */
 public interface InvocationHandler {
-  Invocation.Task invokeAction(EntityID eid, EntityDescriptor entityDescriptor, Set<InvocationCallback.Types> callbacks, InvocationCallback<byte[]> callback, boolean requiresReplication, byte[] payload);
+  Invocation.Task invokeAction(EntityID eid, EntityDescriptor entityDescriptor, Set<InvocationCallback.Types> callbacks, SafeInvocationCallback<byte[]> callback, boolean requiresReplication, byte[] payload);
 }

tc-client/src/main/java/com/tc/object/SafeInvocationCallback.java

@@ -0,0 +1,110 @@
+/*
+ *
+ *  The contents of this file are subject to the Terracotta Public License Version
+ *  2.0 (the "License"); You may not use this file except in compliance with the
+ *  License. You may obtain a copy of the License at
+ *
+ *  http://terracotta.org/legal/terracotta-public-license.
+ *
+ *  Software distributed under the License is distributed on an "AS IS" basis,
+ *  WITHOUT WARRANTY OF ANY KIND, either express or implied. See the License for
+ *  the specific language governing rights and limitations under the License.
+ *
+ *  The Covered Software is Terracotta Core.
+ *
+ *  The Initial Developer of the Covered Software is
+ *  Terracotta, Inc., a Software AG company
+ *
+ */
+package com.tc.object;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+import org.terracotta.entity.InvocationCallback;
+
+public interface SafeInvocationCallback<R> extends InvocationCallback<R> {
+
+    static <R> SafeInvocationCallback<R> safe(InvocationCallback<R> callback) {
+        if (callback instanceof SafeInvocationCallback<?>) {
+            return (SafeInvocationCallback<R>) callback;
+        } else {
+            return new Guard<>(callback);
+        }
+    }
+
+    class Guard<R> implements SafeInvocationCallback<R> {
+        private static final Logger LOGGER = LoggerFactory.getLogger(SafeInvocationCallback.class);
+
+        private final InvocationCallback<R> untrustedCallback;
+
+        private Guard(InvocationCallback<R> untrustedCallback) {
+            this.untrustedCallback = untrustedCallback;
+        }
+
+        @Override
+        public void sent() {
+            try {
+                untrustedCallback.sent();
+            } catch (Exception t) {
+                LOGGER.warn("User-provided callback [" + untrustedCallback + "] threw exception", t);
+            } catch (Throwable t) {
+                LOGGER.warn("User-provided callback [" + untrustedCallback + "] threw throwable", t);
+            }
+        }
+
+        @Override
+        public void received() {
+            try {
+                untrustedCallback.received();
+            } catch (Exception t) {
+                LOGGER.warn("User-provided callback [" + untrustedCallback + "] threw exception", t);
+            } catch (Throwable t) {
+                LOGGER.warn("User-provided callback [" + untrustedCallback + "] threw throwable", t);
+            }
+        }
+
+        @Override
+        public void result(R response) {
+            try {
+                untrustedCallback.result(response);
+            } catch (Exception t) {
+                LOGGER.warn("User-provided callback [" + untrustedCallback + "] threw exception", t);
+            } catch (Throwable t) {
+                LOGGER.warn("User-provided callback [" + untrustedCallback + "] threw throwable", t);
+            }
+        }
+
+        @Override
+        public void failure(Throwable failure) {
+            try {
+                untrustedCallback.failure(failure);
+            } catch (Exception t) {
+                LOGGER.warn("User-provided callback [" + untrustedCallback + "] threw exception", t);
+            } catch (Throwable t) {
+                LOGGER.warn("User-provided callback [" + untrustedCallback + "] threw throwable", t);
+            }
+        }
+
+        @Override
+        public void complete() {
+            try {
+                untrustedCallback.complete();
+            } catch (Exception t) {
+                LOGGER.warn("User-provided callback [" + untrustedCallback + "] threw exception", t);
+            } catch (Throwable t) {
+                LOGGER.warn("User-provided callback [" + untrustedCallback + "] threw throwable", t);
+            }
+        }
+
+        @Override
+        public void retired() {
+            try {
+                untrustedCallback.retired();
+            } catch (Exception t) {
+                LOGGER.warn("User-provided callback [" + untrustedCallback + "] threw exception", t);
+            } catch (Throwable t) {
+                LOGGER.warn("User-provided callback [" + untrustedCallback + "] threw throwable", t);
+            }
+        }
+    }
+}

tc-client/src/main/java/com/terracotta/diagnostic/DiagnosticClientEntityManager.java

@@ -27,6 +27,7 @@
 import com.tc.object.EntityDescriptor;
 import com.tc.object.EntityID;
 import com.tc.object.InFlightMessage;
+import com.tc.object.SafeInvocationCallback;
 import com.tc.object.msg.ClientHandshakeMessage;
 import com.tc.object.tx.TransactionID;
 import com.tc.util.Assert;
@@ -154,7 +155,7 @@ public void shutdown() {
   }
 
   @Override
-  public Invocation.Task invokeAction(EntityID eid, EntityDescriptor entityDescriptor, Set<InvocationCallback.Types> callbacks, InvocationCallback<byte[]> callback, boolean requiresReplication, byte[] payload) {
+  public Invocation.Task invokeAction(EntityID eid, EntityDescriptor entityDescriptor, Set<InvocationCallback.Types> callbacks, SafeInvocationCallback<byte[]> callback, boolean requiresReplication, byte[] payload) {
     DiagnosticMessage network = createMessage(payload);
     InFlightMessage message = new InFlightMessage(eid, ()->network, callback);
     waitingForAnswer.put(network.getTransactionID(), message);

tc-client/src/test/java/com/tc/object/ClientEntityManagerTest.java

@@ -549,7 +549,7 @@ public void testSingleInvoke() throws Exception {
     when(channel.createMessage(TCMessageType.VOLTRON_ENTITY_MESSAGE)).thenReturn(message);
 
     CompletableFuture<byte[]> result = new CompletableFuture<>();
-    InvocationCallback<byte[]> callback = new InvocationCallback<byte[]>() {
+    SafeInvocationCallback<byte[]> callback = new SafeInvocationCallback<byte[]>() {
       @Override
       public void result(byte[] response) {
         result.complete(response);
@@ -570,7 +570,7 @@ public void testSingleInvokeTimeout() throws Exception {
     when(channel.createMessage(TCMessageType.VOLTRON_ENTITY_MESSAGE)).thenReturn(message);
 
     CompletableFuture<byte[]> result = new CompletableFuture<>();
-    InvocationCallback<byte[]> callback = new InvocationCallback<byte[]>() {
+    SafeInvocationCallback<byte[]> callback = new SafeInvocationCallback<byte[]>() {
       @Override
       public void result(byte[] response) {
         result.complete(response);

tc-client/src/test/java/com/tc/object/SafeInvocationCallbackTest.java

@@ -0,0 +1,68 @@
+/*
+ *
+ *  The contents of this file are subject to the Terracotta Public License Version
+ *  2.0 (the "License"); You may not use this file except in compliance with the
+ *  License. You may obtain a copy of the License at
+ *
+ *  http://terracotta.org/legal/terracotta-public-license.
+ *
+ *  Software distributed under the License is distributed on an "AS IS" basis,
+ *  WITHOUT WARRANTY OF ANY KIND, either express or implied. See the License for
+ *  the specific language governing rights and limitations under the License.
+ *
+ *  The Covered Software is Terracotta Core.
+ *
+ *  The Initial Developer of the Covered Software is
+ *  Terracotta, Inc., a Software AG company
+ *
+ */
+package com.tc.object;
+
+import org.junit.Test;
+import org.terracotta.entity.InvocationCallback;
+
+import java.lang.reflect.InvocationTargetException;
+import java.lang.reflect.Method;
+
+import static org.hamcrest.MatcherAssert.assertThat;
+import static org.hamcrest.Matchers.equalTo;
+import static org.hamcrest.Matchers.instanceOf;
+import static org.hamcrest.Matchers.is;
+import static org.junit.Assert.fail;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.times;
+import static org.mockito.Mockito.verify;
+
+public class SafeInvocationCallbackTest {
+
+    @Test
+    public void testSafeInvocationCallbackIsTypedCorrectly() {
+        assertThat(SafeInvocationCallback.safe(mock(InvocationCallback.class)), is(instanceOf(SafeInvocationCallback.class)));
+    }
+
+    @Test
+    public void testSafeInvocationCallbackCatchesAll() throws InvocationTargetException, IllegalAccessException {
+        for (Method method : InvocationCallback.class.getDeclaredMethods()) {
+            System.out.println(method);
+            InvocationCallback<Object> callback = mock(InvocationCallback.class, inv -> {
+                if (inv.getMethod().getDeclaringClass().equals(InvocationCallback.class)) {
+                    throw new Throwable();
+                } else {
+                    return null;
+                }
+            });
+
+            SafeInvocationCallback<Object> safe = SafeInvocationCallback.safe(callback);
+
+            Object[] parameters = new Object[method.getParameterCount()];
+            try {
+                method.invoke(callback, parameters);
+                fail("Expected Throwable");
+            } catch (InvocationTargetException t) {
+                assertThat(t.getCause().getClass(), is(equalTo(Throwable.class)));
+            }
+            method.invoke(safe, parameters);
+            method.invoke(verify(callback, times(2)), parameters);
+        }
+    }
+}