/*
 * The contents of this file are subject to the Mozilla Public
 * License Version 1.1 (the "License"); you may not use this file
 * except in compliance with the License. You may obtain a copy of
 * the License at http://www.mozilla.org/MPL/
 * 
 * Software distributed under the License is distributed on an "AS
 * IS" basis, WITHOUT WARRANTY OF ANY KIND, either express or
 * implied. See the License for the specific language governing
 * rights and limitations under the License.
 * 
 * The Original Code is the Netscape security libraries.
 * 
 * The Initial Developer of the Original Code is Netscape
 * Communications Corporation.  Portions created by Netscape are 
 * Copyright (C) 1994-2000 Netscape Communications Corporation.  All
 * Rights Reserved.
 * 
 * Contributor(s):
 * 
 * Alternatively, the contents of this file may be used under the
 * terms of the GNU General Public License Version 2 or later (the
 * "GPL"), in which case the provisions of the GPL are applicable 
 * instead of those above.  If you wish to allow use of your 
 * version of this file only under the terms of the GPL and not to
 * allow others to use your version of this file under the MPL,
 * indicate your decision by deleting the provisions above and
 * replace them with the notice and other provisions required by
 * the GPL.  If you do not delete the provisions above, a recipient
 * may use your version of this file under either the MPL or the
 * GPL.
 */

/*
 * SSL Sample code. XXX
 *
 */

/****************************************************************************
 *  SSL server program listens on a port, accepts client connection, reads  *
 *  request and responds to it                                              *
 ****************************************************************************/

/* Platform specific header files */

#ifdef XP_UNIX
#define _XOPEN_SOURCE 1 /* to get declarations of getopt() and getpass() */

#ifdef AIX
#define _ALL_SOURCE 1	/* to get declaration for uint on AIX */
#endif
#ifdef OSF1
#define _OSF_SOURCE 1 /* to get declaration for uint on HPUX */
#endif
#ifdef SOLARIS
#define __EXTENSIONS__ 1 /* to get declaration for getpass() on solaris */
#endif
#ifdef IRIX
#define _BSD_TYPES 1	/* to get declaration for uint on IRIX 6.5 */
#endif

#include <unistd.h>	/* for getopt() + getpass() */

#ifdef SOLARIS
#include <signal.h>
#endif	
#endif	/* XP_UNIX */

#ifdef WIN32
#include <winsock.h>
#endif

/* Generic header files */

#include <stdio.h>
#include <string.h>

/* NSPR header files */

#include "nspr.h"
#include "prerror.h"
#include "prnetdb.h"

/* NSS header files */

#include "pk11func.h"
#include "secitem.h"
#include "ssl.h"
#include "certt.h"
#include "nss.h"
#include "secrng.h"
#include "secder.h"
#include "key.h"
#include "sslproto.h"

/* Custom header files */

#include "sslerror.h"

#ifdef _WINDOWS
extern int getopt(int argc, char **argv, char *optstring);
extern char *optarg;
#else
#ifdef STRANGE_PLATFORMS
#include <getopt.h>
#endif
#endif

/* Function prototypes */

void 		usage			(char *progname);
void 		exitErr			(char *function);

PRFileDesc *	setup_socket		(PRFileDesc *sslsocket);

SECStatus 	ownHandshakeCallback	(PRFileDesc *socket,
					 void *arg);

char *		ownPasswd		(PK11SlotInfo *slot,
					 PRBool required,
					 void *arg);

SECStatus 	ownAuthCertificate	(void *arg,
					 PRFileDesc *socket,
					 PRBool checksig,
					 PRBool isServer);

SECStatus 	ownBadCertHandler 	(void *arg,
					 PRFileDesc *socket);

SECStatus 	secCmpCertChainWCANames	(CERTCertificate *cert,
					 CERTDistNames *caNames);

int 		readThisMany		(PRFileDesc *fd,
					 char *buffer,
					 int thisMany);

int 		writeThisMany		(PRFileDesc *fd,
					 char * buffer,
					 int thisMany);

void		checkSecureConnStatus	(PRFileDesc *socket);

/* Global Variable */

