io.undertow.conduits.DeflatingStreamSinkConduit Maven / Gradle / Ivy
/*
* JBoss, Home of Professional Open Source.
* Copyright 2014 Red Hat, Inc., and individual contributors
* as indicated by the @author tags.
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package io.undertow.conduits;
import static org.xnio.Bits.allAreClear;
import static org.xnio.Bits.allAreSet;
import static org.xnio.Bits.anyAreSet;
import java.io.IOException;
import java.nio.ByteBuffer;
import java.nio.channels.ClosedChannelException;
import java.nio.channels.FileChannel;
import java.util.concurrent.TimeUnit;
import java.util.zip.Deflater;
import io.undertow.server.Connectors;
import org.xnio.Buffers;
import org.xnio.IoUtils;
import io.undertow.connector.PooledByteBuffer;
import org.xnio.XnioIoThread;
import org.xnio.XnioWorker;
import org.xnio.channels.StreamSourceChannel;
import org.xnio.conduits.ConduitWritableByteChannel;
import org.xnio.conduits.Conduits;
import org.xnio.conduits.StreamSinkConduit;
import org.xnio.conduits.WriteReadyHandler;
import io.undertow.UndertowLogger;
import io.undertow.server.HttpServerExchange;
import io.undertow.util.ConduitFactory;
import io.undertow.util.NewInstanceObjectPool;
import io.undertow.util.ObjectPool;
import io.undertow.util.Headers;
import io.undertow.util.PooledObject;
import io.undertow.util.SimpleObjectPool;
/**
* Channel that handles deflate compression
*
* @author Stuart Douglas
*/
public class DeflatingStreamSinkConduit implements StreamSinkConduit {
private static final ByteBuffer EMPTY = ByteBuffer.allocate(0);
protected volatile Deflater deflater;
protected final PooledObject pooledObject;
private final ConduitFactory conduitFactory;
private final HttpServerExchange exchange;
private StreamSinkConduit next;
private WriteReadyHandler writeReadyHandler;
/**
* The streams buffer. This is freed when the next is shutdown
*/
protected PooledByteBuffer currentBuffer;
/**
* Buffer used to write the trailer if currentBuffer doesn't have enough space.
*/
private ByteBuffer trailerBuffer;
private int state = 0;
private static final int SHUTDOWN = 1;
private static final int NEXT_SHUTDOWN = 1 << 1;
private static final int FLUSHING_BUFFER = 1 << 2;
private static final int WRITES_RESUMED = 1 << 3;
private static final int CLOSED = 1 << 4;
private static final int WRITTEN_TRAILER = 1 << 5;
public DeflatingStreamSinkConduit(final ConduitFactory conduitFactory, final HttpServerExchange exchange) {
this(conduitFactory, exchange, Deflater.DEFLATED);
}
public DeflatingStreamSinkConduit(final ConduitFactory conduitFactory, final HttpServerExchange exchange, int deflateLevel) {
this(conduitFactory, exchange, newInstanceDeflaterPool(deflateLevel));
}
public DeflatingStreamSinkConduit(final ConduitFactory conduitFactory, final HttpServerExchange exchange, ObjectPool deflaterPool) {
this.pooledObject = deflaterPool.allocate();
this.deflater = pooledObject.getObject();
this.currentBuffer = exchange.getConnection().getByteBufferPool().allocate();
this.exchange = exchange;
this.conduitFactory = conduitFactory;
setWriteReadyHandler(new WriteReadyHandler.ChannelListenerHandler<>(Connectors.getConduitSinkChannel(exchange)));
}
public static ObjectPool newInstanceDeflaterPool(int deflateLevel) {
return new NewInstanceObjectPool<>(() -> new Deflater(deflateLevel, true), Deflater::end);
}
public static ObjectPool simpleDeflaterPool(int poolSize, int deflateLevel) {
return new SimpleObjectPool<>(poolSize, () -> new Deflater(deflateLevel, true), Deflater::reset, Deflater::end);
}
@Override
public int write(final ByteBuffer src) throws IOException {
if (anyAreSet(state, SHUTDOWN | CLOSED) || currentBuffer == null) {
throw new ClosedChannelException();
}
try {
if (!performFlushIfRequired()) {
return 0;
}
if (!src.hasRemaining()) {
return 0;
}
int initialSrcPosition = src.position();
int initialRemaining = src.remaining();
deflater.setInput(src);
deflateData(false);
int consumed = initialRemaining - src.remaining();
Connectors.updateResponseBytesSent(exchange, -consumed);
int endSrcPosition = src.position();
int srcLimit = src.limit();
// Reset the buffer to original values with a limit based on what has
// been deflated such that only data that has been compressed is
// represented by the buffer.
src.position(initialSrcPosition);
src.limit(endSrcPosition);
postDeflate(src);
src.limit(srcLimit);
src.position(endSrcPosition);
// Ensure input buffers are not held or clobbered outside expected usage
deflater.setInput(EMPTY);
return consumed;
} catch (IOException | RuntimeException | Error e) {
freeBuffer();
throw e;
}
}
protected void postDeflate(ByteBuffer data) {
}
@Override
public long write(final ByteBuffer[] srcs, final int offset, final int length) throws IOException {
if (anyAreSet(state, SHUTDOWN | CLOSED) || currentBuffer == null) {
throw new ClosedChannelException();
}
try {
int total = 0;
for (int i = offset; i < offset + length; ++i) {
ByteBuffer buf = srcs[i];
if (buf.hasRemaining()) {
int ret = write(buf);
total += ret;
// Must cease iteration after any buffer is not
// fully exhausted, or after bytes cannot be
// written.
if (ret == 0 || buf.hasRemaining()) {
return total;
}
}
}
return total;
} catch (IOException | RuntimeException | Error e) {
freeBuffer();
throw e;
}
}
@Override
public int writeFinal(ByteBuffer src) throws IOException {
return Conduits.writeFinalBasic(this, src);
}
@Override
public long writeFinal(ByteBuffer[] srcs, int offset, int length) throws IOException {
return Conduits.writeFinalBasic(this, srcs, offset, length);
}
@Override
public long transferFrom(final FileChannel src, final long position, final long count) throws IOException {
if (anyAreSet(state, SHUTDOWN | CLOSED)) {
throw new ClosedChannelException();
}
if (!performFlushIfRequired()) {
return 0;
}
return src.transferTo(position, count, new ConduitWritableByteChannel(this));
}
@Override
public long transferFrom(final StreamSourceChannel source, final long count, final ByteBuffer throughBuffer) throws IOException {
if (anyAreSet(state, SHUTDOWN | CLOSED)) {
throw new ClosedChannelException();
}
if (!performFlushIfRequired()) {
return 0;
}
return IoUtils.transfer(source, count, throughBuffer, new ConduitWritableByteChannel(this));
}
@Override
public XnioWorker getWorker() {
return exchange.getConnection().getWorker();
}
@Override
public void suspendWrites() {
if (next == null) {
state = state & ~WRITES_RESUMED;
} else {
next.suspendWrites();
}
}
@Override
public boolean isWriteResumed() {
if (next == null) {
return anyAreSet(state, WRITES_RESUMED);
} else {
return next.isWriteResumed();
}
}
@Override
public void wakeupWrites() {
if (next == null) {
resumeWrites();
} else {
next.wakeupWrites();
}
}
@Override
public void resumeWrites() {
if (next == null) {
state |= WRITES_RESUMED;
queueWriteListener();
} else {
next.resumeWrites();
}
}
private void queueWriteListener() {
exchange.getConnection().getIoThread().execute(new Runnable() {
@Override
public void run() {
if (writeReadyHandler != null) {
try {
writeReadyHandler.writeReady();
} finally {
//if writes are still resumed queue up another one
if (next == null && isWriteResumed()) {
queueWriteListener();
}
}
}
}
});
}
@Override
public void terminateWrites() throws IOException {
if (deflater != null) {
deflater.finish();
}
state |= SHUTDOWN;
}
@Override
public boolean isWriteShutdown() {
return anyAreSet(state, SHUTDOWN);
}
@Override
public void awaitWritable() throws IOException {
if (next == null) {
return;
} else {
next.awaitWritable();
}
}
@Override
public void awaitWritable(final long time, final TimeUnit timeUnit) throws IOException {
if (next == null) {
return;
} else {
next.awaitWritable(time, timeUnit);
}
}
@Override
public XnioIoThread getWriteThread() {
return exchange.getConnection().getIoThread();
}
@Override
public void setWriteReadyHandler(final WriteReadyHandler handler) {
this.writeReadyHandler = handler;
}
@Override
public boolean flush() throws IOException {
if (currentBuffer == null) {
if (anyAreSet(state, NEXT_SHUTDOWN)) {
return next.flush();
} else {
return true;
}
}
try {
boolean nextCreated = false;
try {
if (anyAreSet(state, SHUTDOWN)) {
if (anyAreSet(state, NEXT_SHUTDOWN)) {
return next.flush();
} else {
if (!performFlushIfRequired()) {
return false;
}
//if the deflater has not been fully flushed we need to flush it
if (!deflater.finished()) {
deflateData(false);
//if could not fully flush
if (!deflater.finished()) {
return false;
}
}
final ByteBuffer buffer = currentBuffer.getBuffer();
if (allAreClear(state, WRITTEN_TRAILER)) {
state |= WRITTEN_TRAILER;
byte[] data = getTrailer();
if (data != null) {
Connectors.updateResponseBytesSent(exchange, data.length);
if(trailerBuffer != null) {
throw new IllegalStateException("trailerBuffer is already set");
} else if(anyAreSet(state, FLUSHING_BUFFER) && buffer.capacity() - buffer.remaining() >= data.length) {
buffer.compact();
buffer.put(data);
buffer.flip();
} else if (data.length <= buffer.remaining() && !anyAreSet(state, FLUSHING_BUFFER)) {
buffer.put(data);
} else {
trailerBuffer = ByteBuffer.wrap(data);
}
}
}
//ok the deflater is flushed, now we need to flush the buffer
if (!anyAreSet(state, FLUSHING_BUFFER)) {
buffer.flip();
state |= FLUSHING_BUFFER;
if (next == null) {
nextCreated = true;
this.next = createNextChannel();
}
}
if (performFlushIfRequired()) {
state |= NEXT_SHUTDOWN;
freeBuffer();
next.terminateWrites();
return next.flush();
} else {
return false;
}
}
} else {
if(allAreClear(state, FLUSHING_BUFFER)) {
if (next == null) {
nextCreated = true;
this.next = createNextChannel();
}
deflateData(true);
if(allAreClear(state, FLUSHING_BUFFER)) {
//deflateData can cause this to be change
currentBuffer.getBuffer().flip();
this.state |= FLUSHING_BUFFER;
}
}
if(!performFlushIfRequired()) {
return false;
}
return next.flush();
}
} finally {
if (nextCreated) {
if (anyAreSet(state, WRITES_RESUMED) && !anyAreSet(state ,NEXT_SHUTDOWN)) {
try {
next.resumeWrites();
} catch (Throwable e) {
UndertowLogger.REQUEST_LOGGER.debug("Failed to resume", e);
}
}
}
}
} catch (IOException | RuntimeException | Error e) {
freeBuffer();
throw e;
}
}
/**
* called before the stream is finally flushed.
*/
protected byte[] getTrailer() {
return null;
}
/**
* The we are in the flushing state then we flush to the underlying stream, otherwise just return true
*
* @return false if there is still more to flush
*/
private boolean performFlushIfRequired() throws IOException {
if (anyAreSet(state, FLUSHING_BUFFER)) {
return trailerBuffer == null
? performFlushIfRequiredSingleBuffer()
: performFlushIfRequiredAdditionalBuffer();
}
return true;
}
private boolean performFlushIfRequiredSingleBuffer() throws IOException {
final ByteBuffer buf = currentBuffer.getBuffer();
long totalLength = buf.remaining();
if (totalLength > 0) {
int total = 0;
int res = 0;
do {
res = next.write(buf);
total += res;
if (res == 0) {
return false;
}
} while (total < totalLength);
}
currentBuffer.getBuffer().clear();
state = state & ~FLUSHING_BUFFER;
return true;
}
private boolean performFlushIfRequiredAdditionalBuffer() throws IOException {
final ByteBuffer[] bufs = new ByteBuffer[] {
currentBuffer.getBuffer(),
trailerBuffer};
long totalLength = Buffers.remaining(bufs);
if (totalLength > 0) {
long total = 0;
long res = 0;
do {
res = next.write(bufs, 0, bufs.length);
total += res;
if (res == 0) {
return false;
}
} while (total < totalLength);
}
trailerBuffer = null;
currentBuffer.getBuffer().clear();
state = state & ~FLUSHING_BUFFER;
return true;
}
private StreamSinkConduit createNextChannel() {
if (deflater.finished() && allAreSet(state, WRITTEN_TRAILER)) {
//the deflater was fully flushed before we created the channel. This means that what is in the buffer is
//all there is
int remaining = currentBuffer.getBuffer().remaining();
if (trailerBuffer != null) {
remaining += trailerBuffer.remaining();
}
if(!exchange.getResponseHeaders().contains(Headers.TRANSFER_ENCODING)) {
exchange.getResponseHeaders().put(Headers.CONTENT_LENGTH, Integer.toString(remaining));
}
} else {
exchange.getResponseHeaders().remove(Headers.CONTENT_LENGTH);
}
return conduitFactory.create();
}
/**
* Runs the current data through the deflater. As much as possible this will be buffered in the current output
* stream.
*
* @throws IOException
*/
private void deflateData(boolean force) throws IOException {
//we don't need to flush here, as this should have been called already by the time we get to
//this point
boolean nextCreated = false;
try {
PooledByteBuffer pooled = this.currentBuffer;
final ByteBuffer outputBuffer = pooled.getBuffer();
final boolean shutdown = anyAreSet(state, SHUTDOWN);
while (force || !deflater.needsInput() || (shutdown && !deflater.finished())) {
int count = deflater.deflate(outputBuffer, force ? Deflater.SYNC_FLUSH : Deflater.NO_FLUSH);
if (count != 0) {
Connectors.updateResponseBytesSent(exchange, count);
if (!outputBuffer.hasRemaining()) {
outputBuffer.flip();
this.state |= FLUSHING_BUFFER;
if (next == null) {
nextCreated = true;
this.next = createNextChannel();
}
if (!performFlushIfRequired()) {
return;
}
}
} else {
force = false;
}
}
} finally {
if (nextCreated) {
if (anyAreSet(state, WRITES_RESUMED)) {
next.resumeWrites();
}
}
}
}
@Override
public void truncateWrites() throws IOException {
freeBuffer();
state |= CLOSED;
next.truncateWrites();
}
private void freeBuffer() {
if (currentBuffer != null) {
currentBuffer.close();
currentBuffer = null;
state = state & ~FLUSHING_BUFFER;
}
if (deflater != null) {
deflater = null;
pooledObject.close();
}
trailerBuffer = null;
}
}
© 2015 - 2024 Weber Informatics LLC | Privacy Policy