/* $Id$ */
/* 
 * Copyright (C) 2008-2011 Teluu Inc. (http://www.teluu.com)
 * Copyright (C) 2003-2008 Benny Prijono <benny@prijono.org>
 *
 * This program 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 2 of the License, or
 * (at your option) any later version.
 *
 * This program 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, write to the Free Software
 * Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA 
 */
#include "turn.h"
#include <pj/compat/socket.h>

#if PJ_HAS_TCP

struct accept_op
{
    pj_ioqueue_op_key_t	op_key;
    pj_sock_t		sock;
    pj_sockaddr		src_addr;
    int			src_addr_len;
};

struct tcp_listener
{
    pj_turn_listener	     base;
    pj_ioqueue_key_t	    *key;
    unsigned		     accept_cnt;
    struct accept_op	    *accept_op;	/* Array of accept_op's	*/
};


static void lis_on_accept_complete(pj_ioqueue_key_t *key, 
				   pj_ioqueue_op_key_t *op_key, 
				   pj_sock_t sock, 
				   pj_status_t status);
static pj_status_t lis_destroy(pj_turn_listener *listener);
static void transport_create(pj_sock_t sock, pj_turn_listener *lis,
			     pj_sockaddr_t *src_addr, int src_addr_len);

static void show_err(const char *sender, const char *title, 
		     pj_status_t status)
{
    char errmsg[PJ_ERR_MSG_SIZE];

    pj_strerror(status, errmsg, sizeof(errmsg));
    PJ_LOG(4,(sender, "%s: %s", title, errmsg));
}


/*
 * Create a new listener on the specified port.
 */
PJ_DEF(pj_status_t) pj_turn_listener_create_tcp(pj_turn_srv *srv,
					        int af,
					        const pj_str_t *bound_addr,
					        unsigned port,
						unsigned concurrency_cnt,
						unsigned flags,
						pj_turn_listener **p_listener)
{
    pj_pool_t *pool;
    struct tcp_listener *tcp_lis;
    pj_ioqueue_callback ioqueue_cb;
    unsigned i;
    pj_status_t status;

    /* Create structure */
    pool = pj_pool_create(srv->core.pf, "tcpl%p", 1000, 1000, NULL);
    tcp_lis = PJ_POOL_ZALLOC_T(pool, struct tcp_listener);
    tcp_lis->base.pool = pool;
    tcp_lis->base.obj_name = pool->obj_name;
    tcp_lis->base.server = srv;
    tcp_lis->base.tp_type = PJ_TURN_TP_TCP;
    tcp_lis->base.sock = PJ_INVALID_SOCKET;
    //tcp_lis->base.sendto = &tcp_sendto;
    tcp_lis->base.destroy = &lis_destroy;
    tcp_lis->accept_cnt = concurrency_cnt;
    tcp_lis->base.flags = flags;

    /* Create socket */
    status = pj_sock_socket(af, pj_SOCK_STREAM(), 0, &tcp_lis->base.sock);
    if (status != PJ_SUCCESS)
	goto on_error;

    /* Init bind address */
    status = pj_sockaddr_init(af, &tcp_lis->base.addr, bound_addr, 
			      (pj_uint16_t)port);
    if (status != PJ_SUCCESS) 
	goto on_error;
    
    /* Create info */
    pj_ansi_strcpy(tcp_lis->base.info, "TCP:");
    pj_sockaddr_print(&tcp_lis->base.addr, tcp_lis->base.info+4, 
		      sizeof(tcp_lis->base.info)-4, 3);

    /* Bind socket */
    status = pj_sock_bind(tcp_lis->base.sock, &tcp_lis->base.addr, 
			  pj_sockaddr_get_len(&tcp_lis->base.addr));
    if (status != PJ_SUCCESS)
	goto on_error;

    /* Listen() */
    status = pj_sock_listen(tcp_lis->base.sock, 5);
    if (status != PJ_SUCCESS)
	goto on_error;

    /* Register to ioqueue */
    pj_bzero(&ioqueue_cb, sizeof(ioqueue_cb));
    ioqueue_cb.on_accept_complete = &lis_on_accept_complete;
    status = pj_ioqueue_register_sock(pool, srv->core.ioqueue, tcp_lis->base.sock,
				      tcp_lis, &ioqueue_cb, &tcp_lis->key);

    /* Create op keys */
    tcp_lis->accept_op = (struct accept_op*)pj_pool_calloc(pool, concurrency_cnt,
						    sizeof(struct accept_op));

    /* Create each accept_op and kick off read operation */
    for (i=0; i<concurrency_cnt; ++i) {
	lis_on_accept_complete(tcp_lis->key, &tcp_lis->accept_op[i].op_key, 
			       PJ_INVALID_SOCKET, PJ_EPENDING);
    }

    /* Done */
    PJ_LOG(4,(tcp_lis->base.obj_name, "Listener %s created", 
	   tcp_lis->base.info));

    *p_listener = &tcp_lis->base;
    return PJ_SUCCESS;


on_error:
    lis_destroy(&tcp_lis->base);
    return status;
}


