Skip to content

Commit

Permalink
Merge pull request #686 from bytedance/main-jvm-benchmark
Browse files Browse the repository at this point in the history 
Main jvm benchmark
  • Loading branch information
@yoloyyh
yoloyyh authored Sep 26, 2024
2 parents 7360c20 + 6371497 commit 7f608f8
Show file tree
Hide file tree
Showing 4 changed files with 289 additions and 4 deletions.
13 changes: 12 additions & 1 deletion rasp/Makefile
View file
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,7 @@ DEBUG_SYMBOLS ?= debug
LIB_OUTPUT ?= $(OUTPUT)/lib-$(VERSION)
VCPKG_OVERLAY_PORTS ?= $(abspath overlay-ports)

.PHONY: all help install clean set-version agent-plugin nsenter pangolin jattach JVMAgent JVMProbe python-probe python-loader go-probe go-probe-ebpf node-probe php-probe librasp rasp-server NSMount
.PHONY: all help install clean check_benchmark set-version agent-plugin nsenter pangolin jattach JVMAgent JVMProbe python-probe python-loader go-probe go-probe-ebpf node-probe php-probe librasp rasp-server NSMount

all: rasp-linux-default-x86_64-$(VERSION).tar.gz rasp-linux-default-x86_64-$(VERSION)-debug.tar.gz SHA256SUMS

Expand Down Expand Up @@ -42,6 +42,17 @@ SHA256SUMS: rasp-linux-default-x86_64-$(VERSION).tar.gz
sha256sum $(OUTPUT)/rasp rasp-linux-default-x86_64-$(VERSION).tar.gz > $@


# for benchmark
check_benchmark:
if echo $(VERSION) | grep -q benchmark; then \
echo "Version contains benchmark. Modifying files..."; \
sed -i "s/isBenchMark = false/isBenchMark = true/g" "jvm/JVMAgent/src/main/java/com/security/smithloader/SmithAgent.java" \
sed -i "s/isBenchMark = false/isBenchMark = true/g" "jvm/JVMProbe/src/main/java/com/security/smith/SmithProbe.java" \
sed -i "s/isBenchMark = false/isBenchMark = true/g" "jvm/JVMProbe/src/main/java/com/security/smith/asm/SmithMethodVisitor.java" \
else \
echo "Version does not contain benchmark. Skipping modification."; \
fi

