feat: servlet skip (#220)

* feat: servlet forbidden * feat: servlet skip
arextest · Aug 8, 2023 · 2727905 · 2727905
1 parent 6a78d51
commit 2727905
Show file tree

Hide file tree

Showing 3 changed files with 16 additions and 3 deletions.
diff --git a/arex-instrumentation-api/src/main/java/io/arex/inst/runtime/model/ArexConstants.java b/arex-instrumentation-api/src/main/java/io/arex/inst/runtime/model/ArexConstants.java
@@ -34,4 +34,5 @@ private ArexConstants() {}
     public static final String CONFIG_DEPENDENCY = "arex_replay_prepare_dependency";
     public static final String PREFIX = "arex-";
     public static final String CONFIG_VERSION = "configBatchNo";
+    public static final String SKIP_FLAG = "arex-skip-flag";
 }
diff --git a/.../servlet/arex-httpservlet/src/main/java/io/arex/inst/httpservlet/ServletAdviceHelper.java b/.../servlet/arex-httpservlet/src/main/java/io/arex/inst/httpservlet/ServletAdviceHelper.java
@@ -1,6 +1,7 @@
 package io.arex.inst.httpservlet;
 
 import io.arex.agent.bootstrap.TraceContextManager;
+import io.arex.agent.bootstrap.constants.ConfigConstants;
 import io.arex.agent.bootstrap.internal.Pair;
 import io.arex.agent.bootstrap.model.MockCategoryType;
 import io.arex.agent.bootstrap.util.StringUtil;
@@ -100,6 +101,12 @@ public static <TRequest, TResponse> Pair<TRequest, TResponse> onServiceEnter(
             return null;
         }
 
+        RequestHandlerManager.preHandle(httpServletRequest, MockCategoryType.SERVLET.getName());
+        // skip servlet if attr with arex-skip-flag
+        if (Boolean.TRUE.equals(adapter.getAttribute(httpServletRequest, ArexConstants.SKIP_FLAG))) {
+            return null;
+        }
+
         // 302 Redirect request
         String redirectRecordId = getRedirectRecordId(adapter, httpServletRequest);
         if (StringUtil.isNotEmpty(redirectRecordId)) {
@@ -109,7 +116,7 @@ public static <TRequest, TResponse> Pair<TRequest, TResponse> onServiceEnter(
             String caseId = adapter.getRequestHeader(httpServletRequest, ArexConstants.RECORD_ID);
             String excludeMockTemplate = adapter.getRequestHeader(httpServletRequest, ArexConstants.HEADER_EXCLUDE_MOCK);
             CaseEventDispatcher.onEvent(CaseEvent.ofCreateEvent(EventSource.of(caseId, excludeMockTemplate)));
-            RequestHandlerManager.preHandle(httpServletRequest, MockCategoryType.SERVLET.getName());
+
         }
 
         if (ContextManager.needRecordOrReplay()) {
@@ -198,7 +205,7 @@ private static <TRequest> boolean shouldSkip(ServletAdapter<TRequest, ?> adapter
 
         // Replay scene
         if (StringUtil.isNotEmpty(caseId)) {
-            return Config.get().getBoolean("arex.disable.replay", false);
+            return Config.get().getBoolean(ConfigConstants.DISABLE_REPLAY, false);
         }
 
         String forceRecord = adapter.getRequestHeader(httpServletRequest, ArexConstants.FORCE_RECORD);

diff --git a/...vlet/arex-httpservlet/src/test/java/io/arex/inst/httpservlet/ServletAdviceHelperTest.java b/...vlet/arex-httpservlet/src/test/java/io/arex/inst/httpservlet/ServletAdviceHelperTest.java
@@ -135,6 +135,10 @@ static Stream<Arguments> onServiceEnterCase() {
             Mockito.when(RecordLimiter.acquire(any())).thenReturn(true);
             Mockito.when(ContextManager.needRecordOrReplay()).thenReturn(true);
         };
+        Runnable shouldSkip9 = () -> {
+            Mockito.when(adapter.getRequestHeader(any(), eq(ArexConstants.RECORD_ID))).thenReturn("mock");
+            Mockito.when(adapter.getAttribute(any(), eq(ArexConstants.SKIP_FLAG))).thenReturn(Boolean.TRUE);
+        };
 
         Predicate<Pair<?, ?>> predicate1 = Objects::isNull;
         Predicate<Pair<?, ?>> predicate2 = Objects::nonNull;
@@ -153,7 +157,8 @@ static Stream<Arguments> onServiceEnterCase() {
             arguments("shouldSkip: IgnoreUtils.ignoreOperation returns true", shouldSkip5, predicate1),
             arguments("shouldSkip: adapter.getRequestURI return .png", shouldSkip6, predicate1),
             arguments("shouldSkip: adapter.getContentType return image/", shouldSkip7, predicate1),
-            arguments("ContextManager.needRecordOrReplay is true", shouldSkip8, predicate2)
+            arguments("ContextManager.needRecordOrReplay is true", shouldSkip8, predicate2),
+            arguments("shouldSkip: adapter.getAttribute returns true", shouldSkip9, predicate1)
         );
     }