/*
 * Destroy listener.
 */
static pj_status_t lis_destroy(pj_turn_listener *listener)
{
    struct tcp_listener *tcp_lis = (struct tcp_listener *)listener;
    unsigned i;

    if (tcp_lis->key) {
	pj_ioqueue_unregister(tcp_lis->key);
	tcp_lis->key = NULL;
	tcp_lis->base.sock = PJ_INVALID_SOCKET;
    } else if (tcp_lis->base.sock != PJ_INVALID_SOCKET) {
	pj_sock_close(tcp_lis->base.sock);
	tcp_lis->base.sock = PJ_INVALID_SOCKET;
    }

    for (i=0; i<tcp_lis->accept_cnt; ++i) {
	/* Nothing to do */
    }

    if (tcp_lis->base.pool) {
	pj_pool_t *pool = tcp_lis->base.pool;

	PJ_LOG(4,(tcp_lis->base.obj_name, "Listener %s destroyed", 
		  tcp_lis->base.info));

	tcp_lis->base.pool = NULL;
	pj_pool_release(pool);
    }
    return PJ_SUCCESS;
}


/*
 * Callback on new TCP connection.
 */
static void lis_on_accept_complete(pj_ioqueue_key_t *key, 
				   pj_ioqueue_op_key_t *op_key, 
				   pj_sock_t sock, 
				   pj_status_t status)
{
    struct tcp_listener *tcp_lis;
    struct accept_op *accept_op = (struct accept_op*) op_key;

    tcp_lis = (struct tcp_listener*) pj_ioqueue_get_user_data(key);

    PJ_UNUSED_ARG(sock);

    do {
	/* Report new connection. */
	if (status == PJ_SUCCESS) {
	    char addr[PJ_INET6_ADDRSTRLEN+8];
	    PJ_LOG(5,(tcp_lis->base.obj_name, "Incoming TCP from %s",
		      pj_sockaddr_print(&accept_op->src_addr, addr,
					sizeof(addr), 3)));
	    transport_create(accept_op->sock, &tcp_lis->base,
			     &accept_op->src_addr, accept_op->src_addr_len);
	} else if (status != PJ_EPENDING) {
	    show_err(tcp_lis->base.obj_name, "accept()", status);
	}

	/* Prepare next accept() */
	accept_op->src_addr_len = sizeof(accept_op->src_addr);
	status = pj_ioqueue_accept(key, op_key, &accept_op->sock,
				   NULL,
				   &accept_op->src_addr,
				   &accept_op->src_addr_len);

    } while (status != PJ_EPENDING && status != PJ_ECANCELLED &&
	     status != PJ_STATUS_FROM_OS(PJ_BLOCKING_ERROR_VAL));
}


/****************************************************************************/
/*
 * Transport
 */
enum
{
    TIMER_NONE,
    TIMER_DESTROY
};

/* The delay in seconds to be applied before TCP transport is destroyed when 
 * no allocation is referencing it. This also means the initial time to wait
 * after the initial TCP connection establishment to receive a valid STUN
 * message in the transport.
 */
#define SHUTDOWN_DELAY  10

struct recv_op
{
    pj_ioqueue_op_key_t	op_key;
    pj_turn_pkt		pkt;
};

struct tcp_transport
{
    pj_turn_transport	 base;
    pj_pool_t		*pool;
    pj_timer_entry	 timer;

    pj_turn_allocation	*alloc;
    int			 ref_cnt;

