/*
   Network glue for kexing

   (C) 2004 Jack Lloyd
*/

#include <cutlass/kex.hpp>
#include <assert.h>

/* Kill the C macros, if set */
#undef min
#undef max

static void send_stuff(conn_t* conn, cutlass_t* cut_handle,
                       byte type, const byte out[], u32bit out_size)
   {
   assert(out_size < conn->mtu);

   cutlass_packet_hdr header;

   header.cut_type = type;

   if(cutlass_send_process(cut_handle, conn, &header, out, out_size) < 0)
      {
      cutlass_sysmsg(cut_handle, CUT_ERROR,
                     "cutlass_send_process failed during kex");
      throw Exception("cutlass_send_process failed, nothing sent");
      }
   }

/*************************************************
* Write a handshake message to the network       *
*************************************************/
void Network_Layer::write(const Handshake_Message& msg,
                          cutlass_t* cut_handle,
                          conn_t* connection)
   {
   if(connection->mtu <= 400) /* i mean it */
      {
      cutlass_sysmsg(cut_handle, CUT_ERROR,
                     "Network_Layer::write: connection->mtu is <= 400; "
                     "you are the weakest link, goodbye");
      return;
      }

   SecureVector<byte> msg_data = msg.serialize();
   write(msg.type(), msg_data.begin(), msg_data.size(),
         cut_handle, connection);

   /* In case we have to resend, store it */
   last_msg_type = msg.type();
   last_msg = msg_data;
   }

/*************************************************
* Write a handshake message to the network       *
*************************************************/
void Network_Layer::write(byte type, const byte msg_data[], u32bit msg_len,
                          cutlass_t* cut_handle, conn_t* conn)
   {
   if(msg_len+2 <= conn->mtu)   /* It will fit in a single packet */
      {
      SecureVector<byte> buf;
      buf.append(1); // frag #
      buf.append(1); // frag count
      buf.append(msg_data, msg_len);

      send_stuff(conn, cut_handle, type, buf, buf.size());
      }
   else
      {
      /* How many fragments to send. The +1 is for the type code, the -2 is for
         the frag header. This will usually send stuff < complete MTU, cause
         it's simpler to handle. The decoder can handle any fragment size, so
         this can be fixed later without any compatability problems.
      */

      u32bit frag_into = (msg_len+1) / (conn->mtu-2);
      u32bit frag_size = ((msg_len+1) / frag_into) + 1;

      /*
        This allows up to 2000 bytes, or 16000 bits. If input is larger than
        this, something has probably gone wrong.
      */
      if(frag_into > 5)
         {
         cutlass_sysmsg(cut_handle, CUT_ERROR, 
                        "Network_Layer::write: Input is waaaayyy too big");
         return;
         }

      for(u32bit j = 0; j != frag_into; j++)
         {
         SecureVector<byte> buf;

         buf.append(j+1); // frag #
         buf.append((byte)(frag_into+1)); // frag count

         u32bit offset = j*frag_size; /* offset into the buffer */
         u32bit this_frag_size = std::min(frag_size, msg_len - offset);
         buf.append(msg_data + offset, this_frag_size);

         send_stuff(conn, cut_handle, type, buf, buf.size());
         }
      }
   }

/*************************************************
* Process incoming data, maybe return packet     *
*************************************************/
Handshake_Message* Network_Layer::read(byte type,
                                       const byte input[], u32bit in_len)
   {
   try {
      if(in_len == 0)
         return 0;

      if(in_len < 3)
         throw Decoding_Error("read: Packet way too short, you are the "
                              "weakest link");

      byte this_frag = input[0];

      /*   printf("Currently %d fragments\n", fragments.size()); */

      if(fragments.size())
         {
         if(total_frags != input[1]) /* should always match */
            throw Decoding_Error("Frag count changed on us, bailing out");
         if(fragments[this_frag].size())
            throw Decoding_Error("Duplicate fragment ID"); /* WTF? */
         }
      else
         total_frags = input[1];

      /*   printf("Total frags %d\n", total_frags); */

      fragments[this_frag].set(input + 2, in_len - 2);

      /*   printf("Have %d of %d frags\n", fragments.size(), total_frags); */

      /* got a full packet, do something useful */
      if(fragments.size() == total_frags)
         {
         /* reassemble, checking that the frag #s are 0...n contigous */

         std::map<byte, SecureVector<byte> >::const_iterator i =
            fragments.begin();

         SecureVector<byte> buffer;
         byte last_frag = 0; // 0 is not a valid frag number: 1+ is
         while(i != fragments.end())
            {
            /*        printf("%d\n", i->first); */
            if(i->first != last_frag + 1)
               throw Decoding_Error("Missing a fragment, or bad frag #s");
            buffer.append(i->second);
            last_frag = i->first;
            i++;
            }

         if(buffer.size() < 2)
            throw Decoding_Error("Error: Handshake packet is empty");

         fragments.clear();

         /*
           NOTE: we only use the type in the last invocation, so we can get
           wildly inconsistent results if the type code randomly changes on us
           between fragments. This is not security-cricial, AFAICT, but that
           doesn't mean it's not ugly as shit. It is. -JL (2004/09/05)
         */
         return decode_packet(type, buffer, buffer.size());
         }
      else
         return 0; // nothing yet
      }
   catch(std::exception& e)
      {
      cutlass_sysmsg(NULL, CUT_ERROR, "%s", e.what());
      return 0;
      }
   }

/*************************************************
* Decoding a complete packet                     *
*************************************************/
Handshake_Message* Network_Layer::decode_packet(byte type,
                                                const byte in[],
                                                u16bit in_len)
   {
   SecureVector<byte> buffer(in, in_len);
   if(type == CLIENT_HELLO) return new ClientHello(buffer);
   if(type == SERVER_HELLO) return new ServerHello(buffer);
   if(type == CLIENT_KEX) return new ClientKeyExchange(buffer);
   if(type == SERVER_KEX) return new ServerKeyExchange(buffer);

   cutlass_sysmsg(NULL, CUT_ERROR, "decode_packet(): Unknown packet type");
   return 0;
   }

/*************************************************
* Resend the last message                        *
*************************************************/
void Network_Layer::resend_last(cutlass_t* cut_handle, conn_t* conn)
   {
   if(last_msg.size() == 0)
     cutlass_sysmsg(cut_handle, CUT_INFO,
                    "Network_Layer::resend_last: No message to resend\n");
   else
     write(last_msg_type, last_msg.begin(), last_msg.size(),
           cut_handle, conn);
   }