Merge pull request #670 from bytedance/feat-rasp-swicth

add jvm switch
bytedance · Aug 6, 2024 · 8029b94 · 8029b94
2 parents 51fc7aa + 380f2fd
commit 8029b94
Show file tree

Hide file tree

Showing 25 changed files with 449 additions and 108 deletions.
diff --git a/rasp/jvm/JVMProbe/src/main/java/com/security/smith/SmithProbe.java b/rasp/jvm/JVMProbe/src/main/java/com/security/smith/SmithProbe.java
@@ -149,7 +149,9 @@ 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 Map<String, Set<String>> hookTypes;
  private Disruptor<Trace> disruptor;
+ private Map<String, Boolean> switchConfig;
 
  private Rule_Mgr rulemgr;
  private Rule_Config ruleconfig;
@@ -199,6 +201,8 @@ public void init() {
  filters = new ConcurrentHashMap<>();
  blocks = new ConcurrentHashMap<>();
  limits = new ConcurrentHashMap<>();
+ hookTypes = new ConcurrentHashMap<>();
+ switchConfig = new ConcurrentHashMap<>();
 
  MessageSerializer.initInstance(proberVersion);
  MessageEncoder.initInstance();
@@ -233,16 +237,21 @@ public Trace newInstance() {
  Reader xreader = new InputStreamReader(inputStream);
  YamlReader yamlReader = new YamlReader(xreader);
  for (SmithClass smithClass : yamlReader.read(SmithClass[].class)) {
+ for (SmithMethod smithMethod : smithClass.getMethods()) {
+
+ if (smithMethod.getTypes() != null && !smithMethod.getTypes().isEmpty())
+ hookTypes.put(smithClass.getId() + "-" + smithMethod.getId(), smithMethod.getTypes());
+ }
  smithClasses.put(smithClass.getName(), smithClass);
  }
- } catch (IOException e) {
+ } catch (Throwable e) {
  SmithLogger.exception(e);
  }
  }
  else {
  SmithLogger.logger.info("not find class.yaml");
  }
-
+  
  SmithLogger.logger.info("probe init leave");
  }
  private boolean isBypassHookClass(String className) {
@@ -260,6 +269,23 @@ private boolean isBypassHookClass(String className) {
 
  return false;
  }
+ public boolean isFunctionEnabled(int classId, int methodId) {
+ String key = classId + "-" + methodId;
+ Set<String> types = hookTypes.get(key);
+
+ if (switchConfig == null || switchConfig.isEmpty()) {
+ return true;
+ }
+
+ if (types != null) {
+ for (String type : types) {
+ if (switchConfig.getOrDefault(type, true)) {
+ return true;
+ }
+ }
+ }
+ return false; 
+ }
 
  public void start() {
  SmithLogger.logger.info("probe start");
@@ -1072,6 +1098,16 @@ private void sendByte(byte[] data, String transId) {
  //}
  }
 
+ @Override
+ public void onSwitches(SwitchConfig switches) {
+ if (switches == null || switches.getSwitches() == null) {
+ return;
+ }
+ switchConfig = switches.getSwitches();
+
+ heartbeat.setSwitches(switches.getUUID());
+ }
+
  public Heartbeat getHeartbeat() {
  return heartbeat;
  }
@@ -1102,4 +1138,21 @@ public Disruptor<Trace> getDisruptor() {
  return disruptor;
  }
 
+ public String getFuncTypes(int classId, int methodId) {
+ String types = "";
+ try {
+
+ if (hookTypes.containsKey(classId + "-" + methodId)) {
+ for (String type: hookTypes.get(classId + "-" + methodId)) {
+ types += type + ",";
+ }
+ }
+ if (types.length() > 0) {
+ types = types.substring(0, types.length() - 1);
+ }
+ } catch (Exception e) {
+ SmithLogger.exception(e);
+ }
+ return types;
+ }
 }