/*

Copyright (C) 2015, 2016 Olaf Till <i7tiol@t-online.de>

This program is free software; you can redistribute it and/or modify
it under the terms of the GNU General Public License as published by
the Free Software Foundation; either version 3 of the License, or
(at your option) any later version.

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.  See the
GNU General Public License for more details.

You should have received a copy of the GNU General Public License
along with this program; If not, see <http://www.gnu.org/licenses/>.

*/

#ifndef __OCT_PARALLEL_STREAMS__

#define __OCT_PARALLEL_STREAMS__

#ifdef HAVE_LIBGNUTLS
#include <gnutls/gnutls.h>
#ifdef HAVE_LIBGNUTLS_EXTRA
#include <gnutls/extra.h>
#endif
#endif

#include <sys/socket.h>

#include <netinet/in.h>

#include <netdb.h>

#include <arpa/inet.h>

#include <string.h>

#include <poll.h>


enum encryption_type
  {
    ENCR_RAW,           // bare sockets
    ENCR_TLS_SRP        // gnutls TLS SRP
  };

#define PROTOCOL_VERSION 1


/*
Although we use the standard file format, we can't rely on
gnutls_srp_set_server_credentials_file(), since current gnutls uses
standard C library streams to read the files (external user space
buffers) and own automatic variables which are not zero'ed to hold the
sensitive data.
 */
#ifdef HAVE_LIBGNUTLS
extern "C" int octave_parallel_server_credentials_callback
(gnutls_session_t session, const char *username, gnutls_datum_t *salt,
 gnutls_datum_t *verifier, gnutls_datum_t *g, gnutls_datum_t *n);
#endif


// It is better to have classes for handling errors in input and
// output. For some tasks we can't use default standard C++
// io-classes: file-io with sensitive data should not be buffered,
// sockets have a system file-descriptor in the first place (and on
// some systems 'read' and 'write' may not work for them), and gnutls
// connections have special functions for sending and receiving. On
// the other hand, it is necessary to have a std::ostream and a
// std::istream for the connection, since Octaves functions for saving
// and loading data require them.
//
// To address these related problems, a common parent class, derived
// from std::streambuf, and several derived classes are declared
// here. Also, a class is provided which takes one of these
// streambuf objects and also contains a std::ostream and a
// std::istream which reference the streambuf object.

