1 /*
2  * Copyright (C) 2023 The Android Open Source Project
3  *
4  * Licensed under the Apache License, Version 2.0 (the "License");
5  * you may not use this file except in compliance with the License.
6  * You may obtain a copy of the License at
7  *
8  *      http://www.apache.org/licenses/LICENSE-2.0
9  *
10  * Unless required by applicable law or agreed to in writing, software
11  * distributed under the License is distributed on an "AS IS" BASIS,
12  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13  * See the License for the specific language governing permissions and
14  * limitations under the License.
15  */
16 
17 package com.android.server.pm;
18 
19 import static org.junit.Assert.fail;
20 
21 import android.content.BroadcastReceiver;
22 import android.content.Context;
23 import android.content.Intent;
24 import android.content.IntentFilter;
25 import android.os.UserHandle;
26 import android.util.Log;
27 
28 import java.io.Closeable;
29 import java.io.IOException;
30 import java.util.Map;
31 import java.util.concurrent.ConcurrentHashMap;
32 import java.util.concurrent.CountDownLatch;
33 import java.util.concurrent.TimeUnit;
34 
35 public class UserRemovalWaiter extends BroadcastReceiver implements Closeable {
36 
37     private final Context mContext;
38     private final String mTag;
39     private final long mTimeoutInSeconds;
40     private final Map<Integer, CountDownLatch> mMap = new ConcurrentHashMap<>();
41 
getLatch(final int userId)42     private CountDownLatch getLatch(final int userId) {
43         return mMap.computeIfAbsent(userId, absentKey -> new CountDownLatch(1));
44     }
45 
UserRemovalWaiter(Context context, String tag, int timeoutInSeconds)46     public UserRemovalWaiter(Context context, String tag, int timeoutInSeconds) {
47         mContext = context;
48         mTag = tag;
49         mTimeoutInSeconds = timeoutInSeconds;
50 
51         mContext.registerReceiver(this, new IntentFilter(Intent.ACTION_USER_REMOVED));
52     }
53 
54     @Override
close()55     public void close() throws IOException {
56         mContext.unregisterReceiver(this);
57     }
58 
59     @Override
onReceive(Context context, Intent intent)60     public void onReceive(Context context, Intent intent) {
61         if (Intent.ACTION_USER_REMOVED.equals(intent.getAction())) {
62             int userId = intent.getIntExtra(Intent.EXTRA_USER_HANDLE, UserHandle.USER_NULL);
63             Log.i(mTag, "ACTION_USER_REMOVED received for user " + userId);
64             getLatch(userId).countDown();
65         }
66     }
67 
68     /**
69      * Waits for the removal of the given user, or fails if it times out.
70      */
waitFor(int userId)71     public void waitFor(int userId) {
72         Log.i(mTag, "Waiting for user " + userId + " to be removed");
73 
74         long startTime = System.currentTimeMillis();
75         try {
76             if (getLatch(userId).await(mTimeoutInSeconds, TimeUnit.SECONDS)) {
77                 Log.i(mTag, "User " + userId + " is removed in "
78                         + (System.currentTimeMillis() - startTime) + " ms");
79             } else {
80                 fail("Timeout waiting for user removal. userId = " + userId);
81             }
82         } catch (InterruptedException e) {
83             throw new AssertionError("Thread interrupted unexpectedly.", e);
84         }
85     }
86 }
87