Merge "fastbootd: Support TCP protocol." into rvc-dev am: de933e859e

Change-Id: Iffb6639a5ff40c199c4770fed7e83d4af02b02fd
diff --git a/fastboot/Android.bp b/fastboot/Android.bp
index a7cde68..cf0f1ac 100644
--- a/fastboot/Android.bp
+++ b/fastboot/Android.bp
@@ -117,8 +117,10 @@
         "device/main.cpp",
         "device/usb.cpp",
         "device/usb_client.cpp",
+        "device/tcp_client.cpp",
         "device/utility.cpp",
         "device/variables.cpp",
+        "socket.cpp",
     ],
 
     shared_libs: [
@@ -143,6 +145,7 @@
     ],
 
     static_libs: [
+        "libgtest_prod",
         "libhealthhalutils",
         "libsnapshot_nobinder",
     ],
diff --git a/fastboot/device/fastboot_device.cpp b/fastboot/device/fastboot_device.cpp
index bb085c5..1b0859f 100644
--- a/fastboot/device/fastboot_device.cpp
+++ b/fastboot/device/fastboot_device.cpp
@@ -19,6 +19,7 @@
 #include <algorithm>
 
 #include <android-base/logging.h>
+#include <android-base/properties.h>
 #include <android-base/strings.h>
 #include <android/hardware/boot/1.0/IBootControl.h>
 #include <android/hardware/fastboot/1.0/IFastboot.h>
@@ -28,6 +29,7 @@
 
 #include "constants.h"
 #include "flashing.h"
+#include "tcp_client.h"
 #include "usb_client.h"
 
 using android::fs_mgr::EnsurePathUnmounted;
@@ -60,11 +62,16 @@
               {FB_CMD_GSI, GsiHandler},
               {FB_CMD_SNAPSHOT_UPDATE, SnapshotUpdateHandler},
       }),