// parent class
class
octave_parallel_streambuf : public std::streambuf
{
public:

  octave_parallel_streambuf (void)
  : fid (-1), store_errno (0), readable (false), writeable (false)
  {
    // no character in the one-character-wide buffer at the beginning
    setg (&inbuf, (&inbuf) + 1, (&inbuf) + 1);
  }

  virtual ~octave_parallel_streambuf (void)
  {
    inbuf = '\0';
  }

  virtual bool good (void) = 0;

  virtual bool connection_state_undefined (void)
  {
    return false;
  }

  virtual void set_connection_state_undefined (void)
  {

  }

  void close (void)
  {
    if (fid >= 0)
      {
        int ret;

        // don't repeat interrupted close()
        ret = do_close ();

        if (ret == -1)
          store_errno = errno;

        fid = -1;
      }
  }

  ssize_t write (const void *buf, size_t count)
  {
    if (! good ())
      return -1;

    if (fid < 0)
      return -1;

    if (count == 0)
      return 0; // if nothing needs to be written, there's no need to
                // risk write() still being interrupted by a signal
    else if (count < 0)
      return -1;

    ssize_t ret;

    size_t left = count;

    // If write returns zero though count is positive, the reason must
    // be something different than being interrupted by a signal,
    // since in the latter case write would return -1 if no data was
    // written at all.
    do
      {
        if ((ret = do_write ((char *) buf + count - left, left)) > 0)
          left -= ret;
      }
    while (left && ret > 0); // don't repeat interrupted write()
    // while (left &&
    //        ((ret == -1 && errno == EINTR) ||
    //         ret > 0));

    if (ret == -1)
      store_errno = errno;
    else if (left)
      store_errno = -1;

    if (store_errno)
      return -1;
    else
      return count;
  }

  virtual int gnutls_write_alert_for_flushing (void)
  {
    if (! good () || fid < 0)
      return -1;
    else
      return 0;
  }

  ssize_t read (void *buf, size_t count)
  {
    if (! good ())
      return -1;

    if (fid < 0)
      return -1;

    if (count == 0)
      return 0;
    else if (count < 0)
      return -1;

    int ndiff = 0;
    if (gptr () != egptr ())
      {
        *((char *) buf) = *(gptr ());
        gbump (1);
        buf = (void *)(((char *) buf) + 1);
        count--;
        ndiff = 1;
      }

    int ret = read_directly (buf, count);
    if (ret > -1)
      ret += ndiff;
    return ret;
  }

  bool readable;

  bool writeable;

  int get_fid (void)
  {
    return fid;
  }

  // redefined streambuf virtual functions

  std::streamsize xsputn (const char *buf, std::streamsize count)
  {
    ssize_t ret = write ((const void *) buf, count);

    return ret >= 0 ? ret : 0;
  }

  int overflow (int c = traits_type::eof ())
  {
    if (c != traits_type::eof ())
      {
        traits_type::char_type tc = traits_type::to_char_type (c);

        if (write ((void *) &tc, sizeof (traits_type::char_type)) == -1)
          return traits_type::eof ();
      }

    return traits_type::to_int_type (c);
  }

  std::streamsize xsgetn (char *buf, std::streamsize count)
  {
    ssize_t ret = read ((void *) buf, count);

    return ret >= 0 ? ret : 0;
  }

  // the only member capable of storing something in the input buffer;
  // and the only reading member which never increases gptr()
  int underflow (void)
  {
    if (gptr () != egptr ())
      return traits_type::to_int_type (*(gptr ()));
    else
      {
        if (read_directly ((void *) eback (), 1) == -1)
          return traits_type::to_int_type (traits_type::eof ());
        else
          {
            gbump (-1);
            return traits_type::to_int_type (*(gptr ()));
          }
      }
  }

  int uflow (void)
  {
    if (gptr () != egptr ())
      {
        gbump (1);
        return traits_type::to_int_type (*(gptr () - 1));
      }
    else
      {
        traits_type::char_type ret;

        if (read_directly ((void *) &ret, sizeof (traits_type::char_type))
            == -1)
          return traits_type::to_int_type (traits_type::eof ());
        else
          return traits_type::to_int_type (ret);
      }
  }

  int network_send_4byteint (uint32_t v, bool)
  {
    uint32_t t = htonl (v);

    write ((void *) &t, 4);

    if (good ())
      return 0;
    else
      return -1;
  }

  int network_recv_4byteint (uint32_t &v)
  {
    uint32_t t;

    read ((void *) &t, 4);

    if (! good ())
      return -1;

    v = ntohl (t);

    return 0;
  }

  int network_send_4byteint (int32_t v)
  {
    if (signed_int_rep () && v < 0)
      {
        once_gripe_neg_int_rep ();

        return -1;
      }

    union
    {
      int32_t s;
      uint32_t u;
    } cast;

    cast.s = v;

    return network_send_4byteint (cast.u, true);
  }

  int network_recv_4byteint (int32_t &v)
  {
    uint32_t t;

    union
    {
      int32_t s;
      uint32_t u;
    } cast;


    read ((void *) &t, 4);

    if (! good ())
      return -1;

    cast.u = ntohl (t);

    v = cast.s;

    if (signed_int_rep () && v < 0)
      {
        once_gripe_neg_int_rep ();

        return -1;
      }

    return 0;
  }

  int network_send_string (const char *s)
  {
    uint32_t l = strlen (s);

    uint32_t t = htonl (l);

    write ((void *) &t, 4);

    write ((const void *) s, l);

    if (good ())
      return 0;
    else
      return -1;
  }

  int network_recv_string (char *s, int maxl)
  {
    uint32_t l, t;

    read ((void *) &t, 4);

    if (! good ())
      return -1;

    l = ntohl (t);

    if (l > maxl - 1)
      {
        c_verror ("string to be received would be longer than buffer");

        return -1;
      }

    s[l] = '\0';

    read ((void *) s, l);

    if (good ())
      return 0;
    else
      return -1;
  }

  int network_recv_string (std::string &str)
  {
    uint32_t l, t;

    read ((void *) &t, 4);

    if (! good ())
      return -1;

    l = ntohl (t);

    struct __bufguard
    {
      __bufguard (octave_idx_type __l) : __buf (new char[__l]) {}
      ~__bufguard (void) { delete [] __buf; }
      char* __get (void) { return __buf; }
      char *__buf;
    } __bufg (l);

    read (__bufg.__get (), l);

    std::string tstr (__bufg.__get (), l);

    str = tstr;

    if (good ())
      return 0;
    else
      return -1;
  }    

  static int signed_int_rep (void)
  {
    // I know this is paranoid, and that current Octave itself just
    // relies on twos complement in saving/loading signed integers.

    static int ret = -1;

    if (ret == -1)
      {
        union
        {
          uint32_t u;
          int32_t s;
        } t;

        t.s = -1;

        if (t.u != 0xFFFFFFFF)
          ret = 1;
        else
          ret = 0;
      }

    return ret;
  }

protected:

  int fid;

  int store_errno;

  virtual ssize_t do_write (const void *buf, size_t count) = 0;

  virtual ssize_t do_read (void *buf, size_t count) = 0;

  virtual int do_close (void) = 0;

private:

  ssize_t read_directly (void *buf, size_t count)
  {
    if (! good ())
      return -1;

    if (fid < 0)
      return -1;

    if (count == 0)
      return 0; // if nothing needs to be written, there's no need to
                // risk read() still being interrupted by a signal
    else if (count < 0)
      return -1;

    ssize_t ret;

    size_t left = count;

    // If read returns zero though count is positive, the reason must
    // be something different than being interrupted by a signal,
    // since in the latter case read would return -1 if no data was
    // written at all.
    do
      {
        if ((ret = do_read ((char *) buf + count - left, left)) > 0)
          left -= ret;
      }
    while (left && ret > 0); // don't repeat interrupted read()
    // while (left &&
    //        ((ret == -1 && errno == EINTR) ||
    //         ret > 0));

    if (ret == -1)
      store_errno = errno;
    else if (left)
      store_errno = -1;

    if (store_errno)
      return -1;
    else
      return count;
  }

  // We need to provide an input buffer of at least one character to
  // be able to implement underflow(), needed by sgetc(), possibly
  // needed by istream.getline(). This buffer must be zero'ed at
  // destruction.
  char inbuf;

  void once_gripe_neg_int_rep (void)
  {
    static bool first_time = true;

    if (first_time)
      {
        c_verror ("This machine doesn't seem to use twos complement as negative integer representation. If you want this to be supported, please file a bug report.");

        first_time = false;
      }
  }
}; // class octave_parallel_streambuf

