uvms/pkgs/ch-proxy/proxy.c

#include <errno.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>

#include <sys/socket.h>
#include <sys/un.h>

#include <linux/vm_sockets.h>

struct msghdr mk_msghdr();
int ch_connect(const char*, const char*);
ssize_t send_fd(int, int);

#define _WRITE_CONFIRM(fd, buf, buflen) {if (write((fd), (buf), (buflen)) != (buflen)) { perror("ch-proxy/write/partial write"); exit(EXIT_FAILURE); }}

void print_usage() {
    fprintf(stderr, "%s",
            "Usage:\n"
            "\tch-proxy uvm/$USER_VM_NAME [PORT]\n"
            "\tch-proxy uuvm/$VM_NAME [PORT]\n"
            "\tch-proxy vsock-mux%$PATH [PORT]\n"
            "\tch-proxy vsock/cid[:port]\n");
}

char *extract_vsock_mux(const char *host_string) {
    const char PREFIX[] = "vsock-mux%";
    const ssize_t PREFIX_LEN = sizeof(PREFIX) - 1;
    if (strncmp(host_string, PREFIX, PREFIX_LEN) == 0) {
        return strdup(&host_string[PREFIX_LEN]);
    }
    return NULL;
}

char *extract_uvm(const char *host_string) {
    const char PREFIX[] = "uuvm/";
    const ssize_t PREFIX_LEN = sizeof(PREFIX) - 1;

    const char *home = getenv("HOME") != NULL ? getenv("HOME") : "./.";

    if (strncmp(host_string, PREFIX, PREFIX_LEN) != 0) {
        return NULL;
    }

    char *result;
    if (asprintf(&result, "%s/uvms/%s/vsock.sock", home, &host_string[PREFIX_LEN]) == -1) {
        perror("ch-proxy/extract_uvm");
        exit(EXIT_FAILURE);
    }

    return result;
}

char *extract_muvm(const char *host_string) {
    const char PREFIX[] = "uvm/";
    const ssize_t PREFIX_LEN = sizeof(PREFIX) - 1;


    if (strncmp(host_string, PREFIX, PREFIX_LEN) != 0) {
        return NULL;
    }

    char *result;
    if (asprintf(&result, "/var/lib/microvms/%s/vsock.sock", &host_string[PREFIX_LEN]) == -1) {
        perror("ch-proxy/extract_muvm");
        exit(EXIT_FAILURE);
    }

    return result;
}

int extract_cid(const char *host, int *cid, int *port) {
    if (errno != 0) {
        perror("extract_cid(...): errno unclean");
        exit(EXIT_FAILURE);
    }

    const char PREFIX[] = "vsock/";
    const ssize_t PREFIX_LEN = sizeof(PREFIX) - 1;

    if (strncmp(host, PREFIX, PREFIX_LEN) != 0) {
        return 1;
    }

    char *in = NULL;
    const char *end = host + strlen(host);
    const char *sCid = host + PREFIX_LEN;

    long x;
    x = strtol(sCid, &in, 10);
    if (errno != 0) {
        perror("strtol(cid, ...)");
        exit(EXIT_FAILURE);
    }
    *cid = (int) x;

    if (in == end) {
        *port = 22;
        return 0;
    } else if (in[0] != ':') {
        perror("extract_cid(...): expected a string of the form cid[:port]");
        perror(host);
        exit(EXIT_FAILURE);
    }

    const char *sPort = in + 1;
    x = strtol(sPort, &in, 10);
    if (errno != 0) {
        perror("strtol(port, ...)");
        exit(EXIT_FAILURE);
    }
    *port = (int) x;
    return 0;
}

