/* nsapi_dns.cpp
 * Original work Copyright (c) 2013 Henry Leinen (henry[dot]leinen [at] online [dot] de)
 * Modified work Copyright (c) 2015 ARM Limited
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
#include "nsapi_dns.h"
#include "netsocket/UDPSocket.h"
#include <string.h>
#include <stdlib.h>
#include <stdio.h>
#include "mbed_shared_queues.h"
#include "EventQueue.h"
#include "OnboardNetworkStack.h"
#include "Kernel.h"

#define CLASS_IN 1

#define RR_A 1
#define RR_AAAA 28

// DNS options
#define DNS_BUFFER_SIZE 512
#define DNS_SERVERS_SIZE 5
#define DNS_RESPONSE_MIN_SIZE 12
#define DNS_MAX_TTL 604800
#define DNS_STACK_SERVERS_NUM 5
#define DNS_QUERY_QUEUE_SIZE 5
#define DNS_HOST_NAME_MAX_LEN 255

#ifndef MBED_CONF_NSAPI_DNS_RESPONSE_WAIT_TIME
#define MBED_CONF_NSAPI_DNS_RESPONSE_WAIT_TIME   5000
#endif

#ifndef MBED_CONF_NSAPI_DNS_TOTAL_RETRIES
#define MBED_CONF_NSAPI_DNS_TOTAL_RETRIES        3
#endif

#ifndef MBED_CONF_NSAPI_DNS_RETRIES
#define MBED_CONF_NSAPI_DNS_RETRIES              0
#endif

#ifndef MBED_CONF_NSAPI_DNS_CACHE_SIZE
#define MBED_CONF_NSAPI_DNS_CACHE_SIZE           3
#endif

struct DNS_CACHE {
    nsapi_addr_t address;
    char *host;
    uint64_t expires;      /*!< time to live in milliseconds */
    uint64_t accessed;     /*!< last accessed */
};

struct DNS_QUERY {
    int unique_id;
    NetworkStack *stack;
    char *host;
    NetworkStack::hostbyname_cb_t callback;
    call_in_callback_cb_t call_in_cb;
    nsapi_size_t addr_count;
    nsapi_version_t version;
    UDPSocket *socket;
    nsapi_addr_t *addrs;
    int dns_server;
    int retries;
    int total_retries;
    int dns_message_id;
    int count;
    uint32_t ttl;
};

typedef nsapi_error_t (*nsapi_dns_call_t)(mbed::Callback<void()> func);
typedef nsapi_error_t (*nsapi_dns_call_in_t)(int delay, mbed::Callback<void()> func);

static void nsapi_dns_cache_add(const char *host, nsapi_addr_t *address, uint32_t ttl);
static nsapi_size_or_error_t nsapi_dns_cache_find(const char *host, nsapi_version_t version, nsapi_addr_t *address);

static nsapi_error_t nsapi_dns_get_server_addr(NetworkStack *stack, int *index, int *total_retries, SocketAddress *dns_addr);

static void nsapi_dns_query_async_create(void *ptr);
static nsapi_error_t nsapi_dns_query_async_delete(int unique_id);
static void nsapi_dns_query_async_send(void *ptr);
static void nsapi_dns_query_async_resp(DNS_QUERY *query, nsapi_error_t status, SocketAddress *address);
static void nsapi_dns_query_async_socket_callback(NetworkStack *stack);
static void nsapi_dns_query_async_response(void *ptr);

static nsapi_addr_t dns_servers[DNS_SERVERS_SIZE] = {
    {NSAPI_IPv4, {8, 8, 8, 8}},                             // Google
    {NSAPI_IPv4, {209, 244, 0, 3}},                         // Level 3
    {NSAPI_IPv4, {84, 200, 69, 80}},                        // DNS.WATCH
    {NSAPI_IPv6, {0x20,0x01, 0x48,0x60, 0x48,0x60, 0,0,     // Google
                  0,0, 0,0, 0,0, 0x88,0x88}},
    {NSAPI_IPv6, {0x20,0x01, 0x16,0x08, 0,0x10, 0,0x25,     // DNS.WATCH
                  0,0, 0,0, 0x1c,0x04, 0xb1,0x2f}},
};

