Closes #377 - Prevent recursive method hook execution

inspectit-ocelot-agent/src/system-test/java/rocks/inspectit/ocelot/instrumentation/hook/RecursionInstrumentation.java

@@ -0,0 +1,34 @@
+package rocks.inspectit.ocelot.instrumentation.hook;
+
+import org.junit.jupiter.api.Test;
+import rocks.inspectit.ocelot.instrumentation.InstrumentationSysTestBase;
+import rocks.inspectit.ocelot.utils.TestUtils;
+
+import java.util.concurrent.atomic.AtomicInteger;
+
+import static org.assertj.core.api.Assertions.assertThat;
+
+public class RecursionInstrumentation extends InstrumentationSysTestBase {
+
+    private static AtomicInteger invocationCount = new AtomicInteger(0);
+
+    /**
+     * This method is instrumentation and will be called in the instrumentation itself
+     * which would result in endless loop and StackOverflowError.
+     */
+    public static void helloWorld() {
+        System.out.println("Hello World!");
+        invocationCount.incrementAndGet();
+    }
+
+    @Test
+    public void recursiveInstrumentation() {
+        TestUtils.waitForClassInstrumentations(RecursionInstrumentation.class);
+
+        helloWorld();
+
+        // two calls - one of the application and one (only!) due to the instrumentation
+        assertThat(invocationCount.get()).isEqualTo(2);
+    }
+
+}

inspectit-ocelot-agent/src/system-test/resources/config/recursiveInstrumentation.yml

@@ -0,0 +1,24 @@
+# Configuration used in rocks.inspectit.ocelot.instrumentation.hook.RecursionInstrumentation
+inspectit:
+  instrumentation:
+    actions:
+      'a_recursive_helloWorld':
+        is-void: true
+        value-body: |
+          // Call instrumented method to provoke infinity loop
+          rocks.inspectit.ocelot.instrumentation.hook.RecursionInstrumentation.helloWorld();
+
+    scopes:
+      's_RecursionInstrumentation_helloWorld':
+        type:
+          name: 'rocks.inspectit.ocelot.instrumentation.hook.RecursionInstrumentation'
+        methods:
+          - name: 'helloWorld'
+
+    rules:
+      'r_recursive_test':
+        scopes:
+          's_RecursionInstrumentation_helloWorld': true
+        entry:
+          'c_invoke_hello':
+            action: 'a_recursive_helloWorld'

inspectit-ocelot-core/src/main/java/rocks/inspectit/ocelot/core/instrumentation/hook/HookManager.java

@@ -1,5 +1,6 @@
 package rocks.inspectit.ocelot.core.instrumentation.hook;
 
+import com.google.common.annotations.VisibleForTesting;
 import lombok.extern.slf4j.Slf4j;
 import lombok.val;
 import net.bytebuddy.description.method.MethodDescription;
