From 7bb049f2df8e2027f8ea292e1e0579bce954bdf0 Mon Sep 17 00:00:00 2001 From: Jast Date: Sun, 1 Sep 2024 00:00:56 +0800 Subject: [PATCH] [improve] improve code (#2523) Signed-off-by: tomsun28 Co-authored-by: YuLuo Co-authored-by: tomsun28 Co-authored-by: aias00 --- .../collector/collect/jmx/JmxClassLoader.java | 55 +++++++++++++++++++ .../collector/collect/jmx/JmxCollectImpl.java | 9 ++- 2 files changed, 63 insertions(+), 1 deletion(-) create mode 100644 collector/src/main/java/org/apache/hertzbeat/collector/collect/jmx/JmxClassLoader.java diff --git a/collector/src/main/java/org/apache/hertzbeat/collector/collect/jmx/JmxClassLoader.java b/collector/src/main/java/org/apache/hertzbeat/collector/collect/jmx/JmxClassLoader.java new file mode 100644 index 00000000000..e0fc8a595b6 --- /dev/null +++ b/collector/src/main/java/org/apache/hertzbeat/collector/collect/jmx/JmxClassLoader.java @@ -0,0 +1,55 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.hertzbeat.collector.collect.jmx; + +import lombok.extern.slf4j.Slf4j; + +/** + * custom class loader config for JMX + */ +@Slf4j +public class JmxClassLoader extends ClassLoader { + + private static final String[] WHITE_PRE_LIST = new String[]{ + "java.", + "javax.management.", + "org.apache.hertzbeat.", + "org.springframework.util.", + "com.sun.", + "sun.", + "org.slf4j.", + "jdk.", + "org.w3c.dom." + }; + + public JmxClassLoader(ClassLoader parent) { + super(parent); + } + + @Override + protected Class loadClass(String name, boolean resolve) throws ClassNotFoundException { + for (String whitePre : WHITE_PRE_LIST) { + if (name.startsWith(whitePre)) { + return super.loadClass(name, resolve); + } + } + log.error("Security vulnerability detection in JMX collect: Forbidden class: {}", name); + throw new ClassNotFoundException("Forbidden unsafe collection request content"); + } + +} diff --git a/collector/src/main/java/org/apache/hertzbeat/collector/collect/jmx/JmxCollectImpl.java b/collector/src/main/java/org/apache/hertzbeat/collector/collect/jmx/JmxCollectImpl.java index fb553a7dd28..396428efddd 100644 --- a/collector/src/main/java/org/apache/hertzbeat/collector/collect/jmx/JmxCollectImpl.java +++ b/collector/src/main/java/org/apache/hertzbeat/collector/collect/jmx/JmxCollectImpl.java @@ -69,10 +69,14 @@ public class JmxCollectImpl extends AbstractCollect { private static final String IGNORED_STUB = "/stub/"; private static final String SUB_ATTRIBUTE = "->"; + private final ConnectionCommonCache connectionCommonCache; + private final ClassLoader jmxClassLoader; + public JmxCollectImpl() { connectionCommonCache = new ConnectionCommonCache<>(); + jmxClassLoader = new JmxClassLoader(ClassLoader.getSystemClassLoader()); } @Override @@ -87,7 +91,8 @@ public void preCheck(Metrics metrics) throws IllegalArgumentException { @Override public void collect(CollectRep.MetricsData.Builder builder, long monitorId, String app, Metrics metrics) { - + ClassLoader currentClassLoader = Thread.currentThread().getContextClassLoader(); + Thread.currentThread().setContextClassLoader(jmxClassLoader); try { JmxProtocol jmxProtocol = metrics.getJmx(); @@ -129,6 +134,8 @@ public void collect(CollectRep.MetricsData.Builder builder, long monitorId, Stri log.error("JMX Error :{}", errorMsg); builder.setCode(CollectRep.Code.FAIL); builder.setMsg(errorMsg); + } finally { + Thread.currentThread().setContextClassLoader(currentClassLoader); } }