static DNS_CACHE *dns_cache[MBED_CONF_NSAPI_DNS_CACHE_SIZE];
static uint16_t dns_message_id = 0;
static int dns_unique_id = 1;
static DNS_QUERY *dns_query_queue[DNS_QUERY_QUEUE_SIZE];
// Protects cache shared between blocking and asynchronous calls
static rtos::Mutex dns_cache_mutex;
// Protects from several threads running asynchronous DNS
static rtos::Mutex dns_mutex;
static nsapi_dns_call_in_t dns_call_in = 0;

// DNS server configuration
extern "C" nsapi_error_t nsapi_dns_add_server(nsapi_addr_t addr)
{
    memmove(&dns_servers[1], &dns_servers[0],
            (DNS_SERVERS_SIZE-1)*sizeof(nsapi_addr_t));

    dns_servers[0] = addr;
    return NSAPI_ERROR_OK;
}


// DNS packet parsing
static void dns_append_byte(uint8_t **p, uint8_t byte)
{
    *(*p)++ = byte;
}

static void dns_append_word(uint8_t **p, uint16_t word)
{

    dns_append_byte(p, 0xff & (word >> 8));
    dns_append_byte(p, 0xff & (word >> 0));
}

static void dns_append_name(uint8_t **p, const char *name, uint8_t len)
{
    dns_append_byte(p, len);
    memcpy(*p, name, len);
    *p += len;
}

static uint8_t dns_scan_byte(const uint8_t **p)
{
    return *(*p)++;
}

static uint16_t dns_scan_word(const uint8_t **p)
{
    uint16_t a = dns_scan_byte(p);
    uint16_t b = dns_scan_byte(p);
    return (a << 8) | b;
}

static uint32_t dns_scan_word32(const uint8_t **p)
{
    uint32_t value = dns_scan_byte(p) << 24;
    value |= dns_scan_byte(p) << 16;
    value |= dns_scan_byte(p) << 8;
    value |= dns_scan_byte(p);

    return value;
}

static int dns_append_question(uint8_t *ptr, uint16_t id, const char *host, nsapi_version_t version)
{
    uint8_t *s_ptr = ptr;
    uint8_t **p = &ptr;

    // fill the header
    dns_append_word(p, id);     // id      = 1
    dns_append_word(p, 0x0100); // flags   = recursion required
    dns_append_word(p, 1);      // qdcount = 1
    dns_append_word(p, 0);      // ancount = 0
    dns_append_word(p, 0);      // nscount = 0
    dns_append_word(p, 0);      // arcount = 0

    // fill out the question names
    while (host[0]) {
        size_t label_len = strcspn(host, ".");
        dns_append_name(p, host, label_len);
        host += label_len + (host[label_len] == '.');
    }

    dns_append_byte(p, 0);

    // fill out question footer
    if (version != NSAPI_IPv6) {
        dns_append_word(p, RR_A);       // qtype  = ipv4
    } else {
        dns_append_word(p, RR_AAAA);    // qtype  = ipv6
    }
    dns_append_word(p, CLASS_IN);

    return *p - s_ptr;
}

static int dns_scan_response(const uint8_t *ptr, uint16_t exp_id, uint32_t *ttl, nsapi_addr_t *addr, unsigned addr_count)
{
    const uint8_t **p = &ptr;

    // scan header
    uint16_t id    = dns_scan_word(p);
    uint16_t flags = dns_scan_word(p);
    bool    qr     = 0x1 & (flags >> 15);
    uint8_t opcode = 0xf & (flags >> 11);
    uint8_t rcode  = 0xf & (flags >>  0);

    uint16_t qdcount = dns_scan_word(p); // qdcount
    uint16_t ancount = dns_scan_word(p); // ancount
    dns_scan_word(p);                    // nscount
    dns_scan_word(p);                    // arcount

    // verify header is response to query
    if (!(id == exp_id && qr && opcode == 0 && rcode == 0)) {
        return -1;
    }

    // skip questions
    for (int i = 0; i < qdcount; i++) {
        while (true) {
            uint8_t len = dns_scan_byte(p);
            if (len == 0) {
                break;
            }

            *p += len;
        }

        dns_scan_word(p); // qtype
        dns_scan_word(p); // qclass
    }

    // scan each response
    unsigned count = 0;

    for (int i = 0; i < ancount && count < addr_count; i++) {
        while (true) {
            uint8_t len = dns_scan_byte(p);
            if (len == 0) {
                break;
            } else if (len & 0xc0) { // this is link
                dns_scan_byte(p);
                break;
            }

            *p += len;
        }

        uint16_t rtype    = dns_scan_word(p);    // rtype
        uint16_t rclass   = dns_scan_word(p);    // rclass
        uint32_t ttl_val  = dns_scan_word32(p);  // ttl
        uint16_t rdlength = dns_scan_word(p);    // rdlength

        if (i == 0) {
            // Is interested only on first address that is stored to cache
            if (ttl_val > DNS_MAX_TTL) {
                ttl_val = DNS_MAX_TTL;
            }
            *ttl = ttl_val;
        }

        if (rtype == RR_A && rclass == CLASS_IN && rdlength == NSAPI_IPv4_BYTES) {
            // accept A record
            addr->version = NSAPI_IPv4;
            for (int i = 0; i < NSAPI_IPv4_BYTES; i++) {
                addr->bytes[i] = dns_scan_byte(p);
            }

            addr += 1;
            count += 1;
        } else if (rtype == RR_AAAA && rclass == CLASS_IN && rdlength == NSAPI_IPv6_BYTES) {
            // accept AAAA record
            addr->version = NSAPI_IPv6;
            for (int i = 0; i < NSAPI_IPv6_BYTES; i++) {
                addr->bytes[i] = dns_scan_byte(p);
            }

            addr += 1;
            count += 1;
        } else {
            // skip unrecognized records
            *p += rdlength;
        }
    }

    return count;
}