PRErrorCode certErr = 0;

/* Function implementation */

/* Function: void usage()
 *
 * Purpose: Upon a bogus invocation of the program, print the usage information
 * and exit with a nonzero status.
 */

void usage(char *progname) 
{

    printf("Usage: %s [-d <cert dir>] -p <server port> -n <nickname> -c <password> \n\n", progname);
    exit(1);
}

/* Function void exitErr()
 *
 * Purpose: When a security function returns an error, get the error number
 * using PR_GetError() and find a matching error string using the custom
 * function SSL_Strerror(). Print out all relevant information, shutdown 
 * NSS, clean up NSPR, and then exit with a nonzero status.
 *
 * Note: SSL_Strerror is defined in sslerror.h.
 */

void exitErr(char *function) 
{

    PRErrorCode	errornumber = PR_GetError();
    const char *errorstring = SSL_Strerror(errornumber);

    printf("Error in function %s: %d\n - %s\n",
	    function, errornumber, errorstring);

    NSS_Shutdown();
    PR_Cleanup();
    exit(1);
}

/* Function: PRFileDesc * setup_socket()
 *
 * Purpose: setup_socket takes a PRFileDesc * and configures it for SSL.
 * We use this to configure sockets returned by PR_Accept in this program.
 */

PRFileDesc * setup_socket(PRFileDesc *sslsocket) 
{

    /* Set the appropriate flags. */

    if ((SSL_Enable(sslsocket, SSL_SECURITY, 1)) != SECSuccess) {
	exitErr("SSL_Enable:SSL_SECURITY");
    }

    if ((SSL_Enable(sslsocket, SSL_HANDSHAKE_AS_SERVER, 1)) != SECSuccess) {
	exitErr("SSL_Enable:SSL_HANDSHAKE_AS_SERVER");
    }	

    if ((SSL_Enable(sslsocket, SSL_REQUEST_CERTIFICATE, 1)) != SECSuccess ) {
	exitErr("SSL_Enable:SSL_REQUEST_CERTIFICATE");
    }

    if ((SSL_Enable(sslsocket, SSL_REQUIRE_CERTIFICATE, 1)) != SECSuccess ) {
	exitErr("SSL_Enable:SSL_REQUIRE_CERTIFICATE");
    }	

    /* Set the appropriate callback routines */

    if ((SSL_AuthCertificateHook(sslsocket,
		                 (SSLAuthCertificate)ownAuthCertificate,
		                 CERT_GetDefaultCertDB() ))
	!= SECSuccess ) {
	exitErr("SSL_AuthCertificateHook");
    }

    if ((SSL_BadCertHook(sslsocket, (SSLBadCertHandler)ownBadCertHandler,
			 &certErr)) != SECSuccess ) {
	exitErr("SSL_BadCertHook");
    }

    if ((SSL_HandshakeCallback(sslsocket, 
                               (SSLHandshakeCallback)ownHandshakeCallback,
			       NULL)) != SECSuccess ) {
	exitErr("SSL_HandshakeCallback");
    }

    return sslsocket;
}

/* Function: SECStatus ownHandshakeCallback()
 *
 * Purpose: Used to inform application that the handshake has completed --
 * useful when a non-blocking SSL_RedoHandshake or SSL_ResetHandshake are used
 * to initiate a handshake.
 *
 * A typical scenario would be:
 *
 * 1. Server accepts an SSL connection from the client without client auth.
 * 2. Client sends a request.
 * 3. Server determines that to service request it needs to authenticate the
 * client and initiates another handshake requesting client auth.
 * 4. While handshake is in progress, server can do other work or spin waiting
 * for the handshake to complete.
 * 5. Server is notified that handshake has been successfully completed by
 * the custom handshake callback function and it can service the client's
 * request.
 *
 * Note: This function is unimplemented in this sample, as we are using
 * blocking sockets.
 */

SECStatus ownHandshakeCallback(PRFileDesc *socket, void *arg) 
{
	return SECSuccess;
}


/* Function: SECStatus ownAuthCertificate()
 *
 * Purpose: This function is our custom certificate authentication handler.
 * 
 */

