/* $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 "test.h"
#include "server.h"

#define SRV_DOMAIN	"pjsip.lab.domain"
#define KA_INTERVAL	50

struct test_result
{
    unsigned    state_called;
    unsigned    rx_data_cnt;
};

struct test_session
{
    pj_pool_t		*pool;
    pj_stun_config	*stun_cfg;
    pj_turn_sock	*turn_sock;
    pj_dns_resolver	*resolver;
    test_server		*test_srv;

    pj_bool_t		 destroy_called;
    int			 destroy_on_state;
    struct test_result	 result;
};

struct test_session_cfg
{
    struct {
	pj_bool_t	enable_dns_srv;
	int		destroy_on_state;
    } client;

    struct {
	pj_uint32_t	flags;
	pj_bool_t	respond_allocate;
	pj_bool_t	respond_refresh;
    } srv;
};

static void turn_on_rx_data(pj_turn_sock *turn_sock,
			    void *pkt,
			    unsigned pkt_len,
			    const pj_sockaddr_t *peer_addr,
			    unsigned addr_len);
static void turn_on_state(pj_turn_sock *turn_sock, 
			  pj_turn_state_t old_state,
			  pj_turn_state_t new_state);

static void destroy_session(struct test_session *sess)
{
    if (sess->resolver) {
	pj_dns_resolver_destroy(sess->resolver, PJ_TRUE);
	sess->resolver = NULL;
    }

    if (sess->turn_sock) {
	if (!sess->destroy_called) {
	    sess->destroy_called = PJ_TRUE;
	    pj_turn_sock_destroy(sess->turn_sock);
	}
	sess->turn_sock = NULL;
    }

    if (sess->test_srv) {
	destroy_test_server(sess->test_srv);
	sess->test_srv = NULL;
    }

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

pj_turn_tp_type get_turn_tp_type(pj_uint32_t flag) {
    if (flag & TURN_TCP) {
	return PJ_TURN_TP_TCP;
    } else if (flag & TURN_TLS) {
	return PJ_TURN_TP_TLS;
    } 
    return PJ_TURN_TP_UDP;
}

static int create_test_session(pj_stun_config  *stun_cfg,
			       const struct test_session_cfg *cfg,
			       struct test_session **p_sess)
{
    struct test_session *sess;
    pj_pool_t *pool;
    pj_turn_sock_cb turn_sock_cb;
    pj_turn_alloc_param alloc_param;
    pj_stun_auth_cred cred;
    pj_status_t status;
    pj_bool_t use_ipv6 = cfg->srv.flags & SERVER_IPV6;
    pj_turn_tp_type tp_type = get_turn_tp_type(cfg->srv.flags);

    /* Create client */
    pool = pj_pool_create(mem, "turnclient", 512, 512, NULL);
    sess = PJ_POOL_ZALLOC_T(pool, struct test_session);
    sess->pool = pool;
    sess->stun_cfg = stun_cfg;
    sess->destroy_on_state = cfg->client.destroy_on_state;

    pj_bzero(&turn_sock_cb, sizeof(turn_sock_cb));
    turn_sock_cb.on_rx_data = &turn_on_rx_data;
    turn_sock_cb.on_state = &turn_on_state;
    status = pj_turn_sock_create(sess->stun_cfg,
				 GET_AF(use_ipv6),
				 tp_type, 
				 &turn_sock_cb, 
				 0, 
				 sess, 
				 &sess->turn_sock);
    if (status != PJ_SUCCESS) {
	destroy_session(sess);
	return -20;
    }

    /* Create test server */
    status = create_test_server(sess->stun_cfg, cfg->srv.flags, SRV_DOMAIN, 
				&sess->test_srv);
    if (status != PJ_SUCCESS) {
	destroy_session(sess);
	return -30;
    }

    sess->test_srv->turn_respond_allocate = cfg->srv.respond_allocate;
    sess->test_srv->turn_respond_refresh = cfg->srv.respond_refresh;

    /* Create client resolver */
    status = pj_dns_resolver_create(mem, "resolver", 0, sess->stun_cfg->timer_heap,
				    sess->stun_cfg->ioqueue, &sess->resolver);
    if (status != PJ_SUCCESS) {
	destroy_session(sess);
	return -40;

    } else {
	pj_str_t dns_srv = use_ipv6?pj_str("::1") : pj_str("127.0.0.1");
	pj_uint16_t dns_srv_port = (pj_uint16_t) DNS_SERVER_PORT;
	status = pj_dns_resolver_set_ns(sess->resolver, 1, &dns_srv, &dns_srv_port);

	if (status != PJ_SUCCESS) {
	    destroy_session(sess);
	    return -50;
	}
    }

    /* Init TURN credential */
    pj_bzero(&cred, sizeof(cred));
    cred.type = PJ_STUN_AUTH_CRED_STATIC;
    cred.data.static_cred.realm = pj_str(SRV_DOMAIN);
    cred.data.static_cred.username = pj_str(TURN_USERNAME);
    cred.data.static_cred.data_type = PJ_STUN_PASSWD_PLAIN;
    cred.data.static_cred.data = pj_str(TURN_PASSWD);

    /* Init TURN allocate parameter */
    pj_turn_alloc_param_default(&alloc_param);
    alloc_param.ka_interval = KA_INTERVAL;

    /* Start the client */
    if (cfg->client.enable_dns_srv) {
	/* Use DNS SRV to resolve server, may fallback to DNS A */
	pj_str_t domain = pj_str(SRV_DOMAIN);
	status = pj_turn_sock_alloc(sess->turn_sock, &domain, TURN_SERVER_PORT,
				    sess->resolver, &cred, &alloc_param);

    } else {
	/* Explicitly specify server address */
	pj_str_t host = use_ipv6?pj_str("::1") : pj_str("127.0.0.1");
	status = pj_turn_sock_alloc(sess->turn_sock, &host, TURN_SERVER_PORT,
				    NULL, &cred, &alloc_param);

    }

    if (status != PJ_SUCCESS) {
	if (cfg->client.destroy_on_state >= PJ_TURN_STATE_READY) {
	    destroy_session(sess);
	    return -70;
	}
    }

    *p_sess = sess;
    return 0;
}


