Add childConsumer to RakChannelFactory

CloudburstMC · Apr 4, 2024 · 4d37bf1 · 4d37bf1
1 parent 954e915
commit 4d37bf1
Show file tree

Hide file tree

Showing 3 changed files with 37 additions and 20 deletions.
diff --git a/transport-raknet/src/main/java/org/cloudburstmc/netty/channel/raknet/RakChannelFactory.java b/transport-raknet/src/main/java/org/cloudburstmc/netty/channel/raknet/RakChannelFactory.java
@@ -25,45 +25,50 @@
 import java.lang.reflect.Constructor;
 import java.util.Objects;
 import java.util.function.Consumer;
+import java.util.function.Function;
 
 public class RakChannelFactory<T extends Channel> implements ChannelFactory<T> {
 
-    private final Constructor<? extends T> constructor;
+    private final Class<T> channelClass;
+    private final Function<DatagramChannel, T> constructor;
     private final Constructor<? extends DatagramChannel> datagramConstructor;
     private final Consumer<DatagramChannel> parentConsumer;
 
-    private RakChannelFactory(Class<? extends T> clazz, Class<? extends DatagramChannel> datagramClass, Consumer<DatagramChannel> parentConsumer) {
-        Objects.requireNonNull(clazz, "clazz");
+    private RakChannelFactory(Class<T> channelClass, Function<DatagramChannel, T> constructor, Class<? extends DatagramChannel> datagramClass,
+                              Consumer<DatagramChannel> parentConsumer) {
+        Objects.requireNonNull(channelClass, "channelClass");
         Objects.requireNonNull(datagramClass, "datagramClass");
-        try {
-            this.constructor = clazz.getConstructor(DatagramChannel.class);
-        } catch (NoSuchMethodException e) {
-            throw new IllegalArgumentException("Proxy class " + StringUtil.simpleClassName(clazz) +
-                    " does not have a public non-arg constructor", e);
-        }
+        Objects.requireNonNull(constructor, "constructor");
+        this.channelClass = channelClass;
+        this.constructor = constructor;
+
         try {
             this.datagramConstructor = datagramClass.getConstructor();
         } catch (NoSuchMethodException e) {
-            throw new IllegalArgumentException("Class " + StringUtil.simpleClassName(clazz) +
+            throw new IllegalArgumentException("Class " + StringUtil.simpleClassName(datagramClass) +
                     " does not have a public non-arg constructor", e);
         }
         this.parentConsumer = parentConsumer;
     }
 
     public static RakChannelFactory<RakServerChannel> server(Class<? extends DatagramChannel> clazz) {
-        return new RakChannelFactory<>(RakServerChannel.class, clazz, null);
+        return new RakChannelFactory<>(RakServerChannel.class, RakServerChannel::new, clazz, null);
     }
 
     public static RakChannelFactory<RakServerChannel> server(Class<? extends DatagramChannel> clazz, Consumer<DatagramChannel> parentConsumer) {
-        return new RakChannelFactory<>(RakServerChannel.class, clazz, parentConsumer);
+        return new RakChannelFactory<>(RakServerChannel.class, RakServerChannel::new, clazz, parentConsumer);
+    }
+
+    public static RakChannelFactory<RakServerChannel> server(Class<? extends DatagramChannel> clazz, Consumer<DatagramChannel> parentConsumer, Consumer<RakChannel> childConsumer) {
+        return new RakChannelFactory<>(RakServerChannel.class, ch -> new RakServerChannel(ch, childConsumer), clazz, parentConsumer);
     }
 
     public static RakChannelFactory<RakClientChannel> client(Class<? extends DatagramChannel> clazz) {
-        return new RakChannelFactory<>(RakClientChannel.class, clazz, null);
+        return new RakChannelFactory<>(RakClientChannel.class, RakClientChannel::new, clazz, null);
     }
 
     public static RakChannelFactory<RakClientChannel> client(Class<? extends DatagramChannel> clazz, Consumer<DatagramChannel> parentConsumer) {
-        return new RakChannelFactory<>(RakClientChannel.class, clazz, parentConsumer);
+        return new RakChannelFactory<>(RakClientChannel.class, RakClientChannel::new, clazz, parentConsumer);
     }
 
     @Override
@@ -73,16 +78,16 @@ public T newChannel() {
             if (this.parentConsumer != null) {
                 this.parentConsumer.accept(channel);
             }
-            return constructor.newInstance(channel);
+            return constructor.apply(channel);
         } catch (Throwable t) {
-            throw new ChannelException("Unable to create Channel from class " + constructor.getDeclaringClass(), t);
+            throw new ChannelException("Unable to create Channel from class " + this.channelClass, t);
         }
     }
 
     @Override
     public String toString() {
         return StringUtil.simpleClassName(RakChannelFactory.class) +
-                '(' + StringUtil.simpleClassName(constructor.getDeclaringClass()) + ".class, " +
+                '(' + StringUtil.simpleClassName(this.channelClass) + ".class, " +
                 StringUtil.simpleClassName(datagramConstructor.getDeclaringClass()) + ".class)";
     }
 }
diff --git a/transport-raknet/src/main/java/org/cloudburstmc/netty/channel/raknet/RakChildChannel.java b/transport-raknet/src/main/java/org/cloudburstmc/netty/channel/raknet/RakChildChannel.java
@@ -28,6 +28,7 @@
 import java.net.SocketAddress;
 import java.nio.channels.ClosedChannelException;
 import java.nio.channels.NonWritableChannelException;
+import java.util.function.Consumer;
 
 public class RakChildChannel extends AbstractChannel implements RakChannel {
 
@@ -39,13 +40,18 @@ public class RakChildChannel extends AbstractChannel implements RakChannel {
     private volatile boolean open = true;
     private volatile boolean active;
 
-    public RakChildChannel(InetSocketAddress remoteAddress, RakServerChannel parent, long guid, int version, int mtu) {
+    RakChildChannel(InetSocketAddress remoteAddress, RakServerChannel parent, long guid, int version, int mtu, Consumer<RakChannel> childConsumer) {
         super(parent);
         this.remoteAddress = remoteAddress;
         this.config = new DefaultRakSessionConfig(this);
         this.config.setGuid(guid);
         this.config.setProtocolVersion(version);
         this.config.setMtu(mtu);
+        // Allow user to configure the child channel before we initialize pipeline
+        // This is not the same as bootstrap.childOption() as Bootstrap does not allow setting options per channel
+        if (childConsumer != null) {
+            childConsumer.accept(this);
+        }
         // Create an internal pipeline for RakNet session logic to take place. We use the parent channel to ensure
         // this all occurs on the parent event loop so the connection is not slowed down by any user code.
         // (compression, encryption, etc.)

diff --git a/transport-raknet/src/main/java/org/cloudburstmc/netty/channel/raknet/RakServerChannel.java b/transport-raknet/src/main/java/org/cloudburstmc/netty/channel/raknet/RakServerChannel.java
@@ -37,14 +37,21 @@
 import java.util.Map;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.TimeUnit;
+import java.util.function.Consumer;
 
 public class RakServerChannel extends ProxyChannel<DatagramChannel> implements ServerChannel {
 
     private final RakServerChannelConfig config;
     private final Map<SocketAddress, RakChildChannel> childChannelMap = new ConcurrentHashMap<>();
+    private final Consumer<RakChannel> childConsumer;
 
     public RakServerChannel(DatagramChannel channel) {
+        this(channel, null);
+    }
+
+    public RakServerChannel(DatagramChannel channel, Consumer<RakChannel> childConsumer) {
         super(channel);
+        this.childConsumer = childConsumer;
         this.config = new DefaultRakServerConfig(this);
         // Default common handler of offline phase. Handles only raknet packets, forwards rest.
         this.pipeline().addLast(UnconnectedPongEncoder.NAME, UnconnectedPongEncoder.INSTANCE);
@@ -65,8 +72,7 @@ public RakChildChannel createChildChannel(InetSocketAddress address, long client
             return null;
         }
 
-        RakChildChannel channel = new RakChildChannel(address, this, clientGuid, protocolVersion, mtu);
-
+        RakChildChannel channel = new RakChildChannel(address, this, clientGuid, protocolVersion, mtu, childConsumer);
         channel.closeFuture().addListener((GenericFutureListener<ChannelFuture>) this::onChildClosed);
         // Fire channel thought ServerBootstrap,
         // register to eventLoop, assign default options and attributes