Merge pull request #254 from Ichbinjoe/aes-nocopy

Make AES crypto operations no-copy
This commit is contained in:
Andrew Steinborn 2019-10-20 00:13:32 -04:00 committed by GitHub
commit 523cea529e
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
8 changed files with 64 additions and 101 deletions

View File

@ -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;
}
}

View File

@ -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

View File

@ -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);
}

View File

@ -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).
*/

View File

@ -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");
}
}
}

View File

@ -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();
}

View File

@ -18,11 +18,14 @@ public class MinecraftCipherDecoder extends ByteToMessageDecoder {
@Override
protected void decode(ChannelHandlerContext ctx, ByteBuf in, List<Object> 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;
}
}

View File

@ -20,9 +20,11 @@ public class MinecraftCipherEncoder extends MessageToMessageEncoder<ByteBuf> {
protected void encode(ChannelHandlerContext ctx, ByteBuf msg, List<Object> 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;
}
}