static void turn_on_rx_data(pj_turn_sock *turn_sock,
			    void *pkt,
			    unsigned pkt_len,
			    const pj_sockaddr_t *peer_addr,
			    unsigned addr_len)
{
    struct test_session *sess;

    PJ_UNUSED_ARG(pkt);
    PJ_UNUSED_ARG(pkt_len);
    PJ_UNUSED_ARG(peer_addr);
    PJ_UNUSED_ARG(addr_len);

    sess = (struct test_session*) pj_turn_sock_get_user_data(turn_sock);
    if (sess == NULL)
	return;

    sess->result.rx_data_cnt++;
}


static void turn_on_state(pj_turn_sock *turn_sock, 
			  pj_turn_state_t old_state,
			  pj_turn_state_t new_state)
{
    struct test_session *sess;
    unsigned i, mask;

    PJ_UNUSED_ARG(old_state);

    sess = (struct test_session*) pj_turn_sock_get_user_data(turn_sock);
    if (sess == NULL)
	return;

    /* This state must not be called before */
    pj_assert((sess->result.state_called & (1<<new_state)) == 0);

    /* new_state must be greater than old_state */
    pj_assert(new_state > old_state);

    /* must not call any greater state before */
    mask = 0;
    for (i=new_state+1; i<31; ++i) mask |= (1 << i);

    pj_assert((sess->result.state_called & mask) == 0);

    sess->result.state_called |= (1 << new_state);

    if (new_state >= sess->destroy_on_state && !sess->destroy_called) {
	sess->destroy_called = PJ_TRUE;
	pj_turn_sock_destroy(turn_sock);
    }

    if (new_state >= PJ_TURN_STATE_DESTROYING) {
	pj_turn_sock_set_user_data(sess->turn_sock, NULL);
	sess->turn_sock = NULL;
    }
}


/////////////////////////////////////////////////////////////////////

static void set_server_flag(struct test_session_cfg *test_cfg, 
			    pj_bool_t use_ipv6,
			    pj_turn_tp_type tp_type)
{
    pj_uint32_t flag = TURN_UDP;
    test_cfg->srv.flags &= ~(SERVER_IPV4+SERVER_IPV6+
			     TURN_UDP+TURN_TCP+TURN_TLS);    
    switch (tp_type) {
	case PJ_TURN_TP_TCP:
	    flag = TURN_TCP;
	    break;
	case PJ_TURN_TP_TLS:
	    flag = TURN_TLS;
	    break;
	default:
	    break;
    }
    test_cfg->srv.flags |= ((use_ipv6)?SERVER_IPV6:SERVER_IPV4)+flag;
}

