/* 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: buffer.c,v 1.51.4.3 1998/07/17 19:56:20 wlu Exp $
 */

#include "socks5p.h"
#include "buffer.h"
#include "block.h"
#include "addr.h"
#include "log.h"
#include "msg.h"

#define DATAHEAD(b) ((b)->data + (b)->off)
#define DATASIZE(b) ((b)->len  - (b)->off)

static int S5BufFillPacket(S5Packet *ebuf, char *buffer, int buflen, int ioflags) {
    int len = DATASIZE(ebuf);

    if (len <= 0) return 0;
    if (len > buflen) len = buflen;
    
    memcpy(buffer, DATAHEAD(ebuf), len);
    S5LogUpdate(S5LogDefaultHandle, S5_LOG_DEBUG(10), 0, "S5BufFill: Filled in %d bytes out of buffer", len);

    if (ioflags & MSG_PEEK) return len;
    ebuf->off += len;

    if (DATASIZE(ebuf) > 0) {
	S5LogUpdate(S5LogDefaultHandle, S5_LOG_DEBUG(10), 0, "S5BufFill: Leaving %d bytes in buffer", DATASIZE(ebuf));
    } else {
	free(ebuf->data);

	ebuf->data = NULL;
	ebuf->off  = 0;
	ebuf->len  = 0;
    }

    return len;
}

static int S5BufPutPacket(S5IOHandle fd, char *buffer, int buflen, int ioflags) {
    int n, m = buflen;
    fd_set fds, wfs;
    struct timeval *stm, tm = { 0, 0 };

    FD_ZERO(&wfs);
    FD_SET(fd, &wfs);
    if (ISNBLOCK(fd)) stm = &tm;
    else	      stm = NULL;

    while (m > 0) {
	fds = wfs;
	switch (REAL(select)(fd + 1, NULL, &fds, NULL, stm)) {
	    case -1:
	    	if (ISSOCKETERROR(EINTR)) continue;
		else {
                    SETSOCKETERROR(EBADF);
		    return -1;
	        }
	    case 0:
		if (m == buflen) {

#if defined(sun) && !defined(__svr4__)
                    SETSOCKETERROR(EWOULDBLOCK);
#else
                    SETSOCKETERROR(EAGAIN);
#endif
		    return -1;
	 	}

		stm = NULL;
		continue;
	}

	if ((n = SENDSOCKET(fd, buffer, m, ioflags)) < 0) {
            if (S5IOCheck(fd) >= 0) continue;
            else {
                SETSOCKETERROR(EBADF);
                return -1;
            }
	}

	m -= n;
	buffer += n;
    }

    return 0;
}

static int S5BufGetPacket(S5IOHandle fd, S5IOInfo *cinfo, int block) {
    int nlen, nr;
    S5Packet buf;
    char *nbuf;
    
    if (cinfo->ibuf.data == NULL) {
	cinfo->ibuf.off = 0;
	cinfo->ibuf.len = 0;
    }

    while (1) {
	/* As long as we haven't read in the whole packet, read in more...   */
	buf.data = cinfo->ibuf.data;
	buf.len  = cinfo->ibuf.off;

	if (!(nlen = cinfo->auth.check(&buf, cinfo->auth.context))) {
	    S5LogUpdate(S5LogDefaultHandle, S5_LOG_DEBUG(10), 0, "S5BufGet: Whole packet available...");
	    return 1;
	}

	/* Grow the buffer if necessary...                                   */
	if (cinfo->ibuf.len < cinfo->ibuf.off + nlen) {
	    if (!cinfo->ibuf.data) nbuf = (char *)malloc(cinfo->ibuf.off + nlen);
	    else                   nbuf = (char *)realloc(cinfo->ibuf.data, cinfo->ibuf.off + nlen);

	    if (!nbuf) {
		S5LogUpdate(S5LogDefaultHandle, S5_LOG_WARNING, 0, "S5BufGet: Out of memory enlarging packet buffer");
		SETSOCKETERROR(EBADF);		
		return -1;
	    }

	    cinfo->ibuf.len  = cinfo->ibuf.off + nlen;
	    cinfo->ibuf.data = nbuf;
	}

        if (!block) {
            int na = 0;

            if (S5IOCheck(fd) < 0) return 0;
#ifdef FIONREAD
            if ((nr = IOCTLSOCKET(fd, FIONREAD, &na)) < 0) {
                SETSOCKETERROR(EBADF);
                return -1;
            }

            if (na < nlen) {
                S5LogUpdate(S5LogDefaultHandle, S5_LOG_DEBUG(10), 0, "S5BufGet: Not enough data");
#if defined(sun) && !defined(__svr4__)
                SETSOCKETERROR(EWOULDBLOCK);
#else
                SETSOCKETERROR(EAGAIN);
#endif
                return -2;
	    }
#endif
        }

	switch (nr = RECVSOCKET(fd, cinfo->ibuf.data + cinfo->ibuf.off, nlen, 0)) {
	    case -1:
	 	if (ISSOCKETERROR(EINTR)) continue;
	    case 0:
		return nr;
	}

	/* Update counters...                                                */
	cinfo->ibuf.off += nr;
    }

    SETSOCKETERROR(EBADF);
    return -1;
}
    
