#include "TLS.h"
#include "Exception.h"
#include "Log.h"
#include <iostream>
#include <openssl/rand.h>

namespace core {
   
   static pthread_mutex_t *lockarray;

   static unsigned long thread_id(void) {	
      return ((unsigned long) pthread_self());
   }
   
   static void lock_callback(int mode, int type, const char *file, int line) {
      if(mode & CRYPTO_LOCK)  
        pthread_mutex_lock(&(lockarray[type]));   
      else 
        pthread_mutex_unlock(&(lockarray[type]));
   }

   static int generate_session_id(const SSL *ssl, unsigned char *id, unsigned int *id_len) {
      char *session_id_prefix = (char *)"BARANT";
      unsigned int count = 0;
      coreutils::Log(coreutils::LOG_DEBUG_3) << "Generating unique session id.";
      do {
	 RAND_bytes(id, *id_len);
	 memcpy(id, session_id_prefix, (strlen(session_id_prefix) < *id_len));
      } while(SSL_has_matching_session_id(ssl, id, *id_len) && (++count < 10));
      return 1;
   }
   
   void handshake_complete(const SSL *ssl, int where, int ret) {
      coreutils::Log(coreutils::LOG_DEBUG_3) << "==>" << SSL_state_string_long(ssl) << "<==" << ret;
      if(where & SSL_CB_HANDSHAKE_DONE) {
	 // Only a context configured with SSL_VERIFY_PEER (see
	 // createContext()'s own cACertificate-gated block) ever sends a
	 // client a CertificateRequest in the first place - ordinary
	 // one-way TLS (a public site's normal case, no caCertificate
	 // configured) never asks the client for one, so
	 // SSL_get_peer_certificate() is unconditionally NULL there and is
	 // not a failure. Checking the mode actually negotiated on THIS
	 // connection (inherited from the context) rather than assuming
	 // every handshake is mutual-TLS is what makes plain
	 // browser-facing HTTPS work at all under this callback.
	 if(SSL_get_verify_mode(ssl) & SSL_VERIFY_PEER) {
	    X509 *ssl_client_cert = SSL_get_peer_certificate(ssl);
	    if(!ssl_client_cert)
	      throw std::string("Unable to get peer certificate.");
	    X509_free(ssl_client_cert);
	    if(SSL_get_verify_result(ssl) != X509_V_OK)
	      throw std::string("Certificate verification failed.");
	    coreutils::Log(coreutils::LOG_DEBUG_3) << "Certificate verified successfully.";
	 }
      }
      else
	coreutils::Log(coreutils::LOG_DEBUG_3) << "No client certificate.";
   }

   TLS::TLS() {} 
   
   TLS::TLS(TLSInfo *tlsInfo) {
      createContext(tlsInfo);
   }

   TLS::TLS(SSL_CTX *ctx) : ctx(ctx) {}
   
   TLS::~TLS() {}

   void TLS::createContext(TLSInfo *tlsInfo) {

      if(tlsInfo) {
	 SSL_library_init();
	 SSL_load_error_strings();
	 lockarray = (pthread_mutex_t *)OPENSSL_malloc(CRYPTO_num_locks() * sizeof(pthread_mutex_t));
	 for(int i = 0; i < CRYPTO_num_locks(); ++i)
	   pthread_mutex_init(&(lockarray[i]), NULL);
	 CRYPTO_set_id_callback((unsigned long (*)())thread_id);
	 CRYPTO_set_locking_callback((void (*)(int, int, const char *, int))lock_callback);
	 SSLeay_add_ssl_algorithms();
	 RAND_load_file("/dev/hwrng", 1024);
	 ctx = buildContext(tlsInfo);
      }
   }

