/*
 * 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 "auth.h"

#define MAX_CLIENTS             32
#define MAX_PEERS_PER_CLIENT    8
//#define MAX_HANDLES           (MAX_CLIENTS*MAX_PEERS_PER_CLIENT+MAX_LISTENERS)
#define MAX_HANDLES             PJ_IOQUEUE_MAX_HANDLES
#define MAX_TIMER               (MAX_HANDLES * 2)
#define MIN_PORT                49152
#define MAX_PORT                65535
#define MAX_LISTENERS           16
#define MAX_THREADS             2
#define MAX_NET_EVENTS          1000

/* Prototypes */
static int server_thread_proc(void *arg);
static pj_status_t on_tx_stun_msg( pj_stun_session *sess,
                                   void *token,
                                   const void *pkt,
                                   pj_size_t pkt_size,
                                   const pj_sockaddr_t *dst_addr,
                                   unsigned addr_len);
static pj_status_t on_rx_stun_request(pj_stun_session *sess,
                                      const pj_uint8_t *pkt,
                                      unsigned pkt_len,
                                      const pj_stun_rx_data *rdata,
                                      void *user_data,
                                      const pj_sockaddr_t *src_addr,
                                      unsigned src_addr_len);

struct saved_cred
{
    pj_str_t realm;
    pj_str_t username;
    pj_str_t nonce;
    int      data_type;
    pj_str_t data;
};


/*
 * Get transport type name, normally for logging purpose only.
 */
PJ_DEF(const char*) pj_turn_tp_type_name(int tp_type)
{
    /* Must be 3 characters long! */
    if (tp_type == PJ_TURN_TP_UDP) {
        return "UDP";
    } else if (tp_type == PJ_TURN_TP_TCP) {
        return "TCP";
    } else {
        pj_assert(!"Unsupported transport");
        return "???";
    }
}

/*
 * Create server.
 */
