src/server/util/io.c

/*
 * DO NOT ALTER OR REMOVE COPYRIGHT NOTICES OR THIS HEADER.
 *
 * Copyright 2013 Olaf Wintermann. All rights reserved.
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions are met:
 *
 *   1. Redistributions of source code must retain the above copyright
 *      notice, this list of conditions and the following disclaimer.
 *
 *   2. Redistributions in binary form must reproduce the above copyright
 *      notice, this list of conditions and the following disclaimer in the
 *      documentation and/or other materials provided with the distribution.
 *
 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
 * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
 * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
 * ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
 * LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
 * CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
 * SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
 * INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
 * CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
 * ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
 * POSSIBILITY OF SUCH DAMAGE.
 */

#ifdef __gnu_linux__
#define _GNU_SOURCE
#endif

#include <unistd.h>
#include <stdlib.h>
#include <sys/uio.h>
#include <sys/uio.h>

#ifndef BSD
#include <sys/sendfile.h>
#else
#if defined(__NetBSD__) || defined(__OpenBSD__)
#define WS_NO_SENDFILE
#define net_stream_sendfile net_fallback_sendfile
#endif
#endif

#include <limits.h> /* asprintf */

#include "../daemon/vfs.h"
#include "io.h"
#include "pool.h"

IOStream native_io_funcs = {
    system_write,
    system_writev,
    system_read,
    NULL,
    NULL,
    NULL
};

IOStream net_io_funcs = {
    (io_write_f)net_stream_write,
    (io_writev_f)net_stream_writev,
    (io_read_f)net_stream_read,
    (io_sendfile_f)net_stream_sendfile,
    (io_close_f)net_stream_close,
    (io_finish_f)net_stream_finish
};

IOStream ssl_io_funcs = {
    (io_write_f)net_ssl_write,
    (io_writev_f)net_ssl_writev,
    (io_read_f)net_ssl_read,
    NULL,
    (io_close_f)net_ssl_close,
    (io_finish_f)net_ssl_finish
};


IOStream* stream_new_from_fd(pool_handle_t *pool, int fd) {
    SystemIOStream *st = pool_malloc(pool, sizeof(SystemIOStream));
    st->st = native_io_funcs;
    st->fd = fd;
    return (IOStream*)st;
}

ssize_t system_write(IOStream *st, void *buf, size_t nbytes) {
    return write(((SystemIOStream*)st)->fd, buf, nbytes);
}

ssize_t system_writev(IOStream *st, struct iovec *iovec, int iovcnt) {
    return writev(((SystemIOStream*)st)->fd, iovec, iovcnt);
}

ssize_t system_read(IOStream *st, void *buf, size_t nbytes) {
    return read(((SystemIOStream*)st)->fd, buf, nbytes);
}


IOStream* net_stream_from_fd(pool_handle_t *pool, int fd) {
    NetIOStream *st = pool_malloc(pool, sizeof(NetIOStream));
    st->st = net_io_funcs;
    st->fd = fd;
    st->max_read = 0;
    st->read = 0;
    st->chunked_enc = 0;
    st->buffered = 0;
    return (IOStream*)st;
}

ssize_t net_stream_write(NetIOStream *st, void *buf, size_t nbytes) {
    if(st->chunked_enc) {
        // TODO: on some plattforms iov_len is smaller than size_t
        struct iovec io[2];
        char chunk_len[16];
        io[0].iov_base = chunk_len;
        io[0].iov_len = snprintf(chunk_len, 16, "\r\n%zx\r\n", nbytes);
        io[1].iov_base = buf;
        io[1].iov_len = nbytes;
        ssize_t r = writev(st->fd, io, 2);
        return r - io[0].iov_len;
    } else {
        return write(st->fd, buf, nbytes);
    }
}

ssize_t net_stream_writev(NetIOStream *st, struct iovec *iovec, int iovcnt) {
    if(st->chunked_enc) {
        struct iovec *io = calloc(iovcnt + 1, sizeof(struct iovec));
        char chunk_len[16];
        io[0].iov_base = chunk_len;
        size_t len = 0;
        for(int i=0;i<iovcnt;i++) {
            len += iovec[i].iov_len;
        }
        io[0].iov_len = snprintf(chunk_len, 16, "\r\n%zx\r\n", len);
        memcpy(io + 1, iovec, iovcnt * sizeof(struct iovec));
        ssize_t r = writev(st->fd, io, iovcnt + 1);
        return r - io[0].iov_len;
    } else {
        return writev(st->fd, iovec, iovcnt);
    }
}

ssize_t net_stream_read(NetIOStream *st, void *buf, size_t nbytes) {
    if(st->max_read != 0 && st->read >= st->max_read) {
        return 0;
    }
    ssize_t r = read(st->fd, buf, nbytes);
    st->read += r;
    return r;
}