// To avoid user space buffers, provided by external libraries, which
// can not be zero'ed after use, we have to directly use system calls
// for reading and writing sensitive data from/to files. Luckily, for
// this small amount of file-io we don't need to construct a custom
// user space buffer.
class
octave_parallel_file_streambuf : public octave_parallel_streambuf
{
public:

  octave_parallel_file_streambuf (const char *path, int flags, mode_t mode = 0)
  {
    while ((fid = open (path, flags, mode)) == -1 && errno == EINTR);

    if (fid == -1)
      store_errno = errno;
    else
      {
        // care is needed here, one of the flag constants can
        // represent the value zero (O_RDONLY on my system), so that
        // checking for (flags & flag_constant) does not work

        if ((O_RDONLY == 0 && ! (flags & O_WRONLY)) ||
            (O_RDONLY && flags & (O_RDONLY | O_RDWR)))
          readable = true;

        if ((O_WRONLY == 0 && ! (flags & O_RDONLY)) ||
            (O_WRONLY && flags & (O_WRONLY | O_RDWR)))
          writeable = true;
      }
  }

  ~octave_parallel_file_streambuf (void)
  {
    if (fid >= 0)
      while (::close (fid) == -1 && errno == EINTR);
  }

  bool good (void)
  {
    return (store_errno == 0);
  }

protected:

  ssize_t do_read (void *buf, size_t count)
  {
    return ::read (fid, buf, count);
  }

  ssize_t do_write (const void *buf, size_t count)
  {
    return ::write (fid, buf, count);
  }

  int do_close (void)
  {
    return ::close (fid);
  }
}; // class octave_parallel_file_streambuf