static void nsapi_dns_cache_add(const char *host, nsapi_addr_t *address, uint32_t ttl)
{
    // RFC 1034: if TTL is zero, entry is not added to cache
    if (!ttl) {
        return;
    }

    // Checks if already cached
    if (nsapi_dns_cache_find(host, address->version, NULL) == NSAPI_ERROR_OK) {
        return;
    }

    dns_cache_mutex.lock();

    int index = -1;
    uint64_t accessed = ~0;

    // Finds free or last accessed entry
    for (int i = 0; i < MBED_CONF_NSAPI_DNS_CACHE_SIZE; i++) {
        if (!dns_cache[i]) {
            index = i;
            break;
        } else if (dns_cache[i]->accessed <= accessed) {
            accessed = dns_cache[i]->accessed;
            index = i;
        }
    }

    if (index < 0) {
        dns_cache_mutex.unlock();
        return;
    }

    // Allocates in case entry is free, otherwise reuses
    if (!dns_cache[index]) {
        dns_cache[index] = new (std::nothrow) DNS_CACHE;
    } else {
        free(dns_cache[index]->host);
    }

    if (dns_cache[index]) {
        dns_cache[index]->address = *address;
        dns_cache[index]->host = (char *) malloc(strlen(host) + 1);
        strcpy(dns_cache[index]->host, host);
        uint64_t ms_count = rtos::Kernel::get_ms_count();
        dns_cache[index]->expires = ms_count + ttl * 1000;
        dns_cache[index]->accessed = ms_count;
    }

    dns_cache_mutex.unlock();
}

static nsapi_error_t nsapi_dns_cache_find(const char *host, nsapi_version_t version, nsapi_addr_t *address)
{
    nsapi_error_t ret_val = NSAPI_ERROR_NO_ADDRESS;

    dns_cache_mutex.lock();

    for (int i = 0; i < MBED_CONF_NSAPI_DNS_CACHE_SIZE; i++) {
        if (dns_cache[i]) {
            uint64_t ms_count = rtos::Kernel::get_ms_count();
            // Checks all entries for expired entries
            if (ms_count > dns_cache[i]->expires) {
                free(dns_cache[i]->host);
                delete dns_cache[i];
                dns_cache[i] = NULL;
            } else if ((version == NSAPI_UNSPEC || version == dns_cache[i]->address.version) &&
                strcmp(dns_cache[i]->host, host) == 0) {
                if (address) {
                    *address = dns_cache[i]->address;
                }
                dns_cache[i]->accessed = ms_count;
                ret_val = NSAPI_ERROR_OK;
            }
        }
    }

    dns_cache_mutex.unlock();

    return ret_val;
}

