Merge pull request #65 from apemanzilla/fix/drop-consumer-overflow

Prevent stack overflows when using turtle.place() with a full inventory

src/main/java/dan200/computercraft/ComputerCraft.java

@@ -98,7 +98,7 @@ import java.net.MalformedURLException;
 import java.net.URISyntaxException;
 import java.net.URL;
 import java.util.*;
-import java.util.function.Consumer;
+import java.util.function.Function;
 import java.util.zip.ZipEntry;
 import java.util.zip.ZipFile;
 
@@ -1125,18 +1125,18 @@ public class ComputerCraft
         turtleProxy.addAllUpgradedTurtles( list );
     }
 
-    public static void setDropConsumer( Entity entity, Consumer<ItemStack> consumer )
+    public static void setDropConsumer( Entity entity, Function<ItemStack, ItemStack> consumer )
     {
         turtleProxy.setDropConsumer( entity, consumer );
     }
 
-    public static void setDropConsumer( World world, BlockPos pos, Consumer<ItemStack> consumer )
+    public static void setDropConsumer( World world, BlockPos pos, Function<ItemStack, ItemStack> consumer )
     {
         turtleProxy.setDropConsumer( world, pos, consumer );
     }
 
-    public static void clearDropConsumer( )
+    public static List<ItemStack> clearDropConsumer( )
     {
-        turtleProxy.clearDropConsumer();
+        return turtleProxy.clearDropConsumer();
     }
 }

src/main/java/dan200/computercraft/shared/proxy/CCTurtleProxyCommon.java

@@ -48,17 +48,19 @@ import net.minecraftforge.registries.IForgeRegistry;
 
 import javax.annotation.Nonnull;
 import java.lang.ref.WeakReference;
+import java.util.ArrayList;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
-import java.util.function.Consumer;
+import java.util.function.Function;
 
 public abstract class CCTurtleProxyCommon implements ICCTurtleProxy
 {
     private Map<Integer, ITurtleUpgrade> m_legacyTurtleUpgrades;
     private Map<String, ITurtleUpgrade> m_turtleUpgrades;
 
-    private Consumer<ItemStack> dropConsumer;
+    private Function<ItemStack, ItemStack> dropConsumer;
+    private List<ItemStack> remainingDrops;
     private WeakReference<World> dropWorld;
     private BlockPos dropPos;
     private AxisAlignedBB dropBounds;
@@ -200,9 +202,10 @@ public abstract class CCTurtleProxyCommon implements ICCTurtleProxy
     }
 
     @Override
-    public void setDropConsumer( Entity entity, Consumer<ItemStack> consumer )
+    public void setDropConsumer( Entity entity, Function<ItemStack, ItemStack> consumer )
     {
         dropConsumer = consumer;
+        remainingDrops = new ArrayList<>();
         dropEntity = new WeakReference<>( entity );
         dropWorld = new WeakReference<>( entity.world );
         dropPos = null;
@@ -212,9 +215,10 @@ public abstract class CCTurtleProxyCommon implements ICCTurtleProxy
     }
 
     @Override
-    public void setDropConsumer( World world, BlockPos pos, Consumer<ItemStack> consumer )
+    public void setDropConsumer( World world, BlockPos pos, Function<ItemStack, ItemStack> consumer )
     {
         dropConsumer = consumer;
+        remainingDrops = new ArrayList<>();
         dropEntity = null;
         dropWorld = new WeakReference<>( world );
         dropPos = pos;
@@ -222,7 +226,7 @@ public abstract class CCTurtleProxyCommon implements ICCTurtleProxy
     }
 
     @Override
-    public void clearDropConsumer()
+    public List<ItemStack> clearDropConsumer()
     {
         if( dropEntity != null )
         {
@@ -232,17 +236,22 @@ public abstract class CCTurtleProxyCommon implements ICCTurtleProxy
                 entity.captureDrops = false;
                 if( entity.capturedDrops != null )
                 {
-                    for( EntityItem entityItem : entity.capturedDrops ) dropConsumer.accept( entityItem.getItem() );
+                    for( EntityItem entityItem : entity.capturedDrops ) handleDrops( entityItem.getItem() );
                     entity.capturedDrops.clear();
                 }
             }
         }
 
+        List<ItemStack> remainingStacks = remainingDrops;
+
         dropConsumer = null;
+        remainingDrops = null;
         dropEntity = null;
         dropWorld = null;
         dropPos = null;
         dropBounds = null;
