/*
 * Copyright (C) 2020 The Android Open Source Project
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#include <errno.h>
#include <sys/socket.h>
#include <unistd.h>

#include <linux/vm_sockets.h>

#include <cstdlib>
#include <cstring>
#include <iostream>

static constexpr const char LOG_TAG[] = "VsockTest: ";
static constexpr int VSOCK_TEST_PORT = 4455;
static constexpr const char VSOCK_TEST_MSG[] = "VsockTestMessage\n";

int VsockServer() {
    int server_fd, client_fd, ret;
    struct sockaddr_vm server_sa = (struct sockaddr_vm) {
        .svm_family = AF_VSOCK,
        .svm_port = VSOCK_TEST_PORT,
        .svm_cid = VMADDR_CID_ANY,
    };
    struct sockaddr_vm client_sa;
    socklen_t socklen = sizeof(client_sa);
    size_t sent = 0, len = strlen(VSOCK_TEST_MSG) + 1;

    server_fd = socket(AF_VSOCK, SOCK_STREAM, 0);
    if (server_fd < 0) {
        std::cerr << "ERROR: socket: " << strerror(errno) << std::endl;
        return -1;
    }

    ret = bind(server_fd, (struct sockaddr*)&server_sa, sizeof(server_sa));
    if (ret != 0) {
        std::cerr << "ERROR: bind: " << strerror(errno) << std::endl;
        close(server_fd);
        return -1;
    }

    std::cerr << LOG_TAG << "Waiting for a vsock connection..." << std::endl;

    ret = listen(server_fd, 1);
    if (ret != 0) {
        std::cerr << "ERROR: listen: " << strerror(errno) << std::endl;
        close(server_fd);
        return -1;
    }

    client_fd = accept(server_fd, (struct sockaddr*)&client_sa, &socklen);
    if (client_fd < 0) {
        std::cerr << "ERROR: accept: " << strerror(errno) << std::endl;
        close(server_fd);
        return -1;
    }

    std::cerr << LOG_TAG << "Sending message to host..." << std::endl;

    while (sent < len) {
        ssize_t nbytes = write(client_fd, VSOCK_TEST_MSG + sent, len - sent);
        if (nbytes > 0) {
            sent += nbytes;
        } else if (nbytes < 0 && errno == EAGAIN) {
            continue;
        } else {
            close(client_fd);
            return EXIT_FAILURE;
        }
    }

    close(client_fd);
    return EXIT_SUCCESS;
}