#define BEGIN_INTERRUPT_IMMEDIATELY_IN_FOREIGN_CODE_PSSTREAMBUF \
  BEGIN_INTERRUPT_IMMEDIATELY_IN_FOREIGN_CODE_1; \
  undefined_connection_state = true; \
  octave_rethrow_exception (); \
  BEGIN_INTERRUPT_IMMEDIATELY_IN_FOREIGN_CODE_2

// TCP sockets without authentication or encryption, as an alternative
// to encrypted connections (currently only TLS-SRP), which are (is)
// derived from this.
class
octave_parallel_socket_streambuf : public octave_parallel_streambuf
{
public:

  // this needs c++11
  octave_parallel_socket_streambuf (int sid, bool server)
    : octave_parallel_socket_streambuf (sid, server, ENCR_RAW)
  {

  }

  // This is extra, so that (derived) encrypted connection streambufs
  // can give it their own 'encrtype' argument to communicate and
  // check encryption type before encryption starts.
  octave_parallel_socket_streambuf (int sid, bool server,
                                    encryption_type encrtype);

  ~octave_parallel_socket_streambuf (void)
  {
    if (fid >= 0)
      while (::close (fid) == -1 && errno == EINTR);
  }

  bool good (void)
  {
    return (store_errno == 0);
  }

  bool connection_state_undefined (void)
  {
    return undefined_connection_state;
  }

  void set_connection_state_undefined (void)
  {
    undefined_connection_state = true;
  }

  // redefined streambuf virtual functions

  std::streamsize showmanyc ()
  {
    if (! good () || fid == -1)
      return -1;

    struct pollfd pfd;
    pfd.fd = fid;
    pfd.events = POLLIN;
    pfd.revents = 0;

    std::streamsize ret;
    if ((ret = poll (&pfd, 1, 0)) == 1 && pfd.revents & POLLIN)
      return 1;
    else
      {
        if (ret < 0)
          store_errno = -1;
        return ret;
      }
  }

protected:

  ssize_t do_read (void *buf, size_t count)
  {
    return ::recv (fid, buf, count, 0);
  }

  ssize_t do_write (const void *buf, size_t count)
  {
    return ::send (fid, buf, count, 0);
  }

  int do_close (void)
  {
    return ::close (fid);
  }

  bool undefined_connection_state;
};

// for usage with class octave_parallel_gnutls_streambuf, one parent
// class and one derived class for server and client, respectively;
// objects meant to be allocated; the class storing the pointer should
// check the refcount and care for deallocation
#ifdef HAVE_LIBGNUTLS
class
octave_parallel_gnutls_srp_credentials
{
public:

  octave_parallel_gnutls_srp_credentials (void)
  : cred (NULL), refcount (0)
  {

  }

  virtual ~octave_parallel_gnutls_srp_credentials (void)
  {

  }

  void *get_ref (void)
  {
    refcount++;
    dlprintf ("cred incr its refcount to %i\n", refcount);

    return cred;
  }

  int release_ref (void)
  {
    dlprintf ("cred will decr its refcount to %i\n", refcount - 1);
    return --refcount;
  }

  int check_ref (void)
  {
    dlprintf ("cred refcount was checked, is %i\n", refcount);
    return refcount;
  }

  void *check_cred (void)
  {
    return cred;
  }

protected:

  void *cred;

  virtual void free_credentials (void) = 0;

private:

  int refcount;
};

