/* 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: libproto.c,v 1.48.2.2.2.1 1997/09/16 20:02:05 wlu Exp $
 */

#include "socks5p.h"
#include "buffer.h"
#include "block.h"
#include "addr.h"
#include "protocol.h"
#include "wrap.h"
#include "libproto.h"
#include "cache.h"
#include "conf.h"
#include "msg.h"
#include "log.h"

/* Send a proxy request to the server                                        */
int lsLibSendRequest(lsSocksInfo *pcon, const S5NetAddr *dest, u_char command) {
    lsProxyInfo *pri = pcon->cur?pcon->cur:pcon->pri;

    return lsSendRequest(pri->cinfo.fd, &pri->cinfo, dest, pri->how, command, 0, NULL);
}

/* Read a response back from the server, assuming it is from a library       */
/* call.  Usually this implies filling in some parts of structures that the  */
/* library keeps and the server does not...                                  */
int lsLibReadResponse(lsSocksInfo *pcon) {
    lsProxyInfo *pri = pcon->cur?pcon->cur:pcon->pri;
    u_char errbyte;
    S5NetAddr *resp;
    int rv;
    
    if (pri == NULL) {
	S5LogUpdate(S5LogDefaultHandle, S5_LOG_DEBUG(1), 0, "lsLibReadResponse: pri was NULL");
	return -1;
    }

    S5LogUpdate(S5LogDefaultHandle, S5_LOG_DEBUG(10), 0, "lsLibReadResponse: Reading Response from fd: %d", pri->cinfo.fd);

    if (pcon->cmd == SOCKS_BIND && pcon->status == CON_ACCEPTING) {
	/* find out who connected                                            */
	resp = &pcon->peer;
    } else if (pcon->cmd == SOCKS_UDP) {
	/* find out where the server wants things send to...                 */
	resp = &pri->prxyin;
    } else {
	/* find out the server's out interface...                            */
	resp = &pri->prxyout;
    }

    if ((rv = lsReadResponse(pri->cinfo.fd, &pri->cinfo, resp, pri->how, &errbyte, &pri->reserved)) < 0) {
	pcon->serrno = GETERRNO();
	S5LogUpdate(S5LogDefaultHandle, S5_LOG_DEBUG(10), 0, "lsLibReadResponse storing errno: %d", pcon->serrno);
    } else {
	S5LogUpdate(S5LogDefaultHandle, S5_LOG_DEBUG(10), 0, "lsLibReadResponse: Response (%s:%d) read", ADDRANDPORT(resp));
    }

    if (pri->reserved != 0x00) {
        u_char flags = 0x00;
 
        if (pcon->cmd == SOCKS_UDP) flags |= S5UDP_USECTRL;
        else if (pcon->cmd == SOCKS_TRACER) flags = flags | SOCKS5_FLAG_VERBOSE | SOCKS5_FLAG_NONAME;
 
        if ((pri->reserved & (~flags)) != 0x00) pri->reserved = 0x00;
    }

    SETERRNO(pcon->serrno);
    return rv;
}

/* Not sure when/if this is ever necessary during failed connects...         */
/* XXX What to do in the nt case...where there's no dup?                     */
static S5IOHandle Reset(S5IOHandle fd, int otype, u_short port) {
    S5NetAddr bndAddr;
    int  len = sizeof(S5NetAddr), optval = 1;
    S5IOHandle nfd;

    nfd = socket(AF_INET, otype, 0);
    if (nfd == S5InvalidIOHandle && port != INVALIDPORT) return nfd;

    if (fd != S5InvalidIOHandle && port != INVALIDPORT) {
    /* It should check all the returned errors. Unfortunately Solaris has    */
    /* bugs that make our checking fails...                                  */
	setsockopt(fd, SOL_SOCKET, SO_REUSEADDR, (char *)&optval, sizeof(int));

	memset((char *)&bndAddr, 0, sizeof(S5NetAddr));
	bndAddr.sa.sa_family = AF_INET;
	bndAddr.sin.sin_port = port;

	while ((fd = REAL(dup2)(nfd, fd)) < 0) {
	    if (!ISSOCKETERROR(EINTR)) break;
	}

	if (fd > 0) REAL(bind)(fd, &bndAddr.sa, len);

	REAL(close)(nfd);
	return fd;
    }

    if (fd != S5InvalidIOHandle) REAL(close)(fd);
    return nfd;
}

