/*
 *  lib/comm.c
 *
 *  Communication subsystem
 *
 *  Copyright (C) 2004  Phoenix Grid Computing Project
 *                      ( http://www.logos.ic.i.u-tokyo.ac.jp/phoenix/ )
 *
 *  2004-12-07  Comm_select () added 
 *              by Kaneda Kenji
 *  2004-11-26  Implemented basic functions 
 *              by Kaneda Kenji
 */


#include "common.h"
#include "std.h"
#include "debug.h"
#include "xssl.h"
#include "opt.h"
#include "comm/common.h"
#include <string.h>
#include <sys/types.h>
#include <sys/utsname.h>
#include <unistd.h>
#include <errno.h>



static void
init_ssl_library ( void ) 
{
	/* Register the available ciphers and digests.
	 * [Note] This function must be called before any other action takes place. */
	SSL_library_init ( );

	/* Register the error strings for all libcrypto functions as
	 * well as the libssl error strings. */
	SSL_load_error_strings ( );
}

static SSL_CTX *
create_ctx_with_sslv3 ( void )
{
	SSL_METHOD *meth;
	SSL_CTX *ctx;
	
	/* Constructor for the SSLv3 SSL_METHOD structure for combined client and server. */
	meth = SSLv3_method ();
	
	/* Create a new SSL_CTX object. 
	 * [Note] only understand the SSLv3 protocol. */
	ctx = xSSL_CTX_new ( meth );

	return ctx;
}


/* Store a user password pointed by <userdata> to <buf>. */
static int
password_cb ( char *buf, int num, int rwflag, void *userdata )
{
	size_t n;
	char *password = ( char * ) userdata;

	assert ( buf != NULL );
	assert ( userdata != NULL );
	
	// printf ( "cb: password = \"%s\" \n", password ) ;

	n = strlen ( password );

	if ( num < n + 1 ) /* No enough space for <password> */
		return 0;

	memmove ( buf, password, n + 1 );

	return n;
}

static void
init_password ( SSL_CTX *ctx, const char *password )
{
	char *ud;

	assert ( ctx != NULL );
	assert ( password != NULL );
	
	ud = Strdup ( password );

	/* Set the default password callback called when
	 * loading/storing a PEM certificate with encryption. */
	SSL_CTX_set_default_passwd_cb_userdata ( ctx, ud );

	SSL_CTX_set_default_passwd_cb ( ctx, &password_cb );
}



#if 0

/* ???? */
/* Check that the common name matches the host name*/
void
check_cert_chain ( SSL *ssl, char *host )
{
	X509 *peer;
	char peer_CN[256];
	
	/* Return the result of the verification of the X509
	 * certificate presented by the peer, if any. */
	if ( SSL_get_verify_result (ssl) != X509_V_OK )
		ssl_failure ("check_cert_chain");

	/*Check the cert chain. The chain length
	  is automatically checked by OpenSSL when we
	  set the verify depth in the ctx */
	
	/* Check the common name*/
	peer = SSL_get_peer_certificate ( ssl );

        int i;
        X509_NAME_ENTRY *e;

	/* Process all commonName entries: */

        int loc;


	X509_NAME *name = X509_get_subject_name (peer);

        loc = -1;
	for ( ; ; ) {
		X509_NAME_ENTRY *e;

		loc = X509_NAME_get_index_by_NID( name, NID_commonName, loc );
		if (loc == -1)
			break;
		
		 e = X509_NAME_get_entry(nm, loc);
		/* Do something with <e> */


		 /*
		   xX509_NAME_get_text_by_NID ( name, NID_commonName, peer_CN, 256);
		   if (strcasecmp (peer_CN,host))
		   err_exit ("Common name doesn't match host name");
		 */
	}
}

#endif