@@ -29,6 +30,13 @@
 @Service
 public class HookManager {
 
+    /**
+     * Thread local flag for marking the current thread that it is currently in the execution/scope of agent actions.
+     * This is used to prevent an endless action recursion in case an instrumented action is invoked within another
+     * action. In that case, the instrumentation has to be suppressed.
+     */
+    public static final ThreadLocal<Boolean> RECURSION_GATE = ThreadLocal.withInitial(() -> false);
+
     @Autowired
     private InstrumentationConfigurationResolver configResolver;
 
@@ -63,17 +71,22 @@ void destroy() {
      * Actual implementation for {@link IHookManager#getHook(Class, String)}.
      *
      * @param clazz           the name of the class to which the method to query the hook for belongs
-     * @param methodSignature the signature of the method in the form of name(parametertype,parametertype,..)
-     * @return
+     * @param methodSignature the signature of the method in the form of name(parametertype, parametertype,..)
+     *
+     * @return the method hook for the specified method
      */
-    private IMethodHook getHook(Class<?> clazz, String methodSignature) {
-        Map<String, MethodHook> methodHooks = hooks.get(clazz);
-        if (methodHooks != null) {
-            MethodHook hook = methodHooks.get(methodSignature);
-            if (hook != null) {
-                return hook;
+    @VisibleForTesting
+    IMethodHook getHook(Class<?> clazz, String methodSignature) {
+        if (!RECURSION_GATE.get()) {
+            Map<String, MethodHook> methodHooks = hooks.get(clazz);
+            if (methodHooks != null) {
+                MethodHook hook = methodHooks.get(methodSignature);
+                if (hook != null) {
+                    return hook;
+                }
             }
         }
+
         return NoopMethodHook.INSTANCE;
     }
 
@@ -88,7 +101,6 @@ public HookUpdate startUpdate() {
         return new HookUpdate();
     }
 
-
     /**
      * A {@link HookUpdate} instance represents a (potential) change to apply to the {@link HookManager}.
      * Hooks for any class can be updated by calling {@link #updateHooksForClass(Class)}. This updates
@@ -120,7 +132,8 @@ private HookUpdate() {
                 newHooks = new WeakHashMap<>();
                 for (Map.Entry<Class<?>, Map<String, MethodHook>> existingMethodHooks : hooks.entrySet()) {
                     HashMap<String, MethodHook> newMethodHooks = new HashMap<>();
-                    existingMethodHooks.getValue().forEach((signature, hook) -> newMethodHooks.put(signature, hook.getResettedCopy()));
+                    existingMethodHooks.getValue()
+                            .forEach((signature, hook) -> newMethodHooks.put(signature, hook.getResettedCopy()));
                     newHooks.put(existingMethodHooks.getKey(), newMethodHooks);
                 }
             }
@@ -159,13 +172,11 @@ private void ensureNotCommitted() {
 
         private Optional<MethodHook> getCurrentHook(Class<?> declaringClass, String methodSignature) {
             Map<String, MethodHook> methodHooks = newHooks.get(declaringClass);
-            return Optional.ofNullable(methodHooks)
-                    .map(myHooks -> myHooks.get(methodSignature));
+            return Optional.ofNullable(methodHooks).map(myHooks -> myHooks.get(methodSignature));
         }
 
         private void setHook(Class<?> declaringClass, String methodSignature, MethodHook newHook) {
-            newHooks.computeIfAbsent(declaringClass, (v) -> new HashMap<>())
-                    .put(methodSignature, newHook);
+            newHooks.computeIfAbsent(declaringClass, (v) -> new HashMap<>()).put(methodSignature, newHook);
         }
 
         private void removeHook(Class<?> declaringClass, String methodSignature) {
@@ -193,8 +204,7 @@ private void removeHook(Class<?> declaringClass, String methodSignature) {
         private void addOrReplaceHooks(Class<?> clazz, Map<MethodDescription, MethodHookConfiguration> hookConfigs) {
             hookConfigs.forEach((method, newConfig) -> {
                 String signature = CoreUtils.getSignature(method);
-                MethodHookConfiguration oldConfig = getCurrentHook(clazz, signature)
-                        .map(MethodHook::getSourceConfiguration)
+                MethodHookConfiguration oldConfig = getCurrentHook(clazz, signature).map(MethodHook::getSourceConfiguration)
                         .orElse(null);
                 if (!Objects.equals(newConfig, oldConfig)) {
                     if (log.isDebugEnabled()) {

inspectit-ocelot-core/src/main/java/rocks/inspectit/ocelot/core/instrumentation/hook/MethodHook.java

@@ -13,7 +13,6 @@
 import rocks.inspectit.ocelot.core.selfmonitoring.ActionScopeFactory;
 import rocks.inspectit.ocelot.core.selfmonitoring.IActionScope;
 
-import javax.validation.constraints.NotNull;
 import java.util.ArrayList;
 import java.util.List;
 import java.util.concurrent.CopyOnWriteArrayList;
@@ -65,7 +64,6 @@ public class MethodHook implements IMethodHook {
      */
     private final MethodReflectionInformation methodInformation;
 
-
     /**
      * The factory that creates/spawns new scopes for an {@link IHookAction}
      */
@@ -88,34 +86,50 @@ public MethodHook(MethodHookConfiguration sourceConfiguration, ContextManager in
 
     @Override
     public InternalInspectitContext onEnter(Object[] args, Object thiz) {
-        InspectitContextImpl inspectitContext = inspectitContextManager.enterNewContext();
-        IHookAction.ExecutionContext executionContext = new IHookAction.ExecutionContext(args, thiz, null, null, this, inspectitContext);
-
-        for (IHookAction action : activeEntryActions) {
-            try (IActionScope scope = actionScopeFactory.createScope(action)) {
-                action.execute(executionContext);
-            } catch (Throwable t) {
-                log.error("Entry action {} executed for method {} threw an exception and from now on is disabled!", action, methodInformation.getMethodFQN(), t);
-                activeEntryActions.remove(action);
+        try {
+            // flags the thread that it is now in action execution
+            HookManager.RECURSION_GATE.set(true);
+
+            InspectitContextImpl inspectitContext = inspectitContextManager.enterNewContext();
+            IHookAction.ExecutionContext executionContext = new IHookAction.ExecutionContext(args, thiz, null, null, this, inspectitContext);
+
+            for (IHookAction action : activeEntryActions) {
+                try (IActionScope scope = actionScopeFactory.createScope(action)) {
+                    action.execute(executionContext);
+                } catch (Throwable t) {
+                    log.error("Entry action {} executed for method {} threw an exception and from now on is disabled!", action, methodInformation
+                            .getMethodFQN(), t);
+                    activeEntryActions.remove(action);
+                }
             }
-        }
 
-        inspectitContext.makeActive();
-        return inspectitContext;
+            inspectitContext.makeActive();
+            return inspectitContext;
+        } finally {
+            HookManager.RECURSION_GATE.set(false);
+        }
     }
 
     @Override
     public void onExit(Object[] args, Object thiz, Object returnValue, Throwable thrown, InternalInspectitContext context) {
-        IHookAction.ExecutionContext executionContext = new IHookAction.ExecutionContext(args, thiz, returnValue, thrown, this, (InspectitContextImpl) context);
-        for (IHookAction action : activeExitActions) {
-            try (IActionScope scope = actionScopeFactory.createScope(action)) {
-                action.execute(executionContext);
-            } catch (Throwable t) {
-                log.error("Exit action {} executed for method {} threw an exception and from now on is disabled!", action, methodInformation.getMethodFQN(), t);
-                activeExitActions.remove(action);
+        try {
+            // flags the thread that it is now in action execution
+            HookManager.RECURSION_GATE.set(true);
+
+            IHookAction.ExecutionContext executionContext = new IHookAction.ExecutionContext(args, thiz, returnValue, thrown, this, (InspectitContextImpl) context);
+            for (IHookAction action : activeExitActions) {
+                try (IActionScope scope = actionScopeFactory.createScope(action)) {
+                    action.execute(executionContext);
+                } catch (Throwable t) {
+                    log.error("Exit action {} executed for method {} threw an exception and from now on is disabled!", action, methodInformation
+                            .getMethodFQN(), t);
+                    activeExitActions.remove(action);
+                }
             }
+            context.close();
+        } finally {
+            HookManager.RECURSION_GATE.set(false);
         }
-        context.close();
     }
 
     /**

inspectit-ocelot-core/src/test/java/rocks/inspectit/ocelot/core/instrumentation/hook/HookManagerTest.java

@@ -0,0 +1,86 @@
+package rocks.inspectit.ocelot.core.instrumentation.hook;
+
+import org.junit.jupiter.api.Nested;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.extension.ExtendWith;
+import org.mockito.InjectMocks;
+import org.mockito.Mock;
+import org.mockito.junit.jupiter.MockitoExtension;
+import org.springframework.test.util.ReflectionTestUtils;
+import rocks.inspectit.ocelot.bootstrap.instrumentation.IMethodHook;
+import rocks.inspectit.ocelot.bootstrap.instrumentation.noop.NoopMethodHook;
+import rocks.inspectit.ocelot.core.instrumentation.config.InstrumentationConfigurationResolver;
+import rocks.inspectit.ocelot.core.selfmonitoring.ActionScopeFactory;
+import rocks.inspectit.ocelot.core.selfmonitoring.SelfMonitoringService;
+
+import java.util.HashMap;
+import java.util.Map;
+
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.mockito.Mockito.mock;
+
+@ExtendWith(MockitoExtension.class)
+public class HookManagerTest {
+
+    @InjectMocks
+    private HookManager manager;
+
+    @Mock
+    private InstrumentationConfigurationResolver configResolver;
+
+    @Mock
+    private SelfMonitoringService selfMonitoring;
+
+    @Mock
+    private MethodHookGenerator hookGenerator;
+
+    @Nested
+    public class GetHook {
+
+        @Test
+        public void hookNotExisting() {
+            IMethodHook result = manager.getHook(HookManagerTest.class, "stark");
+
+            assertThat(result).isSameAs(NoopMethodHook.INSTANCE);
+        }
+
+        @Test
+        public void hookExisting() {
+            // create map with mock hook
+            MethodHook hook = MethodHook.builder().actionScopeFactory(mock(ActionScopeFactory.class)).build();
+            Map<String, MethodHook> methodHookMap = new HashMap<>();
+            methodHookMap.put("lannister", hook);
+            Map<Class<?>, Map<String, MethodHook>> hooks = new HashMap<>();
+            hooks.put(HookManagerTest.class, methodHookMap);
+
+            // inject map
+            ReflectionTestUtils.setField(manager, "hooks", hooks);
+
+            IMethodHook result = manager.getHook(HookManagerTest.class, "lannister");
+
+            assertThat(result).isSameAs(hook);
+        }
+
+        @Test
+        public void preventRecursion() {
+            // create map with mock hook
+            MethodHook hook = MethodHook.builder().actionScopeFactory(mock(ActionScopeFactory.class)).build();
+            Map<String, MethodHook> methodHookMap = new HashMap<>();
+            methodHookMap.put("lannister", hook);
+            Map<Class<?>, Map<String, MethodHook>> hooks = new HashMap<>();
+            hooks.put(HookManagerTest.class, methodHookMap);
+
+            // inject map
+            ReflectionTestUtils.setField(manager, "hooks", hooks);
+
+            IMethodHook resultFirst = manager.getHook(HookManagerTest.class, "lannister");
+            assertThat(resultFirst).isSameAs(hook);
+
+            // set recursion gate
+            HookManager.RECURSION_GATE.set(true);
+
+            IMethodHook resultSecond = manager.getHook(HookManagerTest.class, "lannister");
+            assertThat(resultSecond).isSameAs(NoopMethodHook.INSTANCE);
+        }
+    }
+}