/* This example code is placed in the public domain. */

#ifdef HAVE_CONFIG_H

#include <config.h>

#endif

#include <stdio.h>

#include <stdlib.h>

#include <errno.h>

#include <sys/types.h>

#include <sys/socket.h>

#include <arpa/inet.h>

#include <netinet/in.h>

#include <sys/select.h>

#include <netdb.h>

#include <string.h>

#include <unistd.h>

#include <gnutls/gnutls.h>

#include <gnutls/dtls.h>

#define KEYFILE "key.pem"

#define CERTFILE "cert.pem"

#define CAFILE "/etc/ssl/certs/ca-certificates.crt"

#define CRLFILE "crl.pem"

/* This is a sample DTLS echo server, using X.509 authentication.

 * Note that error checking is minimal to simplify the example.

 */

#define LOOP_CHECK(rval, cmd) \

        do { \

                rval = cmd; \

        } while(rval == GNUTLS_E_AGAIN || rval == GNUTLS_E_INTERRUPTED)

#define MAX_BUFFER 1024

#define PORT 5557

typedef struct {

        gnutls_session_t session;

        int fd;

        struct sockaddr *cli_addr;

        socklen_t cli_addr_size;

} priv_data_st;

static int pull_timeout_func(gnutls_transport_ptr_t ptr, unsigned int ms);

static ssize_t push_func(gnutls_transport_ptr_t p, const void *data,

                         size_t size);

static ssize_t pull_func(gnutls_transport_ptr_t p, void *data,

                         size_t size);

static const char *human_addr(const struct sockaddr *sa, socklen_t salen,

                              char *buf, size_t buflen);

static int wait_for_connection(int fd);

/* Use global credentials and parameters to simplify

 * the example. */

static gnutls_certificate_credentials_t x509_cred;

static gnutls_priority_t priority_cache;

int main(void)

