[improve] improve code (#2523)

Signed-off-by: tomsun28 <[email protected]>
Co-authored-by: YuLuo <[email protected]>
Co-authored-by: tomsun28 <[email protected]>
Co-authored-by: aias00 <[email protected]>

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");
+    }
+
+}

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<CacheIdentifier, JmxConnect> 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);
         }
     }