Code/Projects/SmartPlugDetector/src/main/java/edu/uci/iotproject/TcpReassembler.java

   1 package edu.uci.iotproject;
   2
   3 import org.pcap4j.core.PacketListener;
   4 import org.pcap4j.core.PcapPacket;
   5 import org.pcap4j.packet.TcpPacket;
   6
   7 import java.util.*;
   8
   9 /**
  10  * Reassembles TCP conversations (streams).
  11  *
  12  * @author Janus Varmarken {@literal <jvarmark@uci.edu>}
  13  * @author Rahmadi Trimananda {@literal <rtrimana@uci.edu>}
  14  */
  15 public class TcpReassembler implements PacketListener {
  16
  17     /**
  18      * Holds <em>open</em> {@link Conversation}s, i.e., {@code Conversation}s that have <em>not</em> been detected as
  19      * (gracefully) terminated based on the set of packets observed thus far.
  20      * A {@link Conversation} is moved to {@link #mTerminatedConversations} if it can be determined that it is has
  21      * terminated. Termination can be detected by a) observing two {@link FinAckPair}s, one in each direction, (graceful
  22      * termination, see {@link Conversation#isGracefullyShutdown()}) or b) by observing a SYN packet that matches the
  23      * four tuple of an existing {@code Conversation}, but which holds a <em>different</em> sequence number than the
  24      * same-direction SYN packet recorded for the {@code Conversation}.
  25      * <p>
  26      * Note that due to limitations of the {@link Set} interface (specifically, there is no {@code get(T t)} method),
  27      * we have to resort to a {@link Map} (in which keys map to themselves) to "mimic" a set with {@code get(T t)}
  28      * functionality.
  29      *
  30      * @see <a href="https://stackoverflow.com/questions/7283338/getting-an-element-from-a-set">this question on StackOverflow.com</a>
  31      */
  32     private final Map<Conversation, Conversation> mOpenConversations = new HashMap<>();
  33
  34     /**
  35      * Holds <em>terminated</em> {@link Conversation}s.
  36      * TODO: Should turn this into a list to avoid unintentional overwrite of a Conversation in case ephemeral port number is reused at a later stage...?
  37      */
  38     private final Map<Conversation, Conversation> mTerminatedConversations = new HashMap<>();
  39
  40     @Override
  41     public void gotPacket(PcapPacket pcapPacket) {
  42         TcpPacket tcpPacket = pcapPacket.get(TcpPacket.class);
  43         if (tcpPacket == null) {
  44             return;
  45         }
  46         // ... TODO?
  47         processPacket(pcapPacket);
  48     }
  49
  50     /**
  51      * Get the reassembled TCP connections. Note that if this is called while packets are still being processed (by
  52      * calls to {@link #gotPacket(PcapPacket)}), the behavior is undefined and the returned list may be inconsistent.
  53      * @return The reassembled TCP connections.
  54      */
  55     public List<Conversation> getTcpConversations() {
  56         ArrayList<Conversation> combined = new ArrayList<>();
  57         combined.addAll(mTerminatedConversations.values());
  58         combined.addAll(mOpenConversations.values());
  59         return combined;
  60     }
  61
  62     private void processPacket(PcapPacket pcapPacket) {
  63         TcpPacket tcpPacket = pcapPacket.get(TcpPacket.class);
  64         // Handle client connection initiation attempts.
  65         if (tcpPacket.getHeader().getSyn() && !tcpPacket.getHeader().getAck()) {
  66             // A segment with the SYN flag set, but no ACK flag indicates that a client is attempting to initiate a new
  67             // connection.
  68             processNewConnectionRequest(pcapPacket);
  69             return;
  70         }
  71         // Handle server connection initiation acknowledgement
  72         if (tcpPacket.getHeader().getSyn() && tcpPacket.getHeader().getAck()) {
  73             // A segment with both the SYN and ACK flags set indicates that the server has accepted the client's request
  74             // to initiate a new connection.
  75             processNewConnectionAck(pcapPacket);
  76             return;
  77         }
  78         // Handle resets
  79         if (tcpPacket.getHeader().getRst()) {
  80             processRstPacket(pcapPacket);
  81             return;
  82         }
  83         // Handle FINs
  84         if (tcpPacket.getHeader().getFin()) {
  85             // Handle FIN packet.
  86             processFinPacket(pcapPacket);
  87         }
  88         // Handle ACKs (currently only ACKs of FINS)
  89         if (tcpPacket.getHeader().getAck()) {
  90             processAck(pcapPacket);
  91         }
  92         // Handle packets that carry payload (application data).
  93         if (tcpPacket.getPayload() != null) {
  94             processPayloadPacket(pcapPacket);
  95         }
  96     }
  97
  98     private void processNewConnectionRequest(PcapPacket clientSynPacket) {
  99         // A SYN w/o ACK always originates from the client.
 100         Conversation conv = Conversation.fromPcapPacket(clientSynPacket, true);
 101         conv.addSynPacket(clientSynPacket);
 102         // Is there an ongoing conversation for the same four tuple (clientIp, clientPort, serverIp, serverPort) as
 103         // found in the new SYN packet?
 104         Conversation ongoingConv = mOpenConversations.get(conv);
 105         if (ongoingConv != null) {
 106             if (ongoingConv.isRetransmission(clientSynPacket)) {
 107                 // SYN retransmission detected, do nothing.
 108                 return;
 109                 // TODO: the way retransmission detection is implemented may cause a bug for connections where we have
 110                 // not recorded the initial SYN, but only the SYN ACK, as retransmission is determined by comparing the
 111                 // sequence numbers of initial SYNs -- and if no initial SYN is present for the Conversation, the new
 112                 // SYN will be interpreted as a retransmission. Possible fix: let isRentransmission ALWAYS return false
 113                 // when presented with a SYN packet when the Conversation already holds a SYN ACK packet?
 114             } else {
 115                 // New SYN has different sequence number than SYN recorded for ongoingConv, so this must be an attempt
 116                 // to establish a new conversation with the same four tuple as ongoingConv.
 117                 // Mark existing connection as terminated.
 118                 // TODO: is this 100% theoretically correct, e.g., if many connection attempts are made back to back? And RST packets?
 119                 mTerminatedConversations.put(ongoingConv, ongoingConv);
 120                 mOpenConversations.remove(ongoingConv);
 121             }
 122         }
 123         // Finally, update the map of open connections with the new connection.
 124         mOpenConversations.put(conv, conv);
 125     }
 126
 127
 128     /*
 129      * TODO a problem across the board for all processXPacket methods below:
 130      * if we start the capture in the middle of a TCP connection, we will not have an entry for the conversation in the
 131      * map as we have not seen the initial SYN packet.
 132      * Two ways we can address this:
 133      * a) Perform null-checks and ignore packets for which we have not seen SYN
 134      *    + easy to get correct
 135      *    - we discard data (issue for long-lived connections!)
 136      * b) Add a corresponding conversation entry whenever we encounter a packet that does not map to a conversation
 137      *    + we consider all data
 138      *    - not immediately clear if this will introduce bugs (incorrectly mapping packets to wrong conversations?)
 139      *
 140      *  [[[ I went with option b) for now; see getOngoingConversationOrCreateNew(PcapPacket pcapPacket). ]]]
 141      */
 142
 143     private void processNewConnectionAck(PcapPacket srvSynPacket) {
 144         // Find the corresponding ongoing connection, if any (if we start the capture just *after* the initial SYN, no
 145         // ongoing conversation entry will exist, so it must be created in that case).
 146 //        Conversation conv = mOpenConversations.get(Conversation.fromPcapPacket(srvSynPacket, false));
 147         Conversation conv = getOngoingConversationOrCreateNew(srvSynPacket);
 148         // Note: exploits &&'s short-circuit operation: only attempts to add non-retransmissions.
 149         if (!conv.isRetransmission(srvSynPacket) && !conv.addSynPacket(srvSynPacket)) {
 150             // For safety/debugging: if NOT a retransmission and add fails,
 151             // something has gone terribly wrong/invariant is broken.
 152             throw new IllegalStateException("Attempt to add SYN ACK packet that was NOT a retransmission failed." +
 153                     Conversation.class.getSimpleName() + " invariant broken.");
 154         }
 155     }
 156
 157     private void processRstPacket(PcapPacket rstPacket) {
 158         Conversation conv = getOngoingConversationOrCreateNew(rstPacket);
 159         // Move conversation to set of terminated conversations.
 160         mTerminatedConversations.put(conv, conv);
 161         mOpenConversations.remove(conv, conv);
 162     }
 163
 164     private void processFinPacket(PcapPacket finPacket) {
 165 //        getOngoingConversationForPacket(finPacket).addFinPacket(finPacket);
 166         getOngoingConversationOrCreateNew(finPacket).addFinPacket(finPacket);
 167     }
 168
 169     private void processAck(PcapPacket ackPacket) {
 170 //        getOngoingConversationForPacket(ackPacket).attemptAcknowledgementOfFin(ackPacket);
 171         // Note that unlike the style for SYN, FIN, and payload packets, for "ACK only" packets, we want to avoid
 172         // creating a new conversation.
 173         Conversation conv = getOngoingConversationForPacket(ackPacket);
 174         if (conv != null) {
 175             // The ACK may be an ACK of a FIN, so attempt to mark the FIN as ack'ed.
 176             conv.attemptAcknowledgementOfFin(ackPacket);
 177             if (conv.isGracefullyShutdown()) {
 178                 // Move conversation to set of terminated conversations.
 179                 mTerminatedConversations.put(conv, conv);
 180                 mOpenConversations.remove(conv);
 181             }
 182         }
 183         // Note: add (additional) processing of ACKs (that are not ACKs of FINs) as necessary here...
 184     }
 185
 186     private void processPayloadPacket(PcapPacket pcapPacket) {
 187 //        getOngoingConversationForPacket(pcapPacket).addPacket(pcapPacket, true);
 188         getOngoingConversationOrCreateNew(pcapPacket).addPacket(pcapPacket, true);
 189     }
 190
 191     /**
 192      * Locates an ongoing conversation (if any) that {@code pcapPacket} pertains to.
 193      * @param pcapPacket The packet that is to be mapped to an ongoing {@code Conversation}.
 194      * @return The {@code Conversation} matching {@code pcapPacket} or {@code null} if there is no match.
 195      */
 196     private Conversation getOngoingConversationForPacket(PcapPacket pcapPacket) {
 197         // We cannot know if this is a client-to-server or server-to-client packet without trying both options...
 198         Conversation conv = mOpenConversations.get(Conversation.fromPcapPacket(pcapPacket, true));
 199         if (conv == null) {
 200             conv = mOpenConversations.get(Conversation.fromPcapPacket(pcapPacket, false));
 201         }
 202         return conv;
 203     }
 204
 205     /**
 206      * Like {@link #getOngoingConversationForPacket(PcapPacket)}, but creates and inserts a new {@code Conversation}
 207      * into {@link #mOpenConversations} if no open conversation is found (i.e., in the case that
 208      * {@link #getOngoingConversationForPacket(PcapPacket)} returns {@code null}).
 209      *
 210      * @param pcapPacket The packet that is to be mapped to an ongoing {@code Conversation}.
 211      * @return The existing, ongoing {@code Conversation} matching {@code pcapPacket} or the newly created one in case
 212      *         no match was found.
 213      */
 214     private Conversation getOngoingConversationOrCreateNew(PcapPacket pcapPacket) {
 215         Conversation conv = getOngoingConversationForPacket(pcapPacket);
 216         if (conv == null) {
 217             TcpPacket tcpPacket = pcapPacket.get(TcpPacket.class);
 218             if (tcpPacket.getHeader().getSyn() && tcpPacket.getHeader().getAck()) {
 219                 // A SYN ACK packet always originates from the server (it is a reply to the initial SYN packet from the client)
 220                 conv = Conversation.fromPcapPacket(pcapPacket, false);
 221             } else {
 222                 // TODO: can we do anything else but arbitrarily select who is designated as the server in this case?
 223                 conv = Conversation.fromPcapPacket(pcapPacket, false);
 224             }
 225             mOpenConversations.put(conv, conv);
 226         }
 227         return conv;
 228     }
 229 }