// ---------------------------------------------------------------------------
// - TlsRecord.cpp                                                           -
// - afnix:tls service - tls record class implementation                     -
// ---------------------------------------------------------------------------
// - This program is free software;  you can redistribute it  and/or  modify -
// - it provided that this copyright notice is kept intact.                  -
// -                                                                         -
// - This program  is  distributed in  the hope  that it will be useful, but -
// - without  any  warranty;  without  even   the   implied    warranty   of -
// - merchantability or fitness for a particular purpose.  In no event shall -
// - the copyright holder be liable for any  direct, indirect, incidental or -
// - special damages arising in any way out of the use of this software.     -
// ---------------------------------------------------------------------------
// - copyright (c) 1999-2019 amaury darsch                                   -
// ---------------------------------------------------------------------------

#include "Vector.hpp"
#include "Integer.hpp"
#include "TlsTypes.hxx"
#include "Character.hpp"
#include "TlsRecord.hpp"
#include "QuarkZone.hpp"
#include "Exception.hpp"
#include "BlockCipher.hpp"

namespace afnix {
  
  // -------------------------------------------------------------------------
  // - class section                                                         -
  // -------------------------------------------------------------------------

  // create an empty record

  TlsRecord::TlsRecord (void) {
    p_data = new t_byte[TLS_RLEN_MAX];
    reset ();
  }

  // create a record by input stream

  TlsRecord::TlsRecord (InputStream* is) {
    p_data = new t_byte[TLS_RLEN_MAX];
    read (is);
  }

  // create a record by input stream and cipher

  TlsRecord::TlsRecord (InputStream* is, Cipher* dc) {
    p_data = new t_byte[TLS_RLEN_MAX];
    read (is, dc);
  }
  
  // create a record by type and version

  TlsRecord::TlsRecord (const t_byte type, 
			const t_byte vmaj, const t_byte vmin) {
    // validdate the version
    if (tls_vers_isvalid (vmaj, vmin) == false) {
      throw Exception ("tls-error", "unsupported record version",
		       tls_vers_tostring (vmaj, vmin));
    }
    // validate the record type
    if (tls_type_isvalid (type) == false) {
      throw Exception ("tls-error", "invalid record type");
    }
    // allocate and reset
    p_data = new t_byte[TLS_RLEN_MAX];
    reset ();
    // set type and version
    d_type = type;
    d_vmaj = vmaj;
    d_vmin = vmin;
  }

  // copy construct this record

  TlsRecord::TlsRecord (const TlsRecord& that) {
    that.rdlock ();
    try {
      // record header
      d_type = that.d_type;
      d_vmaj = that.d_vmaj;
      d_vmin = that.d_vmin;
      d_rlen = that.d_rlen;
      d_hlen = that.d_hlen;
      // allocate and copy locally
      p_data = new t_byte[TLS_RLEN_MAX];
      for (long k = 0; k < d_rlen; k++) p_data[k] = that.p_data[k];
      for (long k = d_rlen; k < TLS_RLEN_MAX; k++) p_data[k] = nilc;
    }
    catch (...) {
      that.unlock ();
      throw;
    }
	
  }

  // move construct this record

  TlsRecord::TlsRecord (TlsRecord&& that) {
    d_type = that.d_type;
    d_vmaj = that.d_vmaj;
    d_vmin = that.d_vmin;
    d_rlen = that.d_rlen;
    d_hlen = that.d_hlen;
    p_data = that.p_data;
    // reallocate original object
    that.p_data = new t_byte[TLS_RLEN_MAX];
    that.reset ();
  }

  // destroy this record

  TlsRecord::~TlsRecord (void) {
    // nullify the memory
    reset ();
    // delete data record
    delete [] p_data;
  }

  // assign a record to this one

  TlsRecord& TlsRecord::operator = (const TlsRecord& that) {
    // check for self-assignation
    if (this == &that) return *this;
    // lock and assign
    wrlock ();
    that.rdlock ();
    try {
      // record header
      d_type = that.d_type;
      d_vmaj = that.d_vmaj;
      d_vmin = that.d_vmin;
      d_rlen = that.d_rlen;
      d_hlen = that.d_hlen;
      // copy locally
      for (long k = 0; k < d_rlen; k++) p_data[k] = that.p_data[k];
      for (long k = d_rlen; k < TLS_RLEN_MAX; k++) p_data[k] = nilc;
      // unlock and return
      unlock ();
      that.unlock ();
      return *this;
    } catch (...) {
      unlock ();
      that.unlock ();
      throw;
    }
  }

