Refactor AndroidSetWifiDecorator to circumvent credential leak stripping by inverting lab network provisioning via native gRPC loopback.

src/java/com/google/devtools/mobileharness/api/model/error/AndroidErrorId.java

@@ -711,6 +711,9 @@ public enum AndroidErrorId implements ErrorId {
   // AndroidSetWifiDecorator: 129_201 - 129_250
   ANDROID_SET_WIFI_DECORATOR_GET_DEFAULT_SSID_ERROR(129_201, ErrorType.CUSTOMER_ISSUE),
   ANDROID_SET_WIFI_DECORATOR_SSID_NOT_PRESENT_ERROR(129_202, ErrorType.CUSTOMER_ISSUE),
+  ANDROID_SET_WIFI_DECORATOR_WIFI_CONNECT_ERROR(129_203, ErrorType.INFRA_ISSUE),
+
+  // [AndroidSmsDecorator Error: 129301 ~ 129400]
 
   // AndroidDeviceFeaturesCheckDecorator: 129_351 ~ 129_400
   ANDROID_DEVICE_FEATURES_CHECK_DECORATOR_CEHCK_FAILURE(129_351, ErrorType.DEPENDENCY_ISSUE),

src/java/com/google/devtools/mobileharness/api/model/error/InfraErrorId.java

@@ -189,6 +189,8 @@ public enum InfraErrorId implements ErrorId {
   LAB_RPC_DEVICE_OPS_RUN_TROUBLESHOOT_SCRIPT_GRPC_ERROR(40_488, ErrorType.UNDETERMINED),
   LAB_RPC_DEVICE_OPS_RUN_TROUBLESHOOT_SCRIPT_STUBBY_ERROR(40_489, ErrorType.UNDETERMINED),
   LAB_RPC_DEVICE_OPS_TROUBLESHOOT_SCRIPT_PRECONDITION_FAILED(40_490, ErrorType.CUSTOMER_ISSUE),
+  LAB_RPC_DEVICE_OPS_CONNECT_TO_DEFAULT_WIFI_GRPC_ERROR(40_495, ErrorType.UNDETERMINED),
+  LAB_RPC_DEVICE_OPS_CONNECT_TO_DEFAULT_WIFI_STUBBY_ERROR(40_496, ErrorType.UNDETERMINED),
 
   // Test engine/container: 40_701 ~ 40_800
   TE_CREATE_DEVICE_HELPER_CONTAINER_DOES_NOT_HAVE(40_701, ErrorType.INFRA_ISSUE),

src/java/com/google/devtools/mobileharness/infra/client/api/util/stub/BUILD

@@ -29,10 +29,13 @@ java_library(
     srcs = ["GrpcStubManager.java"],
     visibility = [
         "//src/java/com/google/wireless/qa/mobileharness/client/api/util/stub:__pkg__",
+        "//src/java/com/google/wireless/qa/mobileharness/shared/api/decorator:__pkg__",
     ],
     deps = [
+        "//src/java/com/google/devtools/mobileharness/infra/lab/rpc/stub:device_ops",
         "//src/java/com/google/devtools/mobileharness/infra/lab/rpc/stub:exec_test",
         "//src/java/com/google/devtools/mobileharness/infra/lab/rpc/stub:prepare_test",
+        "//src/java/com/google/devtools/mobileharness/infra/lab/rpc/stub/grpc:device_ops_grpc_stub",
         "//src/java/com/google/devtools/mobileharness/infra/lab/rpc/stub/grpc:exec_test_grpc_stub",
         "//src/java/com/google/devtools/mobileharness/infra/lab/rpc/stub/grpc:prepare_test_grpc_stub",
         "//src/java/com/google/devtools/mobileharness/shared/util/comm/stub:channel_factory",

src/java/com/google/devtools/mobileharness/infra/client/api/util/stub/GrpcStubManager.java

@@ -16,8 +16,10 @@
 
 package com.google.devtools.mobileharness.infra.client.api.util.stub;
 
+import com.google.devtools.mobileharness.infra.lab.rpc.stub.DeviceOpsStub;
 import com.google.devtools.mobileharness.infra.lab.rpc.stub.ExecTestStub;
 import com.google.devtools.mobileharness.infra.lab.rpc.stub.PrepareTestStub;
+import com.google.devtools.mobileharness.infra.lab.rpc.stub.grpc.DeviceOpsGrpcStub;
 import com.google.devtools.mobileharness.infra.lab.rpc.stub.grpc.ExecTestGrpcStub;
 import com.google.devtools.mobileharness.infra.lab.rpc.stub.grpc.PrepareTestGrpcStub;
 import com.google.devtools.mobileharness.shared.util.comm.stub.ChannelManager;
@@ -61,4 +63,9 @@ public ExecTestStub getExecTestGrpcStub(String grpcTarget) {
     return channelManager.createStub(
         grpcTarget, () -> managedChannelSupplier.apply(grpcTarget), ExecTestGrpcStub::new);
   }
+
+  public DeviceOpsStub getDeviceOpsStub(String grpcTarget) {
+    return channelManager.createStub(
+        grpcTarget, () -> managedChannelSupplier.apply(grpcTarget), DeviceOpsGrpcStub::new);
+  }
 }

src/java/com/google/devtools/mobileharness/infra/lab/rpc/stub/DeviceOpsStub.java

@@ -19,6 +19,8 @@
 import com.google.common.util.concurrent.ListenableFuture;
 import com.google.devtools.common.metrics.stability.rpc.RpcExceptionWithErrorId;
 import com.google.devtools.mobileharness.shared.constant.closeable.NonThrowingAutoCloseable;
+import com.google.wireless.qa.mobileharness.lab.proto.DeviceOpsServ.ConnectToDefaultWifiRequest;
+import com.google.wireless.qa.mobileharness.lab.proto.DeviceOpsServ.ConnectToDefaultWifiResponse;
 import com.google.wireless.qa.mobileharness.lab.proto.DeviceOpsServ.GetDeviceDebugInfoRequest;
 import com.google.wireless.qa.mobileharness.lab.proto.DeviceOpsServ.GetDeviceDebugInfoResponse;
 import com.google.wireless.qa.mobileharness.lab.proto.DeviceOpsServ.GetDeviceLogRequest;
@@ -98,4 +100,25 @@ default RunTroubleshootScriptResponse runTroubleshootScript(
   /** Runs a troubleshoot script asynchronously. */
   ListenableFuture<RunTroubleshootScriptResponse> runTroubleshootScriptAsync(
       RunTroubleshootScriptRequest request);
+
+  default ConnectToDefaultWifiResponse connectToDefaultWifi(ConnectToDefaultWifiRequest request)
+      throws RpcExceptionWithErrorId {
+    throw new UnsupportedOperationException();
+  }
+
+  default ConnectToDefaultWifiResponse connectToDefaultWifi(
+      ConnectToDefaultWifiRequest request, @Nullable String impersonationUser)
+      throws RpcExceptionWithErrorId {
+    return connectToDefaultWifi(request);
+  }
+
+  default ListenableFuture<ConnectToDefaultWifiResponse> connectToDefaultWifiAsync(
+      ConnectToDefaultWifiRequest request) {
+    throw new UnsupportedOperationException();
+  }
+
+  default ListenableFuture<ConnectToDefaultWifiResponse> connectToDefaultWifiAsync(
+      ConnectToDefaultWifiRequest request, boolean useClientRpcAuthority) {
+    return connectToDefaultWifiAsync(request);
+  }
 }

src/java/com/google/devtools/mobileharness/infra/lab/rpc/stub/grpc/BUILD

@@ -64,6 +64,7 @@ java_library(
         "//src/java/com/google/wireless/qa/mobileharness/lab/proto:device_ops_serv_grpc",
         "//src/java/com/google/wireless/qa/mobileharness/lab/proto:device_ops_serv_java_proto",
         "@grpc-java//core",
+        "@maven//:com_google_code_findbugs_jsr305",
         "@maven//:com_google_guava_guava",
     ],
 )

src/java/com/google/devtools/mobileharness/infra/lab/rpc/stub/grpc/DeviceOpsGrpcStub.java

@@ -22,6 +22,8 @@
 import com.google.devtools.mobileharness.api.model.error.InfraErrorId;
 import com.google.devtools.mobileharness.infra.lab.rpc.stub.DeviceOpsStub;
 import com.google.devtools.mobileharness.shared.util.comm.stub.GrpcDirectTargetConfigures;
+import com.google.wireless.qa.mobileharness.lab.proto.DeviceOpsServ.ConnectToDefaultWifiRequest;
+import com.google.wireless.qa.mobileharness.lab.proto.DeviceOpsServ.ConnectToDefaultWifiResponse;
 import com.google.wireless.qa.mobileharness.lab.proto.DeviceOpsServ.GetDeviceDebugInfoRequest;
 import com.google.wireless.qa.mobileharness.lab.proto.DeviceOpsServ.GetDeviceDebugInfoResponse;
 import com.google.wireless.qa.mobileharness.lab.proto.DeviceOpsServ.GetDeviceLogRequest;
@@ -32,6 +34,7 @@
 import com.google.wireless.qa.mobileharness.lab.proto.DeviceOpsServ.TakeScreenshotResponse;
 import com.google.wireless.qa.mobileharness.lab.proto.DeviceOpsServiceGrpc;
 import io.grpc.Channel;
+import javax.annotation.Nullable;
 
 /** gRPC stub of {@code DeviceOpsService}. */
 public class DeviceOpsGrpcStub implements DeviceOpsStub {
@@ -141,6 +144,39 @@ public ListenableFuture<RunTroubleshootScriptResponse> runTroubleshootScriptAsyn
     return futureStub.runTroubleshootScript(request);
   }
 
+  @Override
+  public ConnectToDefaultWifiResponse connectToDefaultWifi(ConnectToDefaultWifiRequest request)
+      throws RpcExceptionWithErrorId {
+    return GrpcStubUtil.invoke(
+        stub::connectToDefaultWifi,
+        request,
+        InfraErrorId.LAB_RPC_DEVICE_OPS_CONNECT_TO_DEFAULT_WIFI_GRPC_ERROR,
+        "Failed to connect to default wifi");
+  }
+
+  @Override
+  public ConnectToDefaultWifiResponse connectToDefaultWifi(
+      ConnectToDefaultWifiRequest request, @Nullable String impersonationUser)
+      throws RpcExceptionWithErrorId {
+    return connectToDefaultWifi(request);
+  }
+
+  @Override
+  public ListenableFuture<ConnectToDefaultWifiResponse> connectToDefaultWifiAsync(
+      ConnectToDefaultWifiRequest request) {
+    return futureStub.connectToDefaultWifi(request);
+  }
+
+  @Override
+  public ListenableFuture<ConnectToDefaultWifiResponse> connectToDefaultWifiAsync(
+      ConnectToDefaultWifiRequest request, boolean useClientRpcAuthority) {
+    if (useClientRpcAuthority) {
+      throw new UnsupportedOperationException(
+          "useClientRpcAuthority is not supported in gRPC stub");
+    }
+    return connectToDefaultWifiAsync(request);
+  }
+
   @Override
   public void close() {
     // This stub is not responsible for managing lifecycle of the channel.
@@ -155,6 +191,8 @@ public static interface BlockingInterface {
     GetDeviceDebugInfoResponse getDeviceDebugInfo(GetDeviceDebugInfoRequest request);
 
     RunTroubleshootScriptResponse runTroubleshootScript(RunTroubleshootScriptRequest request);
+
+    ConnectToDefaultWifiResponse connectToDefaultWifi(ConnectToDefaultWifiRequest request);
   }
 
   /** Interface for {@link DeviceOpsServiceFutureStub} */
@@ -168,6 +206,9 @@ ListenableFuture<GetDeviceDebugInfoResponse> getDeviceDebugInfo(
 
     ListenableFuture<RunTroubleshootScriptResponse> runTroubleshootScript(
         RunTroubleshootScriptRequest request);
+
+    ListenableFuture<ConnectToDefaultWifiResponse> connectToDefaultWifi(
+        ConnectToDefaultWifiRequest request);
   }
 
   public static BlockingInterface newBlockingInterface(Channel channel) {

src/java/com/google/wireless/qa/mobileharness/lab/proto/device_ops_serv.proto

@@ -93,6 +93,15 @@ message RunTroubleshootScriptResponse {
   optional string stderr = 3;
 }
 
+message ConnectToDefaultWifiRequest {
+  optional string device_id = 1;
+}
+
+message ConnectToDefaultWifiResponse {
+  optional bool success = 1;
+  optional string error_message = 2;
+}
+
 // Lab service for FE to take the screenshot.
 service DeviceOpsService {
   // Sends the screenshot file info to the lab.
@@ -109,4 +118,8 @@ service DeviceOpsService {
   // Runs a troubleshooting script on the lab host.
   rpc RunTroubleshootScript(RunTroubleshootScriptRequest)
       returns (RunTroubleshootScriptResponse) {}
+
+  // Connects the device to the default lab wifi network.
+  rpc ConnectToDefaultWifi(ConnectToDefaultWifiRequest)
+      returns (ConnectToDefaultWifiResponse) {}
 }

src/java/com/google/wireless/qa/mobileharness/shared/api/decorator/AndroidSetWifiDecorator.java

@@ -17,20 +17,24 @@
 package com.google.wireless.qa.mobileharness.shared.api.decorator;
 
 import com.google.common.annotations.VisibleForTesting;
-import com.google.common.base.Ascii;
 import com.google.common.base.Strings;
 import com.google.common.flogger.FluentLogger;
+import com.google.devtools.common.metrics.stability.rpc.RpcExceptionWithErrorId;
 import com.google.devtools.mobileharness.api.model.error.AndroidErrorId;
 import com.google.devtools.mobileharness.api.model.error.MobileHarnessException;
+import com.google.devtools.mobileharness.infra.client.api.util.stub.GrpcStubManager;
+import com.google.devtools.mobileharness.infra.lab.rpc.stub.DeviceOpsStub;
 import com.google.devtools.mobileharness.platform.android.lightning.networkconnector.NetworkConnector;
 import com.google.devtools.mobileharness.platform.android.lightning.networkconnector.WifiConnectArgs;
+import com.google.devtools.mobileharness.shared.util.flags.Flags;
+import com.google.wireless.qa.mobileharness.lab.proto.DeviceOpsServ.ConnectToDefaultWifiRequest;
+import com.google.wireless.qa.mobileharness.lab.proto.DeviceOpsServ.ConnectToDefaultWifiResponse;
 import com.google.wireless.qa.mobileharness.shared.api.annotation.DecoratorAnnotation;
 import com.google.wireless.qa.mobileharness.shared.api.decorator.base.LifecycleDecorator.SetupContext;
 import com.google.wireless.qa.mobileharness.shared.api.decorator.base.LifecycleDecorator.SetupResult;
 import com.google.wireless.qa.mobileharness.shared.api.decorator.base.SetupOnlyDecorator;
 import com.google.wireless.qa.mobileharness.shared.api.device.Device;
 import com.google.wireless.qa.mobileharness.shared.api.driver.Driver;
-import com.google.wireless.qa.mobileharness.shared.constant.PropertyName.Test;
 import com.google.wireless.qa.mobileharness.shared.model.job.TestInfo;
 import com.google.wireless.qa.mobileharness.shared.model.job.in.spec.SpecConfigable;
 import com.google.wireless.qa.mobileharness.shared.proto.spec.decorator.AndroidSetWifiDecoratorSpec;
@@ -61,6 +65,11 @@ public AndroidSetWifiDecorator(Driver decoratedDriver, TestInfo testInfo) {
     this.networkConnector = networkConnector;
   }
 
+  @VisibleForTesting
+  DeviceOpsStub getDeviceOpsStub(String labIp, int rpcPort) {
+    return GrpcStubManager.getInstance().getDeviceOpsStub(labIp + ":" + rpcPort);
+  }
+
   @Override
   protected SetupResult setUp(SetupContext context)
       throws MobileHarnessException, InterruptedException {
@@ -74,39 +83,32 @@ protected SetupResult setUp(SetupContext context)
     int retryNum = spec.getWifiRetryNum();
     boolean wifiSsidOptional = spec.getWifiSsidOptional();
     if (spec.getUseDefaultSsid()) {
-      // Get the wifi config from the device property.
-      wifiSsid =
-          device.getProperty(
-              Ascii.toLowerCase(Test.AndroidSetWifiDecorator.DEFAULT_WIFI_SSID.name()));
-      wifiPsk =
-          device.getProperty(
-              Ascii.toLowerCase(Test.AndroidSetWifiDecorator.DEFAULT_WIFI_PSK.name()));
-      if (Strings.isNullOrEmpty(wifiSsid)) {
-        if (!wifiSsidOptional) {
-          // Failed to get default ssid for the device.
-          testInfo
-              .log()
-              .atInfo()
-              .alsoTo(logger)
-              .log(
-                  "Could not get default ssid for the device %s from device properties. Have you"
-                      + " set the default ssid in the device config? If you have reset the default"
-                      + " wifi recently, it may have not taken effective.",
-                  deviceId);
+      // Direct the lab server to connect to default wifi via gRPC.
+      String labIp = "localhost";
+      int rpcPort = Flags.rpcPort.getNonNull();
+      if (rpcPort > 0 && labIp != null && !labIp.isEmpty()) {
+        try {
+          DeviceOpsStub stub = getDeviceOpsStub(labIp, rpcPort);
+          ConnectToDefaultWifiResponse response =
+              stub.connectToDefaultWifi(
+                  ConnectToDefaultWifiRequest.newBuilder().setDeviceId(deviceId).build());
+          if (response.getSuccess()) {
+            return SetupResult.continueDecorated();
+          } else {
+            throw new MobileHarnessException(
+                AndroidErrorId.ANDROID_SET_WIFI_DECORATOR_WIFI_CONNECT_ERROR,
+                "Failed to connect to default Wi-Fi: Lab Server returned failure.");
+          }
+        } catch (RpcExceptionWithErrorId e) {
           throw new MobileHarnessException(
-              AndroidErrorId.ANDROID_SET_WIFI_DECORATOR_GET_DEFAULT_SSID_ERROR,
-              "Failed to get default SSID for the device " + deviceId);
-        } else {
-          testInfo
-              .log()
-              .atWarning()
-              .alsoTo(logger)
-              .log(
-                  "Could not get default ssid for the device %s from device properties. Skipping"
-                      + " wifi setup.",
-                  deviceId);
-          return SetupResult.continueDecorated();
+              AndroidErrorId.ANDROID_SET_WIFI_DECORATOR_WIFI_CONNECT_ERROR,
+              "Failed to connect to default lab WiFi via gRPC for device " + deviceId,
+              e);
         }
+      } else {
+        throw new MobileHarnessException(
+            AndroidErrorId.ANDROID_SET_WIFI_DECORATOR_WIFI_CONNECT_ERROR,
+            "Lab server IP or RPC port is unavailable for use_default_ssid on device " + deviceId);
       }
     } else {
       // Get the wifi config from the spec.

src/java/com/google/wireless/qa/mobileharness/shared/api/decorator/BUILD

@@ -677,15 +677,19 @@ java_library(
     name = "android_set_wifi_decorator",
     srcs = ["AndroidSetWifiDecorator.java"],
     deps = [
+        "//src/java/com/google/devtools/common/metrics/stability/rpc:exception",
         "//src/java/com/google/devtools/mobileharness/api/model/error",
+        "//src/java/com/google/devtools/mobileharness/infra/client/api/util/stub:grpc_stub_manager",
+        "//src/java/com/google/devtools/mobileharness/infra/lab/rpc/stub:device_ops",
         "//src/java/com/google/devtools/mobileharness/platform/android/lightning/networkconnector",
+        "//src/java/com/google/devtools/mobileharness/shared/util/flags",
         "//src/java/com/google/devtools/mobileharness/shared/util/logging:google_logger",
+        "//src/java/com/google/wireless/qa/mobileharness/lab/proto:device_ops_serv_java_proto",
         "//src/java/com/google/wireless/qa/mobileharness/shared/api/annotation",
         "//src/java/com/google/wireless/qa/mobileharness/shared/api/decorator/base:lifecycle_decorator",
         "//src/java/com/google/wireless/qa/mobileharness/shared/api/decorator/base:setup_only_decorator",
         "//src/java/com/google/wireless/qa/mobileharness/shared/api/device",
         "//src/java/com/google/wireless/qa/mobileharness/shared/api/driver",
-        "//src/java/com/google/wireless/qa/mobileharness/shared/constant:property",
         "//src/java/com/google/wireless/qa/mobileharness/shared/model/job",
         "//src/java/com/google/wireless/qa/mobileharness/shared/model/job/in/spec",
         "//src/java/com/google/wireless/qa/mobileharness/shared/proto/spec:android_set_wifi_decorator_spec_java_proto",

src/javatests/com/google/devtools/mobileharness/shared/size/BinarySizeTest.java

@@ -59,9 +59,9 @@ public class BinarySizeTest {
           "ats_olc_server",
           40_050_000L,
           "ats_olc_server_local_mode",
-          45_250_000L,
+          46_350_000L,
           "lab_server",
-          46_550_000L,
+          46_750_000L,
           "ats_console",
           22_550_000L,
           "persistent_cache_manager",