c61223de295a8962efe74fc8a7c9a1769449e194
[pingpong.git] / 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.IpV4Packet;
6 import org.pcap4j.packet.TcpPacket;
7
8 import java.util.*;
9
10 /**
11  * Reassembles TCP conversations (streams).
12  *
13  * @author Janus Varmarken {@literal <jvarmark@uci.edu>}
14  * @author Rahmadi Trimananda {@literal <rtrimana@uci.edu>}
15  */
16 public class TcpReassembler implements PacketListener {
17
18     /**
19      * Holds <em>open</em> {@link Conversation}s, i.e., {@code Conversation}s that have <em>not</em> been detected as
20      * (gracefully) terminated based on the set of packets observed thus far.
21      * A {@link Conversation} is moved to {@link #mTerminatedConversations} if it can be determined that it is has
22      * terminated. Termination can be detected by a) observing two {@link FinAckPair}s, one in each direction, (graceful
23      * termination, see {@link Conversation#isGracefullyShutdown()}) or b) by observing a SYN packet that matches the
24      * four tuple of an existing {@code Conversation}, but which holds a <em>different</em> sequence number than the
25      * same-direction SYN packet recorded for the {@code Conversation}.
26      * <p>
27      * Note that due to limitations of the {@link Set} interface (specifically, there is no {@code get(T t)} method),
28      * we have to resort to a {@link Map} (in which keys map to themselves) to "mimic" a set with {@code get(T t)}
29      * functionality.
30      *
31      * @see <a href="https://stackoverflow.com/questions/7283338/getting-an-element-from-a-set">this question on StackOverflow.com</a>
32      */
33     private final Map<Conversation, Conversation> mOpenConversations = new HashMap<>();
34
35     /**
36      * Holds <em>terminated</em> {@link Conversation}s.
37      */
38     private final List<Conversation> mTerminatedConversations = new ArrayList<>();
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);
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.add(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 AssertionError("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         // Add RST packet to conversation.
160         conv.addRstPacket(rstPacket);
161         // Move conversation to set of terminated conversations.
162         mTerminatedConversations.add(conv);
163         mOpenConversations.remove(conv, conv);
164     }
165
166     private void processFinPacket(PcapPacket finPacket) {
167 //        getOngoingConversationForPacket(finPacket).addFinPacket(finPacket);
168         getOngoingConversationOrCreateNew(finPacket).addFinPacket(finPacket);
169     }
170
171     private void processAck(PcapPacket ackPacket) {
172 //        getOngoingConversationForPacket(ackPacket).attemptAcknowledgementOfFin(ackPacket);
173         // Note that unlike the style for SYN, FIN, and payload packets, for "ACK only" packets, we want to avoid
174         // creating a new conversation.
175         Conversation conv = getOngoingConversationForPacket(ackPacket);
176         if (conv != null) {
177             // The ACK may be an ACK of a FIN, so attempt to mark the FIN as ack'ed.
178             conv.attemptAcknowledgementOfFin(ackPacket);
179             if (conv.isGracefullyShutdown()) {
180                 // Move conversation to set of terminated conversations.
181                 mTerminatedConversations.add(conv);
182                 mOpenConversations.remove(conv);
183             }
184         }
185         // Note: add (additional) processing of ACKs (that are not ACKs of FINs) as necessary here...
186     }
187
188     private void processPayloadPacket(PcapPacket pcapPacket) {
189 //        getOngoingConversationForPacket(pcapPacket).addPacket(pcapPacket, true);
190         getOngoingConversationOrCreateNew(pcapPacket).addPacket(pcapPacket, true);
191     }
192
193     /**
194      * Locates an ongoing conversation (if any) that {@code pcapPacket} pertains to.
195      * @param pcapPacket The packet that is to be mapped to an ongoing {@code Conversation}.
196      * @return The {@code Conversation} matching {@code pcapPacket} or {@code null} if there is no match.
197      */
198     private Conversation getOngoingConversationForPacket(PcapPacket pcapPacket) {
199         // We cannot know if this is a client-to-server or server-to-client packet without trying both options...
200         Conversation conv = mOpenConversations.get(Conversation.fromPcapPacket(pcapPacket, true));
201         if (conv == null) {
202             conv = mOpenConversations.get(Conversation.fromPcapPacket(pcapPacket, false));
203         }
204         return conv;
205     }
206
207     /**
208      * Like {@link #getOngoingConversationForPacket(PcapPacket)}, but creates and inserts a new {@code Conversation}
209      * into {@link #mOpenConversations} if no open conversation is found (i.e., in the case that
210      * {@link #getOngoingConversationForPacket(PcapPacket)} returns {@code null}).
211      *
212      * @param pcapPacket The packet that is to be mapped to an ongoing {@code Conversation}.
213      * @return The existing, ongoing {@code Conversation} matching {@code pcapPacket} or the newly created one in case
214      *         no match was found.
215      */
216     private Conversation getOngoingConversationOrCreateNew(PcapPacket pcapPacket) {
217         Conversation conv = getOngoingConversationForPacket(pcapPacket);
218         if (conv == null) {
219             TcpPacket tcpPacket = pcapPacket.get(TcpPacket.class);
220             if (tcpPacket.getHeader().getSyn() && tcpPacket.getHeader().getAck()) {
221                 // A SYN ACK packet always originates from the server (it is a reply to the initial SYN packet from the client)
222                 conv = Conversation.fromPcapPacket(pcapPacket, false);
223             } else {
224                 // TODO: can we do anything else but arbitrarily select who is designated as the server in this case?
225                 // We can check if the IP prefix matches a local IP when handling traffic observed inside the local
226                 // network, but that obviously won't be a useful strategy for an observer at the WAN port.
227                 String srcIp = pcapPacket.get(IpV4Packet.class).getHeader().getSrcAddr().getHostAddress();
228                 boolean clientIsSrc = srcIp.startsWith("10.0.1.") || srcIp.startsWith("192.168.1.");
229                 conv = Conversation.fromPcapPacket(pcapPacket, clientIsSrc);
230             }
231             mOpenConversations.put(conv, conv);
232         }
233         return conv;
234     }
235 }