/* Copyright (c) 1995,1996,1997 NEC Corporation.  All rights reserved.       */
/*                                                                           */
/* The redistribution, use and modification in source or binary forms of     */
/* this software is subject to the conditions set forth in the copyright     */
/* document ("Copyright") included with this distribution.                   */

/*
 * $Id: udputil.c,v 1.46.2.1.2.2 1998/07/19 22:57:56 wlu Exp $
 */

/* This file has a bunch of utility (caching, socket finding, etc...) funcs  */
/* for use in udp.c...Most of them are pretty simple and obvious.            */
#include "socks5p.h"
#include "threads.h"
#include "daemon.h"
#include "protocol.h"
#include "udputil.h"
#include "validate.h"
#include "info.h"
#include "log.h"
#include "msg.h"
#include "s2s.h"

/* Find out which socket we're going to be using to get to host dest...      */
/* Allocate a new entry if we make a new socket to go there...               */
S5IOHandle MakeOutSocket(UdpInfo *u, S5NetAddr *addr, u_short port) {
    static u_short udpport = 0;
    u_short sudpport, tmpport = ntohs(lsAddr2Port(addr));
    S5IOHandle sd;
    S5NetAddr tmpaddr;
    int on = 1;

    if ((sd = socket(AF_INET, SOCK_DGRAM, 0)) == S5InvalidIOHandle) {
	S5LogUpdate(S5LogDefaultHandle, S5_LOG_DEBUG(10), 0, "UDP socket failed for address %s:%d: %m", ADDRANDPORT(addr));
	return S5InvalidIOHandle;
    }

    if (setsockopt(sd, SOL_SOCKET, SO_REUSEADDR, (char *)&on, sizeof(int)) < 0) {
	S5LogUpdate(S5LogDefaultHandle, S5_LOG_DEBUG(10), 0, "UDP setsockopt failed %m");
	CLOSESOCKET(sd);
	return S5InvalidIOHandle;
    }

    if (ludpport > 0 && udpport == 0) udpport = ludpport;

    lsAddrCopy(&tmpaddr, addr, lsAddrSize(addr));
    if (tmpport < ludpport || tmpport > hudpport) {
        if (port == (u_short)0) {
	    lsAddrSetPort(&tmpaddr, htons(udpport));
	    udpport++;
	} else {
	    S5LogUpdate(S5LogDefaultHandle, S5_LOG_DEBUG(10), 0, "UDP client port(%d) is out of range (%d:%d)", (int)tmpport, (int)ludpport, (int)hudpport);
	    CLOSESOCKET(sd);
	    return S5InvalidIOHandle;
        }
    }

    if (bind(sd, &tmpaddr.sa, lsAddrSize(&tmpaddr)) == 0) return sd;

    S5LogUpdate(S5LogDefaultHandle, S5_LOG_DEBUG(10), 0, "UDP bind failed for address %s:%d: %m", ADDRANDPORT(&tmpaddr));

    if (port == (u_short)0) {
	if (ludpport == 0) {
            lsAddrSetPort(&tmpaddr, 0);
	    if (bind(sd, &tmpaddr.sa, lsAddrSize(&tmpaddr)) == 0) return sd;
	} else {
	    sudpport = udpport;
	    while (udpport < hudpport) {
                lsAddrSetPort(&tmpaddr, htons(udpport));
		udpport++;
		if (udpport == hudpport) udpport = ludpport;
	        if (bind(sd, &tmpaddr.sa, lsAddrSize(&tmpaddr)) == 0) return sd;
		if (udpport == sudpport) break;
	    }
	}
    }

    CLOSESOCKET(sd);
    S5LogUpdate(S5LogDefaultHandle, S5_LOG_DEBUG(10), 0, "UDP no port is available");
    return S5InvalidIOHandle;
}