int main(int argc, char** argv) {
    if (!(2 <= argc && argc <= 3)) {
        fprintf(stderr, "%s: Wrong # of arguments: %d\n", argv[0], argc);
        print_usage();
        return 1;
    }

    const char PORT_DEFAULT[] = "22";

    const char *ssh_host = argv[1];
    const char *port_string = argc == 3 ? argv[2] : PORT_DEFAULT;

    int cid = -1, port = -1;
    char *path_un = NULL;
    if ((path_un = extract_uvm(ssh_host)) != NULL) {
    } else if ((path_un = extract_muvm(ssh_host)) != NULL) {
    } else if ((path_un = extract_vsock_mux(ssh_host)) != NULL) {
    } else if (extract_cid(ssh_host, &cid, &port) == 0) {
    } else {
        fprintf(stderr, "ch-proxy/main: unexpected host string format: %s\n", ssh_host);
        print_usage();
        return EXIT_FAILURE;
    }

    int s = -1;

    if (path_un != NULL)  {
        s = ch_connect(path_un, port_string);
        if (s == -1) {
            perror("ssh-vsock-proxy/main/ch_connect");
            return EXIT_FAILURE;
        };
    } else if (cid != -1 && port != -1) {
        struct sockaddr_vm sa = {
            .svm_family = AF_VSOCK,
            .svm_reserved1 = 0,
            .svm_port = port,
            .svm_cid = cid,
        };
        memset(sa.svm_zero, 0, sizeof(sa.svm_zero));
        s = socket(AF_VSOCK, SOCK_STREAM, 0);
        if (s == -1) {
            perror("socket(AF_VSOCK, ...)");
            exit(EXIT_FAILURE);
        }
        if (connect(s, (struct sockaddr*)&sa, sizeof(sa)) != 0) {
            perror("connect(socket(AF_VSOCK, ...), ...)");
            exit(EXIT_FAILURE);
        }
    } else {
        perror("Couldn't parse neither uuvm/ strings nor vsock/cid[:port]");
        exit(EXIT_FAILURE);
    }

    if (send_fd(1, s) == -1) {
        perror("ssh-vsock-proxy/main/send_fd");
        return EXIT_FAILURE;
    }
    return 0;
}

struct msghdr mk_msghdr() {
    struct msghdr msg;
    memset(&msg, 0, sizeof(msg));

    return msg;
}

int ch_connect(const char *path, const char *port) {
    int s = socket(AF_UNIX, SOCK_STREAM, 0);
    if (s == -1) {
        perror("ch-proxy/ch_connect/bind");
        return -1;
    }

    struct sockaddr_un a;
    a.sun_family = AF_UNIX;
    strcpy(a.sun_path, path);

    if (connect(s, (struct sockaddr*) &a, sizeof(a)) == -1) {
        char *err_msg;
        if (asprintf(&err_msg, "ch-proxy/ch_connect/connect(\"%s\")", path) == -1) {
            perror("ch-proxy/ch_connect/while printing connect() error");
            exit(EXIT_FAILURE);
        }
        perror(err_msg);
        exit(EXIT_FAILURE);
    }

    const char CMD[] = "CONNECT ";
    const int CMD_LEN = sizeof(CMD) - 1;

    _WRITE_CONFIRM(s, CMD, CMD_LEN);
    _WRITE_CONFIRM(s, port, strlen(port));
    _WRITE_CONFIRM(s, "\n", 1);

    return s;
}

ssize_t send_fd(int dst_fd, int fd) {
    struct msghdr msg = mk_msghdr();

    /* openssh expects to receive a dummy length=1 iovec? */
    char ch;
    struct iovec vec;
    vec.iov_base = &ch;
    vec.iov_len = 1;
    msg.msg_iov = &vec;
    msg.msg_iovlen = 1;

    union {
        struct cmsghdr align;
        char buf[CMSG_SPACE(sizeof(int))];
    } u;

    msg.msg_control = u.buf;
    msg.msg_controllen = sizeof(u.buf);

    struct cmsghdr *cmptr;
    cmptr = CMSG_FIRSTHDR(&msg);

    if (cmptr == NULL) {
        fprintf(stderr, "ch-proxy/send_fd/CMSG_FIRSTHDR: failed to initialize msg_control\n");
        exit(EXIT_FAILURE);
    }

    cmptr->cmsg_len = CMSG_LEN(sizeof(int));
    cmptr->cmsg_level = SOL_SOCKET;
    cmptr->cmsg_type = SCM_RIGHTS;
    *((int*) CMSG_DATA(cmptr)) = fd;

    return (sendmsg(dst_fd, &msg, 0));
}