  // move a record to this one

  TlsRecord& TlsRecord::operator = (TlsRecord&& that) {
    // check for self-assignation
    if (this == &that) return *this;
    // move locally
    d_type = that.d_type;
    d_vmaj = that.d_vmaj;
    d_vmin = that.d_vmin;
    d_rlen = that.d_rlen;
    d_hlen = that.d_hlen;
    p_data = that.p_data;
    // reallocate original object
    that.p_data = new t_byte[TLS_RLEN_MAX];
    that.reset ();
    // here it is
    return *this;
  }

  // return the class name
  
  String TlsRecord::repr (void) const {
    return "TlsRecord";
  }

  // clone this object

  Object* TlsRecord::clone (void) const {
    return new TlsRecord (*this);
  }

  // reset the record

  void TlsRecord::reset (void) {
    wrlock ();
    try {
      // reset record header
      d_type = nilc;
      d_vmaj = nilc;
      d_vmin = nilc;
      d_rlen = 0L;
      d_hlen = 0L;
      // nullify the memory
      for (long k = 0; k < TLS_RLEN_MAX; k++) p_data[k] = nilc;
      unlock ();
    } catch (...) {
      unlock ();
      throw;
    }
  }

  // get the header type

  t_byte TlsRecord::gettype (void) const {
    rdlock ();
    try {
      t_byte result = d_type;
      unlock ();
      return result;
    } catch (...) {
      unlock ();
      throw;
    }
  }

  // get the header major version

  t_byte TlsRecord::getmajor (void) const {
    rdlock ();
    try {
      t_byte result = d_vmaj;
      unlock ();
      return result;
    } catch (...) {
      unlock ();
      throw;
    }
  }

  // get the header minor version

  t_byte TlsRecord::getminor (void) const {
    rdlock ();
    try {
      t_byte result = d_vmin;
      unlock ();
      return result;
    } catch (...) {
      unlock ();
      throw;
    }
  }

  // get the record length

  long TlsRecord::length (void) const {
    rdlock ();
    try {
      long result = d_rlen;
      unlock ();
      return result;
    } catch (...) {
      unlock ();
      throw;
    }
  }

  // read a record byte by position

  t_byte TlsRecord::getbyte (const long pos) const {
    rdlock ();
    try {
      // check for valid position
      if ((pos < 0) || (pos > d_rlen)) {
	throw Exception ("tls-record", "invalid byte position in get");
      }
      // get byte and unlock
      t_byte result = p_data[pos];
      unlock ();
      return result;
    } catch (...) {
      unlock ();
      throw;
    }
  }

  // add a chunk block to the record
  
  void TlsRecord::add (const TlsChunk& chnk) {
    wrlock ();
    try {
      // lock and check
      chnk.rdlock ();
      long rlen = d_rlen + chnk.d_clen;
      if (rlen >= TLS_RLEN_MAX) {
	throw Exception ("tls-error", "record overflow during chunk add");
      }
      for (long k = 0L; k < chnk.d_clen; k++) {
	p_data[d_rlen++] = chnk.p_data[k];
      }
      chnk.unlock ();
      unlock ();
    } catch (...) {
      chnk.unlock ();
      unlock ();
      throw;
    }
  }

  // push the record data to a buffer

  Buffer& TlsRecord::pushb (Buffer& buf) const {
    rdlock ();
    try {
      if (d_rlen > 0L) buf.add ((char*) p_data, d_rlen);
      unlock ();
      return buf;
    } catch (...) {
      unlock ();
      throw;
    }
  }


  // read a record from an input stream

  long TlsRecord::read (InputStream* is) {
    wrlock ();
    try {
      long result = read (is, nullptr);
      unlock ();
      return result;
    } catch (...) {
      unlock ();
      throw;
    }
  }
    
  // read a record from an input stream