   // The per-context cert/key/cipher/verify setup factored out of
   // createContext() above so addSNIContext() (below) can build
   // additional contexts - one per extra hostname - without repeating
   // the ONE-TIME, process-wide OpenSSL library init that stays in
   // createContext() itself (SSL_library_init(), the lockarray, RAND
   // seeding): that part must only ever run once per process, not once
   // per hostname.
   SSL_CTX *TLS::buildContext(TLSInfo *tlsInfo) {
      SSL_CTX *newCtx;
      if(!(newCtx = SSL_CTX_new(SSLv23_server_method())))
	throw coreutils::Exception("Error while setting server method SSLv23.");
      SSL_CTX_set_mode(newCtx, SSL_MODE_RELEASE_BUFFERS | SSL_MODE_ACCEPT_MOVING_WRITE_BUFFER);
      SSL_CTX_set_options(newCtx, SSL_OP_NO_TICKET);
      SSL_CTX_set_session_cache_mode(newCtx, SSL_SESS_CACHE_SERVER);
      SSL_CTX_set_cipher_list(newCtx, "ECDH-ECDSA-AES256-GCM-SHA384:DHE-DSS-AES256-GCM-SHA384:DHE-RSA-AES256-GCM-SHA384:DHE-RSA-AES256-SHA256:DHE-DSS-AES256-SHA256:DHE-RSA-AES256-SHA:DHE-DSS-AES256-SHA:AES256-GCM-SHA384:AES256-SHA256:AES256-SHA:AES128-GCM-SHA256:AES128-SHA256:AES128-SHA");
      // _chain_file (not _file) - loads every certificate in the file,
      // not just the first. A real Let's Encrypt fullchain.pem is the
      // leaf cert followed by the intermediate CA cert concatenated in
      // one file specifically so a server can hand both to the client
      // in one shot; SSL_CTX_use_certificate_file() only ever reads the
      // first PEM block, silently dropping the intermediate - confirmed
      // via a real handshake against this exact fullchain.pem: clients
      // got "unable to get local issuer certificate" because the server
      // only ever sent the leaf, never the intermediate needed to chain
      // up to a root the client already trusts.
      if(SSL_CTX_use_certificate_chain_file(newCtx, tlsInfo->certificate.c_str()) <= 0)
	throw coreutils::Exception("Error looking up certificate.");
      if(SSL_CTX_use_PrivateKey_file(newCtx, tlsInfo->key.c_str(), SSL_FILETYPE_PEM) < 0)
	throw coreutils::Exception("Error with private key.");
      if(SSL_CTX_check_private_key(newCtx) != 1)
	throw coreutils::Exception("Private key does not match certificate.");
      // Client-certificate verification (mutual TLS) is opt-in, not the
      // default - a public-facing site (e.g. a real Let's Encrypt
      // server cert with no corresponding client CA of its own) has no
      // client certificate to hand out to ordinary visitors, so
      // requesting one via SSL_VERIFY_PEER would surface a "select a
      // certificate" prompt real browsers can't meaningfully answer.
      // Only enable it when a caCertificate is actually configured -
      // existing internal/mTLS use of this class (any config that DOES
      // set tls.caCertificate) keeps behaving exactly as before.
      if(tlsInfo->cACertificate.getLength() > 0) {
	 SSL_CTX_set_verify(newCtx, SSL_VERIFY_PEER, NULL);
	 SSL_CTX_set_verify_depth(newCtx, 1);
	 if(!SSL_CTX_load_verify_locations(newCtx, tlsInfo->cACertificate.c_str(), NULL))
	   throw coreutils::Exception("Cannot verify locations.");
	 SSL_CTX_set_client_CA_list(newCtx, SSL_load_client_CA_file(tlsInfo->cACertificate.c_str()));
	 coreutils::Log(coreutils::LOG_INFO) << "Server key authenticated (client certificate required).";
      } else {
	 coreutils::Log(coreutils::LOG_INFO) << "Server key authenticated (no client certificate required).";
      }
      return newCtx;
   }