PJ_DEF(pj_status_t) pj_turn_srv_create(pj_pool_factory *pf,
                                       pj_turn_srv **p_srv)
{
    pj_pool_t *pool;
    pj_stun_session_cb sess_cb;
    pj_turn_srv *srv;
    unsigned i;
    pj_status_t status;

    PJ_ASSERT_RETURN(pf && p_srv, PJ_EINVAL);

    /* Create server and init core settings */
    pool = pj_pool_create(pf, "srv%p", 1000, 1000, NULL);
    srv = PJ_POOL_ZALLOC_T(pool, pj_turn_srv);
    srv->obj_name = pool->obj_name;
    srv->core.pf = pf;
    srv->core.pool = pool;
    srv->core.tls_key = srv->core.tls_data = -1;

    /* Create ioqueue */
    status = pj_ioqueue_create(pool, MAX_HANDLES, &srv->core.ioqueue);
    if (status != PJ_SUCCESS)
        goto on_error;

    /* Server mutex */
    status = pj_lock_create_recursive_mutex(pool, srv->obj_name,
                                            &srv->core.lock);
    if (status != PJ_SUCCESS)
        goto on_error;

    /* Allocate TLS */
    status = pj_thread_local_alloc(&srv->core.tls_key);
    if (status != PJ_SUCCESS)
        goto on_error;

    status = pj_thread_local_alloc(&srv->core.tls_data);
    if (status != PJ_SUCCESS)
        goto on_error;

    /* Create timer heap */
    status = pj_timer_heap_create(pool, MAX_TIMER, &srv->core.timer_heap);
    if (status != PJ_SUCCESS)
        goto on_error;

    /* Configure lock for the timer heap */
    pj_timer_heap_set_lock(srv->core.timer_heap, srv->core.lock, PJ_FALSE);

    /* Array of listeners */
    srv->core.listener = (pj_turn_listener**)
                         pj_pool_calloc(pool, MAX_LISTENERS,
                                        sizeof(srv->core.listener[0]));

    /* Create hash tables */
    srv->tables.alloc = pj_hash_create(pool, MAX_CLIENTS);
    srv->tables.res = pj_hash_create(pool, MAX_CLIENTS);

    /* Init ports settings */
    srv->ports.min_udp = srv->ports.next_udp = MIN_PORT;
    srv->ports.max_udp = MAX_PORT;
    srv->ports.min_tcp = srv->ports.next_tcp = MIN_PORT;
    srv->ports.max_tcp = MAX_PORT;

    /* Init STUN config */
    pj_stun_config_init(&srv->core.stun_cfg, pf, 0, srv->core.ioqueue,
                        srv->core.timer_heap);

    /* Init STUN credential */
    srv->core.cred.type = PJ_STUN_AUTH_CRED_DYNAMIC;
    srv->core.cred.data.dyn_cred.user_data = srv;
    srv->core.cred.data.dyn_cred.get_auth = &pj_turn_get_auth;
    srv->core.cred.data.dyn_cred.get_password = &pj_turn_get_password;
    srv->core.cred.data.dyn_cred.verify_nonce = &pj_turn_verify_nonce;

    /* Create STUN session to handle new allocation */
    pj_bzero(&sess_cb, sizeof(sess_cb));
    sess_cb.on_rx_request = &on_rx_stun_request;
    sess_cb.on_send_msg = &on_tx_stun_msg;

    status = pj_stun_session_create(&srv->core.stun_cfg, srv->obj_name,
                                    &sess_cb, PJ_FALSE, NULL,
                                    &srv->core.stun_sess);
    if (status != PJ_SUCCESS) {
        goto on_error;
    }

    pj_stun_session_set_user_data(srv->core.stun_sess, srv);
    pj_stun_session_set_credential(srv->core.stun_sess, PJ_STUN_AUTH_LONG_TERM,
                                   &srv->core.cred);


    /* Array of worker threads */
    srv->core.thread_cnt = MAX_THREADS;
    srv->core.thread = (pj_thread_t**)
                       pj_pool_calloc(pool, srv->core.thread_cnt,
                                      sizeof(pj_thread_t*));

    /* Start the worker threads */
    for (i=0; i<srv->core.thread_cnt; ++i) {
        status = pj_thread_create(pool, srv->obj_name, &server_thread_proc,
                                  srv, 0, 0, &srv->core.thread[i]);
        if (status != PJ_SUCCESS)
            goto on_error;
    }

    /* We're done. Application should add listeners now */
    PJ_LOG(4,(srv->obj_name, "TURN server v%s is running",
              pj_get_version()));

    *p_srv = srv;
    return PJ_SUCCESS;

on_error:
    pj_turn_srv_destroy(srv);
    return status;
}


/*
 * Handle timer and network events
 */
static void srv_handle_events(pj_turn_srv *srv, const pj_time_val *max_timeout)
{
    /* timeout is 'out' var. This just to make compiler happy. */
    pj_time_val timeout = { 0, 0};
    unsigned net_event_count = 0;
    int c;

    /* Poll the timer. The timer heap has its own mutex for better
     * granularity, so we don't need to lock the server.
     */
    timeout.sec = timeout.msec = 0;
    c = pj_timer_heap_poll( srv->core.timer_heap, &timeout );

    /* timer_heap_poll should never ever returns negative value, or otherwise
     * ioqueue_poll() will block forever!
     */
    pj_assert(timeout.sec >= 0 && timeout.msec >= 0);
    if (timeout.msec >= 1000) timeout.msec = 999;

    /* If caller specifies maximum time to wait, then compare the value with
     * the timeout to wait from timer, and use the minimum value.
     */
    if (max_timeout && PJ_TIME_VAL_GT(timeout, *max_timeout)) {
        timeout = *max_timeout;
    }

    /* Poll ioqueue.
     * Repeat polling the ioqueue while we have immediate events, because
     * timer heap may process more than one events, so if we only process
     * one network events at a time (such as when IOCP backend is used),
     * the ioqueue may have trouble keeping up with the request rate.
     *
     * For example, for each send() request, one network event will be
     *   reported by ioqueue for the send() completion. If we don't poll
     *   the ioqueue often enough, the send() completion will not be
     *   reported in timely manner.
     */
    do {
        c = pj_ioqueue_poll( srv->core.ioqueue, &timeout);
        if (c < 0) {
            pj_thread_sleep(PJ_TIME_VAL_MSEC(timeout));
            return;
        } else if (c == 0) {
            break;
        } else {
            net_event_count += c;
            timeout.sec = timeout.msec = 0;
        }
    } while (c > 0 && net_event_count < MAX_NET_EVENTS);

}