/* あとでどの設定が重要か調べ直す */
static void
init_cert_verification ( SSL_CTX *ctx, const struct opt_t *opt )
{
	int mode = SSL_VERIFY_PEER | SSL_VERIFY_FAIL_IF_NO_PEER_CERT;

	assert ( ctx != NULL );
	assert ( opt != NULL );

	
	/* Set the verification flag. */

	/* < SSL_VERIFY_PEER >
           [Server mode] the server sends a client certificate request to the
           client.  The certificate returned (if any) is checked. If the veri-
           fication process fails, the TLS/SSL handshake is immediately termi-
           nated with an alert message containing the reason for the verifica-
           tion failure.  The behaviour can be controlled by the additional
           SSL_VERIFY_FAIL_IF_NO_PEER_CERT and SSL_VERIFY_CLIENT_ONCE flags.

           [Client mode] the server certificate is verified. If the verifica-
           tion process fails, the TLS/SSL handshake is immediately terminated
           with an alert message containing the reason for the verification
           failure. If no server certificate is sent, because an anonymous
           cipher is used, SSL_VERIFY_PEER is ignored. */

	/* < SSL_VERIFY_FAIL_IF_NO_PEER_CERT >
	   [Server mode] if the client did not return a certificate, the
           TLS/SSL handshake is immediately terminated with a "handshake fail-
           ure" alert.  This flag must be used together with SSL_VERIFY_PEER.

           [Client mode] ignored */
	
	/* サーバ側も，クライアントの認証を行うようにする．*/

	SSL_CTX_set_verify (ctx, mode, NULL /* use the default verification procedure */ );

	/* Set the maximum depth for the certificate chain
	 * verification that shall be allowed for <ctx>. */
	SSL_CTX_set_verify_depth ( ctx, 1 );

	xSSL_CTX_load_verify_locations ( ctx, opt->ca_cert, NULL );


	/* Set list of CAs sent to the client when requesting a client certificate */	
	{
		STACK_OF(X509_NAME) *cert_names;

		cert_names = xSSL_load_client_CA_file ( opt->ca_cert );
		SSL_CTX_set_client_CA_list ( ctx, cert_names );
	}
}

static void
init_rsa ( SSL_CTX *ctx )
{
	RSA *rsa;
		
	assert ( ctx != NULL );

	rsa = xRSA_generate_key ( 512, RSA_F4, NULL, NULL );
	
	xSSL_CTX_set_tmp_rsa ( ctx, rsa );
	
	RSA_free(rsa);	
}

static DH * 
set_dhe1024(void)
{
	enum { BUF_SIZE = 20 };
	DH *dhparams;
	DSA *dsaparams;
	const char *seed[] = { ";-)  :-(  :-)  :-(  ",
			       ";-)  :-(  :-)  :-(  ",
			       "Random String no. 12",
			       ";-)  :-(  :-)  :-(  ",
			       "hackers have even mo", /* from jargon file */
	};
	const int NELM = ( sizeof ( seed ) ) / ( sizeof ( seed[0] ) );
	unsigned char seedbuf[BUF_SIZE];
	int i;
	
	RAND_bytes ( (unsigned char *) &i, sizeof ( i ) );
	if ( i < 0 ) { i = -i; }
	
	i %= NELM;
	assert ( strlen ( seed[i] ) == BUF_SIZE );
	memmove( seedbuf, seed[i], BUF_SIZE );
	
	dsaparams = xDSA_generate_parameters ( 1024, seedbuf, BUF_SIZE, 
					       NULL, NULL, NULL, NULL );
	dhparams = xDSA_dup_DH(dsaparams);

	DSA_free(dsaparams);

	return dhparams;
}

static void
init_dh ( SSL_CTX *ctx )
{
	/* DH (Diffie-Hellman) */
	DH *dh;
	
	assert ( ctx != NULL );

	dh = set_dhe1024 ();
	
	xSSL_CTX_set_tmp_dh ( ctx, dh );
	
	DH_free ( dh );
	
	/* avoid small subgroup attacks: */
	SSL_CTX_set_options ( ctx, SSL_OP_SINGLE_DH_USE );
}


static void
init_rand_seed_uniquely ( void )
{
    struct {
	    pid_t pid;
	    time_t time;
	    void *stack;
    } data;

    data.pid = getpid();
    data.time = time(NULL);
    data.stack = (void *)&data;

    RAND_seed ( ( const void * ) &data, sizeof ( data) );
}

void
init_rand_seed ( void )
{
    struct {
	    struct utsname uname;
	    int uname_1;
	    int uname_2;
	    uid_t uid;
	    uid_t euid;
	    gid_t gid;
	    gid_t egid;
    } data;
    
    data.uname_1 = uname ( &data.uname );
    data.uname_2 = errno;
    
    data.uid = getuid ( );
    data.euid = geteuid ( );
    data.gid = getgid ( );
    data.egid = getegid ( );
    
    RAND_seed ( (const void * ) &data, sizeof ( data) );
    init_rand_seed_uniquely ( );
}

static void
init_random ( SSL_CTX *ctx, const struct opt_t *opt )
{
	assert ( ctx != NULL );
	assert ( opt != NULL );
	
	init_rand_seed ( );
	xRAND_load_file ( opt->priv_key, LONG_MAX );

	init_rsa ( ctx );
	// init_dh ( ctx );
}