// see comment of parent class
class
octave_parallel_gnutls_srp_client_credentials :
  public octave_parallel_gnutls_srp_credentials
{
public:

  octave_parallel_gnutls_srp_client_credentials (const char *user,
                                                 char *apasswd = NULL);

  octave_parallel_gnutls_srp_client_credentials (std::string passwd_file);

  ~octave_parallel_gnutls_srp_client_credentials (void)
  {
    dlprintf ("ccred destr called\n");
    if (password)
      {
        memset ((void *) password, '\0', strlen (password));

        if (pw_allocated)
          delete [] password;
      }

    if (cred)
      {
        dlprintf ("ccred destr will call 'free_credentials()'\n");
        free_credentials ();
      }
  }

  const char *get_passwd (void)
  {
    return password;
  }

protected:

  void free_credentials (void)
  {
    gnutls_srp_free_client_credentials
      ((gnutls_srp_client_credentials_t) cred);
    cred = NULL;
  }

private:

  char *password;

  bool pw_allocated;
};

// see comment of parent class
class
octave_parallel_gnutls_srp_server_credentials :
  public octave_parallel_gnutls_srp_credentials
{
public:

  octave_parallel_gnutls_srp_server_credentials
  (std::string passwd_file)
  {
    if (! gnutls_srp_allocate_server_credentials
        ((gnutls_srp_server_credentials_t *) &cred))
      {
        // initialize filename in callback function
        octave_parallel_server_credentials_callback (NULL, passwd_file.c_str (),
                                                     NULL, NULL, NULL, NULL);

        gnutls_srp_set_server_credentials_function
          ((gnutls_srp_server_credentials_t) cred,
           octave_parallel_server_credentials_callback);
      }
  }

  ~octave_parallel_gnutls_srp_server_credentials (void)
  {
    if (cred)
      free_credentials ();
  }

protected:

  void free_credentials (void)
  {
    gnutls_srp_free_server_credentials
      ((gnutls_srp_server_credentials_t) cred);
    cred = NULL;
  }
};

