1 files changed, 51 insertions, 12 deletions

diff --git a/fastboot/fuzzy_fastboot/fixtures.cpp b/fastboot/fuzzy_fastboot/fixtures.cpp
index bc13a8c4e..bd76ff4ee 100644
--- a/fastboot/fuzzy_fastboot/fixtures.cpp
+++ b/fastboot/fuzzy_fastboot/fixtures.cpp
@@ -48,12 +48,13 @@
 #include <gtest/gtest.h>
 
 #include "fastboot_driver.h"
+#include "tcp.h"
 #include "usb.h"
 
 #include "extensions.h"
 #include "fixtures.h"
 #include "test_utils.h"
-#include "usb_transport_sniffer.h"
+#include "transport_sniffer.h"
 
 using namespace std::literals::chrono_literals;
 
@@ -74,7 +75,14 @@ int FastBootTest::MatchFastboot(usb_ifc_info* info, const std::string& local_ser
     return 0;
 }
 
+bool FastBootTest::IsFastbootOverTcp() {
+    // serial contains ":" is treated as host ip and port number
+    return (device_serial.find(":") != std::string::npos);
+}
+
 bool FastBootTest::UsbStillAvailible() {
+    if (IsFastbootOverTcp()) return true;
+
     // For some reason someone decided to prefix the path with "usb:"
     std::string prefix("usb:");
     if (std::equal(prefix.begin(), prefix.end(), device_path.begin())) {
@@ -113,15 +121,19 @@ void FastBootTest::SetUp() {
         ASSERT_TRUE(UsbStillAvailible());  // The device disconnected
     }
 
-    const auto matcher = [](usb_ifc_info* info) -> int {
-        return MatchFastboot(info, device_serial);
-    };
-    for (int i = 0; i < MAX_USB_TRIES && !transport; i++) {
-        std::unique_ptr<UsbTransport> usb(usb_open(matcher, USB_TIMEOUT));
-        if (usb)
-            transport = std::unique_ptr<UsbTransportSniffer>(
-                    new UsbTransportSniffer(std::move(usb), serial_port));
-        std::this_thread::sleep_for(std::chrono::milliseconds(10));
+    if (IsFastbootOverTcp()) {
+        ConnectTcpFastbootDevice();
+    } else {
+        const auto matcher = [](usb_ifc_info* info) -> int {
+            return MatchFastboot(info, device_serial);
+        };
+        for (int i = 0; i < MAX_USB_TRIES && !transport; i++) {
+            std::unique_ptr<UsbTransport> usb(usb_open(matcher, USB_TIMEOUT));
+            if (usb)
+                transport = std::unique_ptr<TransportSniffer>(
+                        new TransportSniffer(std::move(usb), serial_port));
+            std::this_thread::sleep_for(std::chrono::milliseconds(10));
+        }
     }
 
     ASSERT_TRUE(transport);  // no nullptr
@@ -154,6 +166,8 @@ void FastBootTest::TearDown() {
 
 // TODO, this should eventually be piped to a file instead of stdout
 void FastBootTest::TearDownSerial() {
+    if (IsFastbootOverTcp()) return;
+
     if (!transport) return;
     // One last read from serial
     transport->ProcessSerial();
@@ -167,9 +181,34 @@ void FastBootTest::TearDownSerial() {
     }
 }
 
+void FastBootTest::ConnectTcpFastbootDevice() {
+    std::size_t found = device_serial.find(":");
+    if (found != std::string::npos) {
+        for (int i = 0; i < MAX_TCP_TRIES && !transport; i++) {
+            std::string error;
+            std::unique_ptr<Transport> tcp(
+                    tcp::Connect(device_serial.substr(0, found), tcp::kDefaultPort, &error)
+                            .release());
+            if (tcp)
+                transport =
+                        std::unique_ptr<TransportSniffer>(new TransportSniffer(std::move(tcp), 0));
+            if (transport != nullptr) break;
+            std::this_thread::sleep_for(std::chrono::milliseconds(10));
+        }
+    }
+}
+
 void FastBootTest::ReconnectFastbootDevice() {
     fb.reset();
     transport.reset();
+
+    if (IsFastbootOverTcp()) {
+        ConnectTcpFastbootDevice();
+        device_path = cb_scratch;
+        fb = std::unique_ptr<FastBootDriver>(new FastBootDriver(transport.get(), {}, true));
+        return;
+    }
+
     while (UsbStillAvailible())
         ;
     printf("WAITING FOR DEVICE\n");
@@ -180,8 +219,8 @@ void FastBootTest::ReconnectFastbootDevice() {
     while (!transport) {
         std::unique_ptr<UsbTransport> usb(usb_open(matcher, USB_TIMEOUT));
         if (usb) {
-            transport = std::unique_ptr<UsbTransportSniffer>(
-                    new UsbTransportSniffer(std::move(usb), serial_port));
+            transport = std::unique_ptr<TransportSniffer>(
+                    new TransportSniffer(std::move(usb), serial_port));
         }
         std::this_thread::sleep_for(1s);
     }