{

        int listen_sd;

        int sock, ret;

        struct sockaddr_in sa_serv;

        struct sockaddr_in cli_addr;

        socklen_t cli_addr_size;

        gnutls_session_t session;

        char buffer[MAX_BUFFER];

        priv_data_st priv;

        gnutls_datum_t cookie_key;

        gnutls_dtls_prestate_st prestate;

        int mtu = 1400;

        unsigned char sequence[8];

        /* 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_crl_file(x509_cred, CRLFILE,

                                             GNUTLS_X509_FMT_PEM);

        ret =

            gnutls_certificate_set_x509_key_file(x509_cred, CERTFILE,

                                                 KEYFILE,

                                                 GNUTLS_X509_FMT_PEM);

        if (ret < 0) {

                printf("No certificate or key were found\n");

                exit(1);

        }

        gnutls_certificate_set_known_dh_params(x509_cred, GNUTLS_SEC_PARAM_MEDIUM);

        /* pre-3.6.3 equivalent:

         * gnutls_priority_init(&priority_cache,

         *                      "NORMAL:-VERS-TLS-ALL:+VERS-DTLS1.0:%SERVER_PRECEDENCE",

         *                      NULL);

         */

        gnutls_priority_init2(&priority_cache,

                              "%SERVER_PRECEDENCE",

                              NULL, GNUTLS_PRIORITY_INIT_DEF_APPEND);

        gnutls_key_generate(&cookie_key, GNUTLS_COOKIE_KEY_SIZE);

        /* Socket operations

         */

        listen_sd = socket(AF_INET, SOCK_DGRAM, 0);

        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);

        {                       /* DTLS requires the IP don't fragment (DF) bit to be set */

#if defined(IP_DONTFRAG)

                int optval = 1;

                setsockopt(listen_sd, IPPROTO_IP, IP_DONTFRAG,

                           (const void *) &optval, sizeof(optval));

#elif defined(IP_MTU_DISCOVER)

                int optval = IP_PMTUDISC_DO;

                setsockopt(listen_sd, IPPROTO_IP, IP_MTU_DISCOVER,

                           (const void *) &optval, sizeof(optval));

#endif

        }

        bind(listen_sd, (struct sockaddr *) &sa_serv, sizeof(sa_serv));

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

        for (;;) {

                printf("Waiting for connection...\n");

                sock = wait_for_connection(listen_sd);

                if (sock < 0)

                        continue;

                cli_addr_size = sizeof(cli_addr);

                ret = recvfrom(sock, buffer, sizeof(buffer), MSG_PEEK,

                               (struct sockaddr *) &cli_addr,

                               &cli_addr_size);

                if (ret > 0) {

                        memset(&prestate, 0, sizeof(prestate));

                        ret =

                            gnutls_dtls_cookie_verify(&cookie_key,

                                                      &cli_addr,

                                                      sizeof(cli_addr),

                                                      buffer, ret,

                                                      &prestate);

                        if (ret < 0) {  /* cookie not valid */

                                priv_data_st s;

                                memset(&s, 0, sizeof(s));

                                s.fd = sock;

                                s.cli_addr = (void *) &cli_addr;

                                s.cli_addr_size = sizeof(cli_addr);

                                printf

                                    ("Sending hello verify request to %s\n",

                                     human_addr((struct sockaddr *)

                                                &cli_addr,

                                                sizeof(cli_addr), buffer,

                                                sizeof(buffer)));

                                gnutls_dtls_cookie_send(&cookie_key,

                                                        &cli_addr,

                                                        sizeof(cli_addr),

                                                        &prestate,

                                                        (gnutls_transport_ptr_t)

                                                        & s, push_func);

                                /* discard peeked data */

                                recvfrom(sock, buffer, sizeof(buffer), 0,

                                         (struct sockaddr *) &cli_addr,

                                         &cli_addr_size);

                                usleep(100);

                                continue;

                        }

                        printf("Accepted connection from %s\n",

                               human_addr((struct sockaddr *)

                                          &cli_addr, sizeof(cli_addr),

                                          buffer, sizeof(buffer)));

                } else

                        continue;

                gnutls_init(&session, GNUTLS_SERVER | GNUTLS_DATAGRAM);

                gnutls_priority_set(session, priority_cache);

                gnutls_credentials_set(session, GNUTLS_CRD_CERTIFICATE,

                                       x509_cred);

                gnutls_dtls_prestate_set(session, &prestate);

                gnutls_dtls_set_mtu(session, mtu);

                priv.session = session;

                priv.fd = sock;

                priv.cli_addr = (struct sockaddr *) &cli_addr;

                priv.cli_addr_size = sizeof(cli_addr);

                gnutls_transport_set_ptr(session, &priv;;

                gnutls_transport_set_push_function(session, push_func);

                gnutls_transport_set_pull_function(session, pull_func);

                gnutls_transport_set_pull_timeout_function(session,

                                                           pull_timeout_func);

                LOOP_CHECK(ret, gnutls_handshake(session));

                /* Note that DTLS may also receive GNUTLS_E_LARGE_PACKET.

                 * In that case the MTU should be adjusted.

                 */

                if (ret < 0) {

                        fprintf(stderr, "Error in handshake(): %s\n",

                                gnutls_strerror(ret));

                        gnutls_deinit(session);

                        continue;

                }

                printf("- Handshake was completed\n");

                for (;;) {

                        LOOP_CHECK(ret,

                                    gnutls_record_recv_seq(session, buffer,

                                                           MAX_BUFFER,

                                                           sequence));

                        if (ret < 0 && gnutls_error_is_fatal(ret) == 0) {

                                fprintf(stderr, "*** Warning: %s\n",

                                        gnutls_strerror(ret));

                                continue;

                        } else if (ret < 0) {

                                fprintf(stderr, "Error in recv(): %s\n",

                                        gnutls_strerror(ret));

                                break;

                        }

                        if (ret == 0) {

                                printf("EOF\n\n");

                                break;

                        }

                        buffer[ret] = 0;

                        printf

                            ("received[%.2x%.2x%.2x%.2x%.2x%.2x%.2x%.2x]: %s\n",

                             sequence[0], sequence[1], sequence[2],

                             sequence[3], sequence[4], sequence[5],

                             sequence[6], sequence[7], buffer);

                        /* reply back */

                        LOOP_CHECK(ret, gnutls_record_send(session, buffer, ret));

                        if (ret < 0) {

                                fprintf(stderr, "Error in send(): %s\n",

                                        gnutls_strerror(ret));

                                break;

                        }

                }

                LOOP_CHECK(ret, gnutls_bye(session, GNUTLS_SHUT_WR));

                gnutls_deinit(session);

        }

        close(listen_sd);

        gnutls_certificate_free_credentials(x509_cred);

        gnutls_priority_deinit(priority_cache);

        gnutls_global_deinit();

        return 0;

}

