
#include <stdio.h>
#include <stdlib.h>
#include <errno.h>
#include <sys/types.h>
#include <sys/socket.h>
#include <netinet/in.h>
#include <arpa/inet.h>
#include <string.h>
#include <unistd.h>
#include <gnutls/gnutls.h>

#define KEYFILE "key.pem"
#define CERTFILE "cert.pem"
#define CAFILE "ca.pem"
#define CRLFILE NULL

/* This is a sample TLS 1.0 echo server.
 * Export-grade ciphersuites and session resuming are supported.
 */

#define SA struct sockaddr
#define SOCKET_ERR(err,s) if(err==-1) {perror(s);return(1);}
#define MAX_BUF 1024
#define PORT 5556               /* listen to 5556 port */
#define DH_BITS 1024

/* These are global */
gnutls_certificate_server_credentials x509_cred;

static void wrap_db_init(void);
static void wrap_db_deinit(void);
static int wrap_db_store(void *dbf, gnutls_datum key, gnutls_datum data);
static gnutls_datum wrap_db_fetch(void *dbf, gnutls_datum key);
static int wrap_db_delete(void *dbf, gnutls_datum key);

#define TLS_SESSION_CACHE 50

gnutls_session initialize_tls_session()
{
   gnutls_session session;

   gnutls_init(&session, GNUTLS_SERVER);

   /* Use the default priorities, plus, export cipher suites.
    */
   gnutls_set_default_export_priority(session);

   gnutls_credentials_set(session, GNUTLS_CRD_CERTIFICATE, x509_cred);

   /* request client certificate if any.
    */
   gnutls_certificate_server_set_request(session, GNUTLS_CERT_REQUEST);

   gnutls_dh_set_prime_bits(session, DH_BITS);

   /* some broken clients may require this in order to connect. 
    * This will weaken security though.
    */
   /* gnutls_handshake_set_rsa_pms_check( session, 1); */

   if (TLS_SESSION_CACHE != 0) {
      gnutls_db_set_retrieve_function(session, wrap_db_fetch);
      gnutls_db_set_remove_function(session, wrap_db_delete);
      gnutls_db_set_store_function(session, wrap_db_store);
      gnutls_db_set_ptr(session, NULL);
   }

   return session;
}

gnutls_dh_params dh_params;
/* Export-grade cipher suites require temporary RSA
 * keys.
 */
gnutls_rsa_params rsa_params;

static int generate_dh_params(void)
{
   gnutls_datum prime, generator;

   /* Generate Diffie Hellman parameters - for use with DHE
    * kx algorithms. These should be discarded and regenerated
    * once a day, once a week or once a month. Depends on the
    * security requirements.
    */
   gnutls_dh_params_init(&dh_params);
   gnutls_dh_params_generate(&prime, &generator, DH_BITS);
   gnutls_dh_params_set(dh_params, prime, generator, DH_BITS);

   free(prime.data);
   free(generator.data);
   
   return 0;
}

static int generate_rsa_params(void)
{
   gnutls_datum m, e, d, p, q, u;

   gnutls_rsa_params_init(&rsa_params);

   /* Generate RSA parameters - for use with RSA-export
    * cipher suites. These should be discarded and regenerated
    * once a day, once every 500 transactions etc. Depends on the
    * security requirements.
    */

   gnutls_rsa_params_generate(&m, &e, &d, &p, &q, &u, 512);
   gnutls_rsa_params_set(rsa_params, m, e, d, p, q, u, 512);

   free(m.data);
   free(e.data);
   free(d.data);
   free(p.data);
   free(q.data);
   free(u.data);

   return 0;
}

