Fixed bug on TCPStream when duplicate packets are sniffed.

include/tcp_stream.h

@@ -96,6 +96,17 @@ public:
         return client_payload_;
     }
     
+    /**
+     * \brief Retrieves the client payload.
+     * 
+     * This is the payload that the connection's client has sent so far.
+     * 
+     * \return payload_type& containing the payload.
+     */
+    payload_type &client_payload() {
+        return client_payload_;
+    }
+    
     /**
      * \brief Retrieves the server payload.
      * 
@@ -106,6 +117,17 @@ public:
     const payload_type &server_payload() const {
         return server_payload_;
     }
+    
+    /**
+     * \brief Retrieves the server payload.
+     * 
+     * This is the payload that the connection's server has sent so far.
+     * 
+     * \return payload_type& containing the payload.
+     */
+    payload_type &server_payload() {
+        return server_payload_;
+    }
 
     /**
      * \brief Retrieves this stream's identification number.
@@ -146,16 +168,6 @@ public:
      * any of the stored payloads.
      */
     bool update(IP *ip, TCP *tcp);
-    
-    /**
-     * Clears the client payload.
-     */
-    void clear_client_payload();
-    
-    /**
-     * Clears the server payload.
-     */
-    void clear_server_payload();
 private:
     typedef std::map<uint32_t, RawPDU*> fragments_type;
     
@@ -171,7 +183,7 @@ private:
     uint64_t identifier;
     payload_type client_payload_, server_payload_;
     fragments_type client_frags, server_frags;
-    bool fin_sent;
+    bool syn_ack_sent, fin_sent;
 };
 
 
@@ -199,7 +211,7 @@ public:
      * closed.
      */
     template<typename DataFunctor, typename EndFunctor>
-    void follow_streams(Sniffer &sniffer, DataFunctor data_fun, EndFunctor end_fun);
+    void follow_streams(BaseSniffer &sniffer, DataFunctor data_fun, EndFunctor end_fun);
 private:
     typedef std::map<TCPStream::StreamInfo, TCPStream> sessions_type;
     
@@ -222,7 +234,7 @@ private:
 };
 
 template<typename DataFunctor, typename EndFunctor>
-void TCPStreamFollower::follow_streams(Sniffer &sniffer, DataFunctor data_fun, EndFunctor end_fun) {
+void TCPStreamFollower::follow_streams(BaseSniffer &sniffer, DataFunctor data_fun, EndFunctor end_fun) {
     typedef proxy_caller<DataFunctor, EndFunctor> proxy_type;
     proxy_type proxy = { this, data_fun, end_fun };
     sniffer.sniff_loop(make_sniffer_handler(&proxy, &proxy_type::callback));
@@ -260,8 +272,8 @@ bool TCPStreamFollower::callback(PDU *pdu, const DataFunctor &data_fun, const En
             end_fun(it->second);
             sessions.erase(it);
         }
-        return true;
     }
+    return true;
 }
 }
 

src/tcp_stream.cpp

@@ -43,7 +43,8 @@ TCPStream::StreamInfo::StreamInfo(IPv4Address client,
 
 TCPStream::TCPStream(IP *ip, TCP *tcp, uint64_t identifier) 
 : client_seq(tcp->seq()), info(ip->src_addr(), ip->dst_addr(), 
-  tcp->sport(), tcp->dport()), identifier(identifier), fin_sent(false)
+  tcp->sport(), tcp->dport()), identifier(identifier), 
+  syn_ack_sent(false), fin_sent(false)
 {
     
 }
@@ -57,6 +58,7 @@ TCPStream& TCPStream::operator=(const TCPStream &rhs) {
     server_seq = rhs.server_seq;
     info = rhs.info;
     identifier = rhs.identifier;
+    syn_ack_sent = rhs.syn_ack_sent;
     fin_sent = rhs.fin_sent;
     client_payload_ = rhs.client_payload_;
     server_payload_ = rhs.server_payload_;
@@ -85,23 +87,20 @@ TCPStream::fragments_type TCPStream::clone_fragments(const fragments_type &frags
 bool TCPStream::generic_process(uint32_t &my_seq, uint32_t &other_seq, 
   payload_type &pload, fragments_type &frags, TCP *tcp, RawPDU *raw) 
 {
-    //std::cout << "Entre, my seq: " << std::hex << my_seq << std::endl;
     bool added_some(false);
     if(tcp->get_flag(TCP::SYN))
         other_seq++;
     if(tcp->get_flag(TCP::FIN) || tcp->get_flag(TCP::RST))
         fin_sent = true;
-    if(raw) {
+    if(raw && tcp->seq() >= my_seq) {
         frags[tcp->seq()] = static_cast<RawPDU*>(tcp->release_inner_pdu()); 
         fragments_type::iterator it = frags.begin();
         while(it != frags.end() && it->first == my_seq) {
-            //std::cout << "Consumo: " << my_seq << std::endl;
             pload.insert(
                 pload.end(),
                 it->second->payload().begin(), 
                 it->second->payload().end()
             );
-            //std::cout << "This size: " << it->second->payload_size() << std::endl;
             my_seq += it->second->payload_size();
             delete it->second;
             frags.erase(it);
@@ -114,20 +113,32 @@ bool TCPStream::generic_process(uint32_t &my_seq, uint32_t &other_seq,
 
 bool TCPStream::update(IP *ip, TCP *tcp) {
     RawPDU *raw = tcp->find_pdu<RawPDU>();
-    if(tcp->get_flag(TCP::SYN) && tcp->get_flag(TCP::ACK)) {
+    if(!syn_ack_sent && tcp->get_flag(TCP::SYN) && tcp->get_flag(TCP::ACK)) {
         server_seq = tcp->seq() + 1;
+        syn_ack_sent = true;
+        return false;
+    }
+    else {
+        if(ip->src_addr() == info.client_addr)
+            return generic_process(client_seq, server_seq, client_payload_, client_frags, tcp, raw);
+        else
+            return generic_process(server_seq, client_seq, server_payload_, server_frags, tcp, raw);
+    }
+}
+
+bool TCPStream::StreamInfo::operator<(const StreamInfo &rhs) const {
+    if(client_addr == rhs.client_addr) {
+        if(server_addr == rhs.server_addr) {
+            if(client_port == rhs.client_port) {
+                return server_port < rhs.server_port;
+            }
+            else
+                return client_port < rhs.client_port;
+        }
+        else
+            return server_addr < rhs.server_addr;
     }
-    if(ip->src_addr() == info.client_addr)
-        return generic_process(client_seq, server_seq, client_payload_, client_frags, tcp, raw);
     else
-        return generic_process(server_seq, client_seq, server_payload_, server_frags, tcp, raw);
-}
-
-void TCPStream::clear_client_payload() {
-    client_payload_.clear();
-}
-
-void TCPStream::clear_server_payload() {
-    server_payload_.clear();
+        return client_addr < rhs.client_addr;
 }
 }