From ca70a561e53a4569077952f5b0b368d6401558cd Mon Sep 17 00:00:00 2001
From: Dmitry Muhomor <muhomor.dmitry@gmail.com>
Date: Tue, 29 Oct 2024 10:31:07 +0200
Subject: [PATCH] fixup! add USB-C port and pogo pins security setting

Toggle USB-C port off and on when changing setting value from "Charging-only" to "Charging-only
when locked", "Charging-only when locked except BFU" or "On" to immediately enable the data
connection.
---
 .../UsbPortSecurityPrefController.java        | 51 +++++++++++++------
 1 file changed, 36 insertions(+), 15 deletions(-)

diff --git a/src/com/android/settings/security/UsbPortSecurityPrefController.java b/src/com/android/settings/security/UsbPortSecurityPrefController.java
index 5e902e300b5..ddea6316299 100644
--- a/src/com/android/settings/security/UsbPortSecurityPrefController.java
+++ b/src/com/android/settings/security/UsbPortSecurityPrefController.java
@@ -4,17 +4,20 @@
 import android.content.res.Resources;
 import android.ext.settings.UsbPortSecurity;
 import android.hardware.usb.UsbManager;
-import android.hardware.usb.UsbPort;
 import android.hardware.usb.ext.PortSecurityState;
 import android.os.Bundle;
 import android.os.ResultReceiver;
 
 import androidx.preference.Preference;
 
+import com.android.internal.infra.AndroidFuture;
 import com.android.settings.R;
 import com.android.settings.ext.IntSettingPrefController;
 
-import java.util.List;
+import java.util.Objects;
+import java.util.concurrent.TimeUnit;
+
+import static java.util.Objects.requireNonNull;
 
 public class UsbPortSecurityPrefController extends IntSettingPrefController {
     private final boolean appliesToPogoPins;
@@ -77,25 +80,43 @@ protected void getEntries(Entries entries) {
                 UsbPortSecurity.MODE_ENABLED);
     }
 
-    private void setState(@android.hardware.usb.ext.PortSecurityState int state) {
-        var um = mContext.getSystemService(UsbManager.class);
-        List<UsbPort> ports = um.getPorts();
-
-        for (UsbPort port : ports) {
-            um.setPortSecurityState(port, state, new ResultReceiver(null) {
-                @Override
-                protected void onReceiveResult(int resultCode, Bundle resultData) {
-                    if (resultCode != android.hardware.usb.ext.IUsbExt.NO_ERROR) {
-                        throw new IllegalStateException("setPortSecurityState failed, " +
-                                "resultCode: " + resultCode + ", port: " + port);
+    private static void setSecurityStateForAllPortsSync(UsbManager usbManager, int state) {
+        var future = new AndroidFuture<>();
+        usbManager.setSecurityStateForAllPorts(state, new ResultReceiver(null) {
+            @Override
+            protected void onReceiveResult(int resultCode, Bundle resultData) {
+                if (resultCode != android.hardware.usb.ext.IUsbExt.NO_ERROR) {
+                    String msg = "setPortSecurityState failed, " +
+                            "resultCode: " + resultCode;
+                    if (resultData != null) {
+                        msg += ", resultData: " + resultData.toStringDeep();
                     }
+                    throw new RuntimeException(msg);
                 }
-            });
+                future.complete(null);
+            }
+        });
+        try {
+            future.get(3, TimeUnit.SECONDS);
+        } catch (Exception e) {
+            throw new RuntimeException(e);
+        }
+    }
+
+    private void setState(int prevSetting, @android.hardware.usb.ext.PortSecurityState int state) {
+        var usbManager = requireNonNull(mContext.getSystemService(UsbManager.class));
+        if (prevSetting == UsbPortSecurity.MODE_CHARGING_ONLY && state >= UsbPortSecurity.MODE_CHARGING_ONLY_WHEN_LOCKED) {
+            // Turn USB ports off first to trigger reconnection of devices that were connected
+            // in charging-only state. Simply enabling the data path is not enough in some
+            // advanced scenarios, e.g. when port alt mode or port role switching are used.
+            setSecurityStateForAllPortsSync(usbManager, PortSecurityState.DISABLED);
         }
+        setSecurityStateForAllPortsSync(usbManager, state);
     }
 
     @Override
     protected boolean setValue(int val) {
+        int prevSetting = getCurrentValue();
         boolean res = super.setValue(val);
         if (!res) {
             return false;
@@ -109,7 +130,7 @@ protected boolean setValue(int val) {
             case UsbPortSecurity.MODE_ENABLED -> PortSecurityState.ENABLED;
             default -> throw new IllegalArgumentException(Integer.toString(val));
         };
-        setState(pss);
+        setState(prevSetting, pss);
         return true;
     }