    pj_sock_t		 sock;
    pj_ioqueue_key_t	*key;
    struct recv_op	 recv_op;
    pj_ioqueue_op_key_t	 send_op;
};


static void tcp_on_read_complete(pj_ioqueue_key_t *key, 
				 pj_ioqueue_op_key_t *op_key, 
				 pj_ssize_t bytes_read);

static pj_status_t tcp_sendto(pj_turn_transport *tp,
			      const void *packet,
			      pj_size_t size,
			      unsigned flag,
			      const pj_sockaddr_t *addr,
			      int addr_len);
static void tcp_destroy(struct tcp_transport *tcp);
static void tcp_add_ref(pj_turn_transport *tp,
			pj_turn_allocation *alloc);
static void tcp_dec_ref(pj_turn_transport *tp,
			pj_turn_allocation *alloc);
static void timer_callback(pj_timer_heap_t *timer_heap,
			   pj_timer_entry *entry);

static void transport_create(pj_sock_t sock, pj_turn_listener *lis,
			     pj_sockaddr_t *src_addr, int src_addr_len)
{
    pj_pool_t *pool;
    struct tcp_transport *tcp;
    pj_ioqueue_callback cb;
    pj_status_t status;

    pool = pj_pool_create(lis->server->core.pf, "tcp%p", 1000, 1000, NULL);

    tcp = PJ_POOL_ZALLOC_T(pool, struct tcp_transport);
    tcp->base.obj_name = pool->obj_name;
    tcp->base.listener = lis;
    tcp->base.info = lis->info;
    tcp->base.sendto = &tcp_sendto;
    tcp->base.add_ref = &tcp_add_ref;
    tcp->base.dec_ref = &tcp_dec_ref;
    tcp->pool = pool;
    tcp->sock = sock;

    pj_timer_entry_init(&tcp->timer, TIMER_NONE, tcp, &timer_callback);

    /* Register to ioqueue */
    pj_bzero(&cb, sizeof(cb));
    cb.on_read_complete = &tcp_on_read_complete;
    status = pj_ioqueue_register_sock(pool, lis->server->core.ioqueue, sock,
				      tcp, &cb, &tcp->key);
    if (status != PJ_SUCCESS) {
	tcp_destroy(tcp);
	return;
    }

    /* Init pkt */
    tcp->recv_op.pkt.pool = pj_pool_create(lis->server->core.pf, "tcpkt%p", 
					   1000, 1000, NULL);
    tcp->recv_op.pkt.transport = &tcp->base;
    tcp->recv_op.pkt.src.tp_type = PJ_TURN_TP_TCP;
    tcp->recv_op.pkt.src_addr_len = src_addr_len;
    pj_memcpy(&tcp->recv_op.pkt.src.clt_addr, src_addr, src_addr_len);

    tcp_on_read_complete(tcp->key, &tcp->recv_op.op_key, -PJ_EPENDING);
    /* Should not access transport from now, it may have been destroyed */
}


static void tcp_destroy(struct tcp_transport *tcp)
{
    if (tcp->key) {
	pj_ioqueue_unregister(tcp->key);
	tcp->key = NULL;
	tcp->sock = 0;
    } else if (tcp->sock) {
	pj_sock_close(tcp->sock);
	tcp->sock = 0;
    }

    if (tcp->pool) {
	pj_pool_release(tcp->pool);
    }
}


static void timer_callback(pj_timer_heap_t *timer_heap,
			   pj_timer_entry *entry)
{
    struct tcp_transport *tcp = (struct tcp_transport*) entry->user_data;

    PJ_UNUSED_ARG(timer_heap);

    tcp_destroy(tcp);
}


