firecracker

1
// Copyright 2018 Amazon.com, Inc. or its affiliates. All Rights Reserved.
2
// SPDX-License-Identifier: Apache-2.0
3

4
// This is a vsock helper tool, used by the Firecracker integration tests,
5
// to - well - test the virtio vsock device. It can be used to
6
// run a vsock echo client, that reads data from STDIN, sends it to an
7
// echo server, then forwards the server's reply to STDOUT.
8

9
#include <sys/socket.h>
10
#include <sys/types.h>
11
#include <sys/select.h>
12
#include <sys/wait.h>
13
#include <linux/vm_sockets.h>
14
#include <stdio.h>
15
#include <stdlib.h>
16
#include <stdint.h>
17
#include <string.h>
18
#include <unistd.h>
19
#include <stdio.h>
20
#include <errno.h>
21

22

23
#define BUF_SIZE (16 * 1024)
24
#define SERVER_ACCEPT_BACKLOG 128
25

26

27
int print_usage() {
28
    fprintf(stderr, "Usage: ./vsock-helper echo <cid> <port>\n");
29
    fprintf(stderr, "\n");
30
    fprintf(stderr, "  echo          connect to an echo server, listening on CID:port.\n");
31
    fprintf(stderr, "                STDIN will be piped through to the echo server, and\n");
32
    fprintf(stderr, "                data coming from the server will pe sent to STDOUT.\n");
33
    fprintf(stderr, "\n");
34
    return -1;
35
}
36

37
int xfer(int src_fd, int dst_fd) {
38
    char buf[BUF_SIZE];
39
    int count =  read(src_fd, buf, sizeof(buf));
40

41
    if (!count) return 0;
42
    if (count < 0) return -1;
43

44
    int offset = 0;
45
    do {
46
        int written;
47
        written = write(dst_fd, &buf[offset], count - offset);
48
        if (written <= 0) return -1;
49
        offset += written;
50
    } while (offset < count);
51

52
    return offset;
53
}
54

55

56
int run_echo(uint32_t cid, uint32_t port) {
57

58
    int sock = socket(AF_VSOCK, SOCK_STREAM, 0);
59
    if (sock < 0) {
60
        perror("socket()");
61
        return -1;
62
    }
63

64
    struct sockaddr_vm vsock_addr = {
65
        .svm_family = AF_VSOCK,
66
        .svm_port = port,
67
        .svm_cid = cid
68
    };
69
    if (connect(sock, (struct sockaddr*)&vsock_addr, sizeof(vsock_addr)) < 0) {
70
        perror("connect()");
71
        return -1;
72
    }
73

74
    for (;;) {
75
        int ping_cnt = xfer(STDIN_FILENO, sock);
76
        if (!ping_cnt) break;
77
        if (ping_cnt < 0) return -1;
78

79
        int pong_cnt = 0;
80
        while (pong_cnt < ping_cnt) {
81
            int res = xfer(sock, STDOUT_FILENO);
82
            if (res <= 0) return -1;
83
            pong_cnt += res;
84
        }
85
    }
86

87
    return 0;
88
}
89

90

91
int main(int argc, char **argv) {
92

93
    if (argc < 3) {
94
        return print_usage();
95
    }
96

97
    if (strcmp(argv[1], "echo") == 0) {
98
        if (argc != 4) {
99
            return print_usage();
100
        }
101
        uint32_t cid = atoi(argv[2]);
102
        uint32_t port = atoi(argv[3]);
103
        if (!cid || !port) {
104
            return print_usage();
105
        }
106
        return run_echo(cid, port);
107
    }
108

109
    return print_usage();
110
}
111