+
+        return remainingStacks;
     }
 
     private void registerTurtleUpgradeInternal( ITurtleUpgrade upgrade )
@@ -471,6 +480,12 @@ public abstract class CCTurtleProxyCommon implements ICCTurtleProxy
         MinecraftForge.EVENT_BUS.register( handlers );
     }
 
+    private void handleDrops(ItemStack stack)
+    {
+        ItemStack remaining = dropConsumer.apply(stack);
+        if (!remaining.isEmpty()) remainingDrops.add(remaining);
+    }
+
     private class ForgeHandlers
     {
         @SubscribeEvent(priority = EventPriority.LOWEST)
@@ -480,7 +495,7 @@ public abstract class CCTurtleProxyCommon implements ICCTurtleProxy
             if( dropEntity != null && event.getEntity() == dropEntity.get() )
             {
                 List<EntityItem> drops = event.getDrops();
-                for( EntityItem entityItem : drops ) dropConsumer.accept( entityItem.getItem() );
+                for( EntityItem entityItem : drops ) handleDrops( entityItem.getItem() );
                 drops.clear();
             }
         }
@@ -494,7 +509,7 @@ public abstract class CCTurtleProxyCommon implements ICCTurtleProxy
             {
                 for( ItemStack item : event.getDrops() )
                 {
-                    if( event.getWorld().rand.nextFloat() < event.getDropChance() ) dropConsumer.accept( item );
+                    if( event.getWorld().rand.nextFloat() < event.getDropChance() ) handleDrops( item );
                 }
                 event.getDrops().clear();
             }
@@ -507,7 +522,7 @@ public abstract class CCTurtleProxyCommon implements ICCTurtleProxy
             if( dropWorld != null && dropWorld.get() == event.getWorld() && event.getEntity() instanceof EntityItem
                 && dropBounds.contains( event.getEntity().getPositionVector() ) )
             {
-                dropConsumer.accept( ((EntityItem) event.getEntity()).getItem() );
+                handleDrops( ((EntityItem) event.getEntity()).getItem() );
                 event.setCanceled( true );
             }
         }

src/main/java/dan200/computercraft/shared/proxy/ICCTurtleProxy.java

@@ -14,7 +14,9 @@ import net.minecraft.util.math.BlockPos;
 import net.minecraft.world.World;
 
 import javax.annotation.Nonnull;
+import java.util.List;
 import java.util.function.Consumer;
+import java.util.function.Function;
 
 public interface ICCTurtleProxy
 {
@@ -27,7 +29,7 @@ public interface ICCTurtleProxy
     ITurtleUpgrade getTurtleUpgrade( @Nonnull ItemStack item );
     void addAllUpgradedTurtles( NonNullList<ItemStack> list );
 
-    void setDropConsumer( Entity entity, Consumer<ItemStack> consumer );
+    void setDropConsumer( Entity entity, Function<ItemStack, ItemStack> consumer );
-    void setDropConsumer( World world, BlockPos pos, Consumer<ItemStack> consumer );
+    void setDropConsumer( World world, BlockPos pos, Function<ItemStack, ItemStack> consumer );
-    void clearDropConsumer();
+    List<ItemStack> clearDropConsumer();
 }

src/main/java/dan200/computercraft/shared/turtle/core/TurtlePlaceCommand.java

@@ -30,7 +30,6 @@ import net.minecraft.util.math.BlockPos;
 import net.minecraft.util.math.Vec3d;
 import net.minecraft.util.text.TextComponentString;
 import net.minecraft.world.World;
-import net.minecraft.world.WorldServer;
 import net.minecraftforge.common.ForgeHooks;
 import net.minecraftforge.common.MinecraftForge;
 import net.minecraftforge.event.entity.player.PlayerInteractEvent;
@@ -38,6 +37,7 @@ import net.minecraftforge.fml.common.eventhandler.Event;
 import org.apache.commons.lang3.tuple.Pair;
 
 import javax.annotation.Nonnull;
