From ec3d666f3ba95b6aa65079347c33a6e2369a35b5 Mon Sep 17 00:00:00 2001
From: Raj Bhattacharjea <raj.b@gatech.edu>
Date: Wed, 2 Apr 2008 14:09:36 -0400
Subject: [PATCH] Add IP layer tracing helpers to InternetStackHelper

---
 examples/tcp-large-transfer.cc      |  2 +-
 src/helper/internet-stack-helper.cc | 67 +++++++++++++++++++++++++++++
 src/helper/internet-stack-helper.h  | 24 +++++++++++
 3 files changed, 92 insertions(+), 1 deletion(-)

diff --git a/examples/tcp-large-transfer.cc b/examples/tcp-large-transfer.cc
index 9c92e3dd9..fcab3dfce 100644
--- a/examples/tcp-large-transfer.cc
+++ b/examples/tcp-large-transfer.cc
@@ -182,9 +182,9 @@ int main (int argc, char *argv[])
 
   std::ofstream ascii;
   ascii.open ("tcp-large-transfer.tr");
-  PointToPointHelper::EnablePcap ("tcp-large-transfer");
   PointToPointHelper::EnableAscii (ascii);
 
+  InternetStackHelper::EnablePcap ("tcp-large-transfer");
 
   Simulator::StopAt (Seconds(1000));
   Simulator::Run ();
diff --git a/src/helper/internet-stack-helper.cc b/src/helper/internet-stack-helper.cc
index 1003dceba..badd49786 100644
--- a/src/helper/internet-stack-helper.cc
+++ b/src/helper/internet-stack-helper.cc
@@ -20,9 +20,13 @@
 #include "internet-stack-helper.h"
 #include "ns3/internet-stack.h"
 #include "ns3/packet-socket-factory.h"
+#include "ns3/config.h"
 
 namespace ns3 {
 
+std::vector<InternetStackHelper::Trace> InternetStackHelper::m_traces;
+std::string InternetStackHelper::m_pcapBaseFilename;
+
 void 
 InternetStackHelper::Build (NodeContainer c)
 {
@@ -35,5 +39,68 @@ InternetStackHelper::Build (NodeContainer c)
     }
 }
 
+void
+InternetStackHelper::EnablePcap (std::string filename)
+{
+  InternetStackHelper::m_pcapBaseFilename = filename;
+  Config::Connect ("/NodeList/*/$ns3::Ipv4L3Protocol/Tx",
+                              MakeCallback (&InternetStackHelper::LogTxIp));
+  Config::Connect ("/NodeList/*/$ns3::Ipv4L3Protocol/Rx",
+                              MakeCallback (&InternetStackHelper::LogRxIp));
+}
+
+uint32_t
+InternetStackHelper::GetNodeIndex (std::string context)
+{
+  std::string::size_type pos;
+  pos = context.find ("/NodeList/");
+  NS_ASSERT (pos == 0);
+  std::string::size_type afterNodeIndex = context.find ("/", 11);
+  NS_ASSERT (afterNodeIndex != std::string::npos);
+  std::string index = context.substr (10, afterNodeIndex - 10);
+  std::istringstream iss;
+  iss.str (index);
+  uint32_t nodeIndex;
+  iss >> nodeIndex;
+  return nodeIndex;
+}
+
+void
+InternetStackHelper::LogTxIp (std::string context, Ptr<const Packet> packet, uint32_t interfaceIndex)
+{
+  Ptr<PcapWriter> writer = InternetStackHelper::GetStream (GetNodeIndex (context), interfaceIndex);
+  writer->WritePacket (packet);
+}
+
+void
+InternetStackHelper::LogRxIp (std::string context, Ptr<const Packet> packet, uint32_t interfaceIndex)
+{
+  Ptr<PcapWriter> writer = InternetStackHelper::GetStream (GetNodeIndex (context), interfaceIndex);
+  writer->WritePacket (packet);
+}
+
+Ptr<PcapWriter>
+InternetStackHelper::GetStream (uint32_t nodeId, uint32_t interfaceId)
+{
+  for (std::vector<Trace>::iterator i = m_traces.begin ();
+       i != m_traces.end (); i++)
+  {
+    if (i->nodeId == nodeId &&
+        i->interfaceId == interfaceId)
+    {
+      return i->writer;
+    }
+  }
+  InternetStackHelper::Trace trace;
+  trace.nodeId = nodeId;
+  trace.interfaceId = interfaceId;
+  trace.writer = Create<PcapWriter> ();
+  std::ostringstream oss;
+  oss << m_pcapBaseFilename << ".pcap-" << nodeId << "-" << interfaceId;
+  trace.writer->Open (oss.str ());
+  trace.writer->WriteIpHeader ();
+  m_traces.push_back (trace);
+  return trace.writer;
+}
 
 } // namespace ns3
diff --git a/src/helper/internet-stack-helper.h b/src/helper/internet-stack-helper.h
index 828d11879..48bc6b633 100644
--- a/src/helper/internet-stack-helper.h
+++ b/src/helper/internet-stack-helper.h
@@ -21,6 +21,9 @@
 #define INTERNET_STACK_HELPER_H
 
 #include "node-container.h"
+#include "net-device-container.h"
+#include "ns3/pcap-writer.h"
+#include "ns3/packet.h"
 
 namespace ns3 {
 
@@ -37,6 +40,27 @@ public:
    * of the ns3::Ipv4, ns3::Udp, and, ns3::Tcp classes.
    */
   void Build (NodeContainer c);
+
+  /**
+   * \param filename filename prefix to use for pcap files.
+   *
+   * Enable pcap output on each protocol instance which is of the
+   * ns3::Ipv4L3Protocol type.  Both Tx and Rx events will be logged.
+   */
+  static void EnablePcap (std::string filename);
+
+private:
+  static void LogRxIp (std::string context, Ptr<const Packet> packet, uint32_t deviceId);
+  static void LogTxIp (std::string context, Ptr<const Packet> packet, uint32_t deviceId);
+  static Ptr<PcapWriter> GetStream (uint32_t nodeId, uint32_t interfaceId);
+  struct Trace {
+    uint32_t nodeId;
+    uint32_t interfaceId;
+    Ptr<PcapWriter> writer;
+  };
+  static std::string m_pcapBaseFilename;
+  static uint32_t GetNodeIndex (std::string context);
+  static std::vector<Trace> m_traces;
 };
 
 } // namespace ns3