/* Find the out socket needed to get to the next hop (either sckAddr or      */
/* dstsin).  If we haven't made a socket for getting there yet, make one and */
/* add it to the route cache.                                                */
S5IOHandle FindOutSocket(UdpInfo *u, S5LinkInfo *pri, const S5NetAddr *nextAddr, const char *nextName) {
    int len = sizeof(S5NetAddr);
    UdpRouteCache *ar;
    S5NetAddr route;

    /* XXX IPV6 stuff not finished here                                      */

    GetRoute(nextAddr, nextName, "udp", &route);
    lsAddrSetPort(&route, lsAddr2Port(&pri->srcAddr));

    for (ar = u->rcache; ar; ar = ar->next) {
        if (!lsAddrIsNull(&ar->raddr) && lsAddr2Port(&ar->raddr) == lsAddr2Port(&pri->srcAddr)) break;
	if (!lsAddrAddrComp(&ar->raddr, &route)) break;
    }
	
    if (!ar) {
	if ((ar = (UdpRouteCache *)calloc(1, sizeof(UdpRouteCache))) == NULL) {
            S5LogUpdate(S5LogDefaultHandle, S5_LOG_ERROR, 0, "UDP: Creating route cache failed");
	    return S5InvalidIOHandle;
	}

	if ((ar->sd = MakeOutSocket(u, &route, pri->clientPort)) == S5InvalidIOHandle) {
	    free(ar);
	    return S5InvalidIOHandle;
	}

	ar->raddr = route;
    	if (getsockname(ar->sd, (ss *)&route.sa, &len) < 0) {
	    CLOSESOCKET(ar->sd);
	    free(ar);
	    return S5InvalidIOHandle;
	}
        lsAddrSetPort(&ar->raddr, lsAddr2Port(&route));

	ar->next  = u->rcache;
	u->rcache = ar;

        u->maxfd = (u->maxfd > ar->sd)?u->maxfd:ar->sd;
        FD_SET(ar->sd, &u->myfds);
    }

    u->curr = ar;
    pri->intAddr = ar->raddr;
    S5LogUpdate(S5LogDefaultHandle, S5_LOG_DEBUG(10), 0, "UDP Out interface: %s:%d", ADDRANDPORT(&ar->raddr));
    return ar->sd;
}

/* find out if i've talked to this socks server, if not, set up a a cache    */
/* entry for them too...  This does something like...send message to next    */
/* socks5 server; receive a message back, mark that port as the port to talk */
/* to that server.  Then do the client side of authentication with that      */
/* server...                                                                 */
static int AddProxyEntry(UdpInfo *u, S5LinkInfo *pri) {
    UdpSocksCache *sc;
    S5IOHandle outfd;
    S5NetAddr tmp;

    if (!(sc = (UdpSocksCache *)calloc(1, sizeof(UdpSocksCache)))) {
	S5LogUpdate(S5LogDefaultHandle, S5_LOG_DEBUG(10), 0, "UDP malloc failed when adding a new proxy");
	return -1;
    }

    sc->next   = u->scache;
    sc->tcpsin = pri->sckAddr;
    strcpy(sc->name, pri->sckName);
    InitIdentEntry(sc->idtentry);
    S5BufSetupContext(&sc->cinfo);
    u->scache  = sc;

    if ((sc->cinfo.fd = socket(AF_INET, SOCK_STREAM, 0)) < 0) {
	S5LogUpdate(S5LogDefaultHandle, S5_LOG_DEBUG(10), 0, "UDP socket failed when contacting proxy: %m");
	goto error;
    }

    if ((outfd = FindOutSocket(u, pri, &pri->sckAddr, pri->sckName)) == S5InvalidIOHandle) {
	S5LogUpdate(S5LogDefaultHandle, S5_LOG_DEBUG(10), 0, "UDP couldn't find an output udp socket to use with proxy");
	goto error;
    }

    lsAddrCopy(&tmp, &pri->intAddr, lsAddrSize(&pri->intAddr));

    if (S5SExchangeProtocol(&u->iio, &sc->cinfo, pri, sc->idtentry, &tmp, &sc->udpsin) < 0) {
	S5LogUpdate(S5LogDefaultHandle, S5_LOG_DEBUG(10), 0, "UDP Socks5 Protocol exchange failed");
	goto error;
    }

    sc->reserved = pri->nextReserved;
    u->curs = sc;
    return 0;

  error:
    u->scache = sc->next;
    S5BufCleanContext(&sc->cinfo);
    free(sc);
    
    return -1;
}

