@Test
public void testWriteEmptyFile() throws Exception {
AbstractStreamWriteFilter<M> filter = createFilter();
M message = createMessage(new byte[0]);
WriteRequest writeRequest = new DefaultWriteRequest(message, new DummyWriteFuture());
NextFilter nextFilter = EasyMock.createMock(NextFilter.class);
/*
* Record expectations
*/
nextFilter.messageSent(session, writeRequest);
/*
* Replay.
*/
EasyMock.replay(nextFilter);
filter.filterWrite(nextFilter, session, writeRequest);
/*
* Verify.
*/
EasyMock.verify(nextFilter);
assertTrue(writeRequest.getFuture().isWritten());
}
/**
* Tests when the contents of the file fits into one write buffer.
*
* @throws Exception when something goes wrong
*/
@Test
public void testWriteSingleBufferFile() throws Exception {
byte[] data = new byte[] {1, 2, 3, 4};
AbstractStreamWriteFilter<M> filter = createFilter();
M message = createMessage(data);
WriteRequest writeRequest = new DefaultWriteRequest(message, new DummyWriteFuture());
NextFilter nextFilter = EasyMock.createMock(NextFilter.class);
/*
* Record expectations
*/
nextFilter.filterWrite(
EasyMock.eq(session), eqWriteRequest(new DefaultWriteRequest(IoBuffer.wrap(data))));
nextFilter.messageSent(session, writeRequest);
/*
* Replay.
*/
EasyMock.replay(nextFilter);
filter.filterWrite(nextFilter, session, writeRequest);
filter.messageSent(nextFilter, session, writeRequest);
/*
* Verify.
*/
EasyMock.verify(nextFilter);
assertTrue(writeRequest.getFuture().isWritten());
}
public boolean matches(Object actual) {
if (actual instanceof WriteRequest) {
WriteRequest w2 = (WriteRequest) actual;
return expected.getMessage().equals(w2.getMessage())
&& expected.getFuture().isWritten() == w2.getFuture().isWritten();
}
return false;
}
public WriteFuture flush() {
Queue<Object> bufferQueue = getMessageQueue();
WriteFuture future = null;
for (; ; ) {
Object encodedMessage = bufferQueue.poll();
if (encodedMessage == null) {
break;
}
// Flush only when the buffer has remaining.
if (!(encodedMessage instanceof IoBuffer) || ((IoBuffer) encodedMessage).hasRemaining()) {
future = new DefaultWriteFuture(session);
nextFilter.filterWrite(
session,
new EncodedWriteRequest(encodedMessage, future, writeRequest.getDestination()));
}
}
if (future == null) {
future =
DefaultWriteFuture.newNotWrittenFuture(
session, new NothingWrittenException(writeRequest));
}
return future;
}
private int writeFile(
S session, WriteRequest req, boolean hasFragmentation, int maxLength, long currentTime)
throws Exception {
int localWrittenBytes;
FileRegion region = (FileRegion) req.getMessage();
if (region.getRemainingBytes() > 0) {
int length;
if (hasFragmentation) {
length = (int) Math.min(region.getRemainingBytes(), maxLength);
} else {
length = (int) Math.min(Integer.MAX_VALUE, region.getRemainingBytes());
}
localWrittenBytes = transferFile(session, region, length);
region.update(localWrittenBytes);
} else {
localWrittenBytes = 0;
}
session.increaseWrittenBytes(localWrittenBytes, currentTime);
if ((region.getRemainingBytes() <= 0) || (!hasFragmentation && (localWrittenBytes != 0))) {
fireMessageSent(session, req);
}
return localWrittenBytes;
}
private void clearWriteRequestQueue(S session) {
WriteRequestQueue writeRequestQueue = session.getWriteRequestQueue();
WriteRequest req;
List<WriteRequest> failedRequests = new ArrayList<WriteRequest>();
if ((req = writeRequestQueue.poll(session)) != null) {
Object message = req.getMessage();
if (message instanceof IoBuffer) {
IoBuffer buf = (IoBuffer) message;
// The first unwritten empty buffer must be
// forwarded to the filter chain.
if (buf.hasRemaining()) {
buf.reset();
failedRequests.add(req);
} else {
IoFilterChain filterChain = session.getFilterChain();
filterChain.fireMessageSent(req);
}
} else {
failedRequests.add(req);
}
// Discard others.
while ((req = writeRequestQueue.poll(session)) != null) {
failedRequests.add(req);
}
}
// Create an exception and notify.
if (!failedRequests.isEmpty()) {
WriteToClosedSessionException cause = new WriteToClosedSessionException(failedRequests);
for (WriteRequest r : failedRequests) {
session.decreaseScheduledBytesAndMessages(r);
r.getFuture().setException(cause);
}
IoFilterChain filterChain = session.getFilterChain();
filterChain.fireExceptionCaught(cause);
}
}
/**
* Tests when the contents of the file doesn't fit into one write buffer.
*
* @throws Exception when something goes wrong
*/
@Test
public void testWriteSeveralBuffersStream() throws Exception {
AbstractStreamWriteFilter<M> filter = createFilter();
filter.setWriteBufferSize(4);
byte[] data = new byte[] {1, 2, 3, 4, 5, 6, 7, 8, 9, 10};
byte[] chunk1 = new byte[] {1, 2, 3, 4};
byte[] chunk2 = new byte[] {5, 6, 7, 8};
byte[] chunk3 = new byte[] {9, 10};
M message = createMessage(data);
WriteRequest writeRequest = new DefaultWriteRequest(message, new DummyWriteFuture());
WriteRequest chunk1Request = new DefaultWriteRequest(IoBuffer.wrap(chunk1));
WriteRequest chunk2Request = new DefaultWriteRequest(IoBuffer.wrap(chunk2));
WriteRequest chunk3Request = new DefaultWriteRequest(IoBuffer.wrap(chunk3));
NextFilter nextFilter = EasyMock.createMock(NextFilter.class);
/*
* Record expectations
*/
nextFilter.filterWrite(EasyMock.eq(session), eqWriteRequest(chunk1Request));
nextFilter.filterWrite(EasyMock.eq(session), eqWriteRequest(chunk2Request));
nextFilter.filterWrite(EasyMock.eq(session), eqWriteRequest(chunk3Request));
nextFilter.messageSent(EasyMock.eq(session), eqWriteRequest(writeRequest));
/*
* Replay.
*/
EasyMock.replay(nextFilter);
filter.filterWrite(nextFilter, session, writeRequest);
filter.messageSent(nextFilter, session, chunk1Request);
filter.messageSent(nextFilter, session, chunk2Request);
filter.messageSent(nextFilter, session, chunk3Request);
/*
* Verify.
*/
EasyMock.verify(nextFilter);
assertTrue(writeRequest.getFuture().isWritten());
}
@Override
public void exceptionCaught(NextFilter nextFilter, IoSession session, Throwable cause)
throws Exception {
if (cause instanceof WriteToClosedSessionException) {
// Filter out SSL close notify, which is likely to fail to flush
// due to disconnection.
WriteToClosedSessionException e = (WriteToClosedSessionException) cause;
List<WriteRequest> failedRequests = e.getRequests();
boolean containsCloseNotify = false;
for (WriteRequest r : failedRequests) {
if (isCloseNotify(r.getMessage())) {
containsCloseNotify = true;
break;
}
}
if (containsCloseNotify) {
if (failedRequests.size() == 1) {
// close notify is the only failed request; bail out.
return;
}
List<WriteRequest> newFailedRequests = new ArrayList<>(failedRequests.size() - 1);
for (WriteRequest r : failedRequests) {
if (!isCloseNotify(r.getMessage())) {
newFailedRequests.add(r);
}
}
if (newFailedRequests.isEmpty()) {
// the failedRequests were full with close notify; bail out.
return;
}
cause =
new WriteToClosedSessionException(
newFailedRequests, cause.getMessage(), cause.getCause());
}
}
nextFilter.exceptionCaught(session, cause);
}
public void messageSent(
IoFilter.NextFilter nextFilter, IoSession session, WriteRequest writeRequest)
throws Exception {
Object message = writeRequest.getMessage();
if (this.SENT_LOGGER.isInfoEnabled() && (message instanceof OriginalBuffer)) {
append(SENT_LOGGER, Category.SENT, session, ((OriginalBuffer) message).getMessage());
}
super.messageSent(nextFilter, session, writeRequest);
}
@Override
public void messageSent(NextFilter nextFilter, final IoSession session, WriteRequest writeRequest)
throws Exception {
int maxWriteThroughput = this.maxWriteThroughput;
// process the request if our max is greater than zero
if (maxWriteThroughput > 0) {
final State state = (State) session.getAttribute(STATE);
long currentTime = System.currentTimeMillis();
long suspendTime = 0;
boolean firstWrite = false;
synchronized (state) {
state.writtenBytes += messageSizeEstimator.estimateSize(writeRequest.getMessage());
if (!state.suspendedWrite) {
if (state.writeStartTime == 0) {
firstWrite = true;
state.writeStartTime = currentTime - 1000;
}
long throughput = (state.writtenBytes * 1000 / (currentTime - state.writeStartTime));
if (throughput >= maxWriteThroughput) {
suspendTime =
Math.max(
0,
state.writtenBytes * 1000 / maxWriteThroughput
- (firstWrite ? 0 : currentTime - state.writeStartTime));
state.writtenBytes = 0;
state.writeStartTime = 0;
state.suspendedWrite = suspendTime != 0;
}
}
}
if (suspendTime != 0) {
log.trace("Suspending write");
session.suspendWrite();
scheduledExecutor.schedule(
new Runnable() {
public void run() {
synchronized (state) {
state.suspendedWrite = false;
}
session.resumeWrite();
log.trace("Resuming write");
}
},
suspendTime,
TimeUnit.MILLISECONDS);
}
}
nextFilter.messageSent(session, writeRequest);
}
public void wrap(NextFilter nextFilter, WriteRequest writeRequest, IoBuffer buf)
throws AuthException {
int start = buf.position();
int len = buf.remaining() - LINE_TERMINATOR.length;
if (len == 0) throw new AuthException("Decryption failed");
// HMAC(Ki, {SeqNum, msg})[0..9]
byte[] originalMessage = new byte[len];
buf.get(originalMessage);
byte[] mac = AuthDigestMD5IoFilter.computeMACBlock(session, originalMessage, false);
// Calculate padding
int bs = encCipher.getBlockSize();
byte[] padding;
if (bs > 1) {
int pad = bs - ((len + 10) % bs); // add 10 for HMAC[0..9]
padding = new byte[pad];
for (int i = 0; i < pad; i++) padding[i] = (byte) pad;
} else padding = EMPTY_BYTE_ARRAY;
byte[] toBeEncrypted = new byte[len + padding.length + 10];
// {msg, pad, HMAC(Ki, {SeqNum, msg}[0..9])}
System.arraycopy(originalMessage, start, toBeEncrypted, 0, len);
System.arraycopy(padding, 0, toBeEncrypted, len, padding.length);
System.arraycopy(mac, 0, toBeEncrypted, len + padding.length, 10);
// CIPHER(Kc, {msg, pad, HMAC(Ki, {SeqNum, msg}[0..9])})
byte[] cipherBlock;
try {
// Do CBC (chaining) across packets
cipherBlock = encCipher.update(toBeEncrypted);
// update() can return null
if (cipherBlock == null) throw new IllegalBlockSizeException("" + toBeEncrypted.length);
} catch (IllegalBlockSizeException e) {
throw new AuthException("Invalid block size for cipher", e);
}
IoBuffer out = IoBuffer.allocate(cipherBlock.length + 2 + 4 + LINE_TERMINATOR.length);
out.put(cipherBlock);
out.put(mac, 10, 6); // messageType & sequenceNum
out.put(LINE_TERMINATOR);
out.flip();
if (out.limit()
> ((Integer) session.getAttribute(AuthDigestMD5Command.CLIENT_MAXBUF)).intValue())
throw new AuthException("Data exceeds client maxbuf capability");
nextFilter.filterWrite(session, new DefaultWriteRequest(out, writeRequest.getFuture()));
}
public void flushWithoutFuture() {
Queue<Object> bufferQueue = getMessageQueue();
for (; ; ) {
Object encodedMessage = bufferQueue.poll();
if (encodedMessage == null) {
break;
}
// Flush only when the buffer has remaining.
if (!(encodedMessage instanceof IoBuffer) || ((IoBuffer) encodedMessage).hasRemaining()) {
SocketAddress destination = writeRequest.getDestination();
WriteRequest writeRequest = new EncodedWriteRequest(encodedMessage, null, destination);
nextFilter.filterWrite(session, writeRequest);
}
}
}
@Override
public void filterWrite(NextFilter nextFilter, IoSession session, WriteRequest writeRequest)
throws SSLException {
boolean needsFlush = true;
SslHandler handler = getSslSessionHandler(session);
synchronized (handler) {
if (!isSslStarted(session)) {
handler.scheduleFilterWrite(nextFilter, writeRequest);
}
// Don't encrypt the data if encryption is disabled.
else if (session.containsAttribute(DISABLE_ENCRYPTION_ONCE)) {
// Remove the marker attribute because it is temporary.
session.removeAttribute(DISABLE_ENCRYPTION_ONCE);
handler.scheduleFilterWrite(nextFilter, writeRequest);
} else {
// Otherwise, encrypt the buffer.
IoBuffer buf = (IoBuffer) writeRequest.getMessage();
if (handler.isWritingEncryptedData()) {
// data already encrypted; simply return buffer
handler.scheduleFilterWrite(nextFilter, writeRequest);
} else if (handler.isHandshakeComplete()) {
// SSL encrypt
int pos = buf.position();
handler.encrypt(buf.buf());
buf.position(pos);
IoBuffer encryptedBuffer = handler.fetchOutNetBuffer();
handler.scheduleFilterWrite(
nextFilter, new EncryptedWriteRequest(writeRequest, encryptedBuffer));
} else {
if (session.isConnected()) {
// Handshake not complete yet.
handler.schedulePreHandshakeWriteRequest(nextFilter, writeRequest);
}
needsFlush = false;
}
}
}
if (needsFlush) {
handler.flushScheduledEvents();
}
}
@Override
public void filterWrite(NextFilter nextFilter, IoSession session, WriteRequest writeRequest)
throws Exception {
Object message = writeRequest.getMessage();
// Bypass the encoding if the message is contained in a IoBuffer,
// as it has already been encoded before
if (message instanceof IoBuffer || message instanceof FileRegion) {
nextFilter.filterWrite(session, writeRequest);
return;
}
// Get the encoder in the session
ProtocolEncoder encoder = getEncoder(session);
ProtocolEncoderOutput encoderOut = getEncoderOut(session, nextFilter, writeRequest);
try {
// Now we can try to encode the response
encoder.encode(session, message, encoderOut);
// Send it directly
((ProtocolEncoderOutputImpl) encoderOut).flushWithoutFuture();
// Call the next filter
nextFilter.filterWrite(session, new MessageWriteRequest(writeRequest));
} catch (Throwable t) {
ProtocolEncoderException pee;
// Generate the correct exception
if (t instanceof ProtocolEncoderException) {
pee = (ProtocolEncoderException) t;
} else {
pee = new ProtocolEncoderException(t);
}
throw pee;
}
}
private int writeBuffer(
S session, WriteRequest req, boolean hasFragmentation, int maxLength, long currentTime)
throws Exception {
IoBuffer buf = (IoBuffer) req.getMessage();
int localWrittenBytes = 0;
if (buf.hasRemaining()) {
int length;
if (hasFragmentation) {
length = Math.min(buf.remaining(), maxLength);
} else {
length = buf.remaining();
}
try {
localWrittenBytes = write(session, buf, length);
} catch (IOException ioe) {
// We have had an issue while trying to send data to the
// peer : let's close the session.
session.close(true);
}
}
session.increaseWrittenBytes(localWrittenBytes, currentTime);
if (!buf.hasRemaining() || (!hasFragmentation && (localWrittenBytes != 0))) {
// Buffer has been sent, clear the current request.
int pos = buf.position();
buf.reset();
fireMessageSent(session, req);
// And set it back to its position
buf.position(pos);
}
return localWrittenBytes;
}
/** {@inheritDoc} */
public void write(NioSession session, WriteRequest writeRequest) {
// We will try to write the message directly
long currentTime = System.currentTimeMillis();
final WriteRequestQueue writeRequestQueue = session.getWriteRequestQueue();
final int maxWrittenBytes =
session.getConfig().getMaxReadBufferSize()
+ (session.getConfig().getMaxReadBufferSize() >>> 1);
int writtenBytes = 0;
// Deal with the special case of a Message marker (no bytes in the request)
// We just have to return after having calle dthe messageSent event
IoBuffer buf = (IoBuffer) writeRequest.getMessage();
if (buf.remaining() == 0) {
// Clear and fire event
session.setCurrentWriteRequest(null);
buf.reset();
session.getFilterChain().fireMessageSent(writeRequest);
return;
}
// Now, write the data
try {
for (; ; ) {
if (writeRequest == null) {
writeRequest = writeRequestQueue.poll(session);
if (writeRequest == null) {
setInterestedInWrite(session, false);
break;
}
session.setCurrentWriteRequest(writeRequest);
}
buf = (IoBuffer) writeRequest.getMessage();
if (buf.remaining() == 0) {
// Clear and fire event
session.setCurrentWriteRequest(null);
buf.reset();
session.getFilterChain().fireMessageSent(writeRequest);
continue;
}
SocketAddress destination = writeRequest.getDestination();
if (destination == null) {
destination = session.getRemoteAddress();
}
int localWrittenBytes = send(session, buf, destination);
if ((localWrittenBytes == 0) || (writtenBytes >= maxWrittenBytes)) {
// Kernel buffer is full or wrote too much
setInterestedInWrite(session, true);
session.getWriteRequestQueue().offer(session, writeRequest);
scheduleFlush(session);
} else {
setInterestedInWrite(session, false);
// Clear and fire event
session.setCurrentWriteRequest(null);
writtenBytes += localWrittenBytes;
buf.reset();
session.getFilterChain().fireMessageSent(writeRequest);
break;
}
}
} catch (Exception e) {
session.getFilterChain().fireExceptionCaught(e);
} finally {
session.increaseWrittenBytes(writtenBytes, currentTime);
}
}
private boolean flush(NioSession session, long currentTime) throws Exception {
final WriteRequestQueue writeRequestQueue = session.getWriteRequestQueue();
final int maxWrittenBytes =
session.getConfig().getMaxReadBufferSize()
+ (session.getConfig().getMaxReadBufferSize() >>> 1);
int writtenBytes = 0;
try {
for (; ; ) {
WriteRequest req = session.getCurrentWriteRequest();
if (req == null) {
req = writeRequestQueue.poll(session);
if (req == null) {
setInterestedInWrite(session, false);
break;
}
session.setCurrentWriteRequest(req);
}
IoBuffer buf = (IoBuffer) req.getMessage();
if (buf.remaining() == 0) {
// Clear and fire event
session.setCurrentWriteRequest(null);
buf.reset();
session.getFilterChain().fireMessageSent(req);
continue;
}
SocketAddress destination = req.getDestination();
if (destination == null) {
destination = session.getRemoteAddress();
}
int localWrittenBytes = send(session, buf, destination);
if ((localWrittenBytes == 0) || (writtenBytes >= maxWrittenBytes)) {
// Kernel buffer is full or wrote too much
setInterestedInWrite(session, true);
return false;
} else {
setInterestedInWrite(session, false);
// Clear and fire event
session.setCurrentWriteRequest(null);
writtenBytes += localWrittenBytes;
buf.reset();
session.getFilterChain().fireMessageSent(req);
}
}
} finally {
session.increaseWrittenBytes(writtenBytes, currentTime);
}
return true;
}
public void appendTo(StringBuffer buffer) {
buffer
.append("Expected a WriteRequest with the message '")
.append(expected.getMessage())
.append("'");
}
@Override
public void messageSent(NextFilter nextFilter, IoSession session, WriteRequest writeRequest)
throws Exception {
log(messageSentLevel, "SENT: {}", writeRequest.getOriginalRequest().getMessage());
nextFilter.messageSent(session, writeRequest);
}
private boolean flushNow(S session, long currentTime) {
if (!session.isConnected()) {
scheduleRemove(session);
return false;
}
final boolean hasFragmentation = session.getTransportMetadata().hasFragmentation();
final WriteRequestQueue writeRequestQueue = session.getWriteRequestQueue();
// Set limitation for the number of written bytes for read-write
// fairness. I used maxReadBufferSize * 3 / 2, which yields best
// performance in my experience while not breaking fairness much.
final int maxWrittenBytes =
session.getConfig().getMaxReadBufferSize()
+ (session.getConfig().getMaxReadBufferSize() >>> 1);
int writtenBytes = 0;
WriteRequest req = null;
try {
// Clear OP_WRITE
setInterestedInWrite(session, false);
do {
// Check for pending writes.
req = session.getCurrentWriteRequest();
if (req == null) {
req = writeRequestQueue.poll(session);
if (req == null) {
break;
}
session.setCurrentWriteRequest(req);
}
int localWrittenBytes = 0;
Object message = req.getMessage();
if (message instanceof IoBuffer) {
localWrittenBytes =
writeBuffer(
session, req, hasFragmentation, maxWrittenBytes - writtenBytes, currentTime);
if ((localWrittenBytes > 0) && ((IoBuffer) message).hasRemaining()) {
// the buffer isn't empty, we re-interest it in writing
writtenBytes += localWrittenBytes;
setInterestedInWrite(session, true);
return false;
}
} else if (message instanceof FileRegion) {
localWrittenBytes =
writeFile(
session, req, hasFragmentation, maxWrittenBytes - writtenBytes, currentTime);
// Fix for Java bug on Linux
// http://bugs.sun.com/bugdatabase/view_bug.do?bug_id=5103988
// If there's still data to be written in the FileRegion,
// return 0 indicating that we need
// to pause until writing may resume.
if ((localWrittenBytes > 0) && (((FileRegion) message).getRemainingBytes() > 0)) {
writtenBytes += localWrittenBytes;
setInterestedInWrite(session, true);
return false;
}
} else {
throw new IllegalStateException(
"Don't know how to handle message of type '"
+ message.getClass().getName()
+ "'. Are you missing a protocol encoder?");
}
if (localWrittenBytes == 0) {
// Kernel buffer is full.
setInterestedInWrite(session, true);
return false;
}
writtenBytes += localWrittenBytes;
if (writtenBytes >= maxWrittenBytes) {
// Wrote too much
scheduleFlush(session);
return false;
}
} while (writtenBytes < maxWrittenBytes);
} catch (Exception e) {
if (req != null) {
req.getFuture().setException(e);
}
IoFilterChain filterChain = session.getFilterChain();
filterChain.fireExceptionCaught(e);
return false;
}
return true;
}