static int wait_for_connection(int fd)

{

        fd_set rd, wr;

        int n;

        FD_ZERO(&rd);

        FD_ZERO(&wr);

        FD_SET(fd, &rd);

        /* waiting part */

        n = select(fd + 1, &rd, &wr, NULL, NULL);

        if (n == -1 && errno == EINTR)

                return -1;

        if (n < 0) {

                perror("select()");

                exit(1);

        }

        return fd;

}

/* Wait for data to be received within a timeout period in milliseconds

 */

static int pull_timeout_func(gnutls_transport_ptr_t ptr, unsigned int ms)

{

        fd_set rfds;

        struct timeval tv;

        priv_data_st *priv = ptr;

        struct sockaddr_in cli_addr;

        socklen_t cli_addr_size;

        int ret;

        char c;

        FD_ZERO(&rfds);

        FD_SET(priv->fd, &rfds);

        tv.tv_sec = ms / 1000;

        tv.tv_usec = (ms % 1000) * 1000;

        ret = select(priv->fd + 1, &rfds, NULL, NULL, &tv;;

        if (ret <= 0)

                return ret;

        /* only report ok if the next message is from the peer we expect

         * from 

         */

        cli_addr_size = sizeof(cli_addr);

        ret =

            recvfrom(priv->fd, &c, 1, MSG_PEEK,

                     (struct sockaddr *) &cli_addr, &cli_addr_size);

        if (ret > 0) {

                if (cli_addr_size == priv->cli_addr_size

                    && memcmp(&cli_addr, priv->cli_addr,

                              sizeof(cli_addr)) == 0)

                        return 1;

        }

        return 0;

}

static ssize_t

push_func(gnutls_transport_ptr_t p, const void *data, size_t size)

{

        priv_data_st *priv = p;

        return sendto(priv->fd, data, size, 0, priv->cli_addr,

                      priv->cli_addr_size);

}

static ssize_t pull_func(gnutls_transport_ptr_t p, void *data, size_t size)

{

        priv_data_st *priv = p;

        struct sockaddr_in cli_addr;

        socklen_t cli_addr_size;

        char buffer[64];

        int ret;

        cli_addr_size = sizeof(cli_addr);

        ret =

            recvfrom(priv->fd, data, size, 0,

                     (struct sockaddr *) &cli_addr, &cli_addr_size);

        if (ret == -1)

                return ret;

        if (cli_addr_size == priv->cli_addr_size

            && memcmp(&cli_addr, priv->cli_addr, sizeof(cli_addr)) == 0)

                return ret;

        printf("Denied connection from %s\n",

               human_addr((struct sockaddr *)

                          &cli_addr, sizeof(cli_addr), buffer,

                          sizeof(buffer)));

        gnutls_transport_set_errno(priv->session, EAGAIN);

        return -1;

}

static const char *human_addr(const struct sockaddr *sa, socklen_t salen,

                              char *buf, size_t buflen)

{

        const char *save_buf = buf;

        size_t l;

        if (!buf || !buflen)

                return NULL;

        *buf = '\0';

        switch (sa->sa_family) {

#if HAVE_IPV6

        case AF_INET6:

                snprintf(buf, buflen, "IPv6 ");

                break;

#endif

        case AF_INET:

                snprintf(buf, buflen, "IPv4 ");

                break;

        }

        l = strlen(buf);

        buf += l;

        buflen -= l;

        if (getnameinfo(sa, salen, buf, buflen, NULL, 0, NI_NUMERICHOST) !=

            0)

                return NULL;

        l = strlen(buf);

        buf += l;

        buflen -= l;

        strncat(buf, " port ", buflen);

        l = strlen(buf);

        buf += l;

        buflen -= l;

        if (getnameinfo(sa, salen, NULL, 0, buf, buflen, NI_NUMERICSERV) !=

            0)

                return NULL;

        return save_buf;

}