  long TlsRecord::read (InputStream* is, Cipher* dc) {
    // check for nil
    if (is == nullptr) return 0L;
    // lock and read
    wrlock ();
    try {
      // reset the record
      reset ();
      // read the type byte
      d_type = is->read ();
      // read the major version
      d_vmaj = is->read ();
      d_vmin = is->read ();
      // read the length
      t_byte hbyt = is->read ();
      t_byte lbyt = is->read ();
      t_word rlen = (((t_word) hbyt) << 8) + ((t_word) lbyt);
      d_rlen = rlen;
      if (d_rlen > TLS_RLEN_MAX) {
	throw Exception ("tls-error", "record size exceed maximum size");
      }
      // read the record
      long result = is->copy ((char*) p_data, d_rlen);
      if (result != d_rlen) {
	throw Exception ("tls-error", "cannot read record from stream");
      }
      // decode the block with a cipher
      if (dc != nullptr) {
	Buffer ib (d_rlen); ib.add ((const char*) p_data, d_rlen);
	// decode the stream into a buffer
	Buffer ob;
	result = dc->stream (ob, ib);
	if (result != ob.length ()) {
	  throw Exception ("record-error", "inconsistent decode size");
	}
	// remove the padding if any with block cipher
	auto* bc = dynamic_cast<BlockCipher*> (dc);
	if ((bc != nullptr) && (result > 0)) {
	  long plen = (long) ob.get(result - 1);
	  if ((plen + 1) >= result) {
	    throw Exception ("record-error", "inconsistent record padding");
	  }
	  result -= (plen + 1);
	}
	// copy the buffer into the record
	d_rlen = ob.copy ((char*) p_data, result);
	if (result != d_rlen) {
	  throw Exception ("record-error", "inconsistent decode copy size");
	}
      }
      unlock ();
      return result;
    } catch (...) {
      // reset everything
      reset ();
      // unlock and throw
      unlock ();
      throw;
    }
  }
  
  // write a record to an output stream

  long TlsRecord::write (OutputStream* os) const {
    rdlock ();
    try {
      long result = write (os, nullptr);
      unlock ();
      return result;
    } catch (...) {
      unlock ();
      throw;
    }
  }
  
  // write a record to an output stream

  long TlsRecord::write (OutputStream* os, Cipher* ec) const {
    // check for nil
    if (os == nullptr) return 0L;
    rdlock ();
    try {
      // prepare the record length
      t_word rlen = (t_word) (d_rlen + d_hlen);
      // eventually encode the data
      Buffer ob;
      if (ec != nullptr) {
	// push the data in the input buffer
	Buffer ib (rlen); ib.add ((const char*) p_data, rlen);
	// eventually pad the data
	auto* bc = dynamic_cast<BlockCipher*> (ec);
	if (bc != nullptr) {
	  long cbsz = bc->getcbsz ();
	  long psiz = cbsz - ((rlen + 1) % cbsz);
	  if (psiz > 255) {
	    throw Exception ("tls-error", "invalid record padding size");
	  }
	  if ((rlen + psiz) >= TLS_RLEN_MAX) {
	    throw Exception ("tls-error", "invalid padded record block size");
	  }
	  for (long k = 0; k <= psiz; k++) {
	    ib.add ((char) psiz);
	  }
	  if (ib.length() != rlen + psiz + 1) {
	    throw Exception ("tls-error", "inconsistent padded record size");
	  }
	}
	rlen = ec->stream (ob, ib);
      } else {
	rlen = ob.add ((const char*) p_data, d_rlen);
      }
      // initialize result
      long result = 0L;
      // write the record header
      char rhrd[5];
      rhrd[0] = (char) d_type;
      rhrd[1] = (char) d_vmaj;
      rhrd[2] = (char) d_vmin;
      rhrd[3] = (char) ((rlen & 0xFF00) >> 8);
      rhrd[4] = (char) (rlen & 0x00FF);
      result += os->write (rhrd, sizeof(rhrd));
      // write the record data
      result += os->write (ob);
      // here it is
      unlock ();
      return result;
    } catch (...) {
      unlock ();
      throw;
    }
  }

  // check the record hmac