/* Check if whole packet is available including data in the I/O buffer       */
/* Return 1 if packet is available, -2 if not available and 0/-1 if I/O      */
/* fails...                                                                  */
int S5BufCheckPacket(S5IOHandle fd, S5IOInfo *cinfo) {
    /* If there's no buffering, or if we've read the whole thing, we're done */
    if (!cinfo || !cinfo->auth.check) return 1;

    return S5BufGetPacket(fd, cinfo, 0);
}

/* Check if whole packet is available in the buffer. Return 1 if packet is   */
/* available, 0 if not.                                                      */
int S5BufCheckData(S5IOHandle fd, S5IOInfo *cinfo) {
    if (cinfo && cinfo->obuf.data) return 1;
    return 0;
}

int S5BufUnreadPacket(S5IOInfo *cinfo, char *ibuf, int ilen) {
    S5Packet nbuf;
    
    if (!cinfo || ilen < 0) return -1;
    if (ilen == 0) return 0;

    if (!cinfo->obuf.data) {
	cinfo->obuf.len  = 0;
	cinfo->obuf.off  = 0;
    }

    /* Try to put the data back into the buffer...                           */
    if (cinfo->obuf.off >= ilen) {
	memcpy(DATAHEAD(&cinfo->obuf) - ilen, ibuf, ilen);
	cinfo->obuf.off -= ilen;
	return 0;
    }

    /* Make a new buffer which can hold what was already there and this.     */
    nbuf.len  = DATASIZE(&cinfo->obuf) + ilen;
    nbuf.off  = 0;
    
    if (!(nbuf.data = (char *)malloc(nbuf.len * sizeof(char)))) {
	S5LogUpdate(S5LogDefaultHandle, S5_LOG_DEBUG(0), 0, "S5BufUnread: Couldn't allocate buffer");
	return -1;
    }

    /* Put the already read data at the beginning of the buffer...           */
    memcpy(DATAHEAD(&nbuf), ibuf, ilen);
    nbuf.off  = ilen;

    /* There was already some data there, so copy it in...                   */
    if (cinfo->obuf.data) {
	memcpy(DATAHEAD(&nbuf), DATAHEAD(&cinfo->obuf), DATASIZE(&cinfo->obuf));
	nbuf.off += DATASIZE(&cinfo->obuf);
    }

    free(cinfo->obuf.data);
    cinfo->obuf = nbuf;
    return 0;
}