static nsapi_error_t nsapi_dns_get_server_addr(NetworkStack *stack, int *index, int *total_retries, SocketAddress *dns_addr)
{
    bool dns_addr_set = false;

    if (*total_retries == 0) {
        return NSAPI_ERROR_NO_ADDRESS;
    }

    if (*index >= DNS_SERVERS_SIZE + DNS_STACK_SERVERS_NUM) {
        if (*total_retries) {
            *index = 0;
        } else {
            return NSAPI_ERROR_NO_ADDRESS;
        }
    }

    if (*index < DNS_STACK_SERVERS_NUM) {
        nsapi_error_t ret = stack->get_dns_server(*index, dns_addr);
        if (ret < 0) {
            *index = DNS_STACK_SERVERS_NUM;
        } else {
            dns_addr_set = true;
        }
    }

    if (!dns_addr_set) {
        dns_addr->set_addr(dns_servers[*index - DNS_STACK_SERVERS_NUM]);
    }

    dns_addr->set_port(53);

    return NSAPI_ERROR_OK;
}

// core query function
static nsapi_size_or_error_t nsapi_dns_query_multiple(NetworkStack *stack, const char *host,
        nsapi_addr_t *addr, unsigned addr_count, nsapi_version_t version)
{
    // check for valid host name
    int host_len = host ? strlen(host) : 0;
    if (host_len > DNS_HOST_NAME_MAX_LEN || host_len == 0) {
        return NSAPI_ERROR_PARAMETER;
    }

    // check cache
    if (nsapi_dns_cache_find(host, version, addr) == NSAPI_ERROR_OK) {
        return 1;
    }

    // create a udp socket
    UDPSocket socket;
    int err = socket.open(stack);
    if (err) {
        return err;
    }

    socket.set_timeout(MBED_CONF_NSAPI_DNS_RESPONSE_WAIT_TIME);

    // create network packet
    uint8_t * const packet = (uint8_t *)malloc(DNS_BUFFER_SIZE);
    if (!packet) {
        return NSAPI_ERROR_NO_MEMORY;
    }

    nsapi_size_or_error_t result = NSAPI_ERROR_DNS_FAILURE;

    int retries = MBED_CONF_NSAPI_DNS_RETRIES;

    int index = 0;
    int total_retries = MBED_CONF_NSAPI_DNS_TOTAL_RETRIES;

    // check against each dns server
    while (true) {
        SocketAddress dns_addr;
        err = nsapi_dns_get_server_addr(stack, &index, &total_retries, &dns_addr);
        if (err != NSAPI_ERROR_OK) {
            break;
        }

        // send the question
        int len = dns_append_question(packet, 1, host, version);

        err = socket.sendto(dns_addr, packet, len);
        // send may fail for various reasons, including wrong address type - move on
        if (err < 0) {
            // goes to next dns server
            retries = MBED_CONF_NSAPI_DNS_RETRIES;
            index++;
            continue;
        }

        if (total_retries) {
            total_retries--;
        }

        // recv the response
        err = socket.recvfrom(NULL, packet, DNS_BUFFER_SIZE);
        if (err == NSAPI_ERROR_WOULD_BLOCK) {
            if (retries) {
                // retries
                retries--;
            } else {
                // goes to next dns server
                retries = MBED_CONF_NSAPI_DNS_RETRIES;
                index++;
            }
            continue;
        } else if (err < 0) {
            result = err;
            break;
        }

        const uint8_t *response = packet;
        uint32_t ttl;
        int resp = dns_scan_response(response, 1, &ttl, addr, addr_count);
        if (resp > 0) {
            nsapi_dns_cache_add(host, addr, ttl);
            result = resp;
        } else if (resp < 0) {
            continue;
        }

        /* The DNS response is final, no need to check other servers */
        break;
    }

    // clean up packet
    free(packet);

    // clean up udp
    err = socket.close();
    if (err) {
        return err;
    }

    // return result
    return result;
}

// convenience functions for other forms of queries
extern "C" nsapi_size_or_error_t nsapi_dns_query_multiple(nsapi_stack_t *stack, const char *host,
        nsapi_addr_t *addr, nsapi_size_t addr_count, nsapi_version_t version)
{
    NetworkStack *nstack = nsapi_create_stack(stack);
    return nsapi_dns_query_multiple(nstack, host, addr, addr_count, version);
}

nsapi_size_or_error_t nsapi_dns_query_multiple(NetworkStack *stack, const char *host,
        SocketAddress *addresses, nsapi_size_t addr_count, nsapi_version_t version)
{
    nsapi_addr_t *addrs = new (std::nothrow) nsapi_addr_t[addr_count];
    nsapi_size_or_error_t result = nsapi_dns_query_multiple(stack, host, addrs, addr_count, version);

    if (result > 0) {
        for (int i = 0; i < result; i++) {
            addresses[i].set_addr(addrs[i]);
        }
    }

    delete[] addrs;
    return result;
}