/*
 * Server worker thread proc.
 */
static int server_thread_proc(void *arg)
{
    pj_turn_srv *srv = (pj_turn_srv*)arg;

    while (!srv->core.quit) {
        pj_time_val timeout_max = {0, 100};
        srv_handle_events(srv, &timeout_max);
    }

    return 0;
}

/*
 * Destroy the server.
 */
PJ_DEF(pj_status_t) pj_turn_srv_destroy(pj_turn_srv *srv)
{
    pj_hash_iterator_t itbuf, *it;
    unsigned i;

    /* Stop all worker threads */
    srv->core.quit = PJ_TRUE;
    for (i=0; i<srv->core.thread_cnt; ++i) {
        if (srv->core.thread[i]) {
            pj_thread_join(srv->core.thread[i]);
            pj_thread_destroy(srv->core.thread[i]);
            srv->core.thread[i] = NULL;
        }
    }

    /* Destroy all allocations FIRST */
    if (srv->tables.alloc) {
        it = pj_hash_first(srv->tables.alloc, &itbuf);
        while (it != NULL) {
            pj_turn_allocation *alloc = (pj_turn_allocation*)
                                        pj_hash_this(srv->tables.alloc, it);
            pj_hash_iterator_t *next = pj_hash_next(srv->tables.alloc, it);
            pj_turn_allocation_destroy(alloc);
            it = next;
        }
    }

    /* Destroy all listeners. */
    for (i=0; i<srv->core.lis_cnt; ++i) {
        if (srv->core.listener[i]) {
            pj_turn_listener_destroy(srv->core.listener[i]);
            srv->core.listener[i] = NULL;
        }
    }

    /* Destroy STUN session */
    if (srv->core.stun_sess) {
        pj_stun_session_destroy(srv->core.stun_sess);
        srv->core.stun_sess = NULL;
    }

    /* Destroy hash tables (well, sort of) */
    if (srv->tables.alloc) {
        srv->tables.alloc = NULL;
        srv->tables.res = NULL;
    }

    /* Destroy timer heap */
    if (srv->core.timer_heap) {
        pj_timer_heap_destroy(srv->core.timer_heap);
        srv->core.timer_heap = NULL;
    }

    /* Destroy ioqueue */
    if (srv->core.ioqueue) {
        pj_ioqueue_destroy(srv->core.ioqueue);
        srv->core.ioqueue = NULL;
    }

    /* Destroy thread local IDs */
    if (srv->core.tls_key != -1) {
        pj_thread_local_free(srv->core.tls_key);
        srv->core.tls_key = -1;
    }
    if (srv->core.tls_data != -1) {
        pj_thread_local_free(srv->core.tls_data);
        srv->core.tls_data = -1;
    }

    /* Destroy server lock */
    if (srv->core.lock) {
        pj_lock_destroy(srv->core.lock);
        srv->core.lock = NULL;
    }

    /* Release pool */
    if (srv->core.pool) {
        pj_pool_t *pool = srv->core.pool;
        srv->core.pool = NULL;
        pj_pool_release(pool);
    }

    /* Done */
    return PJ_SUCCESS;
}


/*
 * Add listener.
 */