static int state_progression_test(pj_stun_config  *stun_cfg, 
				  pj_bool_t use_ipv6,
				  pj_turn_tp_type tp_type)
{
    struct test_session_cfg test_cfg = 
    {
	{   /* Client cfg */			
	    PJ_TRUE,	    /* DNS SRV */       
	    0xFFFF	    /* Destroy on state */
	},
	{   /* Server cfg */
	    0xFFFFFFFF,	    /* flags */
	    PJ_TRUE,	    /* respond to allocate  */
	    PJ_TRUE	    /* respond to refresh   */
	}
    };
    struct test_session *sess;
    unsigned i;
    int rc = 0;

    PJ_LOG(3,("", "  state progression tests - (%s) (%s)",
	      use_ipv6?"IPv6":"IPv4",
	      (tp_type==PJ_TURN_TP_UDP)?"UDP":
		(tp_type==PJ_TURN_TP_TCP)?"TCP":"TLS"));
    
    set_server_flag(&test_cfg, use_ipv6, tp_type);
    for (i=0; i<=1; ++i) {
	enum { TIMEOUT = 60 };
	pjlib_state pjlib_state;
	pj_turn_session_info info;
	struct test_result result;
	pj_time_val tstart;

	PJ_LOG(3,("", "   %s DNS SRV resolution",
	              (i==0? "without" : "with")));

	capture_pjlib_state(stun_cfg, &pjlib_state);

	test_cfg.client.enable_dns_srv = i;

	rc = create_test_session(stun_cfg, &test_cfg, &sess);
	if (rc != 0)
	    return rc;

	pj_bzero(&info, sizeof(info));

	/* Wait until state is READY */
	pj_gettimeofday(&tstart);
	while (sess->turn_sock) {
	    pj_time_val now;

	    poll_events(stun_cfg, 10, PJ_FALSE);
	    if (sess->turn_sock == NULL) {
		break;
	    }
	    rc = pj_turn_sock_get_info(sess->turn_sock, &info);
	    if (rc!=PJ_SUCCESS)
		break;

	    if (info.state >= PJ_TURN_STATE_READY)
		break;

	    pj_gettimeofday(&now);
	    if (now.sec - tstart.sec > TIMEOUT) {
		PJ_LOG(3,("", "    timed-out"));
		break;
	    }
	}

	if (info.state != PJ_TURN_STATE_READY) {
	    PJ_LOG(3,("", "    error: state is not READY"));
	    destroy_session(sess);
	    return -130;
	}

	/* Deallocate */
	pj_turn_sock_destroy(sess->turn_sock);

	/* Wait for couple of seconds.
	 * We can't poll the session info since the session may have
	 * been destroyed
	 */
	poll_events(stun_cfg, 2000, PJ_FALSE);
	sess->turn_sock = NULL;
	pj_memcpy(&result, &sess->result, sizeof(result));
	destroy_session(sess);

	/* Check the result */
	if ((result.state_called & (1<<PJ_TURN_STATE_RESOLVING)) == 0) {
	    PJ_LOG(3,("", "    error: PJ_TURN_STATE_RESOLVING is not called"));
	    return -140;
	}

	if ((result.state_called & (1<<PJ_TURN_STATE_RESOLVED)) == 0) {
	    PJ_LOG(3,("", "    error: PJ_TURN_STATE_RESOLVED is not called"));
	    return -150;
	}

	if ((result.state_called & (1<<PJ_TURN_STATE_ALLOCATING)) == 0) {
	    PJ_LOG(3,("", "    error: PJ_TURN_STATE_ALLOCATING is not called"));
	    return -155;
	}

	if ((result.state_called & (1<<PJ_TURN_STATE_READY)) == 0) {
	    PJ_LOG(3,("", "    error: PJ_TURN_STATE_READY is not called"));
	    return -160;
	}

	if ((result.state_called & (1<<PJ_TURN_STATE_DEALLOCATING)) == 0) {
	    PJ_LOG(3,("", "    error: PJ_TURN_STATE_DEALLOCATING is not called"));
	    return -170;
	}

	if ((result.state_called & (1<<PJ_TURN_STATE_DEALLOCATED)) == 0) {
	    PJ_LOG(3,("", "    error: PJ_TURN_STATE_DEALLOCATED is not called"));
	    return -180;
	}

	if ((result.state_called & (1<<PJ_TURN_STATE_DESTROYING)) == 0) {
	    PJ_LOG(3,("", "    error: PJ_TURN_STATE_DESTROYING is not called"));
	    return -190;
	}

	poll_events(stun_cfg, 500, PJ_FALSE);
	rc = check_pjlib_state(stun_cfg, &pjlib_state);
	if (rc != 0) {
	    PJ_LOG(3,("", "    error: memory/timer-heap leak detected"));
	    return rc;
	}
    }

    if (use_ipv6)
	rc = state_progression_test(stun_cfg, 0, tp_type);

    return rc;
}