#ifndef WS_NO_SENDFILE
ssize_t net_stream_sendfile(NetIOStream *st, sendfiledata *sfd) {  
    ssize_t ret = 0;
    off_t fileoffset = sfd->offset;
    if(sfd->fd->fd != -1) {
#ifdef BSD
        struct iovec hdvec;
        hdvec.iov_base = (void*)sfd->header;
        hdvec.iov_len = sfd->hlen;
        struct iovec trvec;
        trvec.iov_base = (void*)sfd->trailer;
        trvec.iov_len = sfd->tlen;
        struct sf_hdtr hdtr;
        hdtr.headers = &hdvec;
        hdtr.hdr_cnt = 1;
        hdtr.trailers = &trvec;
        hdtr.trl_cnt = 1;
        
        off_t len = sfd->len;
#ifdef OSX
        ret = sendfile(sfd->fd->fd, st->fd, fileoffset, &len, &hdtr, 0);
#else // BSD
        ret = sendfile(
                sfd->fd->fd,
                st->fd,
                fileoffset,
                sfd->len,
                &hdtr,
                NULL,
                0);
#endif
#else // Solaris/Linux
        if(sfd->header) {
            ret += write(st->fd, sfd->header, sfd->hlen);
        }
        ret += sendfile(st->fd, sfd->fd->fd, &fileoffset, sfd->len);
        if(sfd->trailer) {
            ret += write(st->fd, sfd->trailer, sfd->tlen);
        }
#endif
    } else {
        // TODO: regular copy
        fprintf(stderr, "sendfile not implemented for SYS_FILE\n");
    }
    
    return ret;
}
#endif

void net_stream_close(NetIOStream *st) {
    close(st->fd);
}

void net_stream_finish(NetIOStream *st) {
    if(st->chunked_enc) {
        write(st->fd, "0\r\n\r\n", 5);
    }
}


/* ssl stream */
IOStream* net_ssl_stream(pool_handle_t *pool, SSL *ssl) {
    SSLStream *st = pool_malloc(pool, sizeof(SSLStream));
    st->st = ssl_io_funcs;
    st->ssl = ssl;
    return (IOStream*)st;
}

ssize_t net_ssl_write(SSLStream *st, void *buf, size_t nbytes) {
    return SSL_write(st->ssl, buf, nbytes);
}

ssize_t net_ssl_writev(SSLStream *st, struct iovec *iovec, int iovcnt) {
    ssize_t r = 0;
    for(int i=0;i<iovcnt;i++) {
        int ret = SSL_write(st->ssl, iovec[i].iov_base, iovec[i].iov_len);
        if(ret <= 0) {
            return 0;
        }
        r += ret;
    }
    return r;
}

ssize_t net_ssl_read(SSLStream *st, void *buf, size_t nbytes) {
    return SSL_read(st->ssl, buf, nbytes);
}

void net_ssl_close(SSLStream *st) {
    SSL_shutdown(st->ssl);
    close(SSL_get_fd(st->ssl));
}

void net_ssl_finish(SSLStream *st) {
    
}


ssize_t net_read(SYS_NETFD fd, void *buf, size_t nbytes) {
    ssize_t r = ((IOStream*)fd)->read(fd, buf, nbytes);
    if(r == 0) {
        return IO_EOF;
    }
    return r;
}

ssize_t net_write(SYS_NETFD fd, void *buf, size_t nbytes) {
    ssize_t r = ((IOStream*)fd)->write(fd, buf, nbytes);
    if(r < 0) {
        return IO_ERROR;
    }  
    return r;
}

ssize_t net_writev(SYS_NETFD fd, struct iovec *iovec, int iovcnt) {
    ssize_t r = ((IOStream*)fd)->writev(fd, iovec, iovcnt);
    if(r < 0) {
        return IO_ERROR;
    }
    return r;
}

ssize_t net_printf(SYS_NETFD fd, char *format, ...) {
    va_list arg;
    va_start(arg, format);
    char *buf;
    ssize_t len = vasprintf(&buf, format, arg);
    ssize_t r = net_write(fd, buf, len);
    free(buf);
    return r;
}

ssize_t net_sendfile(SYS_NETFD fd, sendfiledata *sfd) {
    NetIOStream *out = fd;
    if(out->st.sendfile && sfd->fd && sfd->fd->fd != -1) {
        ssize_t r = ((IOStream*)fd)->sendfile(fd, sfd);
        if(r < 0) {
            return IO_ERROR;
        }
    } else {
        // stream/file does not support sendfile
        // do regular copy
        return net_fallback_sendfile(fd, sfd);
    }
    return IO_ERROR;
}

int net_flush(SYS_NETFD sd) {
    // TODO: implement
    return 0;
}

void net_close(SYS_NETFD fd) {
    ((IOStream*)fd)->close(fd);
}

void net_finish(SYS_NETFD fd) {
    ((IOStream*)fd)->finish(fd);
}