PJ_DEF(pj_status_t) pj_turn_srv_add_listener(pj_turn_srv *srv,
                                             pj_turn_listener *lis)
{
    unsigned index;

    PJ_ASSERT_RETURN(srv && lis, PJ_EINVAL);
    PJ_ASSERT_RETURN(srv->core.lis_cnt < MAX_LISTENERS, PJ_ETOOMANY);

    /* Add to array */
    index = srv->core.lis_cnt;
    srv->core.listener[index] = lis;
    lis->server = srv;
    lis->id = index;
    srv->core.lis_cnt++;

    PJ_LOG(4,(srv->obj_name, "Listener %s/%s added at index %d",
              lis->obj_name, lis->info, lis->id));

    return PJ_SUCCESS;
}


/*
 * Destroy listener.
 */
PJ_DEF(pj_status_t) pj_turn_listener_destroy(pj_turn_listener *listener)
{
    pj_turn_srv *srv = listener->server;
    unsigned i;

    /* Remove from our listener list */
    pj_lock_acquire(srv->core.lock);
    for (i=0; i<srv->core.lis_cnt; ++i) {
        if (srv->core.listener[i] == listener) {
            srv->core.listener[i] = NULL;
            srv->core.lis_cnt--;
            listener->id = PJ_TURN_INVALID_LIS_ID;
            break;
        }
    }
    pj_lock_release(srv->core.lock);

    /* Destroy */
    return listener->destroy(listener);
}


/**
 * Add a reference to a transport.
 */
PJ_DEF(void) pj_turn_transport_add_ref( pj_turn_transport *transport,
                                        pj_turn_allocation *alloc)
{
    transport->add_ref(transport, alloc);
}


/**
 * Decrement transport reference counter.
 */
PJ_DEF(void) pj_turn_transport_dec_ref( pj_turn_transport *transport,
                                        pj_turn_allocation *alloc)
{
    transport->dec_ref(transport, alloc);
}


/*
 * Register an allocation to the hash tables.
 */
PJ_DEF(pj_status_t) pj_turn_srv_register_allocation(pj_turn_srv *srv,
                                                    pj_turn_allocation *alloc)
{
    /* Add to hash tables */
    pj_lock_acquire(srv->core.lock);
    pj_hash_set(alloc->pool, srv->tables.alloc,
                &alloc->hkey, sizeof(alloc->hkey), 0, alloc);
    pj_hash_set(alloc->pool, srv->tables.res,
                &alloc->relay.hkey, sizeof(alloc->relay.hkey), 0,
                &alloc->relay);
    pj_lock_release(srv->core.lock);

    return PJ_SUCCESS;
}


/*
 * Unregister an allocation from the hash tables.
 */
PJ_DEF(pj_status_t) pj_turn_srv_unregister_allocation(pj_turn_srv *srv,
                                                     pj_turn_allocation *alloc)
{
    /* Unregister from hash tables */
    pj_lock_acquire(srv->core.lock);
    pj_hash_set(alloc->pool, srv->tables.alloc,
                &alloc->hkey, sizeof(alloc->hkey), 0, NULL);
    pj_hash_set(alloc->pool, srv->tables.res,
                &alloc->relay.hkey, sizeof(alloc->relay.hkey), 0, NULL);
    pj_lock_release(srv->core.lock);

    return PJ_SUCCESS;
}


/* Callback from our own STUN session whenever it needs to send
 * outgoing STUN packet.
 */
static pj_status_t on_tx_stun_msg( pj_stun_session *sess,
                                   void *token,
                                   const void *pdu,
                                   pj_size_t pdu_size,
                                   const pj_sockaddr_t *dst_addr,
                                   unsigned addr_len)
{
    pj_turn_transport *transport = (pj_turn_transport*) token;

    PJ_ASSERT_RETURN(transport!=NULL, PJ_EINVALIDOP);

    PJ_UNUSED_ARG(sess);

    return transport->sendto(transport, pdu, pdu_size, 0,
                             dst_addr, addr_len);
}


