From b9c202f356cea79ad85f2c0b7b2ff0ff364fbb34 Mon Sep 17 00:00:00 2001 From: LiangliangSui Date: Tue, 16 Apr 2024 15:06:01 +0800 Subject: [PATCH 1/2] fix(java): ThreadPoolFury and ThreadLocalFury concurrency security issues Signed-off-by: LiangliangSui --- .../fury/pool/ClassLoaderFuryPooled.java | 16 ++ .../org/apache/fury/pool/ThreadPoolFury.java | 2 +- .../org/apache/fury/util/LoaderBinding.java | 144 +++++++++++------- 3 files changed, 110 insertions(+), 52 deletions(-) diff --git a/java/fury-core/src/main/java/org/apache/fury/pool/ClassLoaderFuryPooled.java b/java/fury-core/src/main/java/org/apache/fury/pool/ClassLoaderFuryPooled.java index 7f02fe02c6..ae7c2054de 100644 --- a/java/fury-core/src/main/java/org/apache/fury/pool/ClassLoaderFuryPooled.java +++ b/java/fury-core/src/main/java/org/apache/fury/pool/ClassLoaderFuryPooled.java @@ -29,11 +29,13 @@ import java.util.concurrent.locks.ReentrantLock; import java.util.function.Consumer; import java.util.function.Function; +import javax.annotation.concurrent.ThreadSafe; import org.apache.fury.Fury; import org.apache.fury.logging.Logger; import org.apache.fury.logging.LoggerFactory; /** A thread-safe object pool of {@link Fury}. */ +@ThreadSafe public class ClassLoaderFuryPooled { private static final Logger LOG = LoggerFactory.getLogger(ClassLoaderFuryPooled.class); @@ -122,6 +124,20 @@ private void addFury() { } void setFactoryCallback(Consumer factoryCallback) { + lock.lock(); this.factoryCallback = factoryCallback; + lock.unlock(); + } + + void traversalAllFury(Consumer callback) { + try { + lock.lock(); + allFury.keySet().forEach(callback); + } catch (Exception e) { + LOG.error(e.getMessage(), e); + throw new RuntimeException(e); + } finally { + lock.unlock(); + } } } diff --git a/java/fury-core/src/main/java/org/apache/fury/pool/ThreadPoolFury.java b/java/fury-core/src/main/java/org/apache/fury/pool/ThreadPoolFury.java index e9367dcfc6..4f22e80b7e 100644 --- a/java/fury-core/src/main/java/org/apache/fury/pool/ThreadPoolFury.java +++ b/java/fury-core/src/main/java/org/apache/fury/pool/ThreadPoolFury.java @@ -55,7 +55,7 @@ protected void processCallback(Consumer callback) { factoryCallback = factoryCallback.andThen(callback); for (ClassLoaderFuryPooled furyPooled : furyPooledObjectFactory.classLoaderFuryPooledCache.asMap().values()) { - furyPooled.allFury.keySet().forEach(callback); + furyPooled.traversalAllFury(callback); furyPooled.setFactoryCallback(factoryCallback); } } diff --git a/java/fury-core/src/main/java/org/apache/fury/util/LoaderBinding.java b/java/fury-core/src/main/java/org/apache/fury/util/LoaderBinding.java index 79a1b6118d..95d5b399fa 100644 --- a/java/fury-core/src/main/java/org/apache/fury/util/LoaderBinding.java +++ b/java/fury-core/src/main/java/org/apache/fury/util/LoaderBinding.java @@ -25,10 +25,14 @@ import java.util.HashSet; import java.util.Set; import java.util.WeakHashMap; +import java.util.concurrent.locks.Lock; +import java.util.concurrent.locks.ReentrantLock; import java.util.function.Consumer; import java.util.function.Function; import org.apache.fury.Fury; import org.apache.fury.annotation.Internal; +import org.apache.fury.logging.Logger; +import org.apache.fury.logging.LoggerFactory; /** * An util to bind {@link Fury} with {@link ClassLoader}. If {@link ClassLoader} are changed, the @@ -37,6 +41,7 @@ */ @Internal public final class LoaderBinding { + private static final Logger LOG = LoggerFactory.getLogger(LoaderBinding.class); private final Function furyFactory; // `WeakHashMap` won't work here too, since `Fury` hold classes which reference `ClassLoader`, // which cause @@ -46,6 +51,7 @@ public final class LoaderBinding { private Consumer bindingCallback = f -> {}; private ClassLoader loader; private Fury fury; + private final Lock lock = new ReentrantLock(); public LoaderBinding(Function furyFactory) { this.furyFactory = furyFactory; @@ -56,29 +62,37 @@ public Fury get() { } public void visitAllFury(Consumer consumer) { - if (furySoftMap.isEmpty()) { - for (Fury f : furyMap.values()) { - consumer.accept(f); - } - } else if (furyMap.isEmpty()) { - for (SoftReference ref : furySoftMap.values()) { - Fury f = ref.get(); - if (f != null) { + try { + lock.lock(); + if (furySoftMap.isEmpty()) { + for (Fury f : furyMap.values()) { consumer.accept(f); } - } - } else { - Set furySet = new HashSet<>(furyMap.size()); - Collections.addAll(furyMap.values()); - for (SoftReference ref : furySoftMap.values()) { - Fury f = ref.get(); - if (f != null) { - furySet.add(f); + } else if (furyMap.isEmpty()) { + for (SoftReference ref : furySoftMap.values()) { + Fury f = ref.get(); + if (f != null) { + consumer.accept(f); + } + } + } else { + Set furySet = new HashSet<>(furyMap.size()); + Collections.addAll(furyMap.values()); + for (SoftReference ref : furySoftMap.values()) { + Fury f = ref.get(); + if (f != null) { + furySet.add(f); + } + } + for (Fury f : furySet) { + consumer.accept(f); } } - for (Fury f : furySet) { - consumer.accept(f); - } + } catch (Exception e) { + LOG.error(e.getMessage(), e); + throw new RuntimeException(e); + } finally { + lock.unlock(); } } @@ -114,36 +128,44 @@ public void setClassLoader(ClassLoader classLoader, StagingType stagingType) { classLoader = Fury.class.getClassLoader(); } this.loader = classLoader; - switch (stagingType) { - case NO_STAGING: - fury = furyFactory.apply(classLoader); - bindingCallback.accept(fury); - break; - case SOFT_STAGING: - { - SoftReference furySoftReference = furySoftMap.get(classLoader); - Fury fury = furySoftReference == null ? null : furySoftReference.get(); - if (fury == null) { - fury = furyFactory.apply(classLoader); - bindingCallback.accept(fury); - furySoftMap.put(classLoader, new SoftReference<>(fury)); - this.fury = fury; - } + try { + lock.lock(); + switch (stagingType) { + case NO_STAGING: + fury = furyFactory.apply(classLoader); + bindingCallback.accept(fury); break; - } - case STRONG_STAGING: - { - Fury fury = furyMap.get(classLoader); - if (fury == null) { - fury = furyFactory.apply(classLoader); - bindingCallback.accept(fury); - furyMap.put(classLoader, fury); - this.fury = fury; + case SOFT_STAGING: + { + SoftReference furySoftReference = furySoftMap.get(classLoader); + Fury fury = furySoftReference == null ? null : furySoftReference.get(); + if (fury == null) { + fury = furyFactory.apply(classLoader); + bindingCallback.accept(fury); + furySoftMap.put(classLoader, new SoftReference<>(fury)); + this.fury = fury; + } + break; } - break; - } - default: - throw new IllegalArgumentException(); + case STRONG_STAGING: + { + Fury fury = furyMap.get(classLoader); + if (fury == null) { + fury = furyFactory.apply(classLoader); + bindingCallback.accept(fury); + furyMap.put(classLoader, fury); + this.fury = fury; + } + break; + } + default: + throw new IllegalArgumentException(); + } + } catch (Exception e) { + LOG.error(e.getMessage(), e); + throw new RuntimeException(e); + } finally { + lock.unlock(); } } } @@ -155,8 +177,10 @@ public void setClassLoader(ClassLoader classLoader, StagingType stagingType) { * referenced by other objects. */ public void clearClassLoader(ClassLoader classLoader) { + lock.lock(); furyMap.remove(classLoader); SoftReference softReference = furySoftMap.remove(classLoader); + lock.unlock(); if (softReference != null) { softReference.clear(); } @@ -167,18 +191,36 @@ public void clearClassLoader(ClassLoader classLoader) { } public void register(Class clz) { - furyMap.values().forEach(fury -> fury.register(clz)); - bindingCallback = bindingCallback.andThen(fury -> fury.register(clz)); + try { + lock.lock(); + furyMap.values().forEach(fury -> fury.register(clz)); + bindingCallback = bindingCallback.andThen(fury -> fury.register(clz)); + } catch (Exception e) { + LOG.error(e.getMessage(), e); + throw new RuntimeException(e); + } finally { + lock.unlock(); + } } public void register(Class clz, int id) { Preconditions.checkArgument(id < Short.MAX_VALUE); - furyMap.values().forEach(fury -> fury.register(clz, (short) id)); - bindingCallback = bindingCallback.andThen(fury -> fury.register(clz, (short) id)); + try { + lock.lock(); + furyMap.values().forEach(fury -> fury.register(clz, (short) id)); + bindingCallback = bindingCallback.andThen(fury -> fury.register(clz, (short) id)); + } catch (Exception e) { + LOG.error(e.getMessage(), e); + throw new RuntimeException(e); + } finally { + lock.unlock(); + } } public void setBindingCallback(Consumer bindingCallback) { + lock.lock(); this.bindingCallback = bindingCallback; + lock.unlock(); } public enum StagingType { From fe8f304a39cb266bf2197d7b03b9190588be4b5d Mon Sep 17 00:00:00 2001 From: LiangliangSui Date: Tue, 16 Apr 2024 15:30:44 +0800 Subject: [PATCH 2/2] fix test Signed-off-by: LiangliangSui --- .../org.apache.fury/fury-core/native-image.properties | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/java/fury-core/src/main/resources/META-INF/native-image/org.apache.fury/fury-core/native-image.properties b/java/fury-core/src/main/resources/META-INF/native-image/org.apache.fury/fury-core/native-image.properties index bcb3392d9f..145a17426b 100644 --- a/java/fury-core/src/main/resources/META-INF/native-image/org.apache.fury/fury-core/native-image.properties +++ b/java/fury-core/src/main/resources/META-INF/native-image/org.apache.fury/fury-core/native-image.properties @@ -169,4 +169,5 @@ Args=--initialize-at-build-time=org.apache.fury.memory.MemoryBuffer,\ org.apache.fury.memory.MemoryUtils,\ org.apache.fury.type.DescriptorGrouper,\ sun.misc.Unsafe,\ - com.google.common.collect.Platform + com.google.common.collect.Platform,\ + org.apache.fury.util.LoaderBinding