diff --git a/src/main/java/com/github/hhhzzzsss/songplayer/Util.java b/src/main/java/com/github/hhhzzzsss/songplayer/Util.java index 1b3d1a7..fa0c620 100644 --- a/src/main/java/com/github/hhhzzzsss/songplayer/Util.java +++ b/src/main/java/com/github/hhhzzzsss/songplayer/Util.java @@ -12,6 +12,7 @@ import net.minecraft.text.Style; import net.minecraft.text.Text; import java.io.IOException; +import java.io.InputStream; import java.nio.file.Files; import java.nio.file.Path; import java.util.ArrayList; @@ -30,6 +31,41 @@ public class Util { catch (IOException e) {} } + public static class LimitedSizeInputStream extends InputStream { + private final InputStream original; + private final long maxSize; + private long total; + + public LimitedSizeInputStream(InputStream original, long maxSize) { + this.original = original; + this.maxSize = maxSize; + } + + @Override + public int read() throws IOException { + int i = original.read(); + if (i>=0) incrementCounter(1); + return i; + } + + @Override + public int read(byte b[]) throws IOException { + return read(b, 0, b.length); + } + + @Override + public int read(byte b[], int off, int len) throws IOException { + int i = original.read(b, off, len); + if (i>=0) incrementCounter(i); + return i; + } + + private void incrementCounter(int size) throws IOException { + total += size; + if (total>maxSize) throw new IOException("Input stream exceeded maximum size of " + maxSize + " bytes"); + } + } + public static String formatTime(long milliseconds) { long temp = Math.abs(milliseconds); temp /= 1000; diff --git a/src/main/java/com/github/hhhzzzsss/songplayer/conversion/SPConverter.java b/src/main/java/com/github/hhhzzzsss/songplayer/conversion/SPConverter.java index 7a3fd7d..41da087 100644 --- a/src/main/java/com/github/hhhzzzsss/songplayer/conversion/SPConverter.java +++ b/src/main/java/com/github/hhhzzzsss/songplayer/conversion/SPConverter.java @@ -1,5 +1,6 @@ package com.github.hhhzzzsss.songplayer.conversion; +import com.github.hhhzzzsss.songplayer.Util; import com.github.hhhzzzsss.songplayer.song.Note; import com.github.hhhzzzsss.songplayer.song.Song; @@ -12,9 +13,10 @@ import java.util.zip.GZIPOutputStream; public class SPConverter { public static final byte[] FILE_TYPE_SIGNATURE = {-53, 123, -51, -124, -122, -46, -35, 38}; + public static final long MAX_UNCOMPRESSED_SIZE = 50*1024*1024; public static Song getSongFromBytes(byte[] bytes, String fileName) throws IOException { - InputStream is = new GZIPInputStream(new ByteArrayInputStream(bytes)); + InputStream is = new Util.LimitedSizeInputStream(new GZIPInputStream(new ByteArrayInputStream(bytes)), MAX_UNCOMPRESSED_SIZE); bytes = is.readAllBytes(); is.close(); diff --git a/src/main/java/com/github/hhhzzzsss/songplayer/item/SongItemConfirmationScreen.java b/src/main/java/com/github/hhhzzzsss/songplayer/item/SongItemConfirmationScreen.java index 29769b4..237ab90 100644 --- a/src/main/java/com/github/hhhzzzsss/songplayer/item/SongItemConfirmationScreen.java +++ b/src/main/java/com/github/hhhzzzsss/songplayer/item/SongItemConfirmationScreen.java @@ -25,7 +25,7 @@ public class SongItemConfirmationScreen extends Screen { private static final Text CONFIRM = Text.literal("Play"); private static final Text CANCEL = Text.literal("Cancel"); - public SongItemConfirmationScreen(ItemStack stack) throws IOException { + public SongItemConfirmationScreen(ItemStack stack) throws IOException, IllegalArgumentException { super(Text.literal("Use song item")); this.stack = stack; this.loaderThread = new SongItemLoaderThread(stack); diff --git a/src/main/java/com/github/hhhzzzsss/songplayer/item/SongItemLoaderThread.java b/src/main/java/com/github/hhhzzzsss/songplayer/item/SongItemLoaderThread.java index 819b58e..7c7b6ed 100644 --- a/src/main/java/com/github/hhhzzzsss/songplayer/item/SongItemLoaderThread.java +++ b/src/main/java/com/github/hhhzzzsss/songplayer/item/SongItemLoaderThread.java @@ -14,7 +14,7 @@ public class SongItemLoaderThread extends SongLoaderThread { public int maxNotesPerSecond = 0; public double avgNotesPerSecond = 0; - public SongItemLoaderThread(ItemStack stack) throws IOException { + public SongItemLoaderThread(ItemStack stack) throws IOException, IllegalArgumentException { songData = SongItemUtils.getSongData(stack); if (songData == null) { throw new IOException("Song data is missing"); diff --git a/src/main/java/com/github/hhhzzzsss/songplayer/item/SongItemUtils.java b/src/main/java/com/github/hhhzzzsss/songplayer/item/SongItemUtils.java index 6e2510f..72b8cc0 100644 --- a/src/main/java/com/github/hhhzzzsss/songplayer/item/SongItemUtils.java +++ b/src/main/java/com/github/hhhzzzsss/songplayer/item/SongItemUtils.java @@ -50,7 +50,7 @@ public class SongItemUtils { return getSongItemTag(stack) != null; } - public static byte[] getSongData(ItemStack stack) { + public static byte[] getSongData(ItemStack stack) throws IllegalArgumentException { NbtCompound songPlayerNbt = getSongItemTag(stack); if (songPlayerNbt == null || !songPlayerNbt.contains(SONG_DATA_KEY, NbtElement.STRING_TYPE)) { return null; diff --git a/src/main/java/com/github/hhhzzzsss/songplayer/mixin/MinecraftClientMixin.java b/src/main/java/com/github/hhhzzzsss/songplayer/mixin/MinecraftClientMixin.java index 573f654..564d7de 100644 --- a/src/main/java/com/github/hhhzzzsss/songplayer/mixin/MinecraftClientMixin.java +++ b/src/main/java/com/github/hhhzzzsss/songplayer/mixin/MinecraftClientMixin.java @@ -32,6 +32,9 @@ public class MinecraftClientMixin { @Shadow public HitResult crosshairTarget; + @Shadow + private int itemUseCooldown; + @Inject(at = @At("HEAD"), method = "render(Z)V") public void onRender(boolean tick, CallbackInfo ci) { if (SongPlayer.MC.world != null && SongPlayer.MC.player != null && SongPlayer.MC.interactionManager != null) { @@ -72,9 +75,10 @@ public class MinecraftClientMixin { if (SongItemUtils.isSongItem(stack)) { try { SongPlayer.MC.setScreen(new SongItemConfirmationScreen(stack)); - } catch (IOException e) { + } catch (Exception e) { SongPlayer.addChatMessage("§cFailed to load song item: §4" + e.getMessage()); } + itemUseCooldown = 4; ci.cancel(); } }