/* Respond to STUN request */
static pj_status_t stun_respond(pj_stun_session *sess,
                                pj_turn_transport *transport,
                                const pj_stun_rx_data *rdata,
                                unsigned code,
                                const char *errmsg,
                                pj_bool_t cache,
                                const pj_sockaddr_t *dst_addr,
                                unsigned addr_len)
{
    pj_status_t status;
    pj_str_t reason;
    pj_stun_tx_data *tdata;

    /* Create response */
    status = pj_stun_session_create_res(sess, rdata, code,
                                        (errmsg?pj_cstr(&reason,errmsg):NULL),
                                        &tdata);
    if (status != PJ_SUCCESS)
        return status;

    /* Send the response */
    return pj_stun_session_send_msg(sess, transport, cache, PJ_FALSE,
                                    dst_addr,  addr_len, tdata);
}


/* Callback from our own STUN session when incoming request arrives.
 * This function is triggered by pj_stun_session_on_rx_pkt() call in
 * pj_turn_srv_on_rx_pkt() function below.
 */
static pj_status_t on_rx_stun_request(pj_stun_session *sess,
                                      const pj_uint8_t *pdu,
                                      unsigned pdu_len,
                                      const pj_stun_rx_data *rdata,
                                      void *token,
                                      const pj_sockaddr_t *src_addr,
                                      unsigned src_addr_len)
{
    pj_turn_transport *transport;
    const pj_stun_msg *msg = rdata->msg;
    pj_turn_allocation *alloc;
    pj_status_t status;

    PJ_UNUSED_ARG(pdu);
    PJ_UNUSED_ARG(pdu_len);

    transport = (pj_turn_transport*) token;

    /* Respond any requests other than ALLOCATE with 437 response */
    if (msg->hdr.type != PJ_STUN_ALLOCATE_REQUEST) {
        stun_respond(sess, transport, rdata, PJ_STUN_SC_ALLOCATION_MISMATCH,
                     NULL, PJ_FALSE, src_addr, src_addr_len);
        return PJ_SUCCESS;
    }

    /* Create new allocation. The relay resource will be allocated
     * in this function.
     */
    status = pj_turn_allocation_create(transport, src_addr, src_addr_len,
                                       rdata, sess, &alloc);
    if (status != PJ_SUCCESS) {
        /* STUN response has been sent, no need to reply here */
        return PJ_SUCCESS;
    }

    /* Done. */
    return PJ_SUCCESS;
}

/* Handle STUN Binding request */
static void handle_binding_request(pj_turn_pkt *pkt,
                                   unsigned options)
{
    pj_stun_msg *request, *response;
    pj_uint8_t pdu[200];
    pj_size_t len;
    pj_status_t status;

    /* Decode request */
    status = pj_stun_msg_decode(pkt->pool, pkt->pkt, pkt->len, options,
                                &request, NULL, NULL);
    if (status != PJ_SUCCESS)
        return;

    /* Create response */
    status = pj_stun_msg_create_response(pkt->pool, request, 0, NULL,
                                         &response);
    if (status != PJ_SUCCESS)
        return;

    /* Add XOR-MAPPED-ADDRESS */
    pj_stun_msg_add_sockaddr_attr(pkt->pool, response,
                                  PJ_STUN_ATTR_XOR_MAPPED_ADDR,
                                  PJ_TRUE,
                                  &pkt->src.clt_addr,
                                  pkt->src_addr_len);

    /* Encode */
    status = pj_stun_msg_encode(response, pdu, sizeof(pdu), 0, NULL, &len);
    if (status != PJ_SUCCESS)
        return;

    /* Send response */
    pkt->transport->sendto(pkt->transport, pdu, len, 0,
                           &pkt->src.clt_addr, pkt->src_addr_len);
}

/*
 * This callback is called by UDP listener on incoming packet. This is
 * the first entry for incoming packet (from client) to the server. From
 * here, the packet may be handed over to an allocation if an allocation
 * is found for the client address, or handed over to owned STUN session
 * if an allocation is not found.
 */