SSL_CTX *
init_ctx ( const struct opt_t *opt )
{
	SSL_CTX *ctx;;
	const int mode = SSL_MODE_ENABLE_PARTIAL_WRITE | SSL_MODE_AUTO_RETRY;

	assert ( opt != NULL );

	init_ssl_library ( );

	ctx = create_ctx_with_sslv3 ( );
	
	/* Allow SSL_write(..., n) to return r with 0 < r < n
	 * (i.e. report success when just a single record has been
	 * written). */
	SSL_CTX_set_mode( ctx, mode );


	init_password ( ctx, opt->password );
	xSSL_CTX_use_PrivateKey_file ( ctx, opt->priv_key, SSL_FILETYPE_PEM );

	xSSL_CTX_use_certificate_chain_file ( ctx, opt->certificate );
	xSSL_CTX_check_private_key ( ctx );
	
	init_cert_verification ( ctx, opt );

	init_random ( ctx, opt );

	return ctx;
}

void 
destroy_ctx ( SSL_CTX *ctx )
{
	assert ( ctx != NULL );

	/* Decrement the reference count of <ctx>, and removes the
	 * SSL_CTX object pointed to by <ctx> and frees up the
	 * allocated memory if the reference count has reached 0. */
	SSL_CTX_free (ctx);	
}

struct comm_t *
Comm_new ( const struct opt_t *opt )
{
	struct comm_t *x;

	x = ( struct comm_t *) Malloc ( sizeof ( struct comm_t ) );

	x->ctx = init_ctx ( opt );
	x->conns = ConnList_new ( );

	return x;
}

void 
Comm_destroy ( struct comm_t *x )
{
	assert ( x != NULL );

	destroy_ctx ( x->ctx );
	ConnList_destroy ( x->conns );

	Free ( x );
}

/* Add a new connection */
struct conn_t *
find_conn ( struct comm_t *x, int sockfd )
{
	assert ( x != NULL );

	return ConnList_find ( x->conns, sockfd );
}

/* Add a new connection */
void
add_conn ( struct comm_t *x, struct conn_t *conn )
{
	assert ( x != NULL );
	assert ( conn != NULL );

	ConnList_add ( x->conns, conn );
}

/* Delete a connection of which sockfd is equal to <sockfd> */
void
detele_conn ( struct comm_t *x, int sockfd )
{
	assert ( x != NULL );

	ConnList_delete ( x->conns, sockfd );
}

int
Comm_connect ( struct comm_t *comm, int sockfd, const struct sockaddr *serv_addr, socklen_t addrlen )
{
	int ret;
	struct conn_t *conn;

	assert ( comm != NULL );

	conn = Conn_connect ( comm->ctx, sockfd, serv_addr, addrlen, &ret );

	if ( conn != NULL )
		add_conn ( comm, conn );

	return ret;
}

int
Comm_accept ( struct comm_t *comm, int s, struct sockaddr *addr, socklen_t *addrlen )
{
	int ret;
	struct conn_t *conn;

	assert ( comm != NULL );

	conn = Conn_accept ( comm->ctx, s, addr, addrlen, &ret );

	if ( conn != NULL )
		add_conn ( comm, conn );

	return ret;
}

ssize_t
Comm_send ( struct comm_t *comm, int s, const void *buf, size_t len, int flags )
{
	struct conn_t *conn;

	assert ( comm != NULL );

	conn = find_conn ( comm, s );

	if ( conn == NULL )
		return -1;

	return Conn_send ( conn, buf, len, flags );
}

ssize_t
Comm_recv ( struct comm_t *comm, int s, void *buf, size_t len, int flags )
{
	struct conn_t *conn;

	assert ( comm != NULL );

	conn = find_conn ( comm, s );
	if ( conn == NULL )
		return -1;

	return Conn_recv ( conn, buf, len, flags );
}

/* コネクションを table から取り除くのを，どのタイミングで行うようにするか */
int
Comm_shutdown ( struct comm_t *comm, int s, int how )
{
	struct conn_t *conn;

	assert ( comm != NULL );

	conn = find_conn ( comm, s );

	if ( conn == NULL )
		return -1;

	return Conn_shutdown ( conn, how );
}

/**********************/

static int
check_pending_msgs ( struct comm_t *comm, 
		     int max_fd, fd_set *readfds, fd_set *writefds, fd_set *exceptfds )
{
	int ret;
	int i;
	fd_set ret_fds;

	assert ( comm != NULL );
	assert ( readfds != NULL );
	
	FD_ZERO ( &ret_fds );
	
	ret = 0;
	for ( i = 0; i < max_fd; i++ ) {
		struct conn_t *conn;
		
		if ( ! FD_ISSET ( i, readfds ) ) 
			continue;	
		
		conn = find_conn ( comm, i );
		
		if ( ( conn == NULL ) || ( Conn_pending ( conn ) == 0 ) )
			continue;
		
		ret++;
		FD_SET ( i, &ret_fds );
	}
	
	if ( ret == 0 )
		return 0;

	/* overwrite the fd sets */
	*readfds = ret_fds;
	if ( writefds != NULL )
		FD_ZERO ( writefds );
	if ( exceptfds != NULL )
		FD_ZERO ( exceptfds );
	
	return ret;
}

