/*
 * CDDL HEADER START
 *
 * The contents of this file are subject to the terms of the
 * Common Development and Distribution License (the "Licence").
 * You may not use this file except in compliance with the Licence.
 *
 * You can obtain a copy of the licence at RISC OS path @.^.LICENCE
 * or  http://www.riscosdev.com/lanman98/LICENCE.CDDL
 * See the Licence for the specific language governing permissions
 * and limitations under the Licence.
 *
 * CDDL HEADER END
 */

/*
 *   Portions Copyright RISC OS Developments 2019+, credited to the RISC OS One Project.
 */

#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <ctype.h>
#include <time.h>

#include "kernel.h"
#include "swis.h"
#include "LanMan98BaseLib/strext.h"
#include "LanMan98BaseLib/error.h"
#include "LanMan98BaseLib/memory.h"
#include "LanMan98BaseLib/tcp.h"
#include "LanMan98BaseLib/thread.h"
#include "discover.h"

#define MDNS_PORT   (5353)
#define BUF_SIZE    (2048)
#define NAME_MAX_   (96)
#define PER_PACKET  (12)        /* records of each kind kept per answer */

/*
    How many answers one call to DiscoverPoll will unpack.  It is called
    from a callback and from a foreground loop; both come back for more, so
    a cap costs nothing but the guarantee that neither sits here.
*/
#define DISCOVER_POLL_MAX (8)

#define TYPE_A      (1)
#define TYPE_PTR    (12)
#define TYPE_SRV    (33)

typedef struct
{
    char name[NAME_MAX_];
    char target[NAME_MAX_];
    int  port;
} srv_rec;

typedef struct
{
    char name[NAME_MAX_];
    unsigned char addr[4];
} a_rec;

/* Kept off the stack: a module has little of it and this is a lot */
typedef struct
{
    srv_rec srv[PER_PACKET];
    a_rec   a[PER_PACKET];
    char    ptr[PER_PACKET][NAME_MAX_];
    int     nsrv, na, nptr;
} work_t;

static unsigned int get16(const unsigned char *p)
{
    return (((unsigned int) p[0]) << 8) | p[1];
}

/*
    Read a name, and say where the record carries on.

    A name is a run of parts, each with its length in front.  A part whose
    length has both top bits set is not a part at all but a pointer to
    somewhere earlier in the message holding the rest - which is how the
    same ending gets written once and used by every record after it.
    Following one must not be able to go round for ever, so the jumps are
    counted, and where the record carries on is where the first pointer
    was, not where the name ended up.
*/
static int read_name(const unsigned char *d, int len, int pos,
                     char *out, int out_max)
{
    int n, jumps, after, i, l;

    n = 0;
    jumps = 0;
    after = -1;
    while((pos >= 0) && (pos < len))
    {
        l = d[pos];
        if(l == 0)
        {
            pos++;
            break;
        }
        if((l & 0xC0) == 0xC0)
        {
            if(pos + 1 >= len) break;
            if(after < 0) after = pos + 2;
            pos = (int) (((l & 0x3F) << 8) | d[pos + 1]);
            if(++jumps > 16) break;
            continue;
        }
        if(pos + 1 + l > len) break;
        if((n > 0) && (n < out_max - 1)) out[n++] = '.';
        for(i = 0; i < l; i++)
            if(n < out_max - 1) out[n++] = (char) d[pos + 1 + i];
        pos += 1 + l;
    }
    out[n] = 0;
    return (after >= 0) ? after : pos;
}

/* Already have this one?  A machine with two network cards answers twice */
static int already(found_t list, const char *addr)
{
    while(list)
    {
        if(strcmp(list->addr, addr) == 0) return 1;
        list = list->next;
    }
    return 0;
}

static void add_found(found_t *head, const char *instance,
                      const unsigned char *quad, int port)
{
    found_t f, t;
    char addr[20];
    char name[NAME_MAX_];
    char *dot;

    sprintf(addr, "%d.%d.%d.%d", quad[0], quad[1], quad[2], quad[3]);
    if(already(*head, addr)) return;

    /* The instance is "what it calls itself" followed by the name of the
       service; only the first part is worth showing. */
    strncpy(name, instance, sizeof(name) - 1);
    name[sizeof(name) - 1] = 0;
    dot = strstr(name, "._smb");
    if(dot) *dot = 0;
    if(name[0] == 0) strcpy(name, addr);

    f = Malloc(sizeof(*f));
    memset(f, 0, sizeof(*f));
    ExceptTry
    {
        f->name = strdup(name);
    }
    ExceptCatch
    {
        Free(f);
        ExceptRethrow();
    }
    strcpy(f->addr, addr);
    f->port = port;

    if(*head == NULL)
    {
        *head = f;
    }
    else
    {
        for(t = *head; t->next; t = t->next)
            ;
        t->next = f;
    }
}

