/* 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: select.c,v 1.34.4.3 1998/02/13 23:12:12 wlu Exp $
 */

#define HIDEORIG
#include "socks5p.h"
#include "buffer.h"
#include "wrap.h"
#include "addr.h"
#include "protocol.h"
#include "libproto.h"
#include "cache.h"
#include "msg.h"
#include "log.h"

#define NBCONN(x)   (((x)->cmd == SOCKS_CONNECT) && ((x)->status == CON_INPROGRESS))
#define BUFFERED(x) (((x)->cmd != SOCKS_UDP) && ((x)->pri) && ((x)->pri->cinfo.auth.encode))

/* wrapper around the select system call.                                    */
int LIBPREFIX(select)(S5IOHandle width, fd_set *rfdsp, fd_set *wfdsp, fd_set *efdsp, struct timeval *timeout) {
    fd_set rfs, wfs, efs, w2rfs;
    int n, nbc = 0, nours = 0, size = sizeof(fd_set);
    struct timeval start, end;
    lsSocksInfo *pcon;

#ifdef FOR_SHARED_LIBRARY
    if (lsInRLDFunctions || lsInWrapFunction) return REAL(select)(width, rfdsp, wfdsp, efdsp, timeout);
#endif

    lsInWrapFunction = 1;
    LIBPREFIX2(init)("libsocks5");
    S5LogUpdate(S5LogDefaultHandle, S5_LOG_DEBUG(12), 0, "SOCKS select: FAKE");

    /* If we're not looking at read or write readiness, we can't really do   */
    /* anything, so just return the real select...                           */
    if (!rfdsp && !wfdsp) {
	S5LogUpdate(S5LogDefaultHandle, S5_LOG_DEBUG(12), 0, "SCKS select: REAL: No read/write file descriptors");
    	lsInWrapFunction = 0;
	return REAL(select)(width, NULL, NULL, efdsp, timeout);
    }

    /* Setup the back up file descriptor sets in case we end up looping...   */
    /* Make our best effort not to look at or modify *fdsp beyond size bytes */
    /* out.                                                                  */
    FD_ZERO(&rfs); 
    FD_ZERO(&wfs); 
    FD_ZERO(&efs);
    FD_ZERO(&w2rfs);

    if (rfdsp) memcpy(&rfs, rfdsp, size);
    if (wfdsp) memcpy(&wfs, wfdsp, size);
    if (efdsp) memcpy(&efs, efdsp, size);

    for (pcon = lsConList; pcon; pcon = pcon->next) {
	if (pcon->fd >= width) continue;
	if (pcon->cmd == SOCKS_UDP) continue;
	if (!pcon->pri || pcon->pri->how == DIRECT) continue;
	
	/* The closest thing to anything that makes any sense for            */
	/* non-blocking connects is to connect to the server, and wait for   */
	/* it to connect to the end machine (presumably, thats where the     */
	/* time is spent).  So, we're really looking for read-readiness in   */
	/* the form of a response from the server, not write-readiness in    */
	/* the form of a connect...so, switch them.                          */
	if (wfdsp && FD_ISSET(pcon->fd, &wfs) && NBCONN(pcon)) {
	    S5LogUpdate(S5LogDefaultHandle, S5_LOG_DEBUG(12), 0, "Moving fd: %d from write to read set", pcon->fd);
	    FD_SET(pcon->fd, &rfs);
	    FD_SET(pcon->fd, &w2rfs);
	    FD_CLR(pcon->fd, &wfs);
	    nbc++;
	}

	/* Find out how many of "our" file descriptors we're going to be     */
	/* looking at.  If we find this number is 0, we can just return the  */
	/* real select.                                                      */
	if (FD_ISSET(pcon->fd, &rfs) || FD_ISSET(pcon->fd, &wfs)) {
	    nours++;
	}
    }

    /* If we didn't find any of our file descriptors being examined, just    */
    /* return the real select, we won't be doing anything...                 */
    if (nours == 0) {
	S5LogUpdate(S5LogDefaultHandle, S5_LOG_DEBUG(12), 0, "No file descriptors used by socks; calling real select");
    	lsInWrapFunction = 0;
	return REAL(select)(width, rfdsp, wfdsp, efdsp, timeout);
    }

    if (!rfdsp && nbc) {
	S5LogUpdate(S5LogDefaultHandle, S5_LOG_DEBUG(12), 0, "Using back up read set because we weren't originally interested in reads");
	rfdsp = &rfs;
    }

    /* Before calling the select with timeout, we will check the readable    */
    /* fdsets whether their buffer has something ready for them.             */
    if (rfdsp && rfdsp != &rfs) {
        n = 0;
        FD_ZERO(rfdsp);
        if (wfdsp) FD_ZERO(wfdsp);
        if (efdsp) FD_ZERO(efdsp);

        for (pcon = lsConList; pcon; pcon = pcon->next) {
            if (pcon->fd >= width) continue;
	    if (pcon->cmd == SOCKS_UDP) continue;
	    if (!pcon->pri || pcon->pri->how == DIRECT) continue;
            if (NBCONN(pcon)) continue;
	    if (!FD_ISSET(pcon->fd, &rfs)) continue;

            if (BUFFERED(pcon) && S5BufCheckData(pcon->fd, &pcon->pri->cinfo)) {
                FD_SET(pcon->fd, rfdsp);
                n++;
            }
        }

        if (n > 0) {
            lsInWrapFunction = 0;
            return n;
        }
    }

    if (timeout) gettimeofday(&start, NULL);

    for (;;) {
	if (rfdsp) memcpy(rfdsp, &rfs, size);
	if (wfdsp) memcpy(wfdsp, &wfs, size);
	if (efdsp) memcpy(efdsp, &efs, size);

	S5LogUpdate(S5LogDefaultHandle, S5_LOG_DEBUG(12), 0, "Select: Selecting");
	if ((n = REAL(select)(width, rfdsp, wfdsp, efdsp, timeout)) <= 0) {
    	    lsInWrapFunction = 0;
	    return n;
	}

	/* If we had a non-blocking connecting file descriptor which become  */
	/* read-ready, take care of it...If it didn't really become          */
	/* read-ready, mark it as an error...                                */
	/*                                                                   */
	/* For read file descriptors that became read-ready and have special */
	/* (i.e. buffered) read fucntions, we want to make sure that the     */
	/* whole message is there before we read.  We also don't want to     */
	/* spin forever, so we'll buffer the data that we do read...         */
	/*                                                                   */
	/* We also mark file descriptors which have previously buffered data */
	/* that was already ready to be read but hasn't been.                */
	/*                                                                   */
        /* For non-blocking connecting file descriptor, we have to handle    */
        /* the socks reply.                                                  */
	/*                                                                   */
        /* For non-blocking bind/accept file descriptor, we also have to     */
        /* the socks reply.                                                  */
	if (rfdsp) {
	    for (pcon = lsConList; pcon; pcon = pcon->next) {
		if (pcon->fd >= width) continue;
	        if (pcon->cmd == SOCKS_UDP) continue;
	        if (!pcon->pri || pcon->pri->how == DIRECT) continue;
		if (!FD_ISSET(pcon->fd, rfdsp)) continue;
		if (NBCONN(pcon)) {
		    S5LogUpdate(S5LogDefaultHandle, S5_LOG_DEBUG(12), 0, "Select: Nonblocking connect is read ready");

		    if (lsLibReadResponse(pcon) < 0) {
		    	S5LogUpdate(S5LogDefaultHandle, S5_LOG_DEBUG(12), 0, "Select: Failed to read response: %m");
			SETSOCKETERROR(pcon->serrno);
		        pcon->status = CON_NOTESTABLISHED;
                    } else {
                        pcon->status = CON_ESTABLISHED;
                        lsLastCon = pcon;
                    }

                    if (FD_ISSET(pcon->fd, &w2rfs)) {
                        FD_CLR(pcon->fd, rfdsp);
                        FD_SET(pcon->fd, wfdsp);
                    }
		} else if (BUFFERED(pcon)) {
		    if (S5BufCheckPacket(pcon->fd, &pcon->pri->cinfo) == -2) {
			FD_CLR(pcon->fd, rfdsp);
			n--;
		    }
		}
	    }
	}

	if (n > 0) {
    	    lsInWrapFunction = 0;
	    return n;
	}

	if (!timeout) continue;
	gettimeofday(&end, NULL);

	timeout->tv_sec  -= (end.tv_sec  - start.tv_sec);
	timeout->tv_usec -= (end.tv_usec - start.tv_usec);
	
	while (timeout->tv_usec < 0 && timeout->tv_sec > 0) {
	    timeout->tv_sec  -= 1;
	    timeout->tv_usec += 1000000;
	}
	
	if (timeout->tv_sec < 0 || timeout->tv_usec < 0) {
    	    lsInWrapFunction = 0;
	    return 0;
	}
    }
}