#if 1

int
Comm_select ( struct comm_t *comm, 
	      int n, fd_set *readfds, fd_set *writefds, fd_set *exceptfds,
	      struct timeval *timeout )
{
	int ret;

	assert ( comm != NULL );

	if ( readfds == NULL )
		return select ( n, readfds, writefds, exceptfds, timeout );

	ret = check_pending_msgs ( comm, n, readfds, writefds, exceptfds );
	if ( ret > 0 )
		return ret;

	ret = select ( n, readfds, writefds, exceptfds, timeout );

	/* [TODO]
	 * For some file descriptor, 
	 * the entire record has not yet arrived even if it is set in <readfds>.
	 * ( Solve this problem by non blocking socket??? )
	 */

	return ret;
}

#else

static void
try_save_fd_set ( fd_set *bkup, fd_set *orig )
{
	assert ( bkup != NULL );

	if ( orig != NULL ) {
		*bkup = *orig;
	} else {
		FD_ZERO ( bkup );
	}
}

static void
save_fd_sets ( fd_set *saved_readfds, fd_set *saved_writefds, fd_set *saved_exceptfds,
	       fd_set *readfds, fd_set *writefds, fd_set *exceptfds )
{
	assert ( saved_readfds != NULL );
	assert ( saved_writefds != NULL );
	assert ( saved_exceptfds != NULL );
	assert ( readfds != NULL );

	*saved_readfds = *readfds;
	try_save_fd_set ( saved_writefds, writefds );
	try_save_fd_set ( saved_exceptfds, exceptfds );
}

static void
try_restore_fd_set ( fd_set *bkup, fd_set *orig )
{
	assert ( bkup != NULL );

	if ( orig != NULL )
		*orig = *bkup;
}

static void
restore_fd_sets ( fd_set *saved_readfds, fd_set *saved_writefds, fd_set *saved_exceptfds,
		  fd_set *readfds, fd_set *writefds, fd_set *exceptfds )
{
	assert ( saved_readfds != NULL );
	assert ( saved_writefds != NULL );
	assert ( saved_exceptfds != NULL );
	assert ( readfds != NULL );

	*readfds = *saved_readfds;
	try_restore_fd_set ( saved_writefds, writefds );
	try_restore_fd_set ( saved_exceptfds, exceptfds );
}

static int
remove_unpending_fds ( struct comm_t *comm, int max_fd, fd_set *readfds )
{
	int ret;
	int i;

	assert ( comm != NULL );
	assert ( readfds != NULL );

	ret = 0;
	for ( i = 0; i < max_fd; i++ ) {
		struct conn_t *conn;
			
		if ( ! FD_ISSET ( i, readfds ) ) 
			continue;
		
		conn = find_conn ( comm, i );
		
		if ( ( conn == NULL ) || ( Conn_pending ( conn ) > 0 ) )
			continue;
		
		/* There is no pending messages in <i> even if it is set in <readfds>. */
		FD_CLR ( i,  readfds );
		ret++;
	}

	return ret;
}

static bool_t
is_timeout ( struct timeval *timeout )
{
	return ( ( timeout != NULL) &&
		 ( timeout->tv_sec <= 0 ) && 
		 ( timeout->tv_usec <= 0 ) );	
}

static bool_t
select_has_finished ( int ret, struct timeval *timeout )
{		
	return ( ( ret != 0 ) || ( is_timeout ( timeout ) ) );
}

int
Comm_select ( struct comm_t *comm, 
	      int n, fd_set *readfds, fd_set *writefds, fd_set *exceptfds,
	      struct timeval *timeout )
{
	int ret;

	assert ( comm != NULL );

	if ( readfds == NULL )
		return select ( n, readfds, writefds, exceptfds, timeout );

	ret = check_pending_msgs ( comm, n, readfds, writefds, exceptfds );
	if ( ret > 0 )
		return ret;


	fd_set saved_readfds, saved_writefds, saved_exceptfds;

	save_fd_sets ( &saved_readfds, &saved_writefds, &saved_exceptfds,
		       readfds, writefds, exceptfds );

	for ( ; ; ) {
		ret = select ( n, readfds, writefds, exceptfds, timeout );
		ret -= remove_unpending_fds ( comm,  n, readfds );

		if ( select_has_finished ( ret, timeout ) )
			break;
		
		restore_fd_sets ( &saved_readfds, &saved_writefds, &saved_exceptfds,
				  readfds, writefds, exceptfds );
	}
	return ret;
}

#endif