extern "C" nsapi_error_t nsapi_dns_query(nsapi_stack_t *stack, const char *host,
        nsapi_addr_t *addr, nsapi_version_t version)
{
    NetworkStack *nstack = nsapi_create_stack(stack);
    nsapi_size_or_error_t result = nsapi_dns_query_multiple(nstack, host, addr, 1, version);
    return (nsapi_error_t)((result > 0) ? 0 : result);
}

nsapi_error_t nsapi_dns_query(NetworkStack *stack, const char *host,
        SocketAddress *address, nsapi_version_t version)
{
    nsapi_addr_t addr;
    nsapi_size_or_error_t result = nsapi_dns_query_multiple(stack, host, &addr, 1, version);
    address->set_addr(addr);
    return (nsapi_error_t)((result > 0) ? 0 : result);
}

nsapi_error_t nsapi_dns_query_async(NetworkStack *stack, const char *host,
        NetworkStack::hostbyname_cb_t callback, call_in_callback_cb_t call_in_cb,
        nsapi_version_t version)
{
    return nsapi_dns_query_multiple_async(stack, host, callback, 0, call_in_cb, version);
}

void nsapi_dns_call_in_set(nsapi_dns_call_in_t callback)
{
    dns_call_in = callback;
}

nsapi_error_t nsapi_dns_call_in(call_in_callback_cb_t cb, int delay, mbed::Callback<void()> func)
{
    if (dns_call_in) {
        dns_call_in(delay, func);
    } else {
        return cb(delay, func);
    }
    return NSAPI_ERROR_OK;
}

nsapi_error_t nsapi_dns_query_multiple_async(NetworkStack *stack, const char *host,
    NetworkStack::hostbyname_cb_t callback, nsapi_size_t addr_count,
    call_in_callback_cb_t call_in_cb, nsapi_version_t version)
{
    if (!stack) {
        return NSAPI_ERROR_PARAMETER;
    }

    // check for valid host name
    int host_len = host ? strlen(host) : 0;
    if (host_len > DNS_HOST_NAME_MAX_LEN || host_len == 0) {
        return NSAPI_ERROR_PARAMETER;
    }

    nsapi_addr address;
    if (nsapi_dns_cache_find(host, version, &address) == NSAPI_ERROR_OK) {
        SocketAddress addr(address);
        callback(NSAPI_ERROR_OK, &addr);
        return NSAPI_ERROR_OK;
    }

    dns_mutex.lock();

    int index = -1;

    for (int i = 0; i < DNS_QUERY_QUEUE_SIZE; i++) {
        if (!dns_query_queue[i]) {
            index = i;
            break;
        }
    }

    if (index < 0) {
        dns_mutex.unlock();
        return NSAPI_ERROR_NO_MEMORY;
    }

    DNS_QUERY *query = new (std::nothrow) DNS_QUERY;

    if (!query) {
        dns_mutex.unlock();
        return NSAPI_ERROR_NO_MEMORY;
    }

    query->host = (char *) malloc(host_len + 1);
    strcpy(query->host, host);
    query->callback = callback;
    query->call_in_cb = call_in_cb;
    query->stack = stack;
    query->addr_count = addr_count;
    query->version = version;
    query->socket = NULL;
    query->addrs = NULL;
    query->dns_server = 0;
    query->retries = MBED_CONF_NSAPI_DNS_RETRIES + 1;
    query->total_retries =  MBED_CONF_NSAPI_DNS_TOTAL_RETRIES;
    query->dns_message_id = -1;

    query->unique_id = dns_unique_id++;
    if (query->unique_id > 0x7FFF) {
        query->unique_id = 1;
    }

    dns_query_queue[index] = query;

    if (nsapi_dns_call_in(query->call_in_cb, 0, mbed::callback(nsapi_dns_query_async_create, reinterpret_cast<void *>(query->unique_id))) != NSAPI_ERROR_OK) {
        delete query;
        dns_mutex.unlock();
        return NSAPI_ERROR_NO_MEMORY;
    }

    dns_mutex.unlock();

    return query->unique_id;
}

