From 4f01bea25b8db846f2f4740be7fba76a04647f02 Mon Sep 17 00:00:00 2001
From: Natale Patriciello <natale.patriciello@gmail.com>
Date: Tue, 8 Mar 2016 10:43:54 -0800
Subject: [PATCH] internet: Register callbacks on TCLayer and NetDevices, RX
 side

---
 src/internet/model/ipv4-l3-protocol.cc | 20 +++++++++++++++-----
 src/internet/model/ipv6-l3-protocol.cc | 13 +++++++++++--
 2 files changed, 26 insertions(+), 7 deletions(-)

diff --git a/src/internet/model/ipv4-l3-protocol.cc b/src/internet/model/ipv4-l3-protocol.cc
index 837651503..0069d650e 100644
--- a/src/internet/model/ipv4-l3-protocol.cc
+++ b/src/internet/model/ipv4-l3-protocol.cc
@@ -32,6 +32,7 @@
 #include "ns3/ipv4-header.h"
 #include "ns3/boolean.h"
 #include "ns3/ipv4-routing-table-entry.h"
+#include "ns3/traffic-control-layer.h"
 
 #include "loopback-net-device.h"
 #include "arp-l3-protocol.h"
@@ -371,12 +372,21 @@ uint32_t
 Ipv4L3Protocol::AddInterface (Ptr<NetDevice> device)
 {
   NS_LOG_FUNCTION (this << device);
+  NS_ASSERT (m_node != 0);
 
-  Ptr<Node> node = GetObject<Node> ();
-  node->RegisterProtocolHandler (MakeCallback (&Ipv4L3Protocol::Receive, this), 
-                                 Ipv4L3Protocol::PROT_NUMBER, device);
-  node->RegisterProtocolHandler (MakeCallback (&ArpL3Protocol::Receive, PeekPointer (GetObject<ArpL3Protocol> ())),
-                                 ArpL3Protocol::PROT_NUMBER, device);
+  Ptr<TrafficControlLayer> tc = m_node->GetObject<TrafficControlLayer> ();
+
+  NS_ASSERT (tc != 0);
+
+  m_node->RegisterProtocolHandler (MakeCallback (&TrafficControlLayer::Receive, tc),
+                                   Ipv4L3Protocol::PROT_NUMBER, device);
+  m_node->RegisterProtocolHandler (MakeCallback (&TrafficControlLayer::Receive, tc),
+                                   ArpL3Protocol::PROT_NUMBER, device);
+
+  tc->RegisterProtocolHandler (MakeCallback (&Ipv4L3Protocol::Receive, this),
+                               Ipv4L3Protocol::PROT_NUMBER, device);
+  tc->RegisterProtocolHandler (MakeCallback (&ArpL3Protocol::Receive, PeekPointer (GetObject<ArpL3Protocol> ())),
+                               ArpL3Protocol::PROT_NUMBER, device);
 
   Ptr<Ipv4Interface> interface = CreateObject<Ipv4Interface> ();
   interface->SetNode (m_node);
diff --git a/src/internet/model/ipv6-l3-protocol.cc b/src/internet/model/ipv6-l3-protocol.cc
index 46735ce8f..653889f80 100644
--- a/src/internet/model/ipv6-l3-protocol.cc
+++ b/src/internet/model/ipv6-l3-protocol.cc
@@ -30,6 +30,7 @@
 #include "ns3/ipv6-route.h"
 #include "ns3/mac16-address.h"
 #include "ns3/mac64-address.h"
+#include "ns3/traffic-control-layer.h"
 
 #include "loopback-net-device.h"
 #include "ipv6-l3-protocol.h"
@@ -190,10 +191,18 @@ Ptr<Ipv6RoutingProtocol> Ipv6L3Protocol::GetRoutingProtocol () const
 uint32_t Ipv6L3Protocol::AddInterface (Ptr<NetDevice> device)
 {
   NS_LOG_FUNCTION (this << device);
-  Ptr<Node> node = GetObject<Node> ();
   Ptr<Ipv6Interface> interface = CreateObject<Ipv6Interface> ();
 
-  node->RegisterProtocolHandler (MakeCallback (&Ipv6L3Protocol::Receive, this), Ipv6L3Protocol::PROT_NUMBER, device);
+  Ptr<TrafficControlLayer> tc = m_node->GetObject<TrafficControlLayer> ();
+
+  NS_ASSERT (tc != 0);
+
+  m_node->RegisterProtocolHandler (MakeCallback (&TrafficControlLayer::Receive, tc),
+                                   Ipv6L3Protocol::PROT_NUMBER, device);
+
+  tc->RegisterProtocolHandler (MakeCallback (&Ipv6L3Protocol::Receive, this),
+                               Ipv6L3Protocol::PROT_NUMBER, device);
+
   interface->SetNode (m_node);
   interface->SetDevice (device);
   interface->SetForwarding (m_ipForward);