/* -*- Mode:C++; c-file-style:''gnu''; indent-tabs-mode:nil; -*- */

#include "ns3/core-module.h"
#include "ns3/simulator-module.h"
#include "ns3/node-module.h"
#include "ns3/helper-module.h"
#include "ns3/log.h"
#include "ns3/point-to-point-channel.h"
#include "ns3/point-to-point-net-device.h"

using namespace ns3;

NS_LOG_COMPONENT_DEFINE ("TcpFinWait2Test");

/**
 * Test case for suspected bug in TcpSocketImpl's handling of data attached to a FIN packet received while in the FIN-WAIT-2 state.
 *
 * Demonstrated here using two nodes which follow the following sequence of events:
 *
 * 1. Acceptor node opens socket listening for connections
 * 2. Initiator node connects to acceptor.
 * 3. Initiator node has nothing to send, calls Socket::Close()
 * 4. Acceptor node queues five 100-byte packets.
 * 5. Acceptor node calls Socket::Close()
 * 6. Acceptor's TcpSocketImpl sends the 5 packets, the last one of which contains both data and the FIN flag.
 * 7. Initiator node's TcpSocketImpl receives the first four packets and forwards them to the application's Receive callback properly.
 * 8. Initiator node's TcpSocketImpl (which is in FIN-WAIT-2 state) receives the last packet with 100 bytes of data and a FIN flag.
 *      This fully closes the connection, but the data never makes it to the application; the receive callback is called with a NULL packet.
 *
 *
 * The problem in TcpSocketImpl:
 * ForwardUp() calls ProcessEvent() before it calls ProcessPacketAction(). ProcessEvent() sees that we're in FIN-WAIT-2, sees an incoming
 * FIN (doesn't care that it has data attached to it), goes to the CLOSE-WAIT state and calls NotifyDataRecv(), which calls my RecvCallback
 * with a NULL packet because there is nothing in m_pendingData. m_pendingData WOULD have been filled by ProcessPacketAction()... but that
 * hasn't been called yet.
 *
 * The corresponding code in ns-2's full TCP doesn't have this problem - events occur in an order which allows the data to make it to the application.
 */

class Worker : public Application
{
public:

  Worker(int bytesToSend, int bytesToRecv) :
    m_bytesToSend(bytesToSend), m_bytesToRecv(bytesToRecv)
  {
  }

  void
  Receive(Ptr<Socket> socket)
  {
    NS_LOG_FUNCTION(this << socket);
    Ptr<Packet> packet = socket->Recv();
    if (packet == NULL)
      {
        // This means we have already half-closed our socket, and we just received a FIN; therefore the connection is now fully closed.
        // We should be out of bytes to receive.
        NS_ASSERT_MSG(m_bytesToRecv == 0, "Received FIN, but there are still " << m_bytesToRecv << " bytes to receive.");//
        NS_LOG_LOGIC("Received FIN");
        return;
      } //

    m_bytesToRecv -= packet->GetSize();
    NS_LOG_LOGIC("Received a packet with " << packet->GetSize() << " bytes of payload. " << m_bytesToRecv << " bytes left to receive.");
  }

  void
  Closed(Ptr<Socket> socket)
  {
    NS_LOG_FUNCTION(this << socket);
    // This means that either:
    // A) we just received a FIN+ACK and the other side has half-closed their socket,
    // B) or option A has already happened, we have closed our side as well, and we just received the other side's ACK of our FIN+ACK of their FIN+ACK.

    // In either case, we should either be out of ADUs to receive, or we should have just the explicit FIN ADU remaining.
    NS_ASSERT_MSG(m_bytesToRecv == 0, "Received FIN, but there are still " << m_bytesToRecv << " bytes to receive.");
  }

  void
  ScheduleNextSend(Ptr<Socket> socket)
  {
    const int packetSize = 100;
    if (m_bytesToSend)
      {
        // Interestingly, the problem doesn't occur when the time between sends here is a tenth of a second or greater.
        Simulator::Schedule(Seconds(0.01), &Worker::DoSend, this, socket,
            std::min(packetSize, m_bytesToSend));
      }
    else
      {
        NS_LOG_LOGIC("Nothing to send, closing.");
        socket->Close();
      }
  }

  void
  DoSend(Ptr<Socket> socket, int bytes)
  {
    int bytesSent = socket->Send(Create<Packet> (bytes));
    NS_ASSERT_MSG(bytesSent == bytes, "bytesSent(" << bytesSent << ") != bytes(" << bytes << ")");
    m_bytesToSend -= bytesSent;
    NS_LOG_LOGIC(bytesSent << " bytes accepted for sending. " << m_bytesToSend << " bytes remain to be sent.");

    ScheduleNextSend(socket);
  }

  int m_bytesToSend;
  int m_bytesToRecv;
};

