ssh_driver.cpp

#include <errno.h>
#include <fcntl.h>
#include <sys/select.h>
#include <sys/stat.h>
#include <sys/time.h>
#include <sys/types.h>
#include <sys/wait.h>
#include <unistd.h>
#include <pty.h>
#include <stdlib.h>
#include <stdio.h>
#include <string.h>


void xperror(const char *s) {
    perror(s);
    exit(1);
}

void make_nonblock(int fd) {
    fcntl(fd, F_SETFL, fcntl(fd, F_GETFL) | O_NONBLOCK);
}

ssize_t sr_read(int fd, void* buf, size_t count) {
    ssize_t result;

    do {
        result = read(fd, buf, count);
    } while (-1 == result && EINTR == errno);

    return result;
}

ssize_t sr_write(int fd, const void* buf, size_t count) {
    ssize_t result;

    do {
        result = write(fd, buf, count);
    } while (-1 == result && EINTR == errno);

    return result;
}

const size_t BUFF_SIZE = 1024;

struct relay_context {
    bool reof, werr;
    size_t read_pos, write_pos;
    int read_fd, write_fd;
    char buff[BUFF_SIZE];

    relay_context(int read_fd, int write_fd)
        : reof(false)
        , werr(false)
        , read_pos(0)
        , write_pos(0)
        , read_fd(read_fd)
        , write_fd(write_fd)
    { }

    void fill_fdsets(fd_set* readfds, fd_set* writefds) {
        if (finished()) { return; }

        if (0 == read_pos) {
            FD_SET(read_fd, readfds);
            FD_CLR(write_fd, writefds);
        }
        else {
            FD_CLR(read_fd, readfds);
            FD_SET(write_fd, writefds);
        }
    }

    void perform(fd_set* readfds, fd_set* writefds) {
        if (FD_ISSET(read_fd, readfds)) {
            ssize_t bytes_read = sr_read(read_fd, buff, BUFF_SIZE);
            if (-1 == bytes_read) {
                reof = true;

                char msg[32];
                sprintf(msg, "read fd=%d", read_fd);
                perror(msg);
            }
            else if (0 == bytes_read) {
                reof = true;
            }
            else {
                read_pos = bytes_read;
            }
        }

        if (FD_ISSET(write_fd, writefds)) {
            ssize_t bytes_written = sr_write(write_fd, &buff[write_pos], read_pos - write_pos);
            if (-1 == bytes_written) {
                werr = true;

                char msg[32];
                sprintf(msg, "write fd=%d", write_fd);
                perror(msg);
            }
            else {
                write_pos += bytes_written;

                if (write_pos >= read_pos) {
                    write_pos = read_pos = 0;
                }
            }
        }
    }

    bool finished() const {
        return werr || reof && (0 == read_pos);
    }
};

char* format_fdset(char* buf, int nfds, fd_set* fds) {
    bool need_comma = false;
    for (int i = 0; i < nfds; ++i) {
        if (FD_ISSET(i, fds)) {
            buf += sprintf(buf, need_comma ? ", %d" : "%d", i);
            need_comma = true;
        }
    }
    return buf;
}

void log_dump_fdsets(const char* descr, int nfds, fd_set* readfds, fd_set* writefds) {
    char buf[BUFF_SIZE];

    char* ptr = buf;
    ptr += sprintf(ptr, "r=[");
    ptr = format_fdset(ptr, nfds, readfds);
    ptr += sprintf(ptr, "], w=[");
    ptr = format_fdset(ptr, nfds, writefds);
    ptr += sprintf(ptr, "]");

    fprintf(stderr, "%s: %s\n", descr, buf);
}

void run_child(int amaster) {
    make_nonblock(STDIN_FILENO);
    make_nonblock(STDOUT_FILENO);
    make_nonblock(amaster);

    const int nfds = amaster + 1;

    relay_context from_child(amaster, STDOUT_FILENO), to_child(STDIN_FILENO, amaster);

    fd_set readfds, writefds;

    while (!from_child.finished() && !to_child.finished()) {
        FD_ZERO(&readfds);
        FD_ZERO(&writefds);

        from_child.fill_fdsets(&readfds, &writefds);
        to_child.fill_fdsets(&readfds, &writefds);

        log_dump_fdsets("before select", nfds, &readfds, &writefds);

        int r = select(nfds, &readfds, &writefds, nullptr, nullptr);
        if (-1 == r) {
            if (EINTR == errno) { continue; }
            xperror("select");
        }

        log_dump_fdsets("after select", nfds, &readfds, &writefds);

        from_child.perform(&readfds, &writefds);
        to_child.perform(&readfds, &writefds);
    }
}

int main(int argc, char **argv) {
    int amaster;
    pid_t child = forkpty(&amaster, nullptr, nullptr, nullptr);
    if (-1 == child) {
        xperror("forkpty");
    }

    if (0 == child) {
        char* child_argv[] = { "ssh"
            , "-N", "-2", "-v", "-o", "StrictHostKeyChecking=no"
            , "-o", "UserKnownHostsFile=/dev/null", "-o", "PreferredAuthentications=password,keyboard-interactive"
            , "-o", "ExitOnForwardFailure=yes", "-l", "SSH_USER", "-D", "127.0.0.1:9999", "SSH_HOST"
            , nullptr };

        execvp("ssh", child_argv);
    }

    run_child(amaster);

    int status;
    if (-1 == waitpid(child, &status, 0)) {
        xperror("waitpid");
    }

    if (WIFEXITED(status)) {
        printf("Child exited with %d\n", WEXITSTATUS(status));
        exit(0);
    }
    if (WIFSIGNALED(status)) {
        printf("Child was killed with %d\n", WTERMSIG(status));
        exit(0);
    }

    fprintf(stderr, "Unrecognized waitpid() status\n");
    exit(2);
}