/* Find out if this person is a socks server we've talked to or not.         */
/* If we're recv'ing, we're just checking, so we return our result if it is  */
/* still alive...If we're sending, we need to initialize it if it does not   */
/* exist in our cache, so we do exactly that...                              */
int FindProxyEntry(UdpInfo *u, S5LinkInfo *pri, S5NetAddr *sckAddr, int tcp) {
    int found = 0;
    UdpSocksCache *sc, *prev = NULL;

    for (sc = u->scache; sc; sc = sc->next) { 
	if (ADDRCOMP((tcp?&sc->tcpsin.sin:&sc->udpsin.sin), &sckAddr->sin)) break;
	prev = sc;
    }

    if (sc) {
	found = 1;
	if (S5IOCheck(sc->cinfo.fd) >= 0) {
	    if (!tcp) strcpy(pri->sckName, sc->name);

	    u->oiop          = &sc->cinfo;
	    pri->sckAddr     = sc->udpsin;
	    pri->nextVersion = SOCKS5_VERSION;
	    u->curs = sc;
	    
	    return 0;
	} else {
	    if (prev != NULL) prev->next = sc->next;
	    else u->scache = sc->next;

	    S5BufCleanContext(&sc->cinfo);
	    free(sc);
	}
    }

    if (!tcp) {
	if (found) return -1;
	memset(&pri->sckAddr, 0, sizeof(S5NetAddr));
	pri->nextVersion = 0;
	u->oiop          = NULL;
    } else {
	if (AddProxyEntry(u, pri) < 0) return -1;

	pri->nextVersion = SOCKS5_VERSION;
	pri->sckAddr     = u->scache->udpsin;
	u->oiop          = &u->scache->cinfo;
    }

    return 0;
}

int CheckIfCached(UdpInfo *u, S5LinkInfo *pri, int checkport) {
    UdpAuthCache *ac;
    UdpRouteCache *ar;
    UdpSocksCache *sc;
    S5NetAddr route;
    u_short port;

    u->curs = NULL;
    u->cura = NULL;
    u->curr = NULL;

    for (ac = u->acache; ac; ac = ac->next) {
	if (lsAddrAddrComp(&ac->addr, &pri->dstAddr)) continue;

	if (checkport) {
	    port = lsAddr2Port(&ac->addr);
	    if (port != (u_short)0 && port != lsAddr2Port(&pri->dstAddr)) continue;
	}

	u->cura = ac;

	if (pri->nextVersion) {
    	    for (sc = u->scache; sc; sc = sc->next) {
	        if (!lsAddrComp((checkport)?&sc->tcpsin:&sc->udpsin, &pri->sckAddr)) break;
    	    }

	    if ((u->curs = sc) == NULL) {
		S5LogUpdate(S5LogDefaultHandle, S5_LOG_DEBUG(10), 0, "CheckCache: The dst (%s:%d) come from wrong server", ADDRANDPORT(&pri->dstAddr));
		return -1;
	    } else {
    	    	GetRoute(&sc->udpsin, sc->name, "udp", &route);
    	    	for (ar = u->rcache; ar; ar = ar->next) {
                    if (!lsAddrIsNull(&ar->raddr) && lsAddr2Port(&ar->raddr) == lsAddr2Port(&pri->srcAddr)) break;
	            if (!lsAddrAddrComp(&ar->raddr, &route)) break;
    	    	}

	    	if ((u->curr = ar) == NULL) {
		    S5LogUpdate(S5LogDefaultHandle, S5_LOG_DEBUG(10), 0, "CheckCache: The dst (%s:%d) come from wrong route", ADDRANDPORT(&sc->udpsin));
		    return -1;
	    	} else return 0;
	    }
	} else {
    	    GetRoute(&pri->dstAddr, pri->dstName, "udp", &route);
    	    for (ar = u->rcache; ar; ar = ar->next) {
                if (!lsAddrIsNull(&ar->raddr) && lsAddr2Port(&ar->raddr) == lsAddr2Port(&pri->srcAddr)) break;

	        if (!lsAddrAddrComp(&ar->raddr, &route)) break;
    	    }

	    if ((u->curr = ar) == NULL) {
		S5LogUpdate(S5LogDefaultHandle, S5_LOG_DEBUG(10), 0, "CheckCache: The dst (%s:%d) come from wrong route", ADDRANDPORT(&pri->dstAddr));
		return -1;
	    } else return 0;
	}
    }

    S5LogUpdate(S5LogDefaultHandle, S5_LOG_DEBUG(10), 0, "CheckCache: The dst (%s:%d) is not cached", ADDRANDPORT(&pri->dstAddr));
    return -1;
}

