JNI native zlib compression 🔥

This commit is contained in:
Andrew Steinborn 2018-08-01 22:22:09 -04:00
parent 22dd4bbb99
commit fb4e6fd8f1
22 changed files with 610 additions and 3 deletions

17
native/README.md Normal file
View File

@ -0,0 +1,17 @@
# velocity-natives
This directory contains native acceleration code for Velocity, along with
traditional Java fallbacks.
## Compression
* **Supported platforms**: macOS 10.13, Linux amd64 (precompiled binary is built on Debian 9 with JDK 8)
* **Rationale**: Using a native zlib wrapper, we can avoid multiple trips into Java just to copy memory around.
## Encryption
* No natives available yet, this will use the support inside your Java install.
## OS support
If you are on Alpine Linux, `apk add libc6-compat` will enable native support.

View File

@ -5,4 +5,6 @@ plugins {
dependencies {
compile "com.google.guava:guava:${guavaVersion}"
compile "io.netty:netty-buffer:${nettyVersion}"
testCompile "org.junit.jupiter:junit-jupiter-api:${junitVersion}"
testCompile "org.junit.jupiter:junit-jupiter-engine:${junitVersion}"
}

4
native/compile-linux.sh Executable file
View File

@ -0,0 +1,4 @@
#!/bin/bash
# Modify as you need.
gcc -I$JAVA_HOME/include/ -I$JAVA_HOME/include/linux/ -shared -lz src/main/c/*.c -o src/main/resources/linux_x64/velocity-compress.so

5
native/compile-osx.sh Executable file
View File

@ -0,0 +1,5 @@
#!/bin/bash
# Modify as you need.
export JAVA_HOME=/Library/Java/JavaVirtualMachines/jdk1.8.0_144.jdk/Contents/Home
clang -I$JAVA_HOME/include/ -I$JAVA_HOME/include/darwin/ -shared -lz src/main/c/*.c -o src/main/resources/macosx/velocity-compress.dylib

View File

@ -0,0 +1,12 @@
#include <jni.h>
#include "jni_util.h"
JNIEXPORT void JNICALL
throwException(JNIEnv *env, const char *type, const char *msg)
{
jclass klazz = (*env)->FindClass(env, type);
if (klazz != 0) {
(*env)->ThrowNew(env, klazz, msg);
}
}

View File

@ -0,0 +1,4 @@
#include <jni.h>
JNIEXPORT void JNICALL
throwException(JNIEnv *env, const char *type, const char *msg);

View File

@ -0,0 +1,124 @@
#include <assert.h>
#include <jni.h>
#include <stdlib.h>
#include <zlib.h>
#include "jni_util.h"
static jfieldID finishedID;
static jfieldID consumedID;
JNIEXPORT void JNICALL
Java_com_velocitypowered_natives_compression_NativeZlibDeflate_initIDs(JNIEnv *env, jclass cls)
{
finishedID = (*env)->GetFieldID(env, cls, "finished", "Z");
consumedID = (*env)->GetFieldID(env, cls, "consumed", "I");
}
JNIEXPORT jlong JNICALL
Java_com_velocitypowered_natives_compression_NativeZlibDeflate_init(JNIEnv *env,
jobject obj,
jint level)
{
z_stream* stream = calloc(1, sizeof(z_stream));
if (stream == 0) {
// Out of memory!
throwException(env, "java/lang/OutOfMemoryError", "zlib allocate stream");
return 0;
}
char *msg;
int ret = deflateInit(stream, level);
switch (ret) {
case Z_OK:
return (jlong) stream;
case Z_MEM_ERROR:
free(stream);
throwException(env, "java/lang/OutOfMemoryError", "zlib init");
return 0;
case Z_STREAM_ERROR:
free(stream);
char message[32];
snprintf(message, 32, "invalid level %d", level);
throwException(env, "java/lang/IllegalArgumentException", message);
return 0;
default:
msg = stream->msg;
free(stream);
throwException(env, "java/util/zip/DataFormatException", msg);
return 0;
}
}
JNIEXPORT void JNICALL
Java_com_velocitypowered_natives_compression_NativeZlibDeflate_free(JNIEnv *env,
jobject obj,
jlong ctx)
{
z_stream* stream = (z_stream*) ctx;
int ret = deflateEnd(stream);
char *msg = stream->msg;
free((void*) ctx);
switch (ret) {
case Z_OK:
break;
case Z_STREAM_ERROR:
if (msg == NULL) {
msg = "stream state inconsistent";
}
case Z_DATA_ERROR:
if (msg == NULL) {
msg = "data was discarded";
}
throwException(env, "java/lang/IllegalArgumentException", msg);
break;
}
}
JNIEXPORT int JNICALL
Java_com_velocitypowered_natives_compression_NativeZlibDeflate_process(JNIEnv *env,
jobject obj,
jlong ctx,
jlong sourceAddress,
jint sourceLength,
jlong destinationAddress,
jint destinationLength,
jboolean flush)
{
z_stream* stream = (z_stream*) ctx;
stream->next_in = (Bytef *) sourceAddress;
stream->next_out = (Bytef *) destinationAddress;
stream->avail_in = sourceLength;
stream->avail_out = destinationLength;
int res = deflate(stream, flush ? Z_FINISH : Z_NO_FLUSH);
switch (res) {
case Z_STREAM_END:
// The stream has ended.
(*env)->SetBooleanField(env, obj, finishedID, JNI_TRUE);
// fall-through
case Z_OK:
// Not yet completed, but progress has been made. Tell Java how many bytes we've processed.
(*env)->SetIntField(env, obj, consumedID, sourceLength - stream->avail_in);
return destinationLength - stream->avail_out;
case Z_BUF_ERROR:
// This is not fatal. Just say we need more data. Usually this applies to the next_out buffer,
// which NativeVelocityCompressor will notice and will expand the buffer.
return 0;
default:
throwException(env, "java/util/zip/DataFormatException", stream->msg);
return 0;
}
}
JNIEXPORT void JNICALL
Java_com_velocitypowered_natives_compression_NativeZlibDeflate_reset(JNIEnv *env,
jobject obj,
jlong ctx)
{
z_stream* stream = (z_stream*) ctx;
int ret = deflateReset(stream);
assert(ret == Z_OK);
}

View File

@ -0,0 +1,120 @@
#include <assert.h>
#include <jni.h>
#include <stdlib.h>
#include <zlib.h>
#include "jni_util.h"
static jfieldID finishedID;
static jfieldID consumedID;
JNIEXPORT void JNICALL
Java_com_velocitypowered_natives_compression_NativeZlibInflate_initIDs(JNIEnv *env, jclass cls)
{
finishedID = (*env)->GetFieldID(env, cls, "finished", "Z");
consumedID = (*env)->GetFieldID(env, cls, "consumed", "I");
}
JNIEXPORT jlong JNICALL
Java_com_velocitypowered_natives_compression_NativeZlibInflate_init(JNIEnv *env,
jobject obj)
{
z_stream* stream = calloc(1, sizeof(z_stream));
if (stream == 0) {
// Out of memory!
throwException(env, "java/lang/OutOfMemoryError", "zlib allocate stream");
return 0;
}
char *msg;
int ret = inflateInit(stream);
switch (ret) {
case Z_OK:
return (jlong) stream;
case Z_MEM_ERROR:
free(stream);
throwException(env, "java/lang/OutOfMemoryError", "zlib init");
return 0;
case Z_STREAM_ERROR:
free(stream);
throwException(env, "java/lang/IllegalArgumentException", "stream clobbered?");
return 0;
default:
msg = stream->msg;
free(stream);
throwException(env, "java/util/zip/DataFormatException", msg);
return 0;
}
}
JNIEXPORT void JNICALL
Java_com_velocitypowered_natives_compression_NativeZlibInflate_free(JNIEnv *env,
jobject obj,
jlong ctx)
{
z_stream* stream = (z_stream*) ctx;
int ret = inflateEnd(stream);
char *msg = stream->msg;
free((void*) ctx);
switch (ret) {
case Z_OK:
break;
case Z_STREAM_ERROR:
if (msg == NULL) {
msg = "stream state inconsistent";
}
case Z_DATA_ERROR:
if (msg == NULL) {
msg = "data was discarded";
}
throwException(env, "java/lang/IllegalArgumentException", msg);
break;
}
}
JNIEXPORT int JNICALL
Java_com_velocitypowered_natives_compression_NativeZlibInflate_process(JNIEnv *env,
jobject obj,
jlong ctx,
jlong sourceAddress,
jint sourceLength,
jlong destinationAddress,
jint destinationLength)
{
z_stream* stream = (z_stream*) ctx;
stream->next_in = (Bytef *) sourceAddress;
stream->next_out = (Bytef *) destinationAddress;
stream->avail_in = sourceLength;
stream->avail_out = destinationLength;
int res = inflate(stream, Z_PARTIAL_FLUSH);
switch (res) {
case Z_STREAM_END:
// The stream has ended
(*env)->SetBooleanField(env, obj, finishedID, JNI_TRUE);
// fall-through
case Z_OK:
// Not yet completed, but progress has been made. Tell Java how many bytes we've processed.
(*env)->SetIntField(env, obj, consumedID, sourceLength - stream->avail_in);
return destinationLength - stream->avail_out;
case Z_BUF_ERROR:
// This is not fatal. Just say we need more data. Usually this applies to the next_out buffer,
// which NativeVelocityCompressor will notice and will expand the buffer.
return 0;
default:
throwException(env, "java/util/zip/DataFormatException", stream->msg);
return 0;
}
}
JNIEXPORT void JNICALL
Java_com_velocitypowered_natives_compression_NativeZlibInflate_reset(JNIEnv *env,
jobject obj,
jlong ctx)
{
z_stream* stream = (z_stream*) ctx;
int ret = inflateReset(stream);
assert(ret == Z_OK);
}

View File

@ -16,7 +16,7 @@ public class JavaVelocityCompressor implements VelocityCompressor {
public JavaVelocityCompressor() {
this.deflater = new Deflater();
this.inflater = new Inflater();
this.buf = new byte[8192];
this.buf = new byte[ZLIB_BUFFER_SIZE];
}
@Override

View File

@ -0,0 +1,67 @@
package com.velocitypowered.natives.compression;
import io.netty.buffer.ByteBuf;
import java.util.zip.DataFormatException;
public class NativeVelocityCompressor implements VelocityCompressor {
private final NativeZlibInflate inflate = new NativeZlibInflate();
private final long inflateCtx;
private final NativeZlibDeflate deflate = new NativeZlibDeflate();
private final long deflateCtx;
private boolean disposed = false;
public NativeVelocityCompressor() {
this.inflateCtx = inflate.init();
this.deflateCtx = deflate.init(7);
}
@Override
public void inflate(ByteBuf source, ByteBuf destination) throws DataFormatException {
source.memoryAddress();
destination.memoryAddress();
while (!inflate.finished && source.isReadable()) {
if (!destination.isWritable()) {
destination.ensureWritable(ZLIB_BUFFER_SIZE);
}
int produced = inflate.process(inflateCtx, source.memoryAddress() + source.readerIndex(), source.readableBytes(),
destination.memoryAddress() + destination.writerIndex(), destination.writableBytes());
source.readerIndex(source.readerIndex() + inflate.consumed);
destination.writerIndex(destination.writerIndex() + produced);
}
inflate.reset(inflateCtx);
inflate.consumed = 0;
inflate.finished = false;
}
@Override
public void deflate(ByteBuf source, ByteBuf destination) throws DataFormatException {
source.memoryAddress();
destination.memoryAddress();
while (!deflate.finished) {
if (!destination.isWritable()) {
destination.ensureWritable(ZLIB_BUFFER_SIZE);
}
int produced = deflate.process(deflateCtx, source.memoryAddress() + source.readerIndex(), source.readableBytes(),
destination.memoryAddress() + destination.writerIndex(), destination.writableBytes(), !source.isReadable());
source.readerIndex(source.readerIndex() + deflate.consumed);
destination.writerIndex(destination.writerIndex() + produced);
}
deflate.reset(deflateCtx);
deflate.consumed = 0;
deflate.finished = false;
}
@Override
public void dispose() {
if (!disposed) {
inflate.free(inflateCtx);
deflate.free(deflateCtx);
}
disposed = true;
}
}

View File

@ -0,0 +1,24 @@
package com.velocitypowered.natives.compression;
/**
* Represents a native interface for zlib's deflate functions.
*/
class NativeZlibDeflate {
boolean finished;
int consumed;
native long init(int level);
native long free(long ctx);
native int process(long ctx, long sourceAddress, int sourceLength, long destinationAddress, int destinationLength,
boolean flush);
native void reset(long ctx);
static {
initIDs();
}
static native void initIDs();
}

View File

@ -0,0 +1,23 @@
package com.velocitypowered.natives.compression;
/**
* Represents a native interface for zlib's inflate functions.
*/
public class NativeZlibInflate {
boolean finished;
int consumed;
native long init();
native long free(long ctx);
native int process(long ctx, long sourceAddress, int sourceLength, long destinationAddress, int destinationLength);
native void reset(long ctx);
static {
initIDs();
}
static native void initIDs();
}

View File

@ -6,6 +6,8 @@ import io.netty.buffer.ByteBuf;
import java.util.zip.DataFormatException;
public interface VelocityCompressor extends Disposable {
int ZLIB_BUFFER_SIZE = 8192;
void inflate(ByteBuf source, ByteBuf destination) throws DataFormatException;
void deflate(ByteBuf source, ByteBuf destination) throws DataFormatException;

View File

@ -0,0 +1,92 @@
package com.velocitypowered.natives.util;
import com.google.common.collect.ImmutableList;
import java.util.List;
import java.util.function.BooleanSupplier;
import java.util.function.Supplier;
public class NativeCodeLoader<T> {
private final List<Variant<T>> variants;
private Variant<T> selected;
public NativeCodeLoader(List<Variant<T>> variants) {
this.variants = ImmutableList.copyOf(variants);
}
public Supplier<T> supply() {
if (selected == null) {
selected = select();
}
return selected.supplier;
}
private Variant<T> select() {
for (Variant<T> variant : variants) {
T got = variant.get();
if (got == null) {
continue;
}
return variant;
}
throw new IllegalArgumentException("Can't find any suitable variants");
}
public String getLoadedVariant() {
for (Variant<T> variant : variants) {
T got = variant.get();
if (got == null) {
continue;
}
return variant.name;
}
throw new IllegalArgumentException("Can't find any suitable variants");
}
static class Variant<T> {
private boolean available;
private final Runnable setup;
private final String name;
private final Supplier<T> supplier;
private boolean hasBeenSetup = false;
Variant(BooleanSupplier available, Runnable setup, String name, Supplier<T> supplier) {
this.available = available.getAsBoolean();
this.setup = setup;
this.name = name;
this.supplier = supplier;
}
public boolean setup() {
if (available && !hasBeenSetup) {
try {
setup.run();
hasBeenSetup = true;
} catch (Exception e) {
//logger.error("Unable to set up {}", name, e);
available = false;
}
}
return hasBeenSetup;
}
public T get() {
if (!hasBeenSetup) {
setup();
}
if (available) {
return supplier.get();
}
return null;
}
}
public static final BooleanSupplier MACOS = () -> System.getProperty("os.name").equalsIgnoreCase("Mac OS X") &&
System.getProperty("os.arch").equals("x86_64");
public static final BooleanSupplier LINUX = () -> System.getProperties().getProperty("os.name").equalsIgnoreCase("Linux") &&
System.getProperty("os.arch").equals("amd64");
public static final BooleanSupplier MAC_AND_LINUX = () -> MACOS.getAsBoolean() || LINUX.getAsBoolean();
public static final BooleanSupplier ALWAYS = () -> true;
}

View File

@ -0,0 +1,48 @@
package com.velocitypowered.natives.util;
import com.google.common.collect.ImmutableList;
import com.velocitypowered.natives.compression.JavaVelocityCompressor;
import com.velocitypowered.natives.compression.NativeVelocityCompressor;
import com.velocitypowered.natives.compression.VelocityCompressor;
import java.io.IOException;
import java.nio.file.Files;
import java.nio.file.Path;
import java.nio.file.StandardCopyOption;
public class Natives {
private Natives() {
throw new AssertionError();
}
private static Runnable copyAndLoadNative(String path) {
return () -> {
try {
Path tempFile = Files.createTempFile("native-", path.substring(path.lastIndexOf('.')));
Files.copy(Natives.class.getResourceAsStream(path), tempFile, StandardCopyOption.REPLACE_EXISTING);
Runtime.getRuntime().addShutdownHook(new Thread(() -> {
try {
Files.deleteIfExists(tempFile);
} catch (IOException ignored) {
// Well, it doesn't matter...
}
}));
System.load(tempFile.toAbsolutePath().toString());
} catch (IOException e) {
throw new RuntimeException(e);
}
};
}
public static final NativeCodeLoader<VelocityCompressor> compressor = new NativeCodeLoader<>(
ImmutableList.of(
new NativeCodeLoader.Variant<>(NativeCodeLoader.MACOS,
copyAndLoadNative("/macosx/velocity-compress.dylib"), "native compression (macOS)",
NativeVelocityCompressor::new),
new NativeCodeLoader.Variant<>(NativeCodeLoader.LINUX,
copyAndLoadNative("/linux_x64/velocity-compress.so"), "native compression (Linux amd64)",
NativeVelocityCompressor::new),
new NativeCodeLoader.Variant<>(NativeCodeLoader.ALWAYS, () -> {}, "Java compression", JavaVelocityCompressor::new)
)
);
}

Binary file not shown.

Binary file not shown.

View File

@ -0,0 +1,56 @@
package com.velocitypowered.natives.compression;
import com.velocitypowered.natives.util.Natives;
import io.netty.buffer.ByteBuf;
import io.netty.buffer.ByteBufUtil;
import io.netty.buffer.Unpooled;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.condition.EnabledOnOs;
import java.util.Random;
import java.util.zip.DataFormatException;
import static org.junit.jupiter.api.Assertions.assertTrue;
import static org.junit.jupiter.api.Assertions.fail;
import static org.junit.jupiter.api.condition.OS.LINUX;
import static org.junit.jupiter.api.condition.OS.MAC;
class VelocityCompressorTest {
@Test
@EnabledOnOs({ MAC, LINUX })
void nativeIntegrityCheck() throws DataFormatException {
VelocityCompressor compressor = Natives.compressor.supply().get();
if (compressor instanceof JavaVelocityCompressor) {
fail("Loaded regular compressor");
}
check(compressor);
}
@Test
void javaIntegrityCheck() throws DataFormatException {
JavaVelocityCompressor compressor = new JavaVelocityCompressor();
check(compressor);
}
private void check(VelocityCompressor compressor) throws DataFormatException {
ByteBuf source = Unpooled.directBuffer();
ByteBuf dest = Unpooled.directBuffer();
ByteBuf decompressed = Unpooled.directBuffer();
Random random = new Random(1);
byte[] randomBytes = new byte[1 << 21];
random.nextBytes(randomBytes);
source.writeBytes(randomBytes);
try {
compressor.deflate(source, dest);
compressor.inflate(dest, decompressed);
source.readerIndex(0);
assertTrue(ByteBufUtil.equals(source, decompressed));
} finally {
source.release();
dest.release();
compressor.dispose();
}
}
}

View File

@ -2,6 +2,7 @@ package com.velocitypowered.proxy;
import com.google.gson.Gson;
import com.google.gson.GsonBuilder;
import com.velocitypowered.natives.util.Natives;
import com.velocitypowered.network.ConnectionManager;
import com.velocitypowered.proxy.config.VelocityConfiguration;
import com.velocitypowered.proxy.connection.http.NettyHttpClient;
@ -52,6 +53,8 @@ public class VelocityServer {
}
public void start() {
logger.info("Using {}", Natives.compressor.getLoadedVariant());
// Create a key pair
logger.info("Booting up Velocity...");
try {

View File

@ -1,6 +1,8 @@
package com.velocitypowered.proxy.connection;
import com.google.common.base.Preconditions;
import com.velocitypowered.natives.compression.VelocityCompressor;
import com.velocitypowered.natives.util.Natives;
import com.velocitypowered.proxy.protocol.PacketWrapper;
import com.velocitypowered.proxy.protocol.StateRegistry;
import com.velocitypowered.natives.compression.JavaVelocityCompressor;
@ -189,7 +191,7 @@ public class MinecraftConnection extends ChannelInboundHandlerAdapter {
return;
}
JavaVelocityCompressor compressor = new JavaVelocityCompressor();
VelocityCompressor compressor = Natives.compressor.supply().get();
MinecraftCompressEncoder encoder = new MinecraftCompressEncoder(threshold, compressor);
MinecraftCompressDecoder decoder = new MinecraftCompressDecoder(threshold, compressor);

View File

@ -31,6 +31,7 @@ public class MinecraftCompressDecoder extends MessageToMessageDecoder<ByteBuf> {
}
Preconditions.checkState(uncompressedSize >= threshold, "Uncompressed size %s doesn't make sense with threshold %s", uncompressedSize, threshold);
// Try to use the uncompressed size, but place a cap if it might be too big (possibly malicious).
ByteBuf uncompressed = ctx.alloc().buffer(Math.min(uncompressedSize, MAXIMUM_INITIAL_BUFFER_SIZE));
try {
compressor.inflate(msg, uncompressed);

View File

@ -25,7 +25,8 @@ public class MinecraftCompressEncoder extends MessageToByteEncoder<ByteBuf> {
return;
}
ByteBuf compressedBuffer = ctx.alloc().buffer();
// in other words, see if a plain 8KiB buffer fits us well
ByteBuf compressedBuffer = ctx.alloc().buffer(8192);
try {
int uncompressed = msg.readableBytes();
compressor.deflate(msg, compressedBuffer);