/*
 * shuffler.c
 * $Id: shuffler.c,v 1.2 1996/04/30 21:49:24 abaird Exp $
 */
    
#include <stdio.h>
#include <sys/types.h>
#include <sys/socket.h>
#include <sys/un.h>
#include <sys/uio.h>
#include <sys/stat.h>
#include <sys/mman.h>
#include <thread.h>
#include <errno.h>
#include <fcntl.h>
#include <signal.h>

#include "ShufflerProtocol.h"

#define D(x) if (trace) { x } else {}
    
#if defined(EAGAIN) && defined(EWOULDBLOCK)
#define ETEST(err) ((err == EAGAIN) || (err == EWOULDBLOCK) )
#else
#if defined(EAGAIN)
#define ETEST(err) (err == EAGAIN) 
#else
#define ETEST(err) (err == EWOULDBLOCK)
#endif
#endif

static int trace     = 0 ;
static int thr_count = 0 ;
static FILE *flog = NULL ;

/*
 * Copy the source file descriptor into the target file descriptor.
 * Whatever happens here, when this function exists, the given file descriptors
 * have been cleaned up (closed), and any reply has been sent back.
 * All hacks permitted, provided this is fast (even though, jigsaw server 
 * really is not faster than the shuffler right now).
 */

int getFileSize (int fd) {
    struct stat buf ;
    if ( isastream (fd) )
	return -1 ;
    if ( fstat (fd, &buf) < 0 ) 
	return -2 ;
    return buf.st_size ;
}
	    
	

typedef struct _CopyInfo {
    int msgid ;			/* ID of the message we are processing */
    int cfd ;			/* client file descriptor */
    int src ;   		/* source file descriptor */
    int src_closed  ;		/* did the copy function closed source ? */
    int dst ;   		/* target file descriptor */
    int dst_closed  ;		/* did the copy function closed dst ? */
    int fsize ; 		/* source file size (when available) */
} CopyInfoRec, *CopyInfo ;

#define CopyInfo_CloseSource(i) \
    ((i)->src_closed ? 0 : ((i)->src_closed = 1, close((i)->src)))
#define CopyInfo_CloseDest(i) \
    ((i)->dst_closed ? 0 : ((i)->dst_closed = 1, close((i)->dst)))  
#define COPY_BUFSIZ 8192

/*
 * Copy a (small) file, but read/write.
 */

static int mem_copy (CopyInfo info) {
    char buffer[COPY_BUFSIZ] ;
    int  got = -1, written = 0 , w = -1 ;
    fd_set wfds ;

    while ((got = read (info->src, buffer, COPY_BUFSIZ)) > 0) {
	while ( written < got ) {
	    FD_ZERO(&wfds) ;
	    FD_SET(info->dst, &wfds);
	    if (select(info->dst+1, NULL, &wfds, NULL, NULL) < 0) 
		return -1 ;
	    w = write (info->dst, buffer+written, got-written) ;
	    if ( w < 0 ) {
		if (ETEST(errno)) 
		    continue ;
		return -1 ;
	    } else {
		written += w ;
	    }
	}
    }
    return info->fsize;
}

/*
 * Copy a file, mmaping it.
 */

static int mmap_copy (CopyInfo info) {
    int written = 0 ;
    int w = -1 ;
    fd_set wfds ;
    char *adr = mmap (0, info->fsize, PROT_READ, MAP_SHARED, info->src, 0) ;

    if ( adr == MAP_FAILED ) {
	return mem_copy (info) ;
    } else {
	CopyInfo_CloseSource(info) ;
	while ( written < info->fsize ) {
	    FD_ZERO(&wfds) ;
	    FD_SET(info->dst, &wfds);
	    if (select(info->dst+1, NULL, &wfds, NULL, NULL) < 0) {
		munmap (adr, info->fsize) ;
		return -1 ;
	    }
	    w = write (info->dst, adr+written, info->fsize-written) ;
	    if ( w < 0 ) {
		if (ETEST(errno)) 
		    continue ;
		munmap (adr, info->fsize) ;
		return -1 ;
	    } else {
		written += w ;
	    }
	}
    }
    munmap (adr, info->fsize) ;
    return info->fsize ;
}

static void *copy_thread (void *pinf) {
    ShufflerMessageRec msg ;
    CopyInfo           info    = (CopyInfo) pinf ;
    int                written = -1 ;

    thr_count++ ;
    D ( fprintf (flog, "shuffler[%d]: do_copy %d to %d (%d threads)\n"
		 , info->cfd, info->src, info->dst, thr_count) ; ) ;
    /* select the appropriate copy function, according to the file size */
    info->fsize = getFileSize (info->src) ;
    if ( (info->fsize < 0) || (info->fsize < COPY_BUFSIZ) ) {
	written = mem_copy (info) ;
    } else {
	written = mmap_copy (info) ;
    }
    /* clean up */
    CopyInfo_CloseSource(info) ;
    CopyInfo_CloseDest(info) ;
    /* send back aknowledgement */
    msg.msg.op     = SHUFFLER_OP_DONE ;
    msg.msg.status = written ;
    msg.msg.msgid  = info->msgid ;
    if ( ShufflerSendMessage (info->cfd, &msg) < 0 ) {
	fprintf (flog, "shuffler[%d]: unable to send aknowledge.\n") ;
	exit (1) ;
    }
    D ( fprintf (flog, "shuffler[%d,fd=%d]: do_copy %d to %d done (%d bytes)\n"
                 , info->msgid, info->cfd, info->src, info->dst, written); ) ;
    free (pinf) ;
    thr_count-- ;
    return NULL ;
}