nsapi_error_t nsapi_dns_query_async_cancel(nsapi_error_t id)
{
    dns_mutex.lock();

    nsapi_error_t ret = nsapi_dns_query_async_delete(id);

    dns_mutex.unlock();

    return ret;
}

static void nsapi_dns_query_async_create(void *ptr)
{
    dns_mutex.lock();

    int unique_id = reinterpret_cast<int>(ptr);

    DNS_QUERY *query = NULL;

    for (int i = 0; i < DNS_QUERY_QUEUE_SIZE; i++) {
        if (dns_query_queue[i] && dns_query_queue[i]->unique_id == unique_id) {
            query = dns_query_queue[i];
            break;
        }
    }

    if (!query) {
        // Cancel has been called
        dns_mutex.unlock();
        return;
    }

    for (int i = 0; i < DNS_QUERY_QUEUE_SIZE; i++) {
        if (dns_query_queue[i] && dns_query_queue[i] != query) {
            if (dns_query_queue[i]->socket && dns_query_queue[i]->stack == query->stack) {
                query->socket = dns_query_queue[i]->socket;
            }
        }
    }

    UDPSocket *socket;

    if (query->socket) {
        socket = query->socket;
    } else {
        socket = new (std::nothrow) UDPSocket;
        if (!socket) {
            nsapi_dns_query_async_resp(query, NSAPI_ERROR_NO_MEMORY, NULL);
            dns_mutex.unlock();
            return;
        }

        int err = socket->open(query->stack);
        if (err) {
            delete socket;
            nsapi_dns_query_async_resp(query, err, NULL);
            dns_mutex.unlock();
            return;
        }

        socket->set_timeout(0);
        socket->sigio(mbed::callback(nsapi_dns_query_async_socket_callback, query->stack));

        query->socket = socket;
    }

    dns_mutex.unlock();

    nsapi_dns_query_async_send(reinterpret_cast<void *>(query->unique_id));
}

static nsapi_error_t nsapi_dns_query_async_delete(int unique_id)
{
    int index = -1;

    for (int i = 0; i < DNS_QUERY_QUEUE_SIZE; i++) {
        if (dns_query_queue[i] && dns_query_queue[i]->unique_id == unique_id) {
            index = i;
            break;
        }
    }

    if (index < 0) {
        return NSAPI_ERROR_DEVICE_ERROR;
    }

    bool close_socket = true;

    for (int i = 0; i < DNS_QUERY_QUEUE_SIZE; i++) {
        if (i != index && dns_query_queue[i] && dns_query_queue[i]->socket &&
            dns_query_queue[i]->stack == dns_query_queue[index]->stack) {
            close_socket = false;
        }
    }

    if (close_socket && dns_query_queue[index]->socket) {
        dns_query_queue[index]->socket->close();
        delete dns_query_queue[index]->socket;
    }

    if (dns_query_queue[index]->addrs) {
        delete[] dns_query_queue[index]->addrs;
    }

    free(dns_query_queue[index]->host);

    delete dns_query_queue[index];
    dns_query_queue[index] = NULL;
    return NSAPI_ERROR_OK;
}

static void nsapi_dns_query_async_resp(DNS_QUERY *query, nsapi_error_t status, SocketAddress *address)
{
    query->callback(status, address);
    nsapi_dns_query_async_delete(query->unique_id);
}

