implement new entity limit logic

Author: xGinko

AnarchyExploitFixesFolia/src/main/java/me/xginko/aef/modules/chunklimits/CustomEntityLimit.java

@@ -6,9 +6,9 @@
 import io.papermc.paper.threadedregions.scheduler.ScheduledTask;
 import me.xginko.aef.modules.AEFModule;
 import me.xginko.aef.utils.ChunkUtil;
+import me.xginko.aef.utils.EntityUtil;
 import me.xginko.aef.utils.LocationUtil;
 import org.bukkit.Chunk;
-import org.bukkit.Material;
 import org.bukkit.World;
 import org.bukkit.entity.Entity;
 import org.bukkit.entity.EntityType;
@@ -17,27 +17,28 @@
 import org.bukkit.event.HandlerList;
 import org.bukkit.event.Listener;
 import org.bukkit.event.entity.EntitySpawnEvent;
-import org.bukkit.event.world.ChunkUnloadEvent;
+import org.bukkit.event.world.EntitiesLoadEvent;
+import org.bukkit.event.world.EntitiesUnloadEvent;
 
 import java.util.EnumMap;
+import java.util.List;
 import java.util.Map;
+import java.util.Objects;
 import java.util.TreeMap;
-import java.util.concurrent.atomic.AtomicInteger;
 import java.util.function.Consumer;
 
 public class CustomEntityLimit extends AEFModule implements Consumer<ScheduledTask>, Listener {
 
     private final Map<EntityType, Integer> entityLimits = new EnumMap<>(EntityType.class);
     private final long checkPeriod, minChunkAge;
-    private final boolean logIsEnabled, enableChunkAgeSkip, forceLoadEntities;
+    private final boolean enableChunkAgeSkip, forceLoadEntities;
 
     private ScheduledTask scheduledTask;
 
     public CustomEntityLimit() {
         super("chunk-limits.entity-limits.custom-limit", false, """
                 Limit specific entity types per chunk.\s
                 Read over the defaults at least once before enabling.""");
-        this.logIsEnabled = config.getBoolean(configPath + ".log-removals", true);
         this.checkPeriod = config.getInt(configPath + ".check-period-in-ticks", 1200, """
                 Check all loaded chunks every x ticks.""");
         this.enableChunkAgeSkip = config.getBoolean(configPath + ".chunk-age-skip.enable", true);
@@ -137,14 +138,15 @@ public CustomEntityLimit() {
                 https://jd.papermc.io/paper/1.20.6/org/bukkit/entity/EntityType.html\s
                 Make sure your minecraft version is matching as well.""");
         for (String configuredEntity : section.getKeys(false)) {
+            String configuredLimit = section.getString(configuredEntity);
             try {
-                Integer maxAmountPerChunk = Integer.parseInt(section.getString(configuredEntity));
-                EntityType limitedEntity = EntityType.valueOf(configuredEntity);
-                entityLimits.put(limitedEntity, maxAmountPerChunk);
-            } catch (NumberFormatException e) {
+                entityLimits.put(
+                        EntityType.valueOf(configuredEntity),
+                        Integer.parseInt(Objects.requireNonNull(configuredLimit)));
+            } catch (NumberFormatException | NullPointerException e) {
                 notRecognized(Integer.class, configuredEntity);
             } catch (IllegalArgumentException e) {
-                notRecognized(Material.class, configuredEntity);
+                notRecognized(EntityType.class, configuredEntity);
             }
         }
     }
@@ -162,94 +164,118 @@ public void disable() {
         if (scheduledTask != null) scheduledTask.cancel();
     }
 
-    @EventHandler(priority = EventPriority.HIGHEST, ignoreCancelled = true)
-    private void onChunkUnload(ChunkUnloadEvent event) {
-        if (ChunkUtil.isRetrievalUnsafe(event.getChunk())) return;
+    private void enforceEntityLimit(Map<EntityType, List<Entity>> entitiesByTypes) {
+        for (Map.Entry<EntityType, List<Entity>> entry : entitiesByTypes.entrySet()) {
+            if (!entityLimits.containsKey(entry.getKey())) {
+                continue;
+            }
 
-        for (Map.Entry<EntityType, Integer> limit : entityLimits.entrySet()) {
-            final int maxAllowedPerChunk = limit.getValue();
+            int entityLimit = entityLimits.get(entry.getKey());
 
-            AtomicInteger count = new AtomicInteger();
+            if (entry.getValue().size() <= entityLimit) {
+                continue;
+            }
 
-            for (Entity entity : event.getChunk().getEntities()) {
-                entity.getScheduler().execute(plugin, () -> {
-                    if (entity.getType() != limit.getKey()) return;
-                    if (count.incrementAndGet() <= maxAllowedPerChunk) return;
+            List<Entity> entitiesToRemove = entry.getValue().subList(entityLimit, entry.getValue().size() - 1);
 
+            for (int i = 1; i < entitiesToRemove.size(); i++) {
+                Entity entity = entitiesToRemove.get(i);
+                int index = i;
+                entity.getScheduler().execute(plugin, () -> {
                     entity.remove();
-                    if (logIsEnabled) info("Removed entity " + entity.getType() + " at " +
-                            LocationUtil.toString(entity.getLocation()) + " because reached limit of " + maxAllowedPerChunk);
+                    logger().info("Removed {} at {} ({}/{} limit={})",
+                            entity.getType(),
+                            LocationUtil.toString(entity.getLocation()),
+                            index,
+                            entitiesToRemove.size(),
+                            entityLimit);
                 }, null, 1L);
             }
-        }
-    }
 
-    // Abort spawning process when chunk limit is reached
-    @EventHandler(priority = EventPriority.HIGHEST, ignoreCancelled = true)
-    private void onEntitySpawn(PreCreatureSpawnEvent event) {
-        if (!entityLimits.containsKey(event.getType())) return;
+            entitiesToRemove.clear();
+            entry.getValue().clear();
+        }
 
-        final int maxAllowedPerChunk = entityLimits.get(event.getType());
-        int entityCount = 1; // Count entity that spawned in this event
+        entitiesByTypes.clear();
+    }
 
-        for (Entity entity : event.getSpawnLocation().getChunk().getEntities()) {
-            if (entity.getType() != event.getType()) continue;
+    @Override
+    public void accept(ScheduledTask task) {
+        for (World world : plugin.getServer().getWorlds()) {
+            for (Chunk chunk : world.getLoadedChunks()) {
+                if (ChunkUtil.isRetrievalUnsafe(chunk)) {
+                    continue;
+                }
 
-            entityCount++;
+                plugin.getServer().getRegionScheduler().execute(plugin, world, chunk.getX(), chunk.getZ(), () -> {
+                    if (!forceLoadEntities && !chunk.isEntitiesLoaded()) return;
+                    if (enableChunkAgeSkip && chunk.getInhabitedTime() < minChunkAge) return;
 
-            if (entityCount > maxAllowedPerChunk) {
-                event.setCancelled(true);
-                event.setShouldAbortSpawn(true);
+                    enforceEntityLimit(EntityUtil.getEntityTypeMap(chunk.getEntities()));
+                });
             }
         }
     }
 
     @EventHandler(priority = EventPriority.HIGHEST, ignoreCancelled = true)
-    private void onSpawn(EntitySpawnEvent event) {
-        if (!entityLimits.containsKey(event.getEntityType())) return;
+    private void onEntitiesLoad(EntitiesLoadEvent event) {
+        if (ChunkUtil.isRetrievalUnsafe(event.getChunk())) return;
 
-        final int maxAllowedPerChunk = entityLimits.get(event.getEntityType());
-        int entityCount = 1; // Count entity that spawned in this event
+        enforceEntityLimit(EntityUtil.getEntityTypeMap(event.getEntities()));
+    }
 
-        for (Entity entity : event.getEntity().getChunk().getEntities()) {
-            if (entity.getType() != event.getEntityType()) continue;
+    @EventHandler(priority = EventPriority.HIGHEST, ignoreCancelled = true)
+    private void onEntitiesUnload(EntitiesUnloadEvent event) {
+        if (ChunkUtil.isRetrievalUnsafe(event.getChunk())) return;
+
+        enforceEntityLimit(EntityUtil.getEntityTypeMap(event.getEntities()));
+    }
+
+    @EventHandler(priority = EventPriority.HIGHEST, ignoreCancelled = true)
+    private void onPreCreatureSpawn(PreCreatureSpawnEvent event) {
+        if (!entityLimits.containsKey(event.getType())) return;
+
+        int entityLimit = entityLimits.get(event.getType());
+        int entityCount = 0;
 
-            entityCount++;
-            if (entityCount <= maxAllowedPerChunk) continue;
+        for (Entity entity : event.getSpawnLocation().getChunk().getEntities()) {
+            if (entity.getType() == event.getType()) {
+                entityCount++;
+            }
+        }
 
-            entity.remove();
-            if (logIsEnabled) info("Removed entity " + entity.getType() + " at " +
-                    LocationUtil.toString(entity.getLocation()) + " because reached limit of " + maxAllowedPerChunk);
+        if (entityCount + 1 > entityLimit) {
+            event.setCancelled(true);
+            logger().info("Cancelled {} of {} at {} => count={}, limit={}",
+                    event.getEventName(),
+                    event.getType(),
+                    LocationUtil.toString(event.getSpawnLocation()),
+                    entityCount,
+                    entityLimit);
         }
     }
 
-    @Override
-    public void accept(ScheduledTask task) {
-        for (World world : plugin.getServer().getWorlds()) {
-            for (Chunk chunk : world.getLoadedChunks()) {
-                if (ChunkUtil.isRetrievalUnsafe(chunk)) continue;
+    @EventHandler(priority = EventPriority.HIGHEST, ignoreCancelled = true)
+    private void onEntitySpawn(EntitySpawnEvent event) {
+        if (!entityLimits.containsKey(event.getEntityType())) return;
 
-                plugin.getServer().getRegionScheduler().execute(plugin, world, chunk.getX(), chunk.getZ(), () -> {
-                    if (!forceLoadEntities && !chunk.isEntitiesLoaded()) return;
-                    if (enableChunkAgeSkip && chunk.getInhabitedTime() < minChunkAge) return;
+        int entityLimit = entityLimits.get(event.getEntityType());
+        int entityCount = 0;
 
-                    Entity[] chunkEntities = chunk.getEntities();
-
-                    for (Map.Entry<EntityType, Integer> limit : entityLimits.entrySet()) {
-                        AtomicInteger count = new AtomicInteger();
-                        for (Entity entity : chunkEntities) {
-                            entity.getScheduler().execute(plugin, () -> {
-                                if (entity.getType() != limit.getKey()) return;
-                                if (count.incrementAndGet() <= limit.getValue()) return;
-
-                                entity.remove();
-                                if (logIsEnabled) info("Removed entity " + entity.getType() + " at " +
-                                        LocationUtil.toString(entity.getLocation()) + " because reached limit of " + limit.getValue());
-                            }, null, 1L);
-                        }
-                    }
-                });
+        for (Entity entity : event.getLocation().getChunk().getEntities()) {
+            if (entity.getType() == event.getEntityType()) {
+                entityCount++;
             }
         }
+
+        if (entityCount + 1 > entityLimit) {
+            event.setCancelled(true);
+            logger().info("Cancelled {} of {} at {} => count={}, limit={}",
+                    event.getEventName(),
+                    event.getEntityType(),
+                    LocationUtil.toString(event.getLocation()),
+                    entityCount,
+                    entityLimit);
+        }
     }
 }