class Initiator : public Worker
{
public:
  Initiator(int bytesToSend, int bytesToRecv) :
    Worker(bytesToSend, bytesToRecv)
  {
  }

  void
  Start(InetSocketAddress peer)
  {
    m_socket = Socket::CreateSocket(GetNode(), TcpSocketFactory::GetTypeId());
    m_socket->SetConnectCallback(MakeCallback(&Initiator::ConnectionComplete,
        this), MakeNullCallback<void, Ptr<Socket> > ());
    m_socket->Connect(peer);
  }

  void
  ConnectionComplete(Ptr<Socket> socket)
  {
    NS_LOG_FUNCTION(this << socket);
    socket->SetRecvCallback(MakeCallback(&Worker::Receive, this));
    socket->SetCloseCallbacks(MakeCallback(&Worker::Closed, this),
        MakeNullCallback<void, Ptr<Socket> > ());

    ScheduleNextSend(socket);
  }

  Ptr<Socket> m_socket;
};

class Acceptor : public Worker
{
public:
  Acceptor(int bytesToSend, int bytesToRecv) :
    Worker(bytesToSend, bytesToRecv)
  {
  }

  void
  Start(InetSocketAddress local)
  {
    m_socket = Socket::CreateSocket(GetNode(), TcpSocketFactory::GetTypeId());
    m_socket->SetAcceptCallback(
        MakeCallback(&Acceptor::ConnectionRequest, this), MakeCallback(
            &Acceptor::NewConnectionCreated, this));
    m_socket->Bind(local);
    m_socket->Listen();
  }

  bool
  ConnectionRequest(Ptr<Socket> socket, const Address& address)
  {
    NS_LOG_FUNCTION(this << socket << address);
    // Always accept the connection.
    return true;
  }

  void
  NewConnectionCreated(Ptr<Socket> socket, const Address& address)
  {
    NS_LOG_FUNCTION(this << socket << address);
    socket->SetRecvCallback(MakeCallback(&Worker::Receive, this));
    socket->SetCloseCallbacks(MakeCallback(&Worker::Closed, this),
        MakeNullCallback<void, Ptr<Socket> > ());
    // Close the socket listening for new connections
    m_socket->Close();

    ScheduleNextSend(socket);
  }

  Ptr<Socket> m_socket;
};

int
main(int argc, char** argv)
{
  LogComponentEnableAll(LOG_PREFIX_FUNC);
  LogComponentEnableAll(LOG_PREFIX_NODE);
  LogComponentEnableAll(LOG_PREFIX_TIME);

  //LogComponentEnable("TcpSocketImpl", LOG_LEVEL_LOGIC);
  LogComponentEnable("TcpFinWait2Test", LOG_LEVEL_LOGIC);

  Ptr<Node> node1 = CreateObject<Node> ();
  Ptr<Node> node2 = CreateObject<Node> ();
    {
      InternetStackHelper internet;
      internet.Install(node1);
      internet.Install(node2);
    }

  Ptr<PointToPointChannel> channel = CreateObject<PointToPointChannel> ();
  Ptr<PointToPointNetDevice> device1 = CreateObject<PointToPointNetDevice> ();
  Ptr<PointToPointNetDevice> device2 = CreateObject<PointToPointNetDevice> ();

  Ipv4Address acceptorAddress, initiatorAddress;
    {
      Ipv4AddressHelper addresses;
      addresses.SetBase("10.0.1.0", "255.255.255.0");

      // Device 2: Acceptor
      device2->SetQueue(CreateObject<DropTailQueue> ());
      device2->Attach(channel);
      device2->SetAddress(Mac48Address::Allocate());
      node2->AddDevice(device2);
      acceptorAddress
          = addresses.Assign(NetDeviceContainer(device2)).GetAddress(0);

      // Device 1: Initiator
      device1->SetQueue(CreateObject<DropTailQueue> ());
      device1->Attach(channel);
      device1->SetAddress(Mac48Address::Allocate());
      node1->AddDevice(device1);
      initiatorAddress
          = addresses.Assign(NetDeviceContainer(device1)).GetAddress(0);
    }

  int port = 5000;
  int initiatorSendBytes = 0;
  int acceptorSendBytes = 500;
  Ptr<Initiator> app1 = CreateObject<Initiator> (initiatorSendBytes,
      acceptorSendBytes);
  Ptr<Acceptor> app2 = CreateObject<Acceptor> (acceptorSendBytes,
      initiatorSendBytes);

  node1->AddApplication(app1);
  node2->AddApplication(app2);

  Simulator::Schedule(Seconds(1), &Acceptor::Start, app2, InetSocketAddress(
      acceptorAddress, port));
  Simulator::Schedule(Seconds(1), &Initiator::Start, app1, InetSocketAddress(
      acceptorAddress, port));

  Simulator::Run();
  Simulator::Destroy();

  return 0;
}