/*
    Pull what is useful out of one answer.

    Three kinds of record matter, and they have to be put together: one
    says a machine offers the service, the next says which host and port
    it is on, and the third says what address that host has.
*/
static void collect(const unsigned char *d, int len, found_t *head,
                    work_t *w)
{
    char nm[NAME_MAX_];
    int qd, total, pos, i, j, k, type, rl, rd;

    if(len < 12) return;
    w->nsrv = 0; w->na = 0; w->nptr = 0;

    qd = (int) get16(d + 4);
    total = (int) (get16(d + 6) + get16(d + 8) + get16(d + 10));
    pos = 12;
    for(i = 0; (i < qd) && (pos < len); i++)
        pos = read_name(d, len, pos, nm, sizeof(nm)) + 4;

    for(i = 0; (i < total) && (pos < len); i++)
    {
        pos = read_name(d, len, pos, nm, sizeof(nm));
        if(pos + 10 > len) break;
        type = (int) get16(d + pos);
        rl = (int) get16(d + pos + 8);
        rd = pos + 10;
        if(rd + rl > len) break;

        if((type == TYPE_PTR) && (w->nptr < PER_PACKET))
        {
            read_name(d, len, rd, w->ptr[w->nptr], NAME_MAX_);
            w->nptr++;
        }
        else if((type == TYPE_SRV) && (w->nsrv < PER_PACKET) && (rl >= 7))
        {
            strncpy(w->srv[w->nsrv].name, nm, NAME_MAX_ - 1);
            w->srv[w->nsrv].name[NAME_MAX_ - 1] = 0;
            w->srv[w->nsrv].port = (int) get16(d + rd + 4);
            read_name(d, len, rd + 6, w->srv[w->nsrv].target, NAME_MAX_);
            w->nsrv++;
        }
        else if((type == TYPE_A) && (w->na < PER_PACKET) && (rl == 4))
        {
            strncpy(w->a[w->na].name, nm, NAME_MAX_ - 1);
            w->a[w->na].name[NAME_MAX_ - 1] = 0;
            memcpy(w->a[w->na].addr, d + rd, 4);
            w->na++;
        }
        pos = rd + rl;
    }

    for(i = 0; i < w->nptr; i++)
    {
        for(j = 0; j < w->nsrv; j++)
            if(strcmp(w->srv[j].name, w->ptr[i]) == 0) break;
        if(j == w->nsrv) continue;
        for(k = 0; k < w->na; k++)
            if(strcmp(w->a[k].name, w->srv[j].target) == 0) break;
        if(k == w->na) continue;
        add_found(head, w->ptr[i], w->a[k].addr, w->srv[j].port);
    }
}

void DiscoverFree(found_t list)
{
    found_t next;

    while(list)
    {
        next = list->next;
        if(list->name) Free(list->name);
        Free(list);
        list = next;
    }
}

/* Case-blind compare: a name comes back however the machine spells it */
static int same_name(const char *a, const char *b)
{
    while(*a && *b)
    {
        if(tolower((unsigned char) *a) != tolower((unsigned char) *b))
            return 0;
        a++; b++;
    }
    return (*a == 0) && (*b == 0);
}

/*
    Write a name out the way a query wants it: each part with its length in
    front, and a zero at the end.  ".local" is added where it is not
    already there - that is the domain these questions are asked in, and a
    plain name means the same thing with it on.
*/
static int encode_name(const char *name, char *full, int full_max,
                       unsigned char *out, int out_max)
{
    const char *p;
    int n, len;

    len = (int) strlen(name);
    if((len > 6) && same_name(name + len - 6, ".local"))
        n = len;
    else
        n = len + 6;
    if((n + 1 > full_max) || (n + 2 > out_max))
        return 0;
    strcpy(full, name);
    if(n != len)
        strcat(full, ".local");

    n = 0;
    p = full;
    while(*p)
    {
        const char *dot = strchr(p, '.');
        int l = dot ? (int) (dot - p) : (int) strlen(p);

        if((l <= 0) || (l > 63) || (n + l + 2 > out_max))
            return 0;
        out[n++] = (unsigned char) l;
        memcpy(out + n, p, l);
        n += l;
        p = dot ? dot + 1 : p + l;
    }
    out[n++] = 0;
    return n;
}

