Add Close method for closing OutputStream and PcapWriteHandler (#10638)

Motivation:
We should implement the Closeable method to properly close `OutputStream` and `PcapWriteHandler`. So whenever `handlerRemoved(ChannelHandlerContext)` is called or the user wants to stop the Pcap writes into `OutputStream`, we have a proper method to close it otherwise writing data to close `OutputStream` will result in `IOException`.

Modification:
Implemented `Closeable` in `PcapWriteHandler` which calls `PcapWriter#close` and closes `OutputStream` and stops Pcap writes.

Result:
Better handling of Pcap writes.
This commit is contained in:
Aayush Atharva 2020-10-08 15:22:17 +05:30 committed by Chris Vest
parent bb7a79452e
commit 09dc15aaa0
3 changed files with 77 additions and 25 deletions

View File

@ -17,8 +17,6 @@ package io.netty.handler.pcap;
import io.netty.buffer.ByteBuf;
import java.util.concurrent.TimeUnit;
final class PcapHeaders {
/**

View File

@ -31,6 +31,7 @@ import io.netty.util.internal.ObjectUtil;
import io.netty.util.internal.logging.InternalLogger;
import io.netty.util.internal.logging.InternalLoggerFactory;
import java.io.Closeable;
import java.io.IOException;
import java.io.OutputStream;
import java.net.Inet4Address;
@ -61,8 +62,11 @@ import java.net.InetSocketAddress;
* </ul>
* </p>
*/
public final class PcapWriteHandler extends ChannelDuplexHandler {
public final class PcapWriteHandler extends ChannelDuplexHandler implements Closeable {
/**
* Logger for logging events
*/
private final InternalLogger logger = InternalLoggerFactory.getInstance(PcapWriteHandler.class);
/**
@ -108,12 +112,18 @@ public final class PcapWriteHandler extends ChannelDuplexHandler {
*/
private InetSocketAddress dstAddr;
/**
* Set to {@code true} if {@link #close()} is called and we should stop writing Pcap.
*/
private boolean isClosed;
/**
* Create new {@link PcapWriteHandler} Instance.
* {@code captureZeroByte} is set to {@code false} and
* {@code writePcapGlobalHeader} is set to {@code true}.
*
* @param outputStream OutputStream where Pcap data will be written
* @param outputStream OutputStream where Pcap data will be written. Call {@link #close()} to close this
* OutputStream.
* @throws NullPointerException If {@link OutputStream} is {@code null} then we'll throw an
* {@link NullPointerException}
*/
@ -124,7 +134,8 @@ public final class PcapWriteHandler extends ChannelDuplexHandler {
/**
* Create new {@link PcapWriteHandler} Instance
*
* @param outputStream OutputStream where Pcap data will be written
* @param outputStream OutputStream where Pcap data will be written. Call {@link #close()} to close this
* OutputStream.
* @param captureZeroByte Set to {@code true} to enable capturing packets with empty (0 bytes) payload.
* Otherwise, if set to {@code false}, empty packets will be filtered out.
* @param writePcapGlobalHeader Set to {@code true} to write Pcap Global Header on initialization.
@ -214,6 +225,7 @@ public final class PcapWriteHandler extends ChannelDuplexHandler {
@Override
public void channelRead(ChannelHandlerContext ctx, Object msg) throws Exception {
if (!isClosed) {
if (ctx.channel() instanceof SocketChannel) {
handleTCP(ctx, msg, false);
} else if (ctx.channel() instanceof DatagramChannel) {
@ -221,11 +233,13 @@ public final class PcapWriteHandler extends ChannelDuplexHandler {
} else {
logger.debug("Discarding Pcap Write for Unknown Channel Type: {}", ctx.channel());
}
}
super.channelRead(ctx, msg);
}
@Override
public void write(ChannelHandlerContext ctx, Object msg, ChannelPromise promise) throws Exception {
if (!isClosed) {
if (ctx.channel() instanceof SocketChannel) {
handleTCP(ctx, msg, true);
} else if (ctx.channel() instanceof DatagramChannel) {
@ -233,6 +247,7 @@ public final class PcapWriteHandler extends ChannelDuplexHandler {
} else {
logger.debug("Discarding Pcap Write for Unknown Channel Type: {}", ctx.channel());
}
}
super.write(ctx, msg, promise);
}
@ -495,7 +510,7 @@ public final class PcapWriteHandler extends ChannelDuplexHandler {
logger.debug("Finished Fake TCP FIN+ACK Flow to close connection");
}
this.pCapWriter.close();
close();
super.handlerRemoved(ctx);
}
@ -517,7 +532,25 @@ public final class PcapWriteHandler extends ChannelDuplexHandler {
logger.debug("Sent Fake TCP RST to close connection");
}
this.pCapWriter.close();
close();
ctx.fireExceptionCaught(cause);
}
/**
* <p> Close {@code PcapWriter} and {@link OutputStream}. </p>
* <p> Note: Calling this method does not close {@link PcapWriteHandler}.
* Only Pcap Writes are closed. </p>
*
* @throws IOException If {@link OutputStream#close()} throws an exception
*/
@Override
public void close() throws IOException {
if (isClosed) {
logger.debug("PcapWriterHandler is already closed");
} else {
isClosed = true;
pCapWriter.close();
logger.debug("PcapWriterHandler is now closed");
}
}
}

View File

@ -16,19 +16,30 @@
package io.netty.handler.pcap;
import io.netty.buffer.ByteBuf;
import io.netty.util.internal.logging.InternalLogger;
import io.netty.util.internal.logging.InternalLoggerFactory;
import java.io.Closeable;
import java.io.IOException;
import java.io.OutputStream;
import java.util.concurrent.TimeUnit;
final class PcapWriter implements Closeable {
/**
* Logger
*/
private static final InternalLogger logger = InternalLoggerFactory.getInstance(PcapWriter.class);
/**
* {@link OutputStream} where we'll write Pcap data.
*/
private final OutputStream outputStream;
/**
* Set to {@code true} if {@link #outputStream} is closed.
*/
private boolean isClosed;
/**
* This uses {@link OutputStream} for writing Pcap.
* Pcap Global Header is not written on construction.
@ -58,12 +69,16 @@ final class PcapWriter implements Closeable {
* @throws IOException If {@link OutputStream#write(byte[])} throws an exception
*/
void writePacket(ByteBuf packetHeaderBuf, ByteBuf packet) throws IOException {
long currentTime = System.currentTimeMillis();
if (isClosed) {
logger.debug("Pcap Write attempted on closed PcapWriter");
}
long timestamp = System.currentTimeMillis();
PcapHeaders.writePacketHeader(
packetHeaderBuf,
(int) TimeUnit.MILLISECONDS.toSeconds(currentTime),
(int) TimeUnit.MILLISECONDS.toMicros(currentTime) % 1000000,
(int) (timestamp / 1000L),
(int) (timestamp % 1000L * 1000L),
packet.readableBytes(),
packet.readableBytes()
);
@ -74,7 +89,13 @@ final class PcapWriter implements Closeable {
@Override
public void close() throws IOException {
if (isClosed) {
logger.debug("PcapWriter is already closed");
} else {
isClosed = true;
outputStream.flush();
outputStream.close();
logger.debug("PcapWriter is now closed");
}
}
}