/* Is the current set of global variables allowed to proxy or not...?  If it */
/* is, set the global variables we use (next_version, out) with the cached   */
/* info...                                                                   */
int CheckIfAllowed(UdpInfo *u, S5LinkInfo *pri) {
    UdpAuthCache *ac;

    /* haven't sent from the client to this host yet I guess we had better     */
    /* make sure we're allowed to...then find out where we're sending it,      */
    /* and resolve the names/reset the version if we need to.                  */
    if (Authorize(pri, 0) != AUTH_OK) {
	S5LogUpdate(S5LogDefaultHandle, S5_LOG_DEBUG(10), 0, "UDP authorization failed");
	return -1;
    }

    /* since we were allowed to talk, we should cache this entry...            */
    if (!(ac = (UdpAuthCache *)calloc(1, sizeof(UdpAuthCache)))) {
	S5LogUpdate(S5LogDefaultHandle, S5_LOG_DEBUG(10), 0, "UDP calloc failed when allocating auth cache entry");
	return -1;
    } 

    lsAddrCopy(&ac->addr, &pri->dstAddr, sizeof(S5NetAddr));
    ac->next  = u->acache;
    u->acache = ac;
    u->cura = ac;

    return 0;
}


int WasResolved(UdpInfo *u, S5NetAddr *dst, char *name) {
#ifdef SENDBACK_NAMES
    UdpHostCache *h;
    char *res;

    if (dst->sa.sa_family != AF_INET) {
	return 0;
    }

    /* Find out if we looked up a name and found this address...             */
    for (h = u->hcache; h; h=h->next) {
	if (h->addr.s_addr == dst->sin_addr.s_addr) break;
    }

    if (h && strcmp(h->name, name)) {
	strncpy(name, h->name, MIN(strlen(h->name)+1, S5_HOSTNAME_LEN));
	if (strlen(h->h_name)+1 > S5_HOSTNAME_SIZE) name[S5_HOSTNAME_SIZE-1] = '\0';
    }

    return h?1:0;
#else
    return 0;
#endif
}

void MarkResolved(UdpInfo *u, S5NetAddr *dst, char *name) {
#ifdef SENDBACK_NAMES
    UdpHostCache *h;

    /* If this host wasn't resolve (we're giving the name to the next        */
    /* server), don't cache it, we won't be using it...                      */
    if (dst->sa.sa_family != AF_INET) {
	return;
    }

    /* Find out if there was an entry already...                             */
    for (h = u->hcache; h; h=h->next) {
	if (strcmp(h->name, name)) break;
    }

    if (h) {
	return;
    }

    if (!(h = (UdpHostCache *)malloc(sizeof(UdpHostCache)))) {
	S5LogUpdate(S5LogDefaultHandle, S5_LOG_DEBUG(10), 0, "Malloc failed when allocating resolved host entry");
	return;
    }

    strcpy(h->name, name);
    h->addr = dst->sin.sin_addr;
    h->next = u->hcache;
    u->hcache = h;
#endif
}