PJ_DEF(void) pj_turn_srv_on_rx_pkt(pj_turn_srv *srv,
                                   pj_turn_pkt *pkt)
{
    pj_turn_allocation *alloc;

    /* Get TURN allocation from the source address */
    pj_lock_acquire(srv->core.lock);
    alloc = (pj_turn_allocation*)
            pj_hash_get(srv->tables.alloc, &pkt->src, sizeof(pkt->src), NULL);
    pj_lock_release(srv->core.lock);

    /* If allocation is found, just hand over the packet to the
     * allocation.
     */
    if (alloc) {
        pj_turn_allocation_on_rx_client_pkt(alloc, pkt);
    } else {
        /* Otherwise this is a new client */
        unsigned options;
        pj_size_t parsed_len;
        pj_status_t status;

        /* Check that this is a STUN message */
        options = PJ_STUN_CHECK_PACKET | PJ_STUN_NO_FINGERPRINT_CHECK;
        if (pkt->transport->listener->tp_type == PJ_TURN_TP_UDP)
            options |= PJ_STUN_IS_DATAGRAM;

        status = pj_stun_msg_check(pkt->pkt, pkt->len, options);
        if (status != PJ_SUCCESS) {
            /* If the first byte are not STUN, drop the packet. First byte
             * of STUN message is always 0x00 or 0x01. Otherwise wait for
             * more data as the data might have come from TCP.
             *
             * Also drop packet if it's unreasonably too big, as this might
             * indicate invalid data that's building up in the buffer.
             *
             * Or if packet is a datagram.
             */
            if ((*pkt->pkt != 0x00 && *pkt->pkt != 0x01) ||
                pkt->len > 1600 ||
                (options & PJ_STUN_IS_DATAGRAM))
            {
                char errmsg[PJ_ERR_MSG_SIZE];
                char ip[PJ_INET6_ADDRSTRLEN+10];

                pkt->len = 0;

                pj_strerror(status, errmsg, sizeof(errmsg));
                PJ_LOG(5,(srv->obj_name,
                          "Non-STUN packet from %s is dropped: %s",
                          pj_sockaddr_print(&pkt->src.clt_addr, ip, sizeof(ip), 3),
                          errmsg));
            }
            return;
        }

        /* Special handling for Binding Request. We won't give it to the
         * STUN session since this request is not authenticated.
         */
        if (pkt->pkt[1] == 1) {
            handle_binding_request(pkt, options);
            return;
        }

        /* Hand over processing to STUN session. This will trigger
         * on_rx_stun_request() callback to be called if the STUN
         * message is a request.
         */
        options &= ~PJ_STUN_CHECK_PACKET;
        parsed_len = 0;
        status = pj_stun_session_on_rx_pkt(srv->core.stun_sess, pkt->pkt,
                                           pkt->len, options, pkt->transport,
                                           &parsed_len, &pkt->src.clt_addr,
                                           pkt->src_addr_len);
        if (status != PJ_SUCCESS) {
            char errmsg[PJ_ERR_MSG_SIZE];
            char ip[PJ_INET6_ADDRSTRLEN+10];

            pj_strerror(status, errmsg, sizeof(errmsg));
            PJ_LOG(5,(srv->obj_name,
                      "Error processing STUN packet from %s: %s",
                      pj_sockaddr_print(&pkt->src.clt_addr, ip, sizeof(ip), 3),
                      errmsg));
        }

        if (pkt->transport->listener->tp_type == PJ_TURN_TP_UDP) {
            pkt->len = 0;
        } else if (parsed_len > 0) {
            if (parsed_len == pkt->len) {
                pkt->len = 0;
            } else {
                pj_memmove(pkt->pkt, pkt->pkt+parsed_len,
                           pkt->len - parsed_len);
                pkt->len -= parsed_len;
            }
        }
    }
}


