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