static void nsapi_dns_query_async_send(void *ptr)
{
    dns_mutex.lock();

    int unique_id = reinterpret_cast<int>(ptr);

    DNS_QUERY *query = NULL;

    for (int i = 0; i < DNS_QUERY_QUEUE_SIZE; i++) {
        if (dns_query_queue[i] && dns_query_queue[i]->unique_id == unique_id) {
            query = dns_query_queue[i];
            break;
        }
    }

    if (!query) {
        // Cancel has been called
        dns_mutex.unlock();
        return;
    }

    if (query->retries) {
        query->retries--;
    } else {
        query->dns_server++;
        query->retries = MBED_CONF_NSAPI_DNS_RETRIES;
    }

    query->dns_message_id = dns_message_id++;

    // create network packet
    uint8_t *packet = (uint8_t *)malloc(DNS_BUFFER_SIZE);
    if (!packet) {
        nsapi_dns_query_async_resp(query, NSAPI_ERROR_NO_MEMORY, NULL);
        dns_mutex.unlock();
        return;
    }

    // send the question
    int len = dns_append_question(packet, query->dns_message_id, query->host, query->version);

    while (true) {
        SocketAddress dns_addr;
        nsapi_size_or_error_t err = nsapi_dns_get_server_addr(query->stack, &(query->dns_server), &(query->total_retries), &dns_addr);
        if (err != NSAPI_ERROR_OK) {
            nsapi_dns_query_async_resp(query, NSAPI_ERROR_DNS_FAILURE, NULL);
            free(packet);
            dns_mutex.unlock();
            return;
        }

        err = query->socket->sendto(dns_addr, packet, len);

        if (err < 0) {
            query->dns_server++;
        } else {
            break;
        }
    }

    if (query->total_retries) {
        query->total_retries--;
    }

    free(packet);

    if (nsapi_dns_call_in(query->call_in_cb, MBED_CONF_NSAPI_DNS_RESPONSE_WAIT_TIME,
        mbed::callback(nsapi_dns_query_async_send, reinterpret_cast<void *>(unique_id))) != NSAPI_ERROR_OK) {
        nsapi_dns_query_async_resp(query, NSAPI_ERROR_NO_MEMORY, NULL);
    }

    dns_mutex.unlock();
}

static void nsapi_dns_query_async_socket_callback(NetworkStack *stack)
{
    UDPSocket *socket = NULL;

    dns_mutex.lock();

    for (int i = 0; i < DNS_QUERY_QUEUE_SIZE; i++) {
        if (dns_query_queue[i] && dns_query_queue[i]->stack == stack) {
            socket = dns_query_queue[i]->socket;
            break;
        }
    }

    if (socket) {
        // create network packet
        uint8_t *packet = (uint8_t *)malloc(DNS_BUFFER_SIZE);
        if (!packet) {
            dns_mutex.unlock();
            return;
        }

        while (true) {
            // recv the response
            nsapi_size_or_error_t size = socket->recvfrom(NULL, packet, DNS_BUFFER_SIZE);

            if (size < DNS_RESPONSE_MIN_SIZE) {
                break;
            }

            // gets id from response to associate with correct query
            uint16_t id = (packet[0] << 8) | packet[1];

            DNS_QUERY *query = NULL;

            for (int i = 0; i < DNS_QUERY_QUEUE_SIZE; i++) {
                if (dns_query_queue[i] && dns_query_queue[i]->dns_message_id == id) {
                    query = dns_query_queue[i];
                    break;
                }
            }

            if (!query) {
                continue;
            }

            int requested_count = 1;
            if (query->addr_count > 1) {
                requested_count = query->addr_count;
            }

            query->addrs = new (std::nothrow) nsapi_addr_t[requested_count];

            int resp = dns_scan_response(packet, id, &(query->ttl), query->addrs, requested_count);

            // Ignore invalid responses
            if (resp < 0) {
                delete[] query->addrs;
                query->addrs = 0;
            } else {
                query->count = resp;
                nsapi_dns_call_in(query->call_in_cb, 0, mbed::callback(nsapi_dns_query_async_response, reinterpret_cast<void *>(query->unique_id)));
            }
        }

        free(packet);
    }

    dns_mutex.unlock();
}

static void nsapi_dns_query_async_response(void *ptr)
{
    dns_mutex.lock();

    int unique_id = reinterpret_cast<int>(ptr);

    DNS_QUERY *query = NULL;

    for (int i = 0; i < DNS_QUERY_QUEUE_SIZE; i++) {
        if (dns_query_queue[i] && dns_query_queue[i]->unique_id == unique_id) {
            query = dns_query_queue[i];
            break;
        }
    }

    if (query) {
        if (query->count > 0) {
            SocketAddress *addresses = new (std::nothrow) SocketAddress[query->count];

            for (int i = 0; i < query->count; i++) {
                addresses[i].set_addr(query->addrs[i]);
            }

            // Adds address to cache
            nsapi_dns_cache_add(query->host, &(query->addrs[0]), query->ttl);

            nsapi_error_t status = NSAPI_ERROR_OK;
            if (query->addr_count > 0) {
                status = query->count;
            }
            nsapi_dns_query_async_resp(query, status, addresses);
            delete[] addresses;
        } else {
            nsapi_dns_query_async_resp(query, NSAPI_ERROR_DNS_FAILURE, NULL);
        }
    }
    dns_mutex.unlock();
}