-      transport_(std::make_unique<ClientUsbTransport>()),
       boot_control_hal_(IBootControl::getService()),
       health_hal_(get_health_service()),
       fastboot_hal_(IFastboot::getService()),
       active_slot_("") {
+    if (android::base::GetProperty("fastbootd.protocol", "usb") == "tcp") {
+        transport_ = std::make_unique<ClientTcpTransport>();
+    } else {
+        transport_ = std::make_unique<ClientUsbTransport>();
+    }
+
     if (boot_control_hal_) {
         boot1_1_ = android::hardware::boot::V1_1::IBootControl::castFrom(boot_control_hal_);
     }
diff --git a/fastboot/device/tcp_client.cpp b/fastboot/device/tcp_client.cpp
new file mode 100644
index 0000000..ec5e1e3
--- /dev/null
+++ b/fastboot/device/tcp_client.cpp
@@ -0,0 +1,176 @@
+/*
+ * 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 "tcp_client.h"
+#include "constants.h"
+
+#include <android-base/errors.h>
+#include <android-base/logging.h>
+#include <android-base/parseint.h>
+#include <android-base/properties.h>
+#include <android-base/stringprintf.h>
+#include <android-base/strings.h>
+
+static constexpr int kDefaultPort = 5554;
+static constexpr int kProtocolVersion = 1;
+static constexpr int kHandshakeTimeoutMs = 2000;
+static constexpr size_t kHandshakeLength = 4;
+
+// 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);
+    }
+}
+
+ClientTcpTransport::ClientTcpTransport() {
+    service_ = Socket::NewServer(Socket::Protocol::kTcp, kDefaultPort);
+
+    // A workaround to notify recovery to continue its work.
+    android::base::SetProperty("sys.usb.ffs.ready", "1");
+}
+
+ssize_t ClientTcpTransport::Read(void* data, size_t len) {
+    if (len > SSIZE_MAX) {
+        return -1;
+    }
+
+    size_t total_read = 0;
+    do {
+        // Read a new message
+        while (message_bytes_left_ == 0) {
+            if (socket_ == nullptr) {
+                ListenFastbootSocket();
+            }
+
+            char buffer[8];
+            if (socket_->ReceiveAll(buffer, 8, 0) == 8) {
+                message_bytes_left_ = ExtractMessageLength(buffer);
+            } else {
+                // If connection is closed by host, Receive will return 0 immediately.
+                socket_.reset(nullptr);
+                // In DATA phase, return error.
+                if (downloading_) {
+                    return -1;
+                }
+            }
+        }
+
+        size_t read_length = len - total_read;
+        if (read_length > message_bytes_left_) {
+            read_length = message_bytes_left_;
+        }
+        ssize_t bytes_read =
+                socket_->ReceiveAll(reinterpret_cast<char*>(data) + total_read, read_length, 0);
+        if (bytes_read == -1) {
+            socket_.reset(nullptr);
+            return -1;
+        } else {
+            message_bytes_left_ -= bytes_read;
+            total_read += bytes_read;
+        }
+    // There are more than one DATA phases if the downloading buffer is too
+    // large, like a very big system image. All of data phases should be
+    // received until the whole buffer is filled in that case.
+    } while (downloading_ && total_read < len);
+
+    return total_read;
+}
+
+ssize_t ClientTcpTransport::Write(const void* data, size_t len) {
+    if (socket_ == nullptr || len > SSIZE_MAX) {
+        return -1;
+    }
+
+    // Use multi-buffer writes for better performance.
+    char header[8];
+    EncodeMessageLength(len, header);
+
+    if (!socket_->Send(std::vector<cutils_socket_buffer_t>{{header, 8}, {data, len}})) {
+        socket_.reset(nullptr);
+        return -1;
+    }
+
+    // In DATA phase
+    if (android::base::StartsWith(reinterpret_cast<const char*>(data), RESPONSE_DATA)) {
+        downloading_ = true;
+    } else {
+        downloading_ = false;
+    }
+
+    return len;
+}
+
+int ClientTcpTransport::Close() {
+    if (socket_ == nullptr) {
+        return -1;
+    }
+    socket_.reset(nullptr);
+
+    return 0;
+}
+
+int ClientTcpTransport::Reset() {
+    return Close();
+}
+
+void ClientTcpTransport::ListenFastbootSocket() {
+    while (true) {
+        socket_ = service_->Accept();
+
+        // Handshake
+        char buffer[kHandshakeLength + 1];
+        buffer[kHandshakeLength] = '\0';
+        if (socket_->ReceiveAll(buffer, kHandshakeLength, kHandshakeTimeoutMs) !=
+            kHandshakeLength) {
+            PLOG(ERROR) << "No Handshake message received";
+            socket_.reset(nullptr);
+            continue;
+        }
+
+        if (memcmp(buffer, "FB", 2) != 0) {
+            PLOG(ERROR) << "Unrecognized initialization message";
+            socket_.reset(nullptr);
+            continue;
+        }
+
+        int version = 0;
+        if (!android::base::ParseInt(buffer + 2, &version) || version < kProtocolVersion) {
+            LOG(ERROR) << "Unknown TCP protocol version " << buffer + 2
+                       << ", our version: " << kProtocolVersion;
+            socket_.reset(nullptr);
+            continue;
+        }
+
+        std::string handshake_message(android::base::StringPrintf("FB%02d", kProtocolVersion));
+        if (!socket_->Send(handshake_message.c_str(), kHandshakeLength)) {
+            PLOG(ERROR) << "Failed to send initialization message";
+            socket_.reset(nullptr);
+            continue;
+        }
+
+        break;
+    }
+}
diff --git a/fastboot/device/tcp_client.h b/fastboot/device/tcp_client.h
new file mode 100644
index 0000000..32e9834
--- /dev/null
+++ b/fastboot/device/tcp_client.h
@@ -0,0 +1,42 @@
+/*
+ * 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.
+ */
+#pragma once
+
+#include <memory>
+
+#include "socket.h"
+#include "transport.h"
+
+class ClientTcpTransport : public Transport {
+  public:
+    ClientTcpTransport();
+    ~ClientTcpTransport() override = default;
+
+    ssize_t Read(void* data, size_t len) override;
+    ssize_t Write(const void* data, size_t len) override;
+    int Close() override;
+    int Reset() override;
+
+  private:
+    void ListenFastbootSocket();
+
+    std::unique_ptr<Socket> service_;
+    std::unique_ptr<Socket> socket_;
+    uint64_t message_bytes_left_ = 0;
+    bool downloading_ = false;
+
+    DISALLOW_COPY_AND_ASSIGN(ClientTcpTransport);
+};
diff --git a/fastboot/socket.cpp b/fastboot/socket.cpp
index e56ffcf..5a14b63 100644
--- a/fastboot/socket.cpp
+++ b/fastboot/socket.cpp
@@ -54,7 +54,9 @@
     while (total < length) {
         ssize_t bytes = Receive(reinterpret_cast<char*>(data) + total, length - total, timeout_ms);
 
-        if (bytes == -1) {
+        // Returns 0 only when the peer has disconnected because our requested length is not 0. So
+        // we return immediately to avoid dead loop here.
+        if (bytes <= 0) {
             if (total == 0) {
                 return -1;
             }