SECStatus ownAuthCertificate( 	void *arg,
				PRFileDesc *socket,
				PRBool checksig,
				PRBool isServer) 
{

    SECCertUsage	certUsage;
    CERTCertificate *	cert;
    void *		PinArg;
    SECStatus		rv 		= SECFailure;


    /* Define how the cert is being used based upon the isServer flag. */

    certUsage = isServer ? certUsageSSLClient : certUsageSSLServer;

    cert = SSL_PeerCertificate(socket);
    PinArg = SSL_RevealPinArg(socket);

    if ( !arg ) {
	exitErr("ownAuthCertificate");
    }

    rv = CERT_VerifyCertNow( (CERTCertDBHandle *)arg,
			     cert,
			     checksig,
			     certUsage,
			     PinArg);

    /* There is no need to check the name field in the cert against */
    /* the desired hostname as this is the server side of an SSL	*/
    /* connection. 							*/

    return rv;
}

/* Function: char * ownPasswd()
 *
 * Purpose: This function is our custom password handler that is called by
 * SSL when retriving private certs and keys from the database. Returns a
 * pointer to a string that with a password for the database. Password pointer
 * should point to dynamically allocated memory that will be freed later.
 */

char * ownPasswd(PK11SlotInfo *slot, PRBool retry, void *arg) 
{

    char *passwd = NULL;

    if ( (!retry) && arg ) {
	passwd = PL_strdup((char *)arg);
    }

    return passwd;
}

/* Function: SECStatus ownBadCertHandler()
 *
 * Purpose: This callback is called when the incoming certificate is not
 * valid. We define a certain set of parameters that still cause the
 * certificate to be "valid" for this session, and return SECSuccess to cause
 * the server to continue processing the request when any of these conditions
 * are met. Otherwise, SECFailure is return and the server rejects the 
 * request.
 */

SECStatus ownBadCertHandler(void *arg, PRFileDesc *socket) 
{

    SECStatus	rv = SECFailure;
    PRErrorCode	err;

    /* log invaild cert here */

    if (!arg) {
	return rv;
    }

    *(PRErrorCode *)arg = err = PORT_GetError();

    /* If any of the cases in the switch are met, then we will proceed   */
    /* with the processing of the request anyway. Otherwise, the default */	
    /* case will be reached and we will reject the request.		     */

    switch (err) {
    case SEC_ERROR_INVALID_AVA:
    case SEC_ERROR_INVALID_TIME:
    case SEC_ERROR_BAD_SIGNATURE:
    case SEC_ERROR_EXPIRED_CERTIFICATE:
    case SEC_ERROR_UNKNOWN_ISSUER:
    case SEC_ERROR_UNTRUSTED_CERT:
    case SEC_ERROR_CERT_VALID:
    case SEC_ERROR_EXPIRED_ISSUER_CERTIFICATE:
    case SEC_ERROR_CRL_EXPIRED:
    case SEC_ERROR_CRL_BAD_SIGNATURE:
    case SEC_ERROR_EXTENSION_VALUE_INVALID:
    case SEC_ERROR_CA_CERT_INVALID:
    case SEC_ERROR_CERT_USAGES_INVALID:
    case SEC_ERROR_UNKNOWN_CRITICAL_EXTENSION:
	rv = SECSuccess;
	break;
    default:
	rv = SECFailure;
	break;
    }

    return rv;
}

/* Function: int readThisMany()
 *
 * Purpose: This is a wrapper function around PR_Read that will read exactly
 * the number of bytes requested and return the number of bytes read.
 */

int readThisMany(PRFileDesc *fd, char *buffer, int thisMany) 
{

    int	total = 0;

    while ( total < thisMany ) {

	int got = PR_Read(fd, buffer+total, thisMany-total);

	if ( got < 0 ) {
	    if ( PR_GetError() != PR_WOULD_BLOCK_ERROR ) {
		break;
	    }
	    continue;
	}
	total += got;
    }

    return total;
}

/* Function: int writeThisMany()
 *
 * Purpose: This is a wrapper function around PR_Write that will write exactly
 * the number of bytes requested and return the number of bytes written.
 */