/////////////////////////////////////////////////////////////////////

static int destroy_test(pj_stun_config  *stun_cfg,
			pj_bool_t with_dns_srv,
			pj_bool_t in_callback,
			pj_bool_t use_ipv6,
			pj_turn_tp_type tp_type)
{
    struct test_session_cfg test_cfg = 
    {
	{   /* Client cfg */	    
	    PJ_TRUE,	    /* DNS SRV */   
	    0xFFFF	    /* Destroy on state */
	},
	{   /* Server cfg */
	    0xFFFFFFFF,	    /* flags */
	    PJ_TRUE,	    /* respond to allocate  */
	    PJ_TRUE	    /* respond to refresh   */
	}
    };
    struct test_session *sess;
    int target_state;
    int rc;

    PJ_LOG(3,("", "  destroy test %s %s (%s)",
	          (in_callback? "in callback" : ""),
		  (with_dns_srv? "with DNS srv" : ""),
		  (tp_type==PJ_TURN_TP_UDP)?"UDP":
		    (tp_type==PJ_TURN_TP_TCP)?"TCP":"TLS"));

    test_cfg.client.enable_dns_srv = with_dns_srv;
    set_server_flag(&test_cfg, use_ipv6, tp_type);    

    for (target_state=PJ_TURN_STATE_RESOLVING; target_state<=PJ_TURN_STATE_READY; ++target_state) {
	enum { TIMEOUT = 60 };
	pjlib_state pjlib_state;
	pj_turn_session_info info;
	pj_time_val tstart;

	capture_pjlib_state(stun_cfg, &pjlib_state);

	PJ_LOG(3,("", "   %s", pj_turn_state_name((pj_turn_state_t)target_state)));

	if (in_callback)
	    test_cfg.client.destroy_on_state = target_state;

	rc = create_test_session(stun_cfg, &test_cfg, &sess);
	if (rc != 0)
	    return rc;

	if (in_callback) {
	    pj_gettimeofday(&tstart);
	    rc = 0;
	    while (sess->turn_sock) {
		pj_time_val now;

		poll_events(stun_cfg, 100, PJ_FALSE);

		pj_gettimeofday(&now);
		if (now.sec - tstart.sec > TIMEOUT) {
		    rc = -7;
		    break;
		}
	    }

	} else {
	    pj_gettimeofday(&tstart);
	    rc = 0;
	    while (sess->turn_sock) {
		pj_time_val now;

		poll_events(stun_cfg, 1, PJ_FALSE);

		pj_turn_sock_get_info(sess->turn_sock, &info);
		
		if (info.state >= target_state) {
		    pj_turn_sock_destroy(sess->turn_sock);
		    break;
		}

		pj_gettimeofday(&now);
		if (now.sec - tstart.sec > TIMEOUT) {
		    rc = -8;
		    break;
		}
	    }
	}


	if (rc != 0) {
	    PJ_LOG(3,("", "    error: timeout"));
	    return rc;
	}

	poll_events(stun_cfg, 1000, PJ_FALSE);
	destroy_session(sess);

	rc = check_pjlib_state(stun_cfg, &pjlib_state);
	if (rc != 0) {
	    PJ_LOG(3,("", "    error: memory/timer-heap leak detected"));
	    return rc;
	}
    }

    return 0;
}


/////////////////////////////////////////////////////////////////////

int turn_sock_test(void)
{
    pj_pool_t *pool;
    pj_stun_config stun_cfg;
    int n, i, rc = 0;

    pool = pj_pool_create(mem, "turntest", 512, 512, NULL);
    rc = create_stun_config(pool, &stun_cfg);
    if (rc != PJ_SUCCESS) {
	pj_pool_release(pool);
	return -2;
    }

    for (n = 0; n <= 2; ++n) {
	pj_turn_tp_type tp_type = PJ_TURN_TP_UDP;

	if ((n == 2) && !USE_TLS)
	    break;

	switch (n) {
	case 1:
	    tp_type = PJ_TURN_TP_TCP;
	    break;
	case 2:
	    tp_type = PJ_TURN_TP_TLS;
	}

	rc = state_progression_test(&stun_cfg, USE_IPV6, tp_type);
	if (rc != 0) 
	    goto on_return;

	for (i=0; i<=1; ++i) {
	    int j;
	    for (j=0; j<=1; ++j) {
		rc = destroy_test(&stun_cfg, i, j, USE_IPV6, tp_type);
		if (rc != 0)
		    goto on_return;
	    }
	}
    }

on_return:
    destroy_stun_config(&stun_cfg);
    pj_pool_release(pool);
    return rc;
}