Merge pull request #308 from Nizernizer/main

fix bug: COOKIE_FLAGS_MISSING
HXSecurity · May 16, 2022 · ddd7b16 · ddd7b16
2 parents 4a6c9bc + 521686f
commit ddd7b16
Showing 1 changed file with 6 additions and 1 deletion.
diff --git a/dongtai-core/src/main/java/io/dongtai/iast/core/bytecode/enhance/plugin/cookie/BaseType.java b/dongtai-core/src/main/java/io/dongtai/iast/core/bytecode/enhance/plugin/cookie/BaseType.java
@@ -3,6 +3,7 @@
 import io.dongtai.iast.core.bytecode.enhance.IastContext;
 import io.dongtai.iast.core.bytecode.enhance.plugin.AbstractClassVisitor;
 import io.dongtai.iast.core.bytecode.enhance.plugin.core.adapter.PropagateAdviceAdapter;
+import io.dongtai.iast.core.bytecode.enhance.plugin.core.adapter.SinkAdviceAdapter;
 import io.dongtai.iast.core.utils.AsmUtils;
 import org.objectweb.asm.ClassVisitor;
 import org.objectweb.asm.MethodVisitor;
@@ -31,7 +32,11 @@ public MethodVisitor visitMethod(final int access, final String name, final Stri
         if (match(name, context.getMatchClassName())) {
             String iastMethodSignature = AsmUtils.buildSignature(context.getMatchClassName(), name, desc);
             String framework = "refType";
-            mv = new PropagateAdviceAdapter(mv, access, name, desc, context, framework, iastMethodSignature);
+            if (iastMethodSignature.contains("setSecure") || iastMethodSignature.contains("<init>")){
+                mv = new SinkAdviceAdapter(mv, access, name, desc, context, framework, iastMethodSignature,false);
+            }else {
+                mv = new PropagateAdviceAdapter(mv, access, name, desc, context, framework, iastMethodSignature);
+            }
             setTransformed();
             DongTaiLog.trace("rewrite method {} for listener[match={},class={}]", iastMethodSignature, context.getMatchClassName(), context.getClassName());
         }