common/frontend/socket_vsock_proxy/main.cpp

/*
 * Copyright (C) 2018 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 <thread>
#include <glog/logging.h>
#include <gflags/gflags.h>

#include "common/libs/fs/shared_fd.h"
#include "common/vsoc/lib/socket_forward_region_view.h"

#ifdef CUTTLEFISH_HOST
#include "host/libs/config/cuttlefish_config.h"
#endif

using vsoc::socket_forward::Packet;

DEFINE_uint32(guest_port, 0,
             "Port on which to forward TCP connections to the guest.");
#ifdef CUTTLEFISH_HOST
DEFINE_uint32(host_port, 0, "Ports on which to run a TCP server on the host.");
DEFINE_uint32(vsock_guest_cid, 0, "Guest identifier");
#endif

namespace {
// Sends packets, Shutdown(SHUT_WR) on destruction
class SocketSender {
 public:
  explicit SocketSender(cvd::SharedFD socket) : socket_{socket} {}

  SocketSender(SocketSender&&) = default;
  SocketSender& operator=(SocketSender&&) = default;

  SocketSender(const SocketSender&&) = delete;
  SocketSender& operator=(const SocketSender&) = delete;

  ~SocketSender() {
    if (socket_.operator->()) {  // check that socket_ was not moved-from
      socket_->Shutdown(SHUT_WR);
    }
  }

  ssize_t SendAll(const Packet& packet) {
    ssize_t written{};
    while (written < static_cast<ssize_t>(packet.payload_length())) {
      if (!socket_->IsOpen()) {
        return -1;
      }
      auto just_written =
          socket_->Send(packet.payload() + written,
                        packet.payload_length() - written, MSG_NOSIGNAL);
      if (just_written <= 0) {
        LOG(INFO) << "Couldn't write to client: "
                  << strerror(socket_->GetErrno());
        return just_written;
      }
      written += just_written;
    }
    return written;
  }

 private:
  cvd::SharedFD socket_;
};

class SocketReceiver {
 public:
  explicit SocketReceiver(cvd::SharedFD socket) : socket_{socket} {}

  SocketReceiver(SocketReceiver&&) = default;
  SocketReceiver& operator=(SocketReceiver&&) = default;

  SocketReceiver(const SocketReceiver&&) = delete;
  SocketReceiver& operator=(const SocketReceiver&) = delete;

  // *packet will be empty if Read returns 0 or error
  void Recv(Packet* packet) {
    auto size = socket_->Read(packet->payload(), sizeof packet->payload());
    if (size < 0) {
      size = 0;
    }
    packet->set_payload_length(size);
  }

 private:
  cvd::SharedFD socket_;
};

void SocketToVsock(SocketReceiver socket_receiver,
                   SocketSender vsock_sender) {
  while (true) {
    auto packet = Packet::MakeData();
    socket_receiver.Recv(&packet);
    if (packet.empty() || vsock_sender.SendAll(packet) < 0) {
      break;
    }
  }
  LOG(INFO) << "Socket to vsock exiting";
}

void VsockToSocket(SocketSender socket_sender,
                   SocketReceiver vsock_receiver) {
  auto packet = Packet::MakeData();
  while (true) {
    vsock_receiver.Recv(&packet);
    CHECK(packet.IsData());
    if (packet.empty()) {
      break;
    }
    if (socket_sender.SendAll(packet) < 0) {
      break;
    }
  }
  LOG(INFO) << "Vsock to socket exiting";
}

// One thread for reading from shm and writing into a socket.
// One thread for reading from a socket and writing into shm.
void HandleConnection(cvd::SharedFD vsock,
                      cvd::SharedFD socket) {
  auto socket_to_vsock =
      std::thread(SocketToVsock, SocketReceiver{socket}, SocketSender{vsock});
  VsockToSocket(SocketSender{socket}, SocketReceiver{vsock});
  socket_to_vsock.join();
}

#ifdef CUTTLEFISH_HOST
[[noreturn]] void host() {
  LOG(INFO) << "starting server on " << FLAGS_host_port << " for guest port "
            << FLAGS_guest_port;
  auto server = cvd::SharedFD::SocketLocalServer(FLAGS_host_port, SOCK_STREAM);
  CHECK(server->IsOpen()) << "Could not start server on " << FLAGS_host_port;
  while (true) {
    LOG(INFO) << "waiting for client connection";
    auto client_socket = cvd::SharedFD::Accept(*server);
    CHECK(client_socket->IsOpen()) << "error creating client socket";
    LOG(INFO) << "client socket accepted";
    auto vsock_socket = cvd::SharedFD::VsockClient(FLAGS_vsock_guest_cid,
                                                   FLAGS_guest_port,
                                                   SOCK_STREAM);
    CHECK(vsock_socket->IsOpen()) << "error connecting over vsock";
    auto thread = std::thread(HandleConnection, std::move(vsock_socket),
                              std::move(client_socket));
    thread.detach();
  }
}

#else
cvd::SharedFD OpenSocketConnection() {
  while (true) {
    auto sock = cvd::SharedFD::SocketLocalClient(FLAGS_guest_port, SOCK_STREAM);
    if (sock->IsOpen()) {
      return sock;
    }
    LOG(WARNING) << "could not connect on port " << FLAGS_guest_port
                 << ". sleeping for 1 second";
    sleep(1);
  }
}

[[noreturn]] void guest() {
  LOG(INFO) << "Starting guest mainloop";
  LOG(INFO) << "starting server on " << FLAGS_guest_port;
  auto vsock = cvd::SharedFD::VsockServer(FLAGS_guest_port, SOCK_STREAM);
  CHECK(vsock->IsOpen()) << "Could not start server on " << FLAGS_guest_port;
  while (true) {
    LOG(INFO) << "waiting for vsock connection";
    auto vsock_client = cvd::SharedFD::Accept(*vsock);
    CHECK(vsock_client->IsOpen()) << "error creating vsock socket";
    LOG(INFO) << "vsock socket accepted";
    auto client = OpenSocketConnection();
    CHECK(client->IsOpen()) << "error connecting to guest client";
    auto thread = std::thread(HandleConnection, std::move(vsock_client),
                              std::move(client));
    thread.detach();
  }
}

#endif
}  // namespace

int main(int argc, char* argv[]) {
  gflags::ParseCommandLineFlags(&argc, &argv, true);

  CHECK(FLAGS_guest_port != 0) << "Must specify --guest_port flag";
#ifdef CUTTLEFISH_HOST
  CHECK(FLAGS_vsock_guest_cid != 0) << "Must specify --vsock_guest_cid flag";
  CHECK(FLAGS_host_port != 0) << "Must specify --host_port flag";
  host();
#else
  guest();
#endif
}