int S5BufReadPacket(S5IOHandle fd, S5IOInfo *cinfo, char *ibuf, int ilen, int ioflags) {
    int rval;
    int block = (ISNBLOCK(fd))?0:1;

    if (!cinfo || !cinfo->auth.encode || ilen <= 0) {
	return (int)RECVSOCKET(fd, ibuf, ilen, ioflags);
    }
    
    /* Find out if there was a prior message we had read in an decoded, but  */
    /* couldn't fit into the requested buffer...(this would be typical of a  */
    /* program that calls do (read(fd, buf, 1) while (*buf++ != '\0'), or    */
    /* the like... Also could occur if EINTR occured and restarting wasn't   */
    /* set in S5IORecv...                                                    */
    if (cinfo && cinfo->obuf.data != NULL) {
	return S5BufFillPacket(&cinfo->obuf, (char *)ibuf, ilen, ioflags);
    }

    S5LogUpdate(S5LogDefaultHandle, S5_LOG_DEBUG(10), 0, "S5BufRead: encapsulated...");
    
    switch ((rval = S5BufGetPacket(fd, cinfo, block))) {
	case -2: rval = -1;
	case -1:
	case  0: return rval;
	default: break;
    }

    S5LogUpdate(S5LogDefaultHandle, S5_LOG_DEBUG(10), 0, "S5BufRead: Decoding message...");

    /* Decode message.  Again, if something goes wrong, we have no choice    */
    /* but tho close the file descriptor (since its state is screwed up),    */
    /* and return EBADF.                                                     */
    if (cinfo->auth.encode(&cinfo->ibuf, &cinfo->obuf, S5_DECODE, cinfo->auth.context) < 0) {
	S5LogUpdate(S5LogDefaultHandle, S5_LOG_WARNING, 0, "S5BufRead: Decoding failed.");
	SETSOCKETERROR(EBADF);
	return -1;
    }

    cinfo->obuf.off = 0;
    free(cinfo->ibuf.data);
    cinfo->ibuf.data = NULL;
    
    S5LogUpdate(S5LogDefaultHandle, S5_LOG_DEBUG(10), 0, "S5BufRead: Done");
    return S5BufFillPacket(&cinfo->obuf, ibuf, ilen, ioflags);
}

int S5BufWritePacket(S5IOHandle fd, S5IOInfo *cinfo, char *obuf, int olen, int ioflags) {
    S5Packet buf[2];
    int elen;

    if (!cinfo || !cinfo->auth.encode || olen <= 0) {
	return SENDSOCKET(fd, obuf, olen, ioflags);
    }
    
    buf[0].data = obuf; buf[0].len = olen; buf[0].off = olen;
    buf[1].data = NULL; buf[1].len = 0;	   buf[1].off = 0;
    
    if ((elen = cinfo->auth.encode(&buf[0], &buf[1], S5_ENCODE, cinfo->auth.context)) < 0) {
	S5LogUpdate(S5LogDefaultHandle, S5_LOG_DEBUG(10), 0, "S5BufWrite: encapsulating packet failed");
	SETSOCKETERROR(EBADF);
	return -1;
    }
    
    S5LogUpdate(S5LogDefaultHandle, S5_LOG_DEBUG(10), 0, "S5BufWrite: Sending encapsulated packet");

    /* An easy thing to do is make sure we can write the whole message by    */
    /* making the file descriptor blocking...ugh.                            */
    elen = S5BufPutPacket(fd, buf[1].data, buf[1].len, ioflags);
    free(buf[1].data);

    if (elen < 0) {
    	S5LogUpdate(S5LogDefaultHandle, S5_LOG_DEBUG(10), 0, "S5BufWrite: Network failure");
	return -1;
    } else {
	S5LogUpdate(S5LogDefaultHandle, S5_LOG_DEBUG(10), 0, "S5BufWrite: Done");
        return olen;
    }
}

void S5BufSetupContext(S5IOInfo *cinfo) {
    cinfo->fd           = S5InvalidIOHandle;
    cinfo->ibuf.data    = NULL;
    cinfo->obuf.data    = NULL;
    cinfo->auth.context = NULL;
    cinfo->auth.clean   = NULL;
    cinfo->auth.encode  = NULL;
    cinfo->auth.check   = NULL;
}

void S5BufCleanContext(S5IOInfo *cinfo) {
    if (!cinfo) return;
    
    if (cinfo->auth.clean) cinfo->auth.clean(cinfo->auth.context);
    if (cinfo->fd != S5InvalidIOHandle) CLOSESOCKET(cinfo->fd);
    if (cinfo->ibuf.data) free(cinfo->ibuf.data);
    if (cinfo->obuf.data) free(cinfo->obuf.data);

    S5BufSetupContext(cinfo);
}