/*

 * Copyright (C) 2011-2012 Free Software Foundation, Inc.

 *

 * This file is part of GnuTLS.

 *

 * GnuTLS 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.

 *

 * GnuTLS 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/>.

 */

#include <config.h>

#include <stdio.h>

#if HAVE_SYS_SOCKET_H

#include <sys/socket.h>

#elif HAVE_WS2TCPIP_H

#include <ws2tcpip.h>

#endif

#include <arpa/inet.h>

#ifndef _WIN32

#include <netinet/in.h>

#endif

#include <sys/select.h>

#include <stdlib.h>

#include <string.h>

#include <unistd.h>

#include <errno.h>

#include <common.h>

#include "udp-serv.h"

#include "list.h"

typedef struct {

	gnutls_session_t session;

	int fd;

	struct sockaddr_storage *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);

#define MAX_BUFFER 255		/* Longest string to echo */

void udp_server(const char *name, int port, int mtu)

{

	int sock, ret;

	struct sockaddr_storage cli_addr;

	socklen_t cli_addr_size;

	char buffer[MAX_BUFFER];

	priv_data_st priv;

	gnutls_session_t session;

	gnutls_datum_t cookie_key;

	gnutls_dtls_prestate_st prestate;

	unsigned char sequence[8];

	ret = gnutls_key_generate(&cookie_key, GNUTLS_COOKIE_KEY_SIZE);

	if (ret < 0) {

		fprintf(stderr, "Cannot generate key\n");

		exit(1);

	}

	ret = listen_socket(name, port, SOCK_DGRAM);

	if (ret < 0) {

		fprintf(stderr, "Cannot listen\n");

		exit(1);

	}

	for (;;) {

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

		sock = wait_for_connection();

		if (sock < 0)

			continue;

		cli_addr_size = sizeof(cli_addr);

		ret =

		    recvfrom(sock, buffer, sizeof(buffer)-1, 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,

						      cli_addr_size,

						      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 = cli_addr_size;

				printf

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

				     human_addr((struct sockaddr *)

						&cli_addr,

						cli_addr_size, buffer,

						sizeof(buffer)-1));

				gnutls_dtls_cookie_send(&cookie_key,

							&cli_addr,

							cli_addr_size,

							&prestate,

							(gnutls_transport_ptr_t)

							&s, push_func);

				/* discard peeked data */

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

					 (struct sockaddr *) &cli_addr,

					 &cli_addr_size);

				continue;

			}

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

			       human_addr((struct sockaddr *)

					  &cli_addr, sizeof(cli_addr),

					  buffer, sizeof(buffer)-1));

		} else

			continue;

		session = initialize_session(1);

		gnutls_dtls_prestate_set(session, &prestate);

		if (mtu)

			gnutls_dtls_set_mtu(session, mtu);

		priv.session = session;

		priv.fd = sock;

		priv.cli_addr = &cli_addr;

		priv.cli_addr_size = cli_addr_size;

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

		do {

			ret = gnutls_handshake(session);

		}

		while (ret == GNUTLS_E_AGAIN

		       || ret == GNUTLS_E_INTERRUPTED);

		if (ret < 0) {

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

				gnutls_strerror(ret));

			gnutls_deinit(session);

			continue;

		}

		for (;;) {

			do {

				ret =

				    gnutls_record_recv_seq(session, buffer,

							   sizeof(buffer)-1,

							   sequence);

				if (ret ==

				    GNUTLS_E_HEARTBEAT_PING_RECEIVED)

					gnutls_heartbeat_pong(session, 0);

			}

			while (ret == GNUTLS_E_INTERRUPTED

			       || ret == GNUTLS_E_AGAIN

			       || ret == GNUTLS_E_HEARTBEAT_PING_RECEIVED);

			if (ret == GNUTLS_E_REHANDSHAKE) {

				fprintf(stderr,

					"*** Received hello message\n");

				do {

					ret = gnutls_handshake(session);

				}

				while (ret == GNUTLS_E_INTERRUPTED ||

				       ret == GNUTLS_E_AGAIN);

				if (ret == 0)

					continue;

			}

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

			if (check_command(session, buffer) == 0) {

				/* reply back */

				ret =

				    gnutls_record_send(session, buffer,

						       ret);

				if (ret < 0) {

					fprintf(stderr,

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

						gnutls_strerror(ret));

					break;

				}

			}

		}

	}

	gnutls_deinit(session);

}

/* 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 = 0;

	tv.tv_usec = ms * 1000;

	while (tv.tv_usec >= 1000000) {

		tv.tv_usec -= 1000000;

		tv.tv_sec++;

	}

	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, (struct sockaddr*)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;

}