   int TLS::sniCallback(SSL *ssl, int *ad, void *arg) {
      TLS *self = (TLS *)arg;
      const char *name = SSL_get_servername(ssl, TLSEXT_NAMETYPE_host_name);
      if(name) {
	 auto it = self->sniContexts.find(name);
	 if(it != self->sniContexts.end())
	   SSL_set_SSL_CTX(ssl, it->second);
	 // No match (an unconfigured hostname, or a typo'd Host that still
	 // resolves here via DNS) - fall through and keep whatever SSL_CTX
	 // this SSL* already has, which is the base `ctx` SSL_new() was
	 // constructed from. Exactly the same cert a pre-SNI client (or
	 // one sending no SNI at all) already gets today.
      }
      return SSL_TLSEXT_ERR_OK;
   }

   void TLS::addSNIContext(const std::string &hostname, TLSInfo *tlsInfo) {
      SSL_CTX *newCtx = buildContext(tlsInfo);
      sniContexts[hostname] = newCtx;
      if(!sniCallbackRegistered) {
	 SSL_CTX_set_tlsext_servername_callback(ctx, sniCallback);
	 SSL_CTX_set_tlsext_servername_arg(ctx, this);
	 sniCallbackRegistered = true;
      }
   }


   void TLS::registerSocket(int fd) {
      int ret;

      coreutils::Log(coreutils::LOG_DEBUG_3) << "TLS socket initializing on socket " << fd << "...";      

      fcntl(fd, F_SETFL, fcntl(fd, F_GETFL, 0) | O_NONBLOCK);
      
      ssl = SSL_new(ctx);
//      if(ssl <= 0)
//	throw std::string("Error creating new TLS socket.");     
      
      SSL_set_info_callback(ssl, handshake_complete);
      
      if((ret = SSL_set_fd(ssl, fd)) == 0)
	throw std::string("Error setting TLS socket descriptor.");

//      if(!SSL_set_generate_session_id(ssl, generate_session_id))
//	throw std::string("Error setting session identifier callback.");

      // Was set unconditionally as the FIRST statement in this
      // function (before `ssl` was even assigned) - a flag meant to
      // mean "ssl is a valid, ready-to-use SSL*" needs to become true
      // only once that's actually true, at the end, not the start.
      // Previously set but never actually READ anywhere either -
      // sendData()/receiveData() below now check it before touching
      // `ssl` at all. See their own doc comments for the real,
      // confirmed crash this fixes: a brand-new TLS session's own
      // greeting (queued in onConnected(), see e.g. IMAPSession.h)
      // gets written via Socket::write() -> resetSocket(), which - deep
      // inside TCPServer::accept(), BEFORE TCPServer::connectionRequest()
      // ever calls this session's own performTLSHandshake() - makes the
      // fd immediately eligible for an EPOLLWRNORM event on a SEPARATE
      // epoll worker thread. That thread's TCPSocket::sendData() takes
      // ioLock (see its own doc comment) and calls straight into
      // SSL_write() on `ssl`, which is genuinely uninitialized garbage
      // (TLS.h explicitly documents `ssl` has no default member
      // initializer) until THIS function runs - ioLock alone doesn't
      // prevent it, since performTLSHandshake() (the only other place
      // that takes ioLock) hasn't even been called yet at that point.
      // Confirmed live via gdb under a mixed-protocol connection-churn
      // stress test: SIGSEGV inside libssl.so's SSL_write(), called from
      // TCPSocket::sendData() with ioGuard already held, while
      // registerSocket()/acceptSocket() for that same session were still
      // pending on the accepting thread.
      initialized = true;
   }
   
   void TLS::acceptSocket() {	

      switch (SSL_get_error(ssl, SSL_accept(ssl))) {
       case SSL_ERROR_SSL:
         coreutils::Log(coreutils::LOG_DEBUG_3) << "ERROR_SSL on ssl_accept. errno=" << errno;
         break;
       case SSL_ERROR_WANT_READ:
         coreutils::Log(coreutils::LOG_DEBUG_3) << "ERROR_WANT_READ on ssl_accept.";
         break;
       case SSL_ERROR_WANT_WRITE:
         coreutils::Log(coreutils::LOG_DEBUG_3) << "ERROR_WANT_WRITE on ssl_accept.";
         break;
       case SSL_ERROR_SYSCALL:
         coreutils::Log(coreutils::LOG_DEBUG_3) << "ERROR_SYSCALL on ssl_accept. errno=" << errno;
//         shutdown();
         break;
       default:
         coreutils::Log(coreutils::LOG_DEBUG_3) << "Unknown ERROR on ssl_accept.";
         break;
      }
      
   }
   