static int do_copy (int msgid, int cfd, int src, int dst) {
    CopyInfo info = (CopyInfo) malloc (sizeof(CopyInfoRec)) ;

    if ( ! info ) {
	fprintf (flog, "shuffler[do_copy]: unable to malloc !\n") ;
	exit (1) ;
    }
    info->msgid      = msgid ;
    info->cfd        = cfd ;
    info->src        = src ;
    info->src_closed = 0 ;
    info->dst        = dst ;
    info->dst_closed = 0 ;
    /* Check flags on dst */
    if (thr_create(NULL
		   , 0
		   , copy_thread, (void *) info
		   , THR_DETACHED, NULL) < 0) {
	fprintf (flog, "shuffler[do_copy]: unable to fork thread !\n") ;
	exit (1) ;
    }
    return 1 ;
}

/*
 * Advanced UNIX programming (Stevens, p 501)
 */

static int shuffler_listen (char *path) {
    int fd, len ;
    struct sockaddr_un unadr ;
    
    if ((fd = socket (AF_UNIX, SOCK_STREAM, 0)) < 0 )
	return -1 ;
    unlink (path) ;
    memset (&unadr, 0, sizeof (unadr)) ;
    unadr.sun_family = AF_UNIX ;
    strcpy (unadr.sun_path, path) ;
    len = strlen (path) + sizeof (unadr.sun_family) ; 
    if ( bind (fd, (struct sockaddr *) &unadr, len) < 0 )
	return -2 ;
    if ( chmod (unadr.sun_path, 0777) < 0 )
	perror ("chomd ") ;
    if ( listen (fd, 5) < 0 )
	return -3 ;
    return fd ;
}

/*
 * Advanced UNIX programming (Stevens, p 504)
 */

static int shuffler_accept (int sfd) {
    struct sockaddr_un unadr ;
    int unadrlen, fd ;

    unadrlen = sizeof (unadr) ;
    if ( (fd = accept (sfd, (struct sockaddr *) &unadr, &unadrlen)) < 0 )
	return -1 ;
    return fd ;
}

static void *client_thread (void *pfd) {
    ShufflerMessageRec msg ;
    int cfd   = (int) pfd ;
    int src, dst, err ;

    while ( 1 ) {
	/* read in the message */
	if ( (err = ShufflerGetNextMessage (cfd, &msg)) < 0 ) {
	    fprintf (flog
		     , "shuffler[%d]: ShufflerGetNextMessage failed (e=%d).\n"
		     , cfd
		     , -err) ;
	    exit (1) ;
	}
	switch (msg.msg.op) {
	  case SHUFFLER_OP_COPY:
	      do_copy (msg.msg.msgid, cfd, msg.fds[0], msg.fds[1]) ;
	      break ;
	  default:
	      fprintf (flog, "shuffler[%d]: invalid op (%d)\n", msg.msg.op) ;
	      exit (1) ;
	}
    }
}

static void usage () {
    fprintf (stderr, "usage: shuffler <socket-path> [options]\n") ;
    exit (1) ;
}

void sig_pipe(int sig) {
    return ;
}

void main (int argc, char **argv) {
    char *path = "shuffler" ;
    char *logname = NULL ;
    int fd     = -1 ;
    int cli    = -1 ;
    int i = 0 ;

    /* Parse command line options */
    if ( argc < 2 ) 
	usage() ;
    path = argv[1] ;
    for (i = 2 ; i < argc ; i++) {
	if ( ! strcmp(argv[i], "-v") ) {
	    trace = 1 ;
	} else if ( ! strcmp(argv[i], "-log") ) {
	    logname = argv[++i] ;
	} else {
	    usage() ;
	}
    }
    /* Open the log (if needed) */
    if ((flog = fopen(logname, "w+")) == NULL)
	flog = stderr ;
    else
	setbuffer(flog, NULL, 0) ;
    /* set signals */
    signal(SIGPIPE, SIG_IGN) ;
    /* Set up the listening socket */
    D (fprintf (stderr, "shuffler: try listening on %s\n", path) ; ) ;
    if ( (fd = shuffler_listen (path)) < 0 ) {
	fprintf (stderr, "shuffler: unable to listen on %s.\n", path) ;
	exit (1) ;
    }
    D ( fprintf (stderr, "shuffler: listening on %s\n", path) ; ) ;
    /* main loop */
    while ( 1 ) {
	/* get an incoming connection */
	D (fprintf (stderr, "shuffler: waiting for connections.\n"); ) ;
	if ((cli = shuffler_accept (fd)) < 0 ) {
	    D (fprintf(stderr, "shuffler: invalid accept.");) ;
	    continue;
	}
	D ( fprintf (stderr, "shuffler: got new client %d\n", cli); ) ;
	/* run this new client */
	if (thr_create(NULL
		       , 0
		       , client_thread, (void *) cli
		       , THR_DETACHED, NULL) < 0) {
	    perror ("shuffler: unable to create thread ") ;
	    exit (1) ;
	}
    }
}