+import java.util.List;
 
 public class TurtlePlaceCommand implements ITurtleCommand
 {
@@ -241,14 +241,10 @@ public class TurtlePlaceCommand implements ITurtleCommand
         // Start claiming entity drops
         Entity hitEntity = hit.getKey();
         Vec3d hitPos = hit.getValue();
-        ComputerCraft.setDropConsumer( hitEntity, ( drop ) ->
+        ComputerCraft.setDropConsumer(
-        {
+            hitEntity,
-            ItemStack remainder = InventoryUtil.storeItems( drop, turtle.getItemHandler(), turtle.getSelectedSlot() );
+            drop -> InventoryUtil.storeItems( drop, turtle.getItemHandler(), turtle.getSelectedSlot() )
-            if( !remainder.isEmpty() )
+        );
-            {
-                WorldUtil.dropItemStack( remainder, world, position, turtle.getDirection().getOpposite() );
-            }
-        } );
 
         // Place on the entity
         boolean placed = false;
@@ -285,7 +281,11 @@ public class TurtlePlaceCommand implements ITurtleCommand
         }
 
         // Stop claiming drops
-        ComputerCraft.clearDropConsumer();
+        List<ItemStack> remainingDrops = ComputerCraft.clearDropConsumer();
+        for( ItemStack remaining : remainingDrops )
+        {
+            WorldUtil.dropItemStack( remaining, world, position, turtle.getDirection().getOpposite() );
+        }
 
         // Put everything we collected into the turtles inventory, then return
         ItemStack remainder = turtlePlayer.unloadInventory( turtle );

src/main/java/dan200/computercraft/shared/turtle/upgrades/TurtleTool.java

@@ -42,9 +42,8 @@ import org.apache.commons.lang3.tuple.Pair;
 
 import javax.annotation.Nonnull;
 import javax.vecmath.Matrix4f;
-import java.util.ArrayList;
 import java.util.List;
-import java.util.function.Consumer;
+import java.util.function.Function;
 
 public class TurtleTool implements ITurtleUpgrade
 {
@@ -187,8 +186,7 @@ public class TurtleTool implements ITurtleUpgrade
             }
 
             // Start claiming entity drops
-            List<ItemStack> extra = new ArrayList<>();
+            ComputerCraft.setDropConsumer( hitEntity, turtleDropConsumer( turtle ) );
-            ComputerCraft.setDropConsumer( hitEntity, turtleDropConsumer( turtle, extra ) );
 
             // Attack the entity
             boolean attacked = false;
@@ -220,7 +218,7 @@ public class TurtleTool implements ITurtleUpgrade
             }
 
             // Stop claiming drops
-            stopConsuming( turtle, extra );
+            stopConsuming( turtle );
 
             // Put everything we collected into the turtles inventory, then return
             if( attacked )
@@ -277,8 +275,7 @@ public class TurtleTool implements ITurtleUpgrade
             }
 
             // Consume the items the block drops
-            List<ItemStack> extra = new ArrayList<>();
+            ComputerCraft.setDropConsumer( world, blockPosition, turtleDropConsumer( turtle ) );
-            ComputerCraft.setDropConsumer( world, blockPosition, turtleDropConsumer( turtle, extra ) );
 
             TileEntity tile = world.getTileEntity( blockPosition );
 
@@ -297,7 +294,7 @@ public class TurtleTool implements ITurtleUpgrade
                 state.getBlock().harvestBlock( world, turtlePlayer, blockPosition, state, tile, turtlePlayer.getHeldItemMainhand() );
             }
 
-            stopConsuming( turtle, extra );
+            stopConsuming( turtle );
 
             // Remember the previous block
             if( turtle instanceof TurtleBrain )
@@ -312,18 +309,14 @@ public class TurtleTool implements ITurtleUpgrade
         return TurtleCommandResult.failure( "Nothing to dig here" );
     }
 
-    private Consumer<ItemStack> turtleDropConsumer( ITurtleAccess turtle, List<ItemStack> extra )
+    private Function<ItemStack, ItemStack> turtleDropConsumer( ITurtleAccess turtle )
     {
-        return ( drop ) ->
+        return drop -> InventoryUtil.storeItems( drop, turtle.getItemHandler(), turtle.getSelectedSlot() );
-        {
-            ItemStack remainder = InventoryUtil.storeItems( drop, turtle.getItemHandler(), turtle.getSelectedSlot() );
-            if( !remainder.isEmpty() ) extra.add( remainder );
-        };
     }
 
-    private void stopConsuming( ITurtleAccess turtle, List<ItemStack> extra )
+    private void stopConsuming( ITurtleAccess turtle )
     {
-        ComputerCraft.clearDropConsumer();
+        List<ItemStack> extra = ComputerCraft.clearDropConsumer();
         for( ItemStack remainder : extra )
         {
             WorldUtil.dropItemStack( remainder, turtle.getWorld(), turtle.getPosition(), turtle.getDirection().getOpposite() );