Skip to content

Commit

Permalink
Merge pull request #695 from bytedance/feat-java-performance
Browse files Browse the repository at this point in the history 
feat hook performance
  • Loading branch information
@yoloyyh
yoloyyh authored Oct 12, 2024
2 parents 02d1f9d + d2db1f2 commit 4c46bc5
Show file tree
Hide file tree
Showing 2 changed files with 35 additions and 31 deletions.
54 changes: 28 additions & 26 deletions rasp/jvm/JVMProbe/src/main/java/com/security/smith/SmithProbe.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -160,6 +160,8 @@ public class SmithProbe implements ClassFileTransformer, MessageHandler, EventHa
private SmithProbe ourInstance = null;
private SmithProbeProxy smithProxy = null;
private int TRACE_BUFFER_SIZE = 1024;
private final int CLASS_MAX_ID = 50;
private final int METHOD_MAX_ID = 20;

private Object xClassLoaderObj;
private Boolean disable;
Expand All @@ -177,7 +179,7 @@ public class SmithProbe implements ClassFileTransformer, MessageHandler, EventHa
private final Map<Pair<Integer, Integer>, List<Long>> recordsTotal;
private final Map<Pair<Integer, Integer>, Long> hooktimeRecords;
private final Map<Pair<Integer, Integer>, Long> runtimeRecords;
private Map<String, Set<String>> hookTypes;
private Set<String> [][] hookTypes;
private Disruptor<Trace> disruptor;
private Map<String, Boolean> switchConfig;

