Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

feat hook performance #695

Merged
merged 1 commit into from
Oct 12, 2024
Merged

feat hook performance #695

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
54 changes: 28 additions & 26 deletions rasp/jvm/JVMProbe/src/main/java/com/security/smith/SmithProbe.java
View file
Open in desktop
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
Open in desktop
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
Loading