3 files changed, 78 insertions, 41 deletions

diff --git a/core/java/android/net/ConnectivityDiagnosticsManager.java b/core/java/android/net/ConnectivityDiagnosticsManager.java
index 704f31d7f773..523449497345 100644
--- a/core/java/android/net/ConnectivityDiagnosticsManager.java
+++ b/core/java/android/net/ConnectivityDiagnosticsManager.java
@@ -623,32 +623,41 @@ public class ConnectivityDiagnosticsManager {
         /** @hide */
         @VisibleForTesting
         public void onConnectivityReportAvailable(@NonNull ConnectivityReport report) {
-            Binder.withCleanCallingIdentity(() -> {
+            final long token = Binder.clearCallingIdentity();
+            try {
                 mExecutor.execute(() -> {
                     mCb.onConnectivityReportAvailable(report);
                 });
-            });
+            } finally {
+                Binder.restoreCallingIdentity(token);
+            }
         }
 
         /** @hide */
         @VisibleForTesting
         public void onDataStallSuspected(@NonNull DataStallReport report) {
-            Binder.withCleanCallingIdentity(() -> {
+            final long token = Binder.clearCallingIdentity();
+            try {
                 mExecutor.execute(() -> {
                     mCb.onDataStallSuspected(report);
                 });
-            });
+            } finally {
+                Binder.restoreCallingIdentity(token);
+            }
         }
 
         /** @hide */
         @VisibleForTesting
         public void onNetworkConnectivityReported(
                 @NonNull Network network, boolean hasConnectivity) {
-            Binder.withCleanCallingIdentity(() -> {
+            final long token = Binder.clearCallingIdentity();
+            try {
                 mExecutor.execute(() -> {
                     mCb.onNetworkConnectivityReported(network, hasConnectivity);
                 });
-            });
+            } finally {
+                Binder.restoreCallingIdentity(token);
+            }
         }
     }
 
diff --git a/core/java/android/net/ConnectivityManager.java b/core/java/android/net/ConnectivityManager.java
index 6cef73dcd089..8ea540cd38cd 100644
--- a/core/java/android/net/ConnectivityManager.java
+++ b/core/java/android/net/ConnectivityManager.java
@@ -1833,30 +1833,42 @@ public class ConnectivityManager {
             mCallback = new ISocketKeepaliveCallback.Stub() {
                 @Override
                 public void onStarted(int slot) {
-                    Binder.withCleanCallingIdentity(() ->
-                            mExecutor.execute(() -> {
-                                mSlot = slot;
-                                callback.onStarted();
-                            }));
+                    final long token = Binder.clearCallingIdentity();
+                    try {
+                        mExecutor.execute(() -> {
+                            mSlot = slot;
+                            callback.onStarted();
+                        });
+                    } finally {
+                        Binder.restoreCallingIdentity(token);
+                    }
                 }
 
                 @Override
                 public void onStopped() {
-                    Binder.withCleanCallingIdentity(() ->
-                            mExecutor.execute(() -> {
-                                mSlot = null;
-                                callback.onStopped();
-                            }));
+                    final long token = Binder.clearCallingIdentity();
+                    try {
+                        mExecutor.execute(() -> {
+                            mSlot = null;
+                            callback.onStopped();
+                        });
+                    } finally {
+                        Binder.restoreCallingIdentity(token);
+                    }
                     mExecutor.shutdown();
                 }
 
                 @Override
                 public void onError(int error) {
-                    Binder.withCleanCallingIdentity(() ->
-                            mExecutor.execute(() -> {
-                                mSlot = null;
-                                callback.onError(error);
-                            }));
+                    final long token = Binder.clearCallingIdentity();
+                    try {
+                        mExecutor.execute(() -> {
+                            mSlot = null;
+                            callback.onError(error);
+                        });
+                    } finally {
+                        Binder.restoreCallingIdentity(token);
+                    }
                     mExecutor.shutdown();
                 }
 
diff --git a/core/java/android/net/SocketKeepalive.java b/core/java/android/net/SocketKeepalive.java
index a7dce18a4ff9..d007a9520cb5 100644
--- a/core/java/android/net/SocketKeepalive.java
+++ b/core/java/android/net/SocketKeepalive.java
@@ -187,38 +187,54 @@ public abstract class SocketKeepalive implements AutoCloseable {
         mCallback = new ISocketKeepaliveCallback.Stub() {
             @Override
             public void onStarted(int slot) {
-                Binder.withCleanCallingIdentity(() ->
-                        mExecutor.execute(() -> {
-                            mSlot = slot;
-                            callback.onStarted();
-                        }));
+                final long token = Binder.clearCallingIdentity();
+                try {
+                    mExecutor.execute(() -> {
+                        mSlot = slot;
+                        callback.onStarted();
+                    });
+                } finally {
+                    Binder.restoreCallingIdentity(token);
+                }
             }
 
             @Override
             public void onStopped() {
-                Binder.withCleanCallingIdentity(() ->
-                        executor.execute(() -> {
-                            mSlot = null;
-                            callback.onStopped();
-                        }));
+                final long token = Binder.clearCallingIdentity();
+                try {
+                    executor.execute(() -> {
+                        mSlot = null;
+                        callback.onStopped();
+                    });
+                } finally {
+                    Binder.restoreCallingIdentity(token);
+                }
             }
 
             @Override
             public void onError(int error) {
-                Binder.withCleanCallingIdentity(() ->
-                        executor.execute(() -> {
-                            mSlot = null;
-                            callback.onError(error);
-                        }));
+                final long token = Binder.clearCallingIdentity();
+                try {
+                    executor.execute(() -> {
+                        mSlot = null;
+                        callback.onError(error);
+                    });
+                } finally {
+                    Binder.restoreCallingIdentity(token);
+                }
             }
 
             @Override
             public void onDataReceived() {
-                Binder.withCleanCallingIdentity(() ->
-                        executor.execute(() -> {
-                            mSlot = null;
-                            callback.onDataReceived();
-                        }));
+                final long token = Binder.clearCallingIdentity();
+                try {
+                    executor.execute(() -> {
+                        mSlot = null;
+                        callback.onDataReceived();
+                    });
+                } finally {
+                    Binder.restoreCallingIdentity(token);
+                }
             }
         };
     }