1 files changed, 196 insertions, 0 deletions
diff --git a/fastboot/tcp.cpp b/fastboot/tcp.cpp
new file mode 100644
index 000000000..da2880a5e
--- /dev/null
+++ b/fastboot/tcp.cpp
@@ -0,0 +1,196 @@
+/*
+ * Copyright (C) 2016 The Android Open Source Project
+ * All rights reserved.
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions
+ * are met:
+ *  * Redistributions of source code must retain the above copyright
+ *    notice, this list of conditions and the following disclaimer.
+ *  * Redistributions in binary form must reproduce the above copyright
+ *    notice, this list of conditions and the following disclaimer in
+ *    the documentation and/or other materials provided with the
+ *    distribution.
+ *
+ * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
+ * "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
+ * LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS
+ * FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE
+ * COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT,
+ * INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING,
+ * BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS
+ * OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED
+ * AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+ * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT
+ * OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF
+ * SUCH DAMAGE.
+ */
+
+#include "tcp.h"
+
+#include <android-base/stringprintf.h>
+
+namespace tcp {
+
+static constexpr int kProtocolVersion = 1;
+static constexpr size_t kHandshakeLength = 4;
+static constexpr int kHandshakeTimeoutMs = 2000;
+
+// Extract the big-endian 8-byte message length into a 64-bit number.
+static uint64_t ExtractMessageLength(const void* buffer) {
+    uint64_t ret = 0;
+    for (int i = 0; i < 8; ++i) {
+        ret |= uint64_t{reinterpret_cast<const uint8_t*>(buffer)[i]} << (56 - i * 8);
+    }
+    return ret;
+}
+
+// Encode the 64-bit number into a big-endian 8-byte message length.
+static void EncodeMessageLength(uint64_t length, void* buffer) {
+    for (int i = 0; i < 8; ++i) {
+        reinterpret_cast<uint8_t*>(buffer)[i] = length >> (56 - i * 8);
+    }
+}
+
+class TcpTransport : public Transport {
+  public:
+    // Factory function so we can return nullptr if initialization fails.
+    static std::unique_ptr<TcpTransport> NewTransport(std::unique_ptr<Socket> socket,
+                                                      std::string* error);
+
+    ~TcpTransport() override = default;
+
+    ssize_t Read(void* data, size_t length) override;
+    ssize_t Write(const void* data, size_t length) override;
+    int Close() override;
+
+  private:
+    TcpTransport(std::unique_ptr<Socket> sock) : socket_(std::move(sock)) {}
+
+    // Connects to the device and performs the initial handshake. Returns false and fills |error|
+    // on failure.
+    bool InitializeProtocol(std::string* error);
+
+    std::unique_ptr<Socket> socket_;
+    uint64_t message_bytes_left_ = 0;
+
+    DISALLOW_COPY_AND_ASSIGN(TcpTransport);
+};
+
+std::unique_ptr<TcpTransport> TcpTransport::NewTransport(std::unique_ptr<Socket> socket,
+                                                         std::string* error) {
+    std::unique_ptr<TcpTransport> transport(new TcpTransport(std::move(socket)));
+
+    if (!transport->InitializeProtocol(error)) {
+        return nullptr;
+    }
+
+    return transport;
+}
+
+// These error strings are checked in tcp_test.cpp and should be kept in sync.
+bool TcpTransport::InitializeProtocol(std::string* error) {
+    std::string handshake_message(android::base::StringPrintf("FB%02d", kProtocolVersion));
+
+    if (!socket_->Send(handshake_message.c_str(), kHandshakeLength)) {
+        *error = android::base::StringPrintf("Failed to send initialization message (%s)",
+                                             Socket::GetErrorMessage().c_str());
+        return false;
+    }
+
+    char buffer[kHandshakeLength];
+    if (socket_->ReceiveAll(buffer, kHandshakeLength, kHandshakeTimeoutMs) != kHandshakeLength) {
+        *error = android::base::StringPrintf(
+                "No initialization message received (%s). Target may not support TCP fastboot",
+                Socket::GetErrorMessage().c_str());
+        return false;
+    }
+
+    if (memcmp(buffer, "FB", 2) != 0) {
+        *error = "Unrecognized initialization message. Target may not support TCP fastboot";
+        return false;
+    }
+
+    if (memcmp(buffer + 2, "01", 2) != 0) {
+        *error = android::base::StringPrintf("Unknown TCP protocol version %s (host version %02d)",
+                                             std::string(buffer + 2, 2).c_str(), kProtocolVersion);
+        return false;
+    }
+
+    error->clear();
+    return true;
+}
+
+ssize_t TcpTransport::Read(void* data, size_t length) {
+    if (socket_ == nullptr) {
+        return -1;
+    }
+
+    // Unless we're mid-message, read the next 8-byte message length.
+    if (message_bytes_left_ == 0) {
+        char buffer[8];
+        if (socket_->ReceiveAll(buffer, 8, 0) != 8) {
+            Close();
+            return -1;
+        }
+        message_bytes_left_ = ExtractMessageLength(buffer);
+    }
+
+    // Now read the message (up to |length| bytes).
+    if (length > message_bytes_left_) {
+        length = message_bytes_left_;
+    }
+    ssize_t bytes_read = socket_->ReceiveAll(data, length, 0);
+    if (bytes_read == -1) {
+        Close();
+    } else {
+        message_bytes_left_ -= bytes_read;
+    }
+    return bytes_read;
+}
+
+ssize_t TcpTransport::Write(const void* data, size_t length) {
+    if (socket_ == nullptr) {
+        return -1;
+    }
+
+    // Use multi-buffer writes for better performance.
+    char header[8];
+    EncodeMessageLength(length, header);
+    if (!socket_->Send(std::vector<cutils_socket_buffer_t>{{header, 8}, {data, length}})) {
+        Close();
+        return -1;
+    }
+
+    return length;
+}
+
+int TcpTransport::Close() {
+    if (socket_ == nullptr) {
+        return 0;
+    }
+
+    int result = socket_->Close();
+    socket_.reset();
+    return result;
+}
+
+std::unique_ptr<Transport> Connect(const std::string& hostname, int port, std::string* error) {
+    return internal::Connect(Socket::NewClient(Socket::Protocol::kTcp, hostname, port, error),
+                             error);
+}
+
+namespace internal {
+
+std::unique_ptr<Transport> Connect(std::unique_ptr<Socket> sock, std::string* error) {
+    if (sock == nullptr) {
+        // If Socket creation failed |error| is already set.
+        return nullptr;
+    }
+
+    return TcpTransport::NewTransport(std::move(sock), error);
+}
+
+}  // namespace internal
+
+}  // namespace tcp