int writeThisMany(PRFileDesc *fd, char *buffer, int thisMany) 
{

    int	total = 0;

    while ( total < thisMany ) {

	int	got;	

	got = PR_Write(fd, buffer+total, thisMany-total);
	if ( got < 0 ) {
	    if ( PR_GetError() != PR_WOULD_BLOCK_ERROR ) {
		break;
	    }
	    continue;
	}
	total += got;
    }

    return total;
}

/* Function: void checkSecureConnStatus()
 *
 * Purpose: To print out some useful information about the connection after
 * a handshake.
 */

void checkSecureConnStatus(PRFileDesc *socket) 
{

    int 	on, keySize, secretKeySize;
    char 	*cipher, *issuer, *subject;
    SECStatus	secstatus = SECFailure;

    secstatus = SSL_SecurityStatus(socket, &on, &cipher,
				    &keySize, &secretKeySize, &issuer,
				    &subject);

    if (secstatus != SECSuccess) {
	exitErr("SSL_SecurityStatus");
    }	

    if (cipher) {
    printf("Cipher used is %s, key size is %d, secret key size is %d\n", 
    cipher, keySize, secretKeySize);
    }

    if (on) {
	printf("Security is ON\n"); 
    }

    if (issuer && subject) {
	printf("Client cert issued by %s,\n for %s\n", issuer, subject);
    }

    if (certErr) {
	printf("Client cert is not trusted: %d: %s\n", certErr, 
	SSL_Strerror(certErr));
    }

    PR_Free(issuer);
    PR_Free(subject);
}


/* Function: int main()
 *
 * Purpose: This is the server's main function.
 */