ssize_t net_fallback_sendfile(NetIOStream *st, sendfiledata *sfd) { 
    SYS_NETFD fd = st;
    
    char *buf = malloc(4096);
    if(!buf) {
        // TODO: out of memory error
        return IO_ERROR;
    }
    char *header = (char*)sfd->header;
    int hlen = sfd->hlen;
    char *trailer = (char*)sfd->trailer;
    int tlen = sfd->tlen;
    if(header == NULL) {
        hlen = 0;
    }
    if(trailer == NULL) {
        tlen = 0;
    }

    ssize_t r;
    while(hlen > 0) {
        r = st->st.write(fd, header, hlen);
        header += r;
        hlen -= r;
        if(r <= 0) {
            free(buf);
            return IO_ERROR;
        }
    }

    if(system_lseek(sfd->fd, sfd->offset, SEEK_SET) == -1) {
        free(buf);
        return IO_ERROR;
    }

    size_t length = sfd->len;
    while(length > 0) {
        if((r = system_fread(sfd->fd, buf, 4096)) <= 0) {
            break;
        }
        char *write_buf = buf;
        while(r > 0) {
            ssize_t w = st->st.write(fd, write_buf, r);
            r -= w;
            length -= w;
            write_buf += w;
        }
    }
    free(buf);
    if(length > 0) {
        return IO_ERROR;
    }

    while(tlen > 0) {
        r = st->st.write(fd, trailer, tlen);
        trailer += r;
        tlen -= r;
        if(r <= 0) {
            return IO_ERROR;
        }
    }

    return sfd->hlen + sfd->len + sfd->tlen;
}


/* iovec buffer */
iovec_buf_t *iovec_buf_create(pool_handle_t *pool) {
    iovec_buf_t *buf = pool_malloc(pool, sizeof(iovec_buf_t));

    buf->pool = pool;
    buf->iov  = pool_calloc(pool, 32, sizeof(struct iovec));
    buf->maxiovec = 32;
    buf->iovctn = 0;

    return buf;
}

void iovec_buf_write(iovec_buf_t *io, void *buf, size_t nbyte) {
    if(io->iovctn >= io->maxiovec) {
        io->iov = pool_realloc(
                io->pool,
                io->iov,
                (io->maxiovec + 16) * sizeof(struct iovec));
    }

    io->iov[io->iovctn].iov_base = buf;
    io->iov[io->iovctn].iov_len = nbyte;
    io->iovctn++;
}

ssize_t iovec_buf_flush(iovec_buf_t *io, int fd) {
    return writev(fd, io->iov, io->iovctn);
}


/* TODO: add asprintf to new file */

/* 
 * asprintf implementation for Solaris 10 
 * source from OpenSolaris
 * file: /onnv/onnv-gate/usr/src/lib/libc/port/print/asprintf.c
 */

/*
 * Copyright 2009 Sun Microsystems, Inc.  All rights reserved.
 * Use is subject to license terms.
 */

/*
 * Copyright (c) 2004 Darren Tucker.
 *
 * Based originally on asprintf.c from OpenBSD:
 * Copyright (c) 1997 Todd C. Miller <Todd.Miller@courtesan.com>
 *
 * Permission to use, copy, modify, and distribute this software for any
 * purpose with or without fee is hereby granted, provided that the above
 * copyright notice and this permission notice appear in all copies.
 *
 * THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES
 * WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF
 * MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR
 * ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES
 * WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN
 * ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF
 * OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.
 */

#ifdef __SunOS_5_10

#define	INIT_SZ	128

/* VARARGS2 */
int
vasprintf(char **str, const char *format, va_list ap)
{
	char string[INIT_SZ];
	char *newstr;
	int ret;
	size_t len;

	*str = NULL;
	ret = vsnprintf(string, INIT_SZ, format, ap);
	if (ret < 0)	/* retain the value of errno from vsnprintf() */
		return (-1);
	if (ret < INIT_SZ) {
		len = ret + 1;
		if ((newstr = malloc(len)) == NULL)
			return (-1);	/* retain errno from malloc() */
		(void) strlcpy(newstr, string, len);
		*str = newstr;
		return (ret);
	}
	/*
	 * We will perform this loop more than once only if some other
	 * thread modifies one of the vasprintf() arguments after our
	 * previous call to vsnprintf().
	 */
	for (;;) {
		if (ret == INT_MAX) {	/* Bad length */
			errno = ENOMEM;
			return (-1);
		}
		len = ret + 1;
		if ((newstr = malloc(len)) == NULL)
			return (-1);	/* retain errno from malloc() */
		ret = vsnprintf(newstr, len, format, ap);
		if (ret < 0) {		/* retain errno from vsnprintf() */
			free(newstr);
			return (-1);
		}
		if (ret < len) {
			*str = newstr;
			return (ret);
		}
		free(newstr);
	}
}

int
asprintf(char **str, const char *format, ...)
{
	va_list ap;
	int ret;

	*str = NULL;
	va_start(ap, format);
	ret = vasprintf(str, format, ap);
	va_end(ap);

	return (ret);
}

#endif