static void tcp_on_read_complete(pj_ioqueue_key_t *key, 
				 pj_ioqueue_op_key_t *op_key, 
				 pj_ssize_t bytes_read)
{
    struct tcp_transport *tcp;
    struct recv_op *recv_op = (struct recv_op*) op_key;
    pj_status_t status;

    tcp = (struct tcp_transport*) pj_ioqueue_get_user_data(key);

    do {
	/* Report to server or allocation, if we have allocation */
	if (bytes_read > 0) {

	    recv_op->pkt.len = bytes_read;
	    pj_gettimeofday(&recv_op->pkt.rx_time);

	    tcp_add_ref(&tcp->base, NULL);

	    if (tcp->alloc) {
		pj_turn_allocation_on_rx_client_pkt(tcp->alloc, &recv_op->pkt);
	    } else {
		pj_turn_srv_on_rx_pkt(tcp->base.listener->server, &recv_op->pkt);
	    }

	    pj_assert(tcp->ref_cnt > 0);
	    tcp_dec_ref(&tcp->base, NULL);

	} else if (bytes_read != -PJ_EPENDING) {
	    /* TCP connection closed/error. Notify client and then destroy 
	     * ourselves.
	     * Note: the -PJ_EPENDING is the value passed during init.
	     */
	    ++tcp->ref_cnt;

	    if (tcp->alloc) {
		if (bytes_read != 0) {
		    show_err(tcp->base.obj_name, "TCP socket error", 
			     -bytes_read);
		} else {
		    PJ_LOG(5,(tcp->base.obj_name, "TCP socket closed"));
		}
		pj_turn_allocation_on_transport_closed(tcp->alloc, &tcp->base);
		tcp->alloc = NULL;
	    }

	    pj_assert(tcp->ref_cnt > 0);
	    if (--tcp->ref_cnt == 0) {
		tcp_destroy(tcp);
		return;
	    }
	}

	/* Reset pool */
	pj_pool_reset(recv_op->pkt.pool);

	/* If packet is full discard it */
	if (recv_op->pkt.len == sizeof(recv_op->pkt.pkt)) {
	    PJ_LOG(4,(tcp->base.obj_name, "Buffer discarded"));
	    recv_op->pkt.len = 0;
	}

	/* Read next packet */
	bytes_read = sizeof(recv_op->pkt.pkt) - recv_op->pkt.len;
	status = pj_ioqueue_recv(tcp->key, op_key,
				 recv_op->pkt.pkt + recv_op->pkt.len, 
				 &bytes_read, 0);

	if (status != PJ_EPENDING && status != PJ_SUCCESS)
	    bytes_read = -status;

    } while (status != PJ_EPENDING && status != PJ_ECANCELLED &&
	     status != PJ_STATUS_FROM_OS(PJ_BLOCKING_ERROR_VAL));

}


static pj_status_t tcp_sendto(pj_turn_transport *tp,
			      const void *packet,
			      pj_size_t size,
			      unsigned flag,
			      const pj_sockaddr_t *addr,
			      int addr_len)
{
    struct tcp_transport *tcp = (struct tcp_transport*) tp;
    pj_ssize_t length = size;

    PJ_UNUSED_ARG(addr);
    PJ_UNUSED_ARG(addr_len);

    return pj_ioqueue_send(tcp->key, &tcp->send_op, packet, &length, flag);
}


static void tcp_add_ref(pj_turn_transport *tp,
			pj_turn_allocation *alloc)
{
    struct tcp_transport *tcp = (struct tcp_transport*) tp;

    ++tcp->ref_cnt;

    if (tcp->alloc == NULL && alloc) {
	tcp->alloc = alloc;
    }

    /* Cancel shutdown timer if it's running */
    if (tcp->timer.id != TIMER_NONE) {
	pj_timer_heap_cancel(tcp->base.listener->server->core.timer_heap,
			     &tcp->timer);
	tcp->timer.id = TIMER_NONE;
    }
}


static void tcp_dec_ref(pj_turn_transport *tp,
			pj_turn_allocation *alloc)
{
    struct tcp_transport *tcp = (struct tcp_transport*) tp;

    --tcp->ref_cnt;

    if (alloc && alloc == tcp->alloc) {
	tcp->alloc = NULL;
    }

    if (tcp->ref_cnt == 0 && tcp->timer.id == TIMER_NONE) {
	pj_time_val delay = { SHUTDOWN_DELAY, 0 };
	tcp->timer.id = TIMER_DESTROY;
	pj_timer_heap_schedule(tcp->base.listener->server->core.timer_heap,
			       &tcp->timer, &delay);
    }
}

#else	/* PJ_HAS_TCP */

/* To avoid empty translation unit warning */
int listener_tcp_dummy = 0;

#endif	/* PJ_HAS_TCP */