   int TLS::receiveData(coreutils::ZString &buffer) {

      // `ssl` is genuinely uninitialized until registerSocket() has
      // fully run (see its own doc comment on `initialized` for the
      // real crash this prevents) - a racing thread landing here first
      // just gets "nothing yet", exactly like a real SSL_ERROR_WANT_READ
      // would report, rather than touching a garbage pointer.
      if (!initialized)
         return -1;

      int len;
      coreutils::Log(coreutils::LOG_DEBUG_3) << "receiveData TLS on socket.";

      if((len = ::SSL_read(ssl, buffer.getData(), buffer.getLength())) >= 0) {
	 coreutils::Log(coreutils::LOG_DEBUG_3) << "receiveData TLS...len=" << len;
	 return len;
      }
      else {
	 switch (SSL_get_error(ssl, len)) {
	  case SSL_ERROR_SSL:
	    coreutils::Log(coreutils::LOG_DEBUG_3) << "ERROR_SSL on ssl_read. error=" << errno;
	    break;
	  case SSL_ERROR_WANT_READ:
	    coreutils::Log(coreutils::LOG_DEBUG_3) << "ERROR_WANT_READ on ssl_read.";
	    break;
	  case SSL_ERROR_WANT_WRITE:
	    coreutils::Log(coreutils::LOG_DEBUG_3) << "ERROR_WANT_WRITE on ssl_read.";
	    break;
	  case SSL_ERROR_SYSCALL:
	    coreutils::Log(coreutils::LOG_DEBUG_3) << "ERROR_SYSCALL on ssl_read. errno=" << errno;
	    break;
	  default:
	    coreutils::Log(coreutils::LOG_DEBUG_3) << "Unknown ERROR on ssl_read.";
	    break;
	 }
	 return -1;
      }

   }

   int TLS::sendData(const char *data, int length) {

      // Same "ssl isn't ready yet" guard as receiveData() - see
      // registerSocket()'s doc comment on `initialized` for the real,
      // confirmed crash this prevents (SIGSEGV in SSL_write() on a
      // still-uninitialized `ssl`, reached via a brand-new session's own
      // greeting racing ahead of its own TLS handshake). Returning <= 0
      // here is already the exact case Socket::writeSocket() treats as
      // "nothing sent this round, leave it queued and retry on the next
      // EPOLLWRNORM" - no special-casing needed at any call site.
      if (!initialized)
         return -1;

      int len;

      if((len = ::SSL_write(ssl, data, length)) >= 0) {
	 coreutils::Log(coreutils::LOG_DEBUG_3) << "sendData TLS...len=" << len;
	 lastSendWantedRead = false;
	 return len;
      }
      else {
	 lastSendWantedRead = false;
	 switch (SSL_get_error(ssl, len)) {
	  case SSL_ERROR_SSL:
	    coreutils::Log(coreutils::LOG_DEBUG_3) << "ERROR_SSL on ssl_write. error=" << errno;
	    break;
	  case SSL_ERROR_WANT_READ:
	    coreutils::Log(coreutils::LOG_DEBUG_3) << "ERROR_WANT_READ on ssl_write.";
	    lastSendWantedRead = true;
	    break;
	  case SSL_ERROR_WANT_WRITE:
	    coreutils::Log(coreutils::LOG_DEBUG_3) << "ERROR_WANT_WRITE on ssl_write.";
	    break;
	  case SSL_ERROR_SYSCALL:
	    coreutils::Log(coreutils::LOG_DEBUG_3) << "ERROR_SYSCALL on ssl_write. errno=" << errno;
	    break;
	  default:
	    coreutils::Log(coreutils::LOG_DEBUG_3) << "Unknown ERROR on ssl_write.";
	    break;
	 }
	 return -1;
      }

   }

}