diff --git a/native/src/main/java/com/velocitypowered/natives/encryption/JavaVelocityCipher.java b/native/src/main/java/com/velocitypowered/natives/encryption/JavaVelocityCipher.java index 941f83cfb..dc3cd2318 100644 --- a/native/src/main/java/com/velocitypowered/natives/encryption/JavaVelocityCipher.java +++ b/native/src/main/java/com/velocitypowered/natives/encryption/JavaVelocityCipher.java @@ -35,50 +35,25 @@ public class JavaVelocityCipher implements VelocityCipher { } @Override - public void process(ByteBuf source, ByteBuf destination) throws ShortBufferException { + public void process(ByteBuf source) { ensureNotDisposed(); + Preconditions.checkArgument(source.hasArray(), "No source array"); int inBytes = source.readableBytes(); - byte[] asBytes = ByteBufUtil.getBytes(source); - int outputSize = cipher.getOutputSize(inBytes); - byte[] outBuf = new byte[outputSize]; - cipher.update(asBytes, 0, inBytes, outBuf); - destination.writeBytes(outBuf); - } - - @Override - public ByteBuf process(ChannelHandlerContext ctx, ByteBuf source) throws ShortBufferException { - ensureNotDisposed(); - - int inBytes = source.readableBytes(); - ByteBuf asHeapBuf = toHeap(source); - ByteBuf out = ctx.alloc().heapBuffer(cipher.getOutputSize(inBytes)); try { - out.writerIndex( - cipher.update(asHeapBuf.array(), asHeapBuf.arrayOffset() + asHeapBuf.readerIndex(), - inBytes, out.array(), out.arrayOffset() + out.writerIndex())); - return out; - } catch (ShortBufferException e) { - out.release(); - throw e; - } finally { - asHeapBuf.release(); + cipher.update(source.array(), source.arrayOffset(), inBytes, source.array(), + source.arrayOffset()); + } catch (ShortBufferException ex) { + /* This _really_ shouldn't happen - AES CFB8 will work in place. + If you run into this, that means that for whatever reason the Java Runtime has determined + that the output buffer needs more bytes than the input buffer. When we are working with + AES-CFB8, the output size is equal to the input size. See the problem? */ + throw new AssertionError("Cipher update did not operate in place and requested a larger " + + "buffer than the source buffer"); } } - private static ByteBuf toHeap(ByteBuf src) { - if (src.hasArray()) { - return src.retain(); - } - - // Copy into a temporary heap buffer. We could use a local buffer, but Netty pools all buffers, - // so we'd lose more than we gain. - ByteBuf asHeapBuf = src.alloc().heapBuffer(src.readableBytes()); - asHeapBuf.writeBytes(src); - return asHeapBuf; - } - @Override public void dispose() { disposed = true; @@ -90,6 +65,6 @@ public class JavaVelocityCipher implements VelocityCipher { @Override public BufferPreference preferredBufferType() { - return BufferPreference.HEAP_PREFERRED; + return BufferPreference.HEAP_REQUIRED; } } diff --git a/native/src/main/java/com/velocitypowered/natives/encryption/NativeVelocityCipher.java b/native/src/main/java/com/velocitypowered/natives/encryption/NativeVelocityCipher.java index 35de5519d..ee7ab1532 100644 --- a/native/src/main/java/com/velocitypowered/natives/encryption/NativeVelocityCipher.java +++ b/native/src/main/java/com/velocitypowered/natives/encryption/NativeVelocityCipher.java @@ -33,40 +33,14 @@ public class NativeVelocityCipher implements VelocityCipher { } @Override - public void process(ByteBuf source, ByteBuf destination) throws ShortBufferException { + public void process(ByteBuf source) { ensureNotDisposed(); source.memoryAddress(); - destination.memoryAddress(); - // The exact amount we read in is also the amount we write out. + long base = source.memoryAddress() + source.readerIndex(); int len = source.readableBytes(); - destination.ensureWritable(len); - impl.process(ctx, source.memoryAddress() + source.readerIndex(), len, - destination.memoryAddress() + destination.writerIndex(), encrypt); - - source.skipBytes(len); - destination.writerIndex(destination.writerIndex() + len); - } - - @Override - public ByteBuf process(ChannelHandlerContext ctx, ByteBuf source) throws ShortBufferException { - ensureNotDisposed(); - source.memoryAddress(); // sanity check - - int len = source.readableBytes(); - ByteBuf out = ctx.alloc().directBuffer(len); - - try { - impl.process(this.ctx, source.memoryAddress() + source.readerIndex(), len, - out.memoryAddress(), encrypt); - source.skipBytes(len); - out.writerIndex(len); - return out; - } catch (Exception e) { - out.release(); - throw e; - } + impl.process(ctx, base, len, base, encrypt); } @Override diff --git a/native/src/main/java/com/velocitypowered/natives/encryption/VelocityCipher.java b/native/src/main/java/com/velocitypowered/natives/encryption/VelocityCipher.java index c52c88737..3f17a6604 100644 --- a/native/src/main/java/com/velocitypowered/natives/encryption/VelocityCipher.java +++ b/native/src/main/java/com/velocitypowered/natives/encryption/VelocityCipher.java @@ -7,8 +7,5 @@ import io.netty.channel.ChannelHandlerContext; import javax.crypto.ShortBufferException; public interface VelocityCipher extends Disposable, Native { - - void process(ByteBuf source, ByteBuf destination) throws ShortBufferException; - - ByteBuf process(ChannelHandlerContext ctx, ByteBuf source) throws ShortBufferException; + void process(ByteBuf source); } diff --git a/native/src/main/java/com/velocitypowered/natives/util/BufferPreference.java b/native/src/main/java/com/velocitypowered/natives/util/BufferPreference.java index 7be4fa74e..6814b11fd 100644 --- a/native/src/main/java/com/velocitypowered/natives/util/BufferPreference.java +++ b/native/src/main/java/com/velocitypowered/natives/util/BufferPreference.java @@ -1,6 +1,10 @@ package com.velocitypowered.natives.util; public enum BufferPreference { + /** + * A heap buffer is required. + */ + HEAP_REQUIRED, /** * A heap buffer is preferred (but not required). */ diff --git a/native/src/main/java/com/velocitypowered/natives/util/MoreByteBufUtils.java b/native/src/main/java/com/velocitypowered/natives/util/MoreByteBufUtils.java index 5bf328396..dd41e3fc7 100644 --- a/native/src/main/java/com/velocitypowered/natives/util/MoreByteBufUtils.java +++ b/native/src/main/java/com/velocitypowered/natives/util/MoreByteBufUtils.java @@ -20,28 +20,30 @@ public class MoreByteBufUtils { * @return a buffer compatible with the native */ public static ByteBuf ensureCompatible(ByteBufAllocator alloc, Native nativeStuff, ByteBuf buf) { - if (nativeStuff.preferredBufferType() != BufferPreference.DIRECT_REQUIRED - || buf.hasMemoryAddress()) { + if (isCompatible(nativeStuff, buf)) { return buf.retain(); } // It's not, so we must make a direct copy. - ByteBuf newBuf = alloc.directBuffer(buf.readableBytes()); + ByteBuf newBuf = preferredBuffer(alloc, nativeStuff, buf.readableBytes()); newBuf.writeBytes(buf); return newBuf; } - /** - * Creates a {@link ByteBuf} that will have the best performance with the specified - * {@code nativeStuff}. - * - * @param alloc the {@link ByteBufAllocator} to use - * @param nativeStuff the native we are working with - * @return a buffer compatible with the native - */ - public static ByteBuf preferredBuffer(ByteBufAllocator alloc, Native nativeStuff) { - return nativeStuff.preferredBufferType() != BufferPreference.HEAP_PREFERRED - ? alloc.directBuffer() : alloc.heapBuffer(); + private static boolean isCompatible(Native nativeStuff, ByteBuf buf) { + BufferPreference preferred = nativeStuff.preferredBufferType(); + switch (preferred) { + case DIRECT_PREFERRED: + case HEAP_PREFERRED: + // The native prefers this type, but doesn't strictly require we provide it. + return true; + case DIRECT_REQUIRED: + return buf.hasMemoryAddress(); + case HEAP_REQUIRED: + return buf.hasArray(); + default: + throw new AssertionError("Preferred buffer type unknown"); + } } /** @@ -55,7 +57,15 @@ public class MoreByteBufUtils { */ public static ByteBuf preferredBuffer(ByteBufAllocator alloc, Native nativeStuff, int initialCapacity) { - return nativeStuff.preferredBufferType() != BufferPreference.HEAP_PREFERRED - ? alloc.directBuffer(initialCapacity) : alloc.heapBuffer(initialCapacity); + switch (nativeStuff.preferredBufferType()) { + case HEAP_REQUIRED: + case HEAP_PREFERRED: + return alloc.heapBuffer(initialCapacity); + case DIRECT_PREFERRED: + case DIRECT_REQUIRED: + return alloc.directBuffer(initialCapacity); + default: + throw new AssertionError("Preferred buffer type unknown"); + } } } diff --git a/native/src/test/java/com/velocitypowered/natives/encryption/VelocityCipherTest.java b/native/src/test/java/com/velocitypowered/natives/encryption/VelocityCipherTest.java index d7328c922..93ab89117 100644 --- a/native/src/test/java/com/velocitypowered/natives/encryption/VelocityCipherTest.java +++ b/native/src/test/java/com/velocitypowered/natives/encryption/VelocityCipherTest.java @@ -56,20 +56,18 @@ class VelocityCipherTest { VelocityCipher encrypt = factory.forEncryption(new SecretKeySpec(AES_KEY, "AES")); ByteBuf source = bufSupplier.get(); - ByteBuf dest = bufSupplier.get(); - ByteBuf decryptionBuf = bufSupplier.get(); source.writeBytes(TEST_DATA); + ByteBuf workingBuf = source.copy(); + try { - encrypt.process(source, dest); - decrypt.process(dest, decryptionBuf); - source.readerIndex(0); - assertTrue(ByteBufUtil.equals(source, decryptionBuf)); + encrypt.process(workingBuf); + decrypt.process(workingBuf); + assertTrue(ByteBufUtil.equals(source, workingBuf)); } finally { source.release(); - dest.release(); - decryptionBuf.release(); + workingBuf.release(); decrypt.dispose(); encrypt.dispose(); } diff --git a/proxy/src/main/java/com/velocitypowered/proxy/protocol/netty/MinecraftCipherDecoder.java b/proxy/src/main/java/com/velocitypowered/proxy/protocol/netty/MinecraftCipherDecoder.java index 4541fc902..2b4d0b502 100644 --- a/proxy/src/main/java/com/velocitypowered/proxy/protocol/netty/MinecraftCipherDecoder.java +++ b/proxy/src/main/java/com/velocitypowered/proxy/protocol/netty/MinecraftCipherDecoder.java @@ -18,11 +18,14 @@ public class MinecraftCipherDecoder extends ByteToMessageDecoder { @Override protected void decode(ChannelHandlerContext ctx, ByteBuf in, List out) throws Exception { - ByteBuf compatible = MoreByteBufUtils.ensureCompatible(ctx.alloc(), cipher, in); + ByteBuf compatible = MoreByteBufUtils.ensureCompatible(ctx.alloc(), cipher, in).slice(); try { - out.add(cipher.process(ctx, compatible)); - } finally { - compatible.release(); + cipher.process(compatible); + out.add(compatible); + in.skipBytes(in.readableBytes()); + } catch (Exception e) { + compatible.release(); // compatible will never be used if we throw an exception + throw e; } } diff --git a/proxy/src/main/java/com/velocitypowered/proxy/protocol/netty/MinecraftCipherEncoder.java b/proxy/src/main/java/com/velocitypowered/proxy/protocol/netty/MinecraftCipherEncoder.java index dd944a19f..ecaf9adb6 100644 --- a/proxy/src/main/java/com/velocitypowered/proxy/protocol/netty/MinecraftCipherEncoder.java +++ b/proxy/src/main/java/com/velocitypowered/proxy/protocol/netty/MinecraftCipherEncoder.java @@ -20,9 +20,11 @@ public class MinecraftCipherEncoder extends MessageToMessageEncoder { protected void encode(ChannelHandlerContext ctx, ByteBuf msg, List out) throws Exception { ByteBuf compatible = MoreByteBufUtils.ensureCompatible(ctx.alloc(), cipher, msg); try { - out.add(cipher.process(ctx, compatible)); - } finally { - compatible.release(); + cipher.process(compatible); + out.add(compatible); + } catch (Exception e) { + compatible.release(); // compatible will never be used if we throw an exception + throw e; } }