NR-272899: Fix for user class detection grpc

newrelic · Sep 17, 2024 · fed109e · fed109e
1 parent 339ee15
commit fed109e
Showing 1 changed file with 16 additions and 10 deletions.
diff --git a/newrelic-security-agent/src/main/java/com/newrelic/api/agent/security/Agent.java b/newrelic-security-agent/src/main/java/com/newrelic/api/agent/security/Agent.java
@@ -496,6 +496,7 @@ private UserClassEntity setUserClassEntityByAnnotation(StackTraceElement[] servi
     }
 
     private UserClassEntity setUserClassEntity(AbstractOperation operation, SecurityMetaData securityMetaData) {
+        boolean frameworkSpecificEntityFound = false;
         UserClassEntity userClassEntity = new UserClassEntity();
         StackTraceElement userStackTraceElement = securityMetaData.getCustomAttribute(GenericHelper.USER_CLASS_ENTITY, StackTraceElement.class);
         if(userStackTraceElement == null && securityMetaData.getMetaData().getServiceTrace() != null && securityMetaData.getMetaData().getServiceTrace().length > 0){
@@ -520,21 +521,18 @@ private UserClassEntity setUserClassEntity(AbstractOperation operation, Security
             switch (framework){
                 case "vertx-web":
                 case "GRPC":
-                    if(i-1 >= 0) {
+                    if(!frameworkSpecificEntityFound && i-1 >= 0) {
                         userClassEntity = setUserClassEntityForVertx(operation, userStackTraceElement, userClassEntity, securityMetaData.getMetaData().isUserLevelServiceMethodEncountered(), i);
                         if(userClassEntity.getUserClassElement() != null){
-                            return userClassEntity;
+                            frameworkSpecificEntityFound = true;
                         }
                     }
                     break;
                 default:
-                    if(userStackTraceElement != null){
-                        if(StringUtils.equals(stackTraceElement.getClassName(), userStackTraceElement.getClassName())
-                                && StringUtils.equals(stackTraceElement.getMethodName(), userStackTraceElement.getMethodName())){
-                            userClassEntity.setUserClassElement(stackTraceElement);
-                            userClassEntity.setCalledByUserCode(securityMetaData.getMetaData().isUserLevelServiceMethodEncountered());
-                            return userClassEntity;
-                        }
+                    if(userStackTraceElement != null && StringUtils.equals(stackTraceElement.getClassName(), userStackTraceElement.getClassName())
+                            && StringUtils.equals(stackTraceElement.getMethodName(), userStackTraceElement.getMethodName())){
+                        userClassEntity.setUserClassElement(stackTraceElement);
+                        userClassEntity.setCalledByUserCode(securityMetaData.getMetaData().isUserLevelServiceMethodEncountered());
                     }
             }
         }
@@ -544,8 +542,16 @@ private UserClassEntity setUserClassEntity(AbstractOperation operation, Security
             return userClassEntity;
         }
 
+        if (frameworkSpecificEntityFound && userClassEntity.getUserClassElement() != null && !securityMetaData.getMetaData().isFoundAnnotedUserLevelServiceMethod()){
+            return userClassEntity;
+        }
+
+        if (userClassEntity.getUserClassElement() != null){
+            return userClassEntity;
+        }
+
         // user class identification using annotations
-        if (userClassEntity.getUserClassElement() == null && securityMetaData.getMetaData().isFoundAnnotedUserLevelServiceMethod()) {
+        if (securityMetaData.getMetaData().isFoundAnnotedUserLevelServiceMethod()) {
             return setUserClassEntityByAnnotation(securityMetaData.getMetaData().getServiceTrace());
         }