int main (int argc, char **argv) 
{

    char *             progName;
    char *             nickname  = NULL;
    char *             certdir   = NULL;
    char *             password  = NULL;
    PRFileDesc *       sock      = NULL;
    PRFileDesc *       respsock  = NULL;
    PRFileDesc *       sslsocket = NULL;
    SECKEYPrivateKey * privKey   = NULL;
    CERTCertificate *  cert      = NULL;

    int                i;
    int                option;
    int                port      = 0;
    SSLKEAType         certKEAType;
    SECStatus          secstatus = SECFailure;
    PRStatus           prstatus  = PR_FAILURE;

    PRNetAddr          addr;
    PRSocketOptionData sockdata;


#ifdef SOLARIS
    sigignore(SIGPIPE);
#endif

    PR_Init(PR_USER_THREAD, PR_PRIORITY_NORMAL, 1);

    progName = PL_strdup(argv[0]);

    while ((option = getopt(argc, argv, "d:p:n:c:")) != -1) {

	switch (option) {
	case 'd': certdir  = optarg; 		break;

	case 'p': port     = atoi(optarg); 	break;

	case 'n': nickname = optarg; 		break;

	case 'c': password = optarg; 		break;

	default:  usage(progName); 		break;
	}
    }

    if (nickname == NULL || port == 0) {
	usage(progName);
    }

    if (certdir == NULL) {
	certdir = PR_smprintf("%s/.netscape", getenv("HOME"));
    }

    /* Set the cert database password callback. */
    PK11_SetPasswordFunc(ownPasswd);  

    /* Initialize NSS. */
    secstatus = NSS_Init(certdir);
    if (secstatus != SECSuccess) {
	exitErr("NSS_Init");
    }

    /* Set the policy for this server. */
    secstatus = NSS_SetDomesticPolicy();
    if (secstatus != SECSuccess) {
	exitErr("NSS_SetDomesticPolicy");
    }

    /* Get own certificate and private key. */
    cert = PK11_FindCertFromNickname(nickname, password);
    if (cert == NULL) {
	exitErr("PK11_FindCertFromNickname");
    }

    privKey = PK11_FindKeyByAnyCert(cert, password);
    if (privKey == NULL) {
	exitErr("PK11_FindKeyByAnyCert");
    }

    /* Now configure the server's cache. */  
    SSL_ConfigServerSessionIDCache(10, 0, 0, ".");


    /* 
     * Create SSL socket
     */

    /* create a new socket */
    sock = PR_NewTCPSocket();
    if (!sock) {
	exitErr("PR_NewTCPSocket");
    }

    /* Set socket to be blocking - 
     * on some platforms the default is nonblocking. 
     */
    sockdata.option = PR_SockOpt_Nonblocking;
    sockdata.value.non_blocking = PR_FALSE;

    if (PR_SetSocketOption(sock, &sockdata) != PR_SUCCESS) {
	exitErr("PR_SetSocketOption"); 
    }

    /* Import the socket into SSL layer. */
    sslsocket = SSL_ImportFD(NULL, sock);

    if (!sslsocket) {
	exitErr("SSL_ImportFD");
    }

    secstatus = SSL_SetPKCS11PinArg(sslsocket, password);

    if (secstatus != SECSuccess) {
	exitErr("SSL_SetPKCS11PinArg");
    }

    certKEAType = NSS_FindCertKEAType(cert); 

    secstatus = SSL_ConfigSecureServer(	sslsocket, cert, 
					    privKey, certKEAType); 

    if (secstatus != SECSuccess) {
	exitErr("SSL_ConfigSecureServer");
    }

    /* Setup the network connection */
    addr.inet.family = AF_INET;
    addr.inet.ip     = INADDR_ANY;
    addr.inet.port   = PR_htons(port);

    prstatus = PR_Bind(sslsocket, &addr);

    if (prstatus != PR_SUCCESS) {
	exitErr("PR_Bind");
    }

    prstatus = PR_Listen(sslsocket, 5);
    if (prstatus != PR_SUCCESS) {
	exitErr("PR_Listen");
    }

    while (1) {
	PRInt32 numbytes = 0, readbytes = 0;
	char buffer[256];
	char reply[]="Server has received:";

	printf("\nReady to accept client connection:\n\n");
	respsock = PR_Accept(sslsocket, &addr, PR_INTERVAL_NO_TIMEOUT);

	if (respsock == NULL) {
	    exitErr("PR_Accept");
	}

	respsock = setup_socket(respsock);

	secstatus = SSL_ForceHandshake(respsock);

	if (secstatus == SECSuccess) {
	    checkSecureConnStatus(respsock);
	} else {
	    printf("SSL handshake was unsuccesful.\n");
	}
	readbytes = numbytes = 0;

	numbytes = readThisMany(respsock, (char *)&readbytes, sizeof readbytes);

	if (numbytes != sizeof readbytes) {
	    exitErr("PR_Read message length");
	}

	readbytes = PR_ntohl(readbytes);
	readbytes = (readbytes > 255) ? 255 : readbytes;

	numbytes = readThisMany(respsock, buffer, readbytes); 

	if (numbytes != readbytes) {
	    exitErr("PR_Read");
	}

	buffer[readbytes] = 0;

	printf("Read from the client *%s*\n", buffer);

	i = PR_htonl(strlen(reply)+readbytes);
	numbytes = writeThisMany(respsock, (char *)&i, sizeof i);

	if (numbytes != sizeof i) {
	    exitErr("PR_Write message length");
	}

	numbytes = writeThisMany(respsock, reply, strlen(reply));

	if (numbytes < strlen(reply)) {
	    exitErr("PR_Write reply");
	}

	numbytes = writeThisMany(respsock, buffer, readbytes);

	if (numbytes < readbytes) {
	    exitErr("PR_Write client data");
		break;
	}

	/* The break in the conditional above is a kludge to get this
	 * code to compile without warnings. The compiler sees a
	 * break and, even though it will never be executed given
	 * the nature of exitErr(), is satistifed that the loop will
	 * be left at some point. A true server would have a shutdown
	 * command that would cause the server to break out of the
	 * "accept loop" and execute the proper shutdown functions.
	 * Those functions are illustrated below.
	 */

	PR_Close(respsock);
    }

    /* PR_Shutdown closes communication channel on the socket
     * PR_Close frees the socket structure
     * NSS_Shutdown closes the certificate and key databases
     */

    if (PR_Close(sock) != PR_SUCCESS) {
	exitErr("PR_Close");
    }

    NSS_Shutdown();
    PR_Cleanup();
    return 0;
}
