From 461b6e0f7b3277fe980b7293a6ab644340af1d7c Mon Sep 17 00:00:00 2001 From: Aayush Atharva Date: Thu, 8 Oct 2020 15:22:17 +0530 Subject: [PATCH] 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. --- .../io/netty/handler/pcap/PcapHeaders.java | 2 - .../netty/handler/pcap/PcapWriteHandler.java | 67 ++++++++++++++----- .../io/netty/handler/pcap/PcapWriter.java | 33 +++++++-- 3 files changed, 77 insertions(+), 25 deletions(-) diff --git a/handler/src/main/java/io/netty/handler/pcap/PcapHeaders.java b/handler/src/main/java/io/netty/handler/pcap/PcapHeaders.java index 96eacd27bf..a460222273 100644 --- a/handler/src/main/java/io/netty/handler/pcap/PcapHeaders.java +++ b/handler/src/main/java/io/netty/handler/pcap/PcapHeaders.java @@ -17,8 +17,6 @@ package io.netty.handler.pcap; import io.netty.buffer.ByteBuf; -import java.util.concurrent.TimeUnit; - final class PcapHeaders { /** diff --git a/handler/src/main/java/io/netty/handler/pcap/PcapWriteHandler.java b/handler/src/main/java/io/netty/handler/pcap/PcapWriteHandler.java index 5c5e270651..f55dabd77d 100644 --- a/handler/src/main/java/io/netty/handler/pcap/PcapWriteHandler.java +++ b/handler/src/main/java/io/netty/handler/pcap/PcapWriteHandler.java @@ -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; * *

*/ -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,24 +225,28 @@ public final class PcapWriteHandler extends ChannelDuplexHandler { @Override public void channelRead(ChannelHandlerContext ctx, Object msg) throws Exception { - if (ctx.channel() instanceof SocketChannel) { - handleTCP(ctx, msg, false); - } else if (ctx.channel() instanceof DatagramChannel) { - handleUDP(ctx, msg); - } else { - logger.debug("Discarding Pcap Write for Unknown Channel Type: {}", ctx.channel()); + if (!isClosed) { + if (ctx.channel() instanceof SocketChannel) { + handleTCP(ctx, msg, false); + } else if (ctx.channel() instanceof DatagramChannel) { + handleUDP(ctx, msg); + } 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 (ctx.channel() instanceof SocketChannel) { - handleTCP(ctx, msg, true); - } else if (ctx.channel() instanceof DatagramChannel) { - handleUDP(ctx, msg); - } else { - logger.debug("Discarding Pcap Write for Unknown Channel Type: {}", ctx.channel()); + if (!isClosed) { + if (ctx.channel() instanceof SocketChannel) { + handleTCP(ctx, msg, true); + } else if (ctx.channel() instanceof DatagramChannel) { + handleUDP(ctx, msg); + } 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); } + + /** + *

Close {@code PcapWriter} and {@link OutputStream}.

+ *

Note: Calling this method does not close {@link PcapWriteHandler}. + * Only Pcap Writes are closed.

+ * + * @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"); + } + } } diff --git a/handler/src/main/java/io/netty/handler/pcap/PcapWriter.java b/handler/src/main/java/io/netty/handler/pcap/PcapWriter.java index bae06b27b3..7933a8881a 100644 --- a/handler/src/main/java/io/netty/handler/pcap/PcapWriter.java +++ b/handler/src/main/java/io/netty/handler/pcap/PcapWriter.java @@ -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 { - outputStream.flush(); - outputStream.close(); + if (isClosed) { + logger.debug("PcapWriter is already closed"); + } else { + isClosed = true; + outputStream.flush(); + outputStream.close(); + logger.debug("PcapWriter is now closed"); + } } }