// TCP sockets with a gnutls connection.
class
octave_parallel_gnutls_streambuf : public octave_parallel_socket_streambuf
{

public:

  octave_parallel_gnutls_streambuf
  (int sid, octave_parallel_gnutls_srp_credentials *accred, bool server);

  ~octave_parallel_gnutls_streambuf (void)
  {
    if (session)
      gnutls_deinit (session);

    if (ccred)
      dlprintf ("gstreambuf destructor will requcest decr of cred refcount\n");
    int tpdebug;
    if (ccred && (tpdebug = ccred->release_ref ()) <= 0)
      {
        dlprintf ("gstreambuf destructor got decr cred refcount of %i and will delete cred\n", tpdebug);
        delete ccred;
      }
  }

  bool good (void)
  {
    return (store_errno == 0);
  }

  const char *server_get_username (void)
  {
    return gnutls_srp_server_get_username (session);
  }

  int gnutls_write_alert_for_flushing (void)
  {
    if (! good () || fid < 0)
      return -1;

    int ret;
    // this is a dummy alert to force flushing
    while ((ret = gnutls_alert_send
            (session, GNUTLS_AL_WARNING, GNUTLS_A_INTERNAL_ERROR)) ==
           GNUTLS_E_AGAIN); // don't repeat if interrupted
    // GNUTLS_E_INTERRUPTED ||
    // ret == GNUTLS_E_AGAIN);

    if (ret != GNUTLS_E_SUCCESS)
      {
        store_errno = -1;
        return -1;
      }
    else
      return 0;
  }

protected:

  ssize_t do_read (void *buf, size_t count)
  {
    ssize_t ret;

    while ((ret = gnutls_record_recv (session, buf, count)) ==
           GNUTLS_E_WARNING_ALERT_RECEIVED)
      gnutls_alert_get (session); // discard the alert

    if (ret < 0)
      {
        // although checking GNUTLS_E_AGAIN is probably not necessary
        // with blocking calls configured
        if (ret == GNUTLS_E_INTERRUPTED || ret == GNUTLS_E_AGAIN)
          errno = EINTR;

        ret = -1;
      }

    return ret;
  }

  ssize_t do_write (const void *buf, size_t count)
  {
    ssize_t ret;

    while ((ret = gnutls_record_send (session, buf, count)) ==
           GNUTLS_E_WARNING_ALERT_RECEIVED)
      gnutls_alert_get (session); // discard the alert

    if (ret < 0)
      {
        // although checking GNUTLS_E_AGAIN is probably not necessary
        // with blocking calls configured
        if (ret == GNUTLS_E_INTERRUPTED || ret == GNUTLS_E_AGAIN)
          errno = EINTR;

        ret = -1;
      }

    return ret;
  }

  int do_close (void)
  {
    if (session)
      {
        gnutls_deinit (session);

        session = NULL;
      }

    if (ccred && ccred->release_ref () <= 0)
      delete ccred;

    ccred = NULL;

    return octave_parallel_socket_streambuf::do_close ();
  }

  octave_parallel_gnutls_srp_credentials *ccred;

  void *cred;

  gnutls_session_t session;
};

#endif // HAVE_LIBGNUTLS

class
octave_parallel_stream
{
public:

  octave_parallel_stream (octave_parallel_streambuf *strbuf_arg)
  : strbuf (strbuf_arg), istr (NULL), ostr (NULL)
  {
    if (strbuf->readable)
      istr = new std::istream (strbuf);

    if (strbuf->writeable)
      ostr = new std::ostream (strbuf);
  }

  ~octave_parallel_stream (void)
  {
    if (istr)
      delete istr;

    if (ostr)
      delete ostr;

    dlprintf ("pstream destructor will delete streambuf\n");
    delete strbuf;
  }

  bool good (void) const
  {
    return strbuf->good ();
  }

  void close (void) const
  {
    if (strbuf)
      strbuf->close ();
  }

  std::istream &get_istream (void) const
  {
    return *istr;
  }

  std::ostream &get_ostream (void) const
  {
    return *ostr;
  }

  int network_send_4byteint (uint32_t v, bool) const
  {
    return strbuf->network_send_4byteint (v, true);
  }

  int network_recv_4byteint (uint32_t &v) const
  {
    return strbuf->network_recv_4byteint (v);
  }

  int network_send_4byteint (int32_t v) const
  {
    return strbuf->network_send_4byteint (v);
  }

  int network_recv_4byteint (int32_t &v) const
  {
    return strbuf->network_recv_4byteint (v);
  }

  int network_send_string (const char *s) const
  {
    return strbuf->network_send_string (s);
  }

  int network_recv_string (char *s, int maxl) const
  {
    return strbuf->network_recv_string (s, maxl);
  }

  int network_recv_string (std::string &str) const
  {
    return strbuf->network_recv_string (str);
  }    

  octave_parallel_streambuf *get_strbuf (void)
  {
    return strbuf;
  }

  static int signed_int_rep (void)
  {
    return octave_parallel_streambuf::signed_int_rep ();
  }

private:

  octave_parallel_streambuf *strbuf;

  std::istream *istr;

  std::ostream *ostr;
};

#endif // __OCT_PARALLEL_STREAMS__


/*
;;; Local Variables: ***
;;; mode: C++ ***
;;; End: ***
*/