char *DiscoverResolve(const char *name, int wait_cs)
{
    udp_conn_t conn;
    tcp_addr_t group;
    unsigned char quad[4];
    unsigned char *buf;
    unsigned char qname[NAME_MAX_ + 8];
    char full[NAME_MAX_];
    char nm[NAME_MAX_];
    char *res;
    volatile int size;
    clock_t stop;
    int qlen, n;

    if((name == NULL) || (*name == 0))
        return NULL;

    res = NULL;
    buf = Malloc(BUF_SIZE);
    conn = NULL;
    group = NULL;

    ExceptTry
    {
        qlen = encode_name(name, full, sizeof(full), qname, sizeof(qname));
        if(qlen > 0)
        {
            char *var;

            quad[0] = 224; quad[1] = 0; quad[2] = 0; quad[3] = 251;
            var = getenv("LanMan98$DiscoverAddr");
            if(var)
            {
                int a, b, c2, d2;

                if(sscanf(var, "%d.%d.%d.%d", &a, &b, &c2, &d2) == 4)
                {
                    quad[0] = (unsigned char) a; quad[1] = (unsigned char) b;
                    quad[2] = (unsigned char) c2; quad[3] = (unsigned char) d2;
                }
            }
            group = TcpAddrFromQuad((char *) quad);
            conn = UdpCreate(TCP_PORT_ANY);
            UdpSetPeer(MDNS_PORT, group, conn);

            memset(buf, 0, 12);
            buf[5] = 1;                 /* one question */
            n = 12;
            memcpy(buf + n, qname, qlen);
            n += qlen;
            buf[n++] = 0; buf[n++] = TYPE_A;
            buf[n++] = 0; buf[n++] = 1; /* on the internet */
            UdpWrite(buf, n, conn);

            stop = clock() + (clock_t) wait_cs * (CLOCKS_PER_SEC / 100);
            while((res == NULL) && (clock() < stop))
            {
                size = 0;
                ThreadSetTimeout(0);
                ExceptTry
                {
                    size = UdpRead(buf, BUF_SIZE, conn);
                }
                ExceptCatch
                {
                    size = 0;
                }
                ThreadResetTimeout();
                if(size >= 12)
                {
                    int qd, total, pos, i, type, rl, rd;

                    qd = (int) get16(buf + 4);
                    total = (int) (get16(buf + 6) + get16(buf + 8) +
                                   get16(buf + 10));
                    pos = 12;
                    for(i = 0; (i < qd) && (pos < size); i++)
                        pos = read_name(buf, size, pos, nm, sizeof(nm)) + 4;
                    for(i = 0; (i < total) && (pos < size); i++)
                    {
                        pos = read_name(buf, size, pos, nm, sizeof(nm));
                        if(pos + 10 > size) break;
                        type = (int) get16(buf + pos);
                        rl = (int) get16(buf + pos + 8);
                        rd = pos + 10;
                        if(rd + rl > size) break;
                        if((type == TYPE_A) && (rl == 4) && same_name(nm, full))
                        {
                            char dotted[20];

                            sprintf(dotted, "%d.%d.%d.%d",
                                    buf[rd], buf[rd+1], buf[rd+2], buf[rd+3]);
                            res = strdup(dotted);
                            break;
                        }
                        pos = rd + rl;
                    }
                }
                else
                {
                    _swix(OS_Byte, _IN(0), 19);
                }
            }
        }
    }
    ExceptCatch
    {
        if(conn) UdpDestruct(conn);
        if(group) TcpAddrDestruct(group);
        Free(buf);
        if(res) Free(res);
        return NULL;        /* nothing answered is not an error */
    }
    if(conn) UdpDestruct(conn);
    if(group) TcpAddrDestruct(group);
    Free(buf);
    return res;
}

/*
    Asking in the background.

    The question is asked once and the answers trickle back over the
    seconds that follow, so the waiting is nearly all of the cost.  A sweep
    driven from a callback cannot sit through that: it has to ask, go away,
    and come back later for whatever has arrived since.

    So the wait is turned inside out.  DiscoverBegin asks and returns
    somewhere to keep the answers; DiscoverPoll takes whatever has turned
    up without waiting at all, and says whether the collecting period still
    has time left in it; DiscoverEnd hands back what answered.

    DiscoverServers is then that loop with a wait in it, which leaves one
    implementation of the protocol rather than a foreground copy and a
    background one that could drift apart.
*/

struct sweep_s
{
    udp_conn_t conn;
    tcp_addr_t group;
    unsigned char *buf;
    work_t *work;
    found_t head;
    clock_t stop;
};