  bool TlsRecord::chkhmc (Hmac* hmac, const t_octa snum) {
    rdlock ();
    try {
      // check for nil hmac
      if (hmac == nullptr) {
	unlock ();
	return true;
      }
      // reset the hmac
      hmac->reset();
      // get the hmac length and record length
      long hlen = hmac->length();
      long rlen = d_rlen - hlen;
      if (rlen < 0) {
	unlock ();
	return false;
      }
      // create the hmac buffer
      Buffer hbuf;
      hbuf.addho (snum);
      hbuf.add ((char) d_type);
      hbuf.add ((char) d_vmaj);
      hbuf.add ((char) d_vmin);
      hbuf.addhw ((t_word) rlen);
      // compute the mac
      Buffer mbuf;
      hmac->process (hbuf);
      hmac->process (p_data, rlen);
      hmac->finish ();
      hmac->pushb (mbuf);
      // check the computed result
      Buffer vbuf (hlen, hlen, (char*) &p_data[rlen]);
      bool result = (mbuf == vbuf);
      // update if valid
      if (result == true) {
	d_hlen  = hlen;
	d_rlen -= hlen;
      }
      unlock ();
      return result;
    } catch (...) {
      unlock ();
      throw;
    }
  }

  // add the record hmac

  bool TlsRecord::addhmc (Hmac* hmac, const t_octa snum) {
    // lock and add
    rdlock ();
    try {
      // check for nil hmac
      if (hmac == nullptr) {
	unlock ();
	return true;
      }
      // check for valid hmac len or record length
      if ((d_rlen == 0) || (d_hlen != 0)) {
	unlock ();
	return false;
      }
      // reset the hmac
      hmac->reset();
      // create the hmac buffer
      Buffer hbuf;
      hbuf.addho (snum);
      hbuf.add ((char) d_type);
      hbuf.add ((char) d_vmaj);
      hbuf.add ((char) d_vmin);
      hbuf.addhw ((t_word) d_rlen);
      // compute the mac
      Buffer mbuf;
      hmac->process (hbuf);
      hmac->process (p_data, d_rlen);
      hmac->finish ();
      hmac->pushb (mbuf);
      // paranoid check
      long hlen = hmac->length();
      if (mbuf.length() != hlen) return false;
      // map the buffer to the record
      mbuf.tomap(&p_data[d_rlen], hlen);
      d_hlen = hlen;
      unlock ();
      return true;
    } catch (...) {
      unlock ();
      throw;
    }
  }
  
  // -------------------------------------------------------------------------
  // - object section                                                        -
  // -------------------------------------------------------------------------

  // the quark zone
  static const long QUARK_ZONE_LENGTH = 1;
  static QuarkZone  zone (QUARK_ZONE_LENGTH);

  // the object supported quarks
  static const long QUARK_READ    = zone.intern ("read");

  // create a new object in a generic way

  Object* TlsRecord::mknew (Vector* argv) {
    // get the number of arguments
    long argc = (argv == nullptr) ? 0 : argv->length ();

    // check for 0 argument
    if (argc == 0) return new TlsRecord;
    // check for 1 argument
    if (argc == 1) {
      Object*     obj = argv->get (0);
      InputStream* is = dynamic_cast<InputStream*> (obj);
      if (is == nullptr) {
	throw Exception ("type-error", "invalid object as tls input stream",
			 Object::repr (obj));
      }
      return new TlsRecord (is);
    }
    // too many arguments
    throw Exception ("argument-error", 
                     "too many argument with tls record constructor");
  }

  // return true if the given quark is defined

  bool TlsRecord::isquark (const long quark, const bool hflg) const {
    rdlock ();
    try {
      if (zone.exists (quark) == true) {
	unlock ();
	return true;
      }
      bool result = hflg ? TlsHeader::isquark (quark, hflg) : false;
      unlock ();
      return result;
    } catch (...) {
      unlock ();
      throw;
    }
  }

  // apply this object with a set of arguments and a quark
  
  Object* TlsRecord::apply (Runnable* robj, Nameset* nset, const long quark,
			    Vector* argv) {
    // get the number of arguments
    long argc = (argv == nullptr) ? 0 : argv->length ();

    // dispatch 1 argument
    if (argc == 1) {
      if (quark == QUARK_READ) {
	Object* obj = argv->get (0);
	InputStream* is = dynamic_cast<InputStream*> (obj);
	if (is == nullptr) {
	  throw Exception ("type-error", "invalid object as tls input stream",
			   Object::repr (obj));
	}
	return new Integer (read (is));
      }
    }
    // call the tle header method
    return TlsHeader::apply (robj, nset, quark, argv);
  }
}
