/*

 * The Qubes OS Project, http://www.qubes-os.org

 *

 * Copyright (C) 2020 Paweł Marczewski <pawel@invisiblethingslab.com>

 *

 * This program is free software; you can redistribute it and/or

 * modify it under the terms of the GNU General Public License

 * as published by the Free Software Foundation; either version 2

 * of the License, or (at your option) any later version.

 *

 * This program is distributed in the hope that it will be useful,

 * but WITHOUT ANY WARRANTY; without even the implied warranty of

 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the

 * GNU General Public License for more details.

 *

 * You should have received a copy of the GNU General Public License

 * along with this program; if not, write to the Free Software

 * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301, USA.

 *

 */

#include <string.h>

#include <stdlib.h>

#include <stdio.h>

#include <inttypes.h>

#include <unistd.h>

#include <errno.h>

#include <sys/socket.h>

#include <limits.h>

#include <assert.h>

#include "libqrexec-utils.h"

int handle_remote_data(

    libvchan_t *data_vchan, int stdin_fd, int *status,

    struct buffer *stdin_buf, int data_protocol_version,

    bool replace_chars_stdout, bool replace_chars_stderr, bool is_service)

{

    struct msg_header hdr;

    const size_t max_len = max_data_chunk_size(data_protocol_version);

    char *buf;

    int rc = REMOTE_ERROR;

    bool msg_data_warned = false;

    /* do not receive any data if we have something already buffered */

    switch (flush_client_data(stdin_fd, stdin_buf)) {

        case WRITE_STDIN_OK:

            break;

        case WRITE_STDIN_BUFFERED:

            return REMOTE_OK;

        case WRITE_STDIN_ERROR:

            PERROR("write");

            return REMOTE_EOF;

    }

    buf = malloc(max_len);

    if (!buf) {

        PERROR("malloc");

        return REMOTE_ERROR;

    }

    while (libvchan_data_ready(data_vchan) > 0) {

        if (libvchan_recv(data_vchan, &hdr, sizeof(hdr)) != sizeof(hdr))

            goto out;

        if (hdr.len > max_len) {

            LOG(ERROR, "Too big data chunk received: %" PRIu32 " > %zu",

                hdr.len, max_len);

            goto out;

        }

        if (!read_vchan_all(data_vchan, buf, hdr.len))

            goto out;

        switch (hdr.type) {

            /* handle both directions because this can be either server or client

             * of VM-VM connection */

            case MSG_DATA_STDIN:

            case MSG_DATA_STDOUT:

                if (hdr.type != (is_service ? MSG_DATA_STDIN : MSG_DATA_STDOUT) &&

                        !msg_data_warned) {

                    LOG(ERROR, is_service ? "client sent MSG_DATA_STDOUT" : "service sent MSG_DATA_STDIN");

                    msg_data_warned = true;

                }

                if (stdin_fd < 0)

                    /* discard the data */

                    continue;

                if (hdr.len == 0) {

                    rc = REMOTE_EOF;

                    goto out;

                } else {

                    if (replace_chars_stdout)

                        do_replace_chars(buf, hdr.len);

                    switch (write_stdin(stdin_fd, buf, hdr.len, stdin_buf)) {

                        case WRITE_STDIN_OK:

                            break;

                        case WRITE_STDIN_BUFFERED:

                            rc = REMOTE_OK;

                            goto out;

                        case WRITE_STDIN_ERROR:

                            if (!(errno == EPIPE || errno == ECONNRESET)) {

                                PERROR("write");

                            }

                            rc = REMOTE_EOF;

                            goto out;

                    }

                }

                break;

            case MSG_DATA_STDERR:

                if (is_service) {

                    LOG(ERROR, "client sent MSG_DATA_STDERR");

                    continue;

                }

                if (replace_chars_stderr)

                    do_replace_chars(buf, hdr.len);

                /* stderr of remote service, log locally */

                if (!write_all(2, buf, hdr.len)) {

                    PERROR("write");

                    /* only log the error */

                }

                break;

            case MSG_DATA_EXIT_CODE:

                if (is_service) {

                    LOG(ERROR, "client sent MSG_DATA_EXIT_CODE");

                    continue;

                }

                /* remote process exited, so there is no sense to send any data

                 * to it */

                if (hdr.len < sizeof(*status)) {

                    LOG(ERROR, "MSG_DATA_EXIT_CODE too short: " PRIu32 " < %zu",

                        hdr.len, sizeof(*status));

                    *status = 255;

                } else

                    memcpy(status, buf, sizeof(*status));

                rc = REMOTE_EXITED;

                goto out;

            default:

                LOG(ERROR, "unknown msg %d", hdr.type);

                rc = REMOTE_ERROR;

                goto out;

        }

    }

    rc = REMOTE_OK;

out:

    free(buf);

    return rc;

}

int handle_input(

    libvchan_t *vchan, int fd, int msg_type,

    int data_protocol_version, struct prefix_data *prefix_data)

{

    const size_t max_len = max_data_chunk_size(data_protocol_version);

    char *buf;

    ssize_t len;

    struct msg_header hdr;

    int rc = REMOTE_ERROR, buf_space;

    buf = malloc(max_len);

    if (!buf) {

        PERROR("malloc");

        return REMOTE_ERROR;

    }

    static_assert(SSIZE_MAX >= INT_MAX, "can't happen on Linux");

    hdr.type = msg_type;

    while ((buf_space = libvchan_buffer_space(vchan)) > (int)sizeof(struct msg_header)) {

        len = buf_space - sizeof(struct msg_header);

        if ((size_t)len > max_len)

            len = max_len;

        if (prefix_data->len) {

            if ((size_t)len > prefix_data->len)

                len = prefix_data->len;

            memcpy(buf, prefix_data->data, len);

            prefix_data->data += len;

            prefix_data->len -= len;

        } else {

            len = read(fd, buf, len);

            /* If the other side of the socket is a process that is already dead,

             * read from such socket could fail with ECONNRESET instead of

             * just 0. */

            if (len < 0 && errno == ECONNRESET)

                len = 0;

            if (len < 0) {

                if (errno == EAGAIN || errno == EWOULDBLOCK)

                    rc = REMOTE_OK;

                /* otherwise keep rc = REMOTE_ERROR */

                goto out;

            }

        }

        hdr.len = (uint32_t)len;

        /* do not fail on sending EOF (think: close()), it will be handled just below */

        if (libvchan_send(vchan, &hdr, sizeof(hdr)) != sizeof(hdr) && hdr.len != 0)

            goto out;

        if (len && !write_vchan_all(vchan, buf, len))

            goto out;

        if (len == 0) {

            rc = REMOTE_EOF;

            goto out;

        }

    }

    rc = REMOTE_OK;

out:

    free(buf);

    return rc;

}

int send_exit_code(libvchan_t *vchan, int status)

{

    struct msg_header hdr;

    hdr.type = MSG_DATA_EXIT_CODE;

    hdr.len = sizeof(int);

    if (libvchan_send(vchan, &hdr, sizeof(hdr)) != sizeof(hdr)) {

        PERROR("send_exit_code hdr");

        return -1;

    }

    if (libvchan_send(vchan, &status, sizeof(status)) != sizeof(status)) {

        PERROR("send_exit_code status");

        return -1;

    }

    return 0;

}