| // Copyright 2017 The Chromium OS Authors. All rights reserved. |
| // Use of this source code is governed by a BSD-style license that can be |
| // found in the LICENSE file. |
| |
| #include "vm_tools/vsh/utils.h" |
| |
| #include <errno.h> |
| #include <stddef.h> |
| #include <stdint.h> |
| #include <unistd.h> |
| |
| #include <arpa/inet.h> |
| #include <sys/types.h> |
| |
| #include <base/bind.h> |
| #include <base/files/file_util.h> |
| #include <base/logging.h> |
| #include <base/posix/eintr_wrapper.h> |
| #include <brillo/message_loops/message_loop.h> |
| |
| using google::protobuf::MessageLite; |
| |
| namespace vm_tools { |
| namespace vsh { |
| namespace { |
| |
| bool SendAllBytes(int sockfd, const uint8_t* buf, uint32_t buf_size) { |
| uint32_t msg_size = htole32(buf_size); |
| |
| if (!base::WriteFileDescriptor( |
| sockfd, reinterpret_cast<const char*>(&msg_size), sizeof(msg_size))) { |
| PLOG(ERROR) << "Failed to write message size to socket"; |
| return false; |
| } |
| |
| if (!base::WriteFileDescriptor(sockfd, reinterpret_cast<const char*>(buf), |
| buf_size)) { |
| PLOG(ERROR) << "Failed to write message to socket"; |
| return false; |
| } |
| |
| return true; |
| } |
| |
| ssize_t RecvAllBytes(int sockfd, uint8_t* buf, uint32_t buf_size) { |
| uint32_t msg_size; |
| |
| if (!base::ReadFromFD(sockfd, reinterpret_cast<char*>(&msg_size), |
| sizeof(msg_size))) { |
| PLOG(ERROR) << "Failed to read message size from socket"; |
| return -1; |
| } |
| msg_size = le32toh(msg_size); |
| |
| if (buf_size < msg_size) { |
| LOG(ERROR) << "Message size of " << msg_size << " exceeds buffer size of " |
| << buf_size; |
| return -1; |
| } |
| |
| if (!base::ReadFromFD(sockfd, reinterpret_cast<char*>(buf), msg_size)) { |
| PLOG(ERROR) << "Failed to read message from socket"; |
| return -1; |
| } |
| |
| return msg_size; |
| } |
| |
| void ShutdownTask() { |
| brillo::MessageLoop::current()->BreakLoop(); |
| } |
| |
| } // namespace |
| |
| bool SendMessage(int sockfd, const MessageLite& message) { |
| size_t msg_size = message.ByteSizeLong(); |
| if (msg_size > kMaxMessageSize) { |
| LOG(ERROR) << "Serialized message too large: " << msg_size; |
| return false; |
| } |
| |
| uint8_t buf[kMaxMessageSize]; |
| |
| if (!message.SerializeToArray(buf, sizeof(buf))) { |
| LOG(ERROR) << "Failed to serialize message"; |
| return false; |
| } |
| |
| if (!SendAllBytes(sockfd, buf, msg_size)) { |
| return false; |
| } |
| |
| return true; |
| } |
| |
| bool RecvMessage(int sockfd, MessageLite* message) { |
| ssize_t count; |
| uint8_t buf[kMaxMessageSize]; |
| |
| count = RecvAllBytes(sockfd, buf, sizeof(buf)); |
| if (count < 0) { |
| return false; |
| } |
| |
| if (!message->ParseFromArray(buf, count)) { |
| LOG(ERROR) << "Failed to parse message"; |
| return false; |
| } |
| |
| return true; |
| } |
| |
| // Posts a shutdown task to the main message loop. |
| void Shutdown() { |
| brillo::MessageLoop::current()->PostTask(FROM_HERE, |
| base::Bind(&ShutdownTask)); |
| } |
| |
| } // namespace vsh |
| } // namespace vm_tools |