set-version:
sed -i "s/1.0.0.1/${VERSION}/g" "librasp/src/settings.rs"
sed -i "s/1.0.0/${VERSION}/g" "node/src/client.js"
Expand Down
24 changes: 24 additions & 0 deletions rasp/jvm/JVMAgent/src/main/java/com/security/smithloader/SmithAgent.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -25,6 +25,8 @@ public class SmithAgent {
private static String checksumStr = null;
private static String proberPath = null;
private static Instrumentation instrumentation = null;
// just fo benchmark test
private static boolean isBenchMark = false;
public static InheritableThreadLocal<Boolean> checkRecursive = new InheritableThreadLocal<Boolean>() {
@Override
protected Boolean initialValue() {
Expand Down Expand Up @@ -99,6 +101,28 @@ public static Object ExceptionProxy(Object MethodNameObj,int classID, int method
return obj;
}

public static void RecordProxy(int classID, int methodID, Long t1, Long t2) {
if (isBenchMark) {
try {
if (checkRecursive != null && checkRecursive.get() == true) {
return;
}
if (checkRecursive != null && checkRecursive.get() == false) {
checkRecursive.set(true);
}
if(SmithProberObj != null) {
Class<?>[] argType = new Class[]{int.class,int.class, Long.class, Long.class};
Reflection.invokeMethod(SmithProberObj,"record",argType,classID,methodID,t1,t2);
}
if (checkRecursive != null && checkRecursive.get() == true) {
checkRecursive.set(false);
}
} catch (Throwable e) {
SmithAgentLogger.exception(e);
}
}
}

private static boolean loadSmithProber(String proberPath, Instrumentation inst) {
boolean bret = false;
boolean bexception = false;
Expand Down
120 changes: 119 additions & 1 deletion 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 @@ -115,6 +115,28 @@ public boolean cancel() {
}
}

class BenchMarkTimerTask extends TimerTask {
private boolean isCancel = false;
private SmithProbe smithProbe = null;

public void setSmithProbe(SmithProbe smithProbe) {
this.smithProbe = smithProbe;
}

@Override
public void run() {
if(!isCancel) {
smithProbe.show();
}
}

@Override
public boolean cancel() {
isCancel = true;
return super.cancel();
}
}

class MatchRulePredicate implements Predicate<MatchRule> {
private final Trace trace;

Expand Down Expand Up @@ -151,23 +173,36 @@ public class SmithProbe implements ClassFileTransformer, MessageHandler, EventHa
private Map<Pair<Integer, Integer>, Filter> filters;
private Map<Pair<Integer, Integer>, Block> blocks;
private Map<Pair<Integer, Integer>, Integer> limits;
private final Map<Pair<Integer, Integer>, List<Long>> records;
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 Disruptor<Trace> disruptor;
private Map<String, Boolean> switchConfig;

private Rule_Mgr rulemgr;
private Rule_Config ruleconfig;
private Timer detectTimer;
private Timer benchMarkTimer;
private Timer smithproxyTimer;
private DetectTimerTask detectTimerTask;
private SmithproxyTimerTask smithproxyTimerTask;
private BenchMarkTimerTask benchMarkTimerTask;
private String proberVersion;
private String proberPath;
private JsRuleEngine jsRuleEngine;
// just for benchmark test
private boolean isBenchMark;

public SmithProbe() {
disable = false;
scanswitch = true;
records = new HashMap<>();
recordsTotal = new HashMap<>();
hooktimeRecords = new HashMap<>();
runtimeRecords = new HashMap<>();
isBenchMark = false;
}

public void setInst(Instrumentation inst) {
Expand Down Expand Up @@ -211,6 +246,7 @@ public void init() {
limits = new ConcurrentHashMap<>();
hookTypes = new ConcurrentHashMap<>();
switchConfig = new ConcurrentHashMap<>();


MessageSerializer.initInstance(proberVersion);
MessageEncoder.initInstance();
Expand Down Expand Up @@ -389,6 +425,19 @@ public void start() {
detectTimerTask,
TimeUnit.MINUTES.toMillis(1)
);

if (isBenchMark) {
benchMarkTimerTask = new BenchMarkTimerTask();
benchMarkTimerTask.setSmithProbe(this);

benchMarkTimer = new Timer(true);
benchMarkTimer.schedule(
benchMarkTimerTask,
TimeUnit.SECONDS.toMillis(5),
TimeUnit.SECONDS.toMillis(10)
);
}

smithproxyTimerTask = new SmithproxyTimerTask();
smithproxyTimerTask.setSmithProxy(smithProxy);

Expand All @@ -413,6 +462,7 @@ public void start() {
reloadClasses();

SmithLogger.logger.info("probe start leave");

}

public void stop() {
Expand All @@ -432,6 +482,11 @@ public void 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");
Expand All @@ -445,6 +500,9 @@ public void stop() {
detectTimerTask = null;
detectTimer =null;

benchMarkTimerTask = null;
benchMarkTimer = null;

smithproxyTimerTask = null;
smithproxyTimer = null;

Expand Down Expand Up @@ -584,6 +642,66 @@ private void reloadClasses(Collection<String> classes) {
}
}

private Long tp(List<Long> times, double percent) {
return times.get((int)(percent / 100 * times.size() - 1));
}

public void show() {
synchronized (records) {
SmithLogger.logger.info("=================== statistics ===================");

records.forEach((k, v) -> {
Collections.sort(v);

List<Long> tv = recordsTotal.get(new ImmutablePair<>(k.getLeft(), k.getRight()));

Collections.sort(tv);

Long hooktime = hooktimeRecords.get(new ImmutablePair<>(k.getLeft(), k.getRight()));
Long runtime = runtimeRecords.get(new ImmutablePair<>(k.getLeft(), k.getRight()));

SmithLogger.logger.info(
String.format(
"class: %d method: %d count: %d tp50: %d tp90: %d tp95: %d tp99: %d tp99.99: %d max: %d total-max:%d hooktime:%d runtime:%d",
k.getLeft(),
k.getRight(),
v.size(),
tp(v, 50),
tp(v, 90),
tp(v, 95),
tp(v, 99),
tp(v, 99.99),
v.get(v.size() - 1),
tv.get(tv.size() - 1),
hooktime,
runtime
)
);
});
}
}

public void record(int classID, int methodID, long time,long totaltime) {
SmithLogger.logger.info("record: " + classID + " " + methodID + " " + time);
synchronized (records) {
records.computeIfAbsent(new ImmutablePair<>(classID, methodID), k -> new ArrayList<>()).add(time);
}

synchronized (recordsTotal) {
recordsTotal.computeIfAbsent(new ImmutablePair<>(classID, methodID), k -> new ArrayList<>()).add(totaltime);
}

synchronized (hooktimeRecords) {
hooktimeRecords.computeIfAbsent(new ImmutablePair<>(classID, methodID), k -> time);
hooktimeRecords.computeIfPresent(new ImmutablePair<>(classID, methodID),(k,v) -> v+time);
}

synchronized (runtimeRecords) {
runtimeRecords.computeIfAbsent(new ImmutablePair<>(classID, methodID), k -> totaltime);
runtimeRecords.computeIfPresent(new ImmutablePair<>(classID, methodID),(k,v) -> v+totaltime);
}
}

@Override
public void onEvent(Trace trace, long sequence, boolean endOfBatch) {
Filter filter = filters.get(new ImmutablePair<>(trace.getClassID(), trace.getMethodID()));
Expand Down Expand Up @@ -835,7 +953,7 @@ public byte[] transform(ClassLoader loader, String className, Class<?> classBein
methodMap
);

classReader.accept(classVisitor, ClassReader.EXPAND_FRAMES);
classReader.accept(classVisitor, ClassReader.EXPAND_FRAMES);

return classWriter.toByteArray();
}
Expand Down
Loading

0 comments on commit 7f608f8

Please sign in to comment.