Expand Down Expand Up @@ -244,10 +246,15 @@ public void init() {
filters = new ConcurrentHashMap<>();
blocks = new ConcurrentHashMap<>();
limits = new ConcurrentHashMap<>();
hookTypes = new ConcurrentHashMap<>();
switchConfig = new ConcurrentHashMap<>();

hookTypes = new Set[CLASS_MAX_ID][METHOD_MAX_ID];

for (int i = 0; i < CLASS_MAX_ID; i++) {
for (int j = 0; j < METHOD_MAX_ID; j++) {
hookTypes[i][j] = new HashSet<>();
}
}
MessageSerializer.initInstance(proberVersion);
MessageEncoder.initInstance();
MessageDecoder.initInstance();
Expand All @@ -271,7 +278,7 @@ public Trace newInstance() {
rulemgr = new Rule_Mgr();
ruleconfig = new Rule_Config(rulemgr);

smithProxy = new SmithProbeProxy();
smithProxy = new SmithProbeProxy(CLASS_MAX_ID, METHOD_MAX_ID);


InputStream inputStream = getResourceAsStream("class.yaml");
Expand All @@ -285,7 +292,7 @@ public Trace newInstance() {
for (SmithMethod smithMethod : smithClass.getMethods()) {

if (smithMethod.getTypes() != null && !smithMethod.getTypes().isEmpty())
hookTypes.put(smithClass.getId() + "-" + smithMethod.getId(), smithMethod.getTypes());
hookTypes[smithClass.getId()][smithMethod.getId()] = smithMethod.getTypes();
}
smithClasses.put(smithClass.getName(), smithClass);
}
Expand All @@ -298,7 +305,7 @@ public Trace newInstance() {
}

try {
SmithLogger.logger.info("jsRuleEngine init");
// SmithLogger.logger.info("jsRuleEngine init");
jsRuleEngine = JsRuleEngine.InitializeEngine();
if (jsRuleEngine != null) {
SmithLogger.logger.info("jsRuleEngine init success");
Expand All @@ -314,16 +321,10 @@ public boolean addJsFile(Path scriptFilePath) {
boolean ret = false;
try {
if (scriptFilePath != null && jsRuleEngine != null) {
SmithLogger.logger.info("add js rule enter");
int result = jsRuleEngine.addJsRule(scriptFilePath);
if (result == 0) {
SmithLogger.logger.info("add js rule success");
ret = true;
} else {
SmithLogger.logger.info("add js rule failed, ret :" + result);
}
} else {
SmithLogger.logger.info("not find js rule path: " + scriptFilePath);
}
}
catch (Throwable e) {
Expand Down Expand Up @@ -384,11 +385,10 @@ private boolean isBypassHookClass(String className) {

public boolean isFunctionEnabled(int classId, int methodId) {

if (switchConfig == null || switchConfig.isEmpty()) {
if (switchConfig == null || switchConfig.isEmpty() || classId >= CLASS_MAX_ID || methodId >= METHOD_MAX_ID || hookTypes == null) {
return true;
}
String key = classId + "-" + methodId;
Set<String> types = hookTypes.get(key);
Set<String> types = hookTypes[classId][methodId];

if (types != null) {
for (String type : types) {
Expand All @@ -404,7 +404,6 @@ public void start() {
SmithLogger.logger.info("probe start");
AttachInfo.info();

SmithLogger.logger.info("init ClassUploadTransformer");
ClassUploadTransformer.getInstance().start(client, inst);


Expand Down Expand Up @@ -451,6 +450,7 @@ public void start() {
smithProxy.setDisruptor(disruptor);
smithProxy.setProbe(this);


try {
addJsRule();
} catch (Exception e) {
Expand All @@ -470,32 +470,25 @@ public void stop() {

inst.removeTransformer(this);
reloadClasses();
SmithLogger.logger.info("Transformer stop");

disable = true;
scanswitch = false;

ClassUploadTransformer.getInstance().stop();

SmithLogger.logger.info("Upload Transformer stop");

detectTimer.cancel();
smithproxyTimer.cancel();
SmithLogger.logger.info("detect Timer stop");

if (isBenchMark) {
benchMarkTimer.cancel();
SmithLogger.logger.info("benchMark Timer stop");
}

client.stop();
SmithLogger.logger.info("client stop");

ruleconfig.destry();
SmithLogger.logger.info("ruleconfig stop");

rulemgr.destry();
SmithLogger.logger.info("rulemgr stop");

detectTimerTask = null;
detectTimer =null;
Expand Down Expand Up @@ -537,11 +530,18 @@ public void uninit() {
value.removeAll();
blocks.remove(key);
}
for (int i = 0; i < hookTypes.length; i++) {
for (int j = 0; j < hookTypes[i].length; j++) {
hookTypes[i][j].clear();
hookTypes[i][j] = null;
}
}
hookTypes = null;
blocks.clear();
blocks = null;
limits.clear();
limits = null;
SmithLogger.logger.info("probe uninit 0");


disruptor = null;
ruleconfig = null;
Expand Down Expand Up @@ -1343,10 +1343,12 @@ public Disruptor<Trace> getDisruptor() {

public String getFuncTypes(int classId, int methodId) {
String types = "";
if (classId < 0 || methodId < 0 || classId >= CLASS_MAX_ID || methodId >= METHOD_MAX_ID ) {
return types;
}
try {

if (hookTypes.containsKey(classId + "-" + methodId)) {
for (String type: hookTypes.get(classId + "-" + methodId)) {
if (hookTypes[classId][methodId] != null) {
for (String type: hookTypes[classId][methodId]) {
types += type + ",";
}
}
Expand Down
12 changes: 7 additions & 5 deletions rasp/jvm/JVMProbe/src/main/java/com/security/smith/SmithProbeProxy.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -26,8 +26,8 @@
import com.google.gson.JsonElement;
import com.google.gson.GsonBuilder;
public class SmithProbeProxy {
private final int CLASS_MAX_ID = 50;
private final int METHOD_MAX_ID = 20;
private final int CLASS_MAX_ID;
private final int METHOD_MAX_ID;
private final int DEFAULT_QUOTA = 12000;

private SmithProbe SmithProbeObj = null;
Expand Down Expand Up @@ -117,9 +117,11 @@ private void RemoveThreadLocalVar() {
}
}

public SmithProbeProxy() {
public SmithProbeProxy(int classMaxID, int methodMaxID) {
stopX = false;

CLASS_MAX_ID = classMaxID;
METHOD_MAX_ID = methodMaxID;
quotas = new AtomicIntegerArray[CLASS_MAX_ID];
for (int i = 0; i < CLASS_MAX_ID; i++) {
quotas[i] = new AtomicIntegerArray(METHOD_MAX_ID);
Expand Down Expand Up @@ -280,7 +282,7 @@ public void sendMetadataClass(Class<?> cla, int classID, int methodID) {

JsRuleResult result = SmithProbeObj.getJsRuleEngine().detect(1,argsX);
if(result != null) {
SmithLogger.logger.info("Js Rule Result +" + result.toString());
// SmithLogger.logger.info("Js Rule Result +" + result.toString());
ClassFilter classFilter = new ClassFilter();
SmithHandler.queryClassFilter(cla, classFilter);
classFilter.setTransId();
Expand Down Expand Up @@ -803,7 +805,7 @@ public void checkWildflyaddFilterPre(int classID, int methodID, Object[] args) {
}

public void handleReflectField(int classID, int methodID, Object[] args, Object ret, boolean blocked) {
if(stopX) {
if(stopX || SmithProbeObj.isFunctionEnabled(classID, methodID) == false) {
return;
}
if (args.length < 2) {
Expand Down

0 comments on commit 4c46bc5

Please sign in to comment.