/* Establish a proxy channel with a SOCKS server...                          */
lsSocksInfo *lsLibProtoExchg(S5IOHandle fd, const S5NetAddr *rsin, u_char command) {
    int len = sizeof(S5NetAddr), i, nproxies, connected = 0;
    S5NetAddr junk, dest, *proxies;
    S5IOHandle cfd = S5InvalidIOHandle;
    lsSocksInfo *pcon;
    lsProxyInfo *pri;
    u_char how, usectl = 0x00;
    u_short port = INVALIDPORT;

    pcon = lsConnectionFind(fd);

    if (pcon && command != SOCKS_UDP && command != SOCKS_BIND) {
	S5LogUpdate(S5LogDefaultHandle, S5_LOG_DEBUG(10), 0, "lsLibProtoExchg: Deleting bogus connection");
	lsConnectionDel(fd);
	pcon = NULL;
    }

    if (pcon) {
	S5LogUpdate(S5LogDefaultHandle, S5_LOG_DEBUG(10), 0, "lsLibProtoExchg: Connection found");
    } else {
	S5LogUpdate(S5LogDefaultHandle, S5_LOG_DEBUG(10), 0, "lsLibProtoExchg: No connection found");
	if ((pcon = lsConnectionAdd(fd)) == NULL) return NULL;
	pcon->cmd = command;
	S5LogUpdate(S5LogDefaultHandle, S5_LOG_DEBUG(10), 0, "lsLibProtoExchg: Connection added");
    } 

    if (rsin == NULL) return pcon;

    S5LogUpdate(S5LogDefaultHandle, S5_LOG_DEBUG(10), 0, "lsLibProtoExchg: rsin is (%s:%d)", ADDRANDPORT(rsin));
    lsAddrCopy(&pcon->peer, rsin, lsAddrSize(rsin));

    if ((how = lsHowToConnect(rsin, command, &proxies, &nproxies, lsEffUser(), &dest)) == (u_char)-1) {
        S5LogUpdate(S5LogDefaultHandle, S5_LOG_DEBUG(10), 0, "lsProtoExchg: Authorization failed");
	if (command != SOCKS_UDP) lsConnectionDel(fd);
        return NULL;
    }

    if (how == DIRECT || nproxies == 0) {
	S5LogUpdate(S5LogDefaultHandle, S5_LOG_DEBUG(10), 0, "lsLibProtoExchg: Direct -- done...");
	pcon->cur = NULL;
	return pcon;
    }

    /* We are going to walk through the proxy cache first.                       */
    for (i = 0; i < nproxies; i++) {
	/* If we've already talked to this server before, it's still alive, and  */
	/* its udp, we're done.  If it's dead (but the rest is true), we need to */
	/* reestablish the connection, so we delete it...                        */
	S5LogUpdate(S5LogDefaultHandle, S5_LOG_DEBUG(10), 0, "lsLibProtoExchg: Checking proxy cache for (%s:%d)", ADDRANDPORT(&proxies[i])); 

	if ((pri = lsProxyCacheFind(pcon, &proxies[i], how, 0)) != NULL && command == SOCKS_UDP) {
	    if (pri->how == DIRECT || S5IOCheck(pri->cinfo.fd) >= 0) {
		S5LogUpdate(S5LogDefaultHandle, S5_LOG_DEBUG(10), 0, "lsLibProtoExchg: Valid proxy cache entry found");
		pcon->cur = pri;
		return pcon;
	    }
	}

	if (pri != NULL) {
	    S5LogUpdate(S5LogDefaultHandle, S5_LOG_DEBUG(10), 0, "lsLibProtoExchg: Deleting stale proxy cache entry...");
	    lsProxyCacheDel(pcon, pri);
	}
    }

    if (command != SOCKS_UDP) {
	cfd = fd;

	memset((char *)&junk, 0 , sizeof(S5NetAddr));
	if (REAL(getsockname)(fd, &junk.sa, &len) == 0) port = lsAddr2Port(&junk);
    } else if ((cfd = socket(AF_INET, SOCK_STREAM, 0)) == S5InvalidIOHandle) return NULL;

    /* Then we connect to proxy server now...                                     */
    for (i = 0; i < nproxies; i++) {
	S5LogUpdate(S5LogDefaultHandle, S5_LOG_DEBUG(10), 0, "lsLibProtoExchg: Adding to proxy cache...");

	if ((pri = lsProxyCacheAdd(pcon, &proxies[i], how)) == NULL) {
	    if (command != SOCKS_UDP) lsConnectionDel(fd);
	    else if (cfd != S5InvalidIOHandle) CLOSESOCKET(cfd);
	    return NULL;
	}

	S5LogUpdate(S5LogDefaultHandle, S5_LOG_DEBUG(10), 0, "lsLibProtoExchg: Connecting to socks server %s:%d", ADDRANDPORT(&pri->prxyin));

	if (REAL(connect)(cfd, &pri->prxyin.sa, lsAddrSize(&pri->prxyin)) < 0) {
	    if (!ISSOCKETERROR(EINPROGRESS) && !ISSOCKETERROR(EINTR) && !ISSOCKETERROR(EWOULDBLOCK)) {
		S5LogUpdate(S5LogDefaultHandle, S5_LOG_WARNING, 0, "Unable to connect to socks server: %s:%d: %m", ADDRANDPORT(&pri->prxyin));
	        lsProxyCacheDel(pcon, pri);
		if ((cfd = Reset(cfd, SOCK_STREAM, port)) == S5InvalidIOHandle) goto error;
		continue;
	    }
	    
	    for (;;) {
		int snalen = sizeof(S5NetAddr);
		S5NetAddr sna;
		fd_set fds;
		
		FD_ZERO(&fds);
		FD_SET(cfd, &fds);
		
		if (REAL(select)(cfd+1, NULL, &fds, NULL, NULL) < 0) {
		    if (ISSOCKETERROR(EINTR)) continue;
		    S5LogUpdate(S5LogDefaultHandle, S5_LOG_WARNING, 0, "lsLibProtoExchg: Unable to connect to socks server: %s:%d: %m", ADDRANDPORT(&pri->prxyin));
		    break;	    
		}
		
		if (REAL(getpeername)(cfd, &sna.sa, &snalen) < 0) {
		    S5LogUpdate(S5LogDefaultHandle, S5_LOG_DEBUG(1), 0, "lsLibProtoExchg: Non-blocking connect error socks server: %s:%d: %m", ADDRANDPORT(&pri->prxyin));
		    break;
		}

	 	connected = 1;
		break;
	    }

	    if (connected) break;
	    lsProxyCacheDel(pcon, pri);
	    if ((cfd = Reset(cfd, SOCK_STREAM, port)) == S5InvalidIOHandle) goto error;
	} else {
	    connected = 1;
	    break;
	}
    }

    if (connected == 0) {
	S5LogUpdate(S5LogDefaultHandle, S5_LOG_DEBUG(10), 0, "lsLibProtoExchg: Unable to Connect to any socks server");
	goto error;	
    }

    /* if there are more than one server in the list and the connected     */
    /* one is not the first in the list, shuffle the list so that the      */
    /* connected one will be the first in the list. Since the list is      */
    /* global, the consequent use of the list will have better luck to     */
    /* connect a live server.                                              */
    /*                                                                     */
    /* Also, we might shuffle the list everytime we used even no server    */
    /* is dead (load balancing with round-robin selection of servers...    */
    if (nproxies > 1 && i > 0) {
        S5NetAddr tmp;

        S5LogUpdate(S5LogDefaultHandle, S5_LOG_DEBUG(10), 0, "lsLibProtoExchg: Switching the server order");
        memcpy((char *)&tmp, (char *)&proxies[0], sizeof(S5NetAddr));
        memcpy((char *)&proxies[0], (char *)&proxies[i], sizeof(S5NetAddr));
        memcpy((char *)&proxies[i], (char *)&tmp, sizeof(S5NetAddr));
    }

    S5LogUpdate(S5LogDefaultHandle, S5_LOG_DEBUG(10), 0, "lsLibProtoExchg: Connected to socks server");

    if (command == SOCKS_UDP) {
	usectl = S5UDP_USECTRL;
	memset((char *)&junk, 0, sizeof(S5NetAddr));

        if (REAL(getsockname)(fd, &junk.sa, &len) < 0 || lsAddr2Port(&junk) == 0) {
       	    junk.sin.sin_family = AF_INET;
	    if (REAL(bind)(fd, &junk.sa, lsAddrSize(&junk)) < 0) goto error;
	    REAL(getsockname)(fd, &junk.sa, &len);
	} else pcon->myport = lsAddr2Port(&junk);

	REAL(getsockname)(cfd, &dest.sa, &len);
        lsAddrSetPort(&dest, lsAddr2Port(&junk));
    }

    if (lsProtoExchg(cfd, &pri->cinfo, &dest, lsEffUser(), pri->how, command, usectl) < 0) {
	S5LogUpdate(S5LogDefaultHandle, S5_LOG_DEBUG(10), 0, "lsLibProtoExchg: lsProtoExchg Failed");
	goto error;
    }     

    S5LogUpdate(S5LogDefaultHandle, S5_LOG_DEBUG(10), 0, "lsLibProtoExchg: Initial protocol exchanged");

    /* call lsLibReadResponse for bind(), even if it is a non-blocking call. */
    /* this way, getsockname will have the correct information filled in.    */
    /* (We want to block on the _second_ message, not the first one)         */
    /*     --  Jonathan Lemon (jlemon@netcom.com)                            */
    if (command != SOCKS_UDP && command != SOCKS_BIND && ISNBLOCK(fd)) {
	pcon->status = CON_INPROGRESS;
        return pcon;
    }

    if (lsLibReadResponse(pcon) < 0) {
	S5LogUpdate(S5LogDefaultHandle, S5_LOG_DEBUG(10), 0, "lsLibProtoExchg: lsLibReadResponse Failed");
	goto error;
    }

    pcon->cur = pri;
    return pcon;

  error:
    if (command != SOCKS_UDP) lsConnectionDel(fd);
    else {
	if (cfd != S5InvalidIOHandle) CLOSESOCKET(cfd);
	pri->cinfo.fd = S5InvalidIOHandle;
	lsProxyCacheDel(pcon, pri);
    }

    return NULL;
}

/* Perform the UDP sub-commands via the TCP control channel...               */
int lsLibExchgUdpCmd(lsSocksInfo *pcon, const S5NetAddr *dest, u_char cmd) {
    u_char err, flags;
    lsProxyInfo *pri;

    if (!pcon || !(pri = pcon->cur) || pri->reserved != S5UDP_USECTRL) return -1;

    if (S5IOCheck(pri->cinfo.fd) < 0) {
	lsProxyCacheDel(pcon, pri);
	return -1;
    }

    if (lsSendRequest(pri->cinfo.fd, &pri->cinfo, dest, SOCKS5_VERSION, cmd, 0, NULL) < 0) return -1;
    return lsReadResponse(pri->cinfo.fd, &pri->cinfo, &pri->prxyout, SOCKS5_VERSION, &err, &flags);
}