void DiscoverAbandon(sweep_t sw)
{
    if(sw == NULL)
        return;
    if(sw->conn)  UdpDestruct(sw->conn);
    if(sw->group) TcpAddrDestruct(sw->group);
    if(sw->work)  Free(sw->work);
    if(sw->buf)   Free(sw->buf);
    DiscoverFree(sw->head);
    Free(sw);
}

sweep_t DiscoverBegin(int wait_cs)
{
    static const char service[] = "\4_smb\4_tcp\5local";
    sweep_t volatile sw;
    unsigned char quad[4];
    int n;

    sw = Malloc(sizeof(*sw));
    sw->conn = NULL;
    sw->group = NULL;
    sw->buf = NULL;
    sw->work = NULL;
    sw->head = NULL;
    sw->stop = 0;

    ExceptTry
    {
        sw->buf = Malloc(BUF_SIZE);
        sw->work = Malloc(sizeof(*sw->work));

        /*
            The address everything on the network listens to for this.

            LanMan98$DiscoverAddr asks one machine instead, which is what
            to do on a network that will not carry a question addressed to
            everybody - and is the only way to try any of this where the
            question cannot get out at all.
        */
        quad[0] = 224; quad[1] = 0; quad[2] = 0; quad[3] = 251;
        {
            char *var = getenv("LanMan98$DiscoverAddr");

            if(var)
            {
                int a, b, c2, d2;

                if(sscanf(var, "%d.%d.%d.%d", &a, &b, &c2, &d2) == 4)
                {
                    quad[0] = (unsigned char) a; quad[1] = (unsigned char) b;
                    quad[2] = (unsigned char) c2; quad[3] = (unsigned char) d2;
                }
            }
        }
        sw->group = TcpAddrFromQuad((char *) quad);

        /*
            Asked from an ordinary port rather than the service's own, so
            the answers come back here by themselves rather than to
            everybody - which is what makes this possible without joining
            anything.
        */
        sw->conn = UdpCreate(TCP_PORT_ANY);
        UdpSetPeer(MDNS_PORT, sw->group, sw->conn);

        n = 0;
        memset(sw->buf, 0, 12);
        sw->buf[5] = 1;                 /* one question */
        n = 12;
        memcpy(sw->buf + n, service, sizeof(service) - 1);
        n += (int) sizeof(service) - 1;
        sw->buf[n++] = 0;               /* end of the name */
        sw->buf[n++] = 0; sw->buf[n++] = TYPE_PTR;
        sw->buf[n++] = 0; sw->buf[n++] = 1;     /* on the internet */
        UdpWrite(sw->buf, n, sw->conn);

        sw->stop = clock() + (clock_t) wait_cs * (CLOCKS_PER_SEC / 100);
    }
    ExceptCatch
    {
        DiscoverAbandon(sw);
        ExceptRethrow();
    }
    return sw;
}

int DiscoverPoll(sweep_t sw)
{
    volatile int size;
    int budget;

    if(sw == NULL)
        return 0;
    /*
        A few of whatever is waiting, and then back to the caller.

        This is called from a callback, which has to return promptly: the
        desktop is not running while it is in here.  Taking everything the
        socket has in one go is fine on a quiet network and is not fine on
        a busy one, where more can arrive while the last lot is being
        unpacked and the loop need never end - which stops the Wimp
        redrawing or answering anything, and looks exactly like a hung
        machine.

        A read that finds nothing is how the socket says so rather than an
        error worth reporting, which is why the exception here is dropped.
    */
    for(budget = DISCOVER_POLL_MAX; budget > 0; budget--)
    {
        size = 0;
        ThreadSetTimeout(0);
        ExceptTry
        {
            size = UdpRead(sw->buf, BUF_SIZE, sw->conn);
        }
        ExceptCatch
        {
            size = 0;
        }
        ThreadResetTimeout();
        if(size < 12)
            break;
        collect(sw->buf, size, &sw->head, sw->work);
    }
    return clock() < sw->stop;
}

found_t DiscoverEnd(sweep_t sw)
{
    found_t head;

    if(sw == NULL)
        return NULL;
    head = sw->head;
    sw->head = NULL;            /* the caller owns it now */
    DiscoverAbandon(sw);
    return head;
}

found_t DiscoverServers(int wait_cs)
{
    sweep_t sw;

    sw = DiscoverBegin(wait_cs);
    while(DiscoverPoll(sw))
        _swix(OS_Byte, _IN(0), 19);     /* wait a frame, not a spin */
    return DiscoverEnd(sw);
}