int main()
{
   int err, listen_sd, i;
   int sd, ret;
   struct sockaddr_in sa_serv;
   struct sockaddr_in sa_cli;
   int client_len;
   char topbuf[512];
   gnutls_session session;
   char buffer[MAX_BUF + 1];
   int optval = 1;
   char name[256];

   strcpy(name, "Echo Server");

   /* this must be called once in the program
    */
   gnutls_global_init();

   gnutls_certificate_allocate_credentials(&x509_cred);

   gnutls_certificate_set_x509_trust_file(x509_cred, CAFILE,
                                          GNUTLS_X509_FMT_PEM);

   gnutls_certificate_set_x509_key_file(x509_cred, CERTFILE, KEYFILE,
                                        GNUTLS_X509_FMT_PEM);

   generate_dh_params();
   generate_rsa_params();

   if (TLS_SESSION_CACHE != 0) {
      wrap_db_init();
   }

   gnutls_certificate_set_dh_params(x509_cred, dh_params);
   gnutls_certificate_set_rsa_params(x509_cred, rsa_params);

   /* Socket operations
    */
   listen_sd = socket(AF_INET, SOCK_STREAM, 0);
   SOCKET_ERR(listen_sd, "socket");

   memset(&sa_serv, '\0', sizeof(sa_serv));
   sa_serv.sin_family = AF_INET;
   sa_serv.sin_addr.s_addr = INADDR_ANY;
   sa_serv.sin_port = htons(PORT);      /* Server Port number */

   setsockopt(listen_sd, SOL_SOCKET, SO_REUSEADDR, &optval, sizeof(int));

   err = bind(listen_sd, (SA *) & sa_serv, sizeof(sa_serv));
   SOCKET_ERR(err, "bind");
   err = listen(listen_sd, 1024);
   SOCKET_ERR(err, "listen");

   printf("%s ready. Listening to port '%d'.\n\n", name, PORT);

   client_len = sizeof(sa_cli);
   for (;;) {
      session = initialize_tls_session();

      sd = accept(listen_sd, (SA *) & sa_cli, &client_len);

      printf("- connection from %s, port %d\n",
             inet_ntop(AF_INET, &sa_cli.sin_addr, topbuf,
                       sizeof(topbuf)), ntohs(sa_cli.sin_port));

      gnutls_transport_set_ptr(session, sd);
      ret = gnutls_handshake(session);
      if (ret < 0) {
         close(sd);
         gnutls_deinit(session);
         fprintf(stderr, "*** Handshake has failed (%s)\n\n",
                 gnutls_strerror(ret));
         continue;
      }
      printf("- Handshake was completed\n");

      /* print_info(session); */

      i = 0;
      for (;;) {
         bzero(buffer, MAX_BUF + 1);
         ret = gnutls_record_recv(session, buffer, MAX_BUF);

         if (ret == 0) {
            printf("\n- Peer has closed the TLS connection\n");
            break;
         } else if (ret < 0) {
            fprintf(stderr,
                    "\n*** Received corrupted data(%d). Closing the connection.\n\n",
                    ret);
            break;
         } else if (ret > 0) {
            /* echo data back to the client
             */
            gnutls_record_send(session, buffer, strlen(buffer));
         }
      }
      printf("\n");
      gnutls_bye(session, GNUTLS_SHUT_WR);      /* do not wait for
                                                   * the peer to close the connection.
                                                 */

      close(sd);
      gnutls_deinit(session);

   }
   close(listen_sd);

   gnutls_certificate_free_credentials(x509_cred);

   gnutls_global_deinit();

   return 0;

}


/* Functions and other stuff needed for session resuming.
 * This is done using a very simple list which holds session ids
 * and session data.
 */

#define MAX_SESSION_ID_SIZE 32
#define MAX_SESSION_DATA_SIZE 512

typedef struct {
   char session_id[MAX_SESSION_ID_SIZE];
   int session_id_size;

   char session_data[MAX_SESSION_DATA_SIZE];
   int session_data_size;
} CACHE;

static CACHE *cache_db;
static int cache_db_ptr = 0;

static void wrap_db_init(void)
{

   /* allocate cache_db */
   cache_db = calloc(1, TLS_SESSION_CACHE * sizeof(CACHE));
}

static void wrap_db_deinit(void)
{
   return;
}

static int wrap_db_store(void *dbf, gnutls_datum key, gnutls_datum data)
{

   if (cache_db == NULL)
      return -1;

   if (key.size > MAX_SESSION_ID_SIZE)
      return -1;
   if (data.size > MAX_SESSION_DATA_SIZE)
      return -1;

   memcpy(cache_db[cache_db_ptr].session_id, key.data, key.size);
   cache_db[cache_db_ptr].session_id_size = key.size;

   memcpy(cache_db[cache_db_ptr].session_data, data.data, data.size);
   cache_db[cache_db_ptr].session_data_size = data.size;

   cache_db_ptr++;
   cache_db_ptr %= TLS_SESSION_CACHE;

   return 0;
}

static gnutls_datum wrap_db_fetch(void *dbf, gnutls_datum key)
{
   gnutls_datum res = { NULL, 0 };
   int i;

   if (cache_db == NULL)
      return res;

   for (i = 0; i < TLS_SESSION_CACHE; i++) {
      if (key.size == cache_db[i].session_id_size &&
          memcmp(key.data, cache_db[i].session_id, key.size) == 0) {


         res.size = cache_db[i].session_data_size;

         res.data = gnutls_malloc(res.size);
         if (res.data == NULL)
            return res;

         memcpy(res.data, cache_db[i].session_data, res.size);

         return res;
      }
   }
   return res;
}

static int wrap_db_delete(void *dbf, gnutls_datum key)
{
   int i;

   if (cache_db == NULL)
      return -1;

   for (i = 0; i < TLS_SESSION_CACHE; i++) {
      if (key.size == cache_db[i].session_id_size &&
          memcmp(key.data, cache_db[i].session_id, key.size) == 0) {

         cache_db[i].session_id_size = 0;
         cache_db[i].session_data_size = 0;

         return 0;
      }
   }

   return -1;

}

