1 /*
2  * Copyright (C) 2021 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.cts.deviceowner;
18 
19 import static com.google.common.truth.Truth.assertWithMessage;
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 androidx.localbroadcastmanager.content.LocalBroadcastManager;
29 
30 import java.util.ArrayList;
31 import java.util.Arrays;
32 import java.util.Collections;
33 import java.util.List;
34 import java.util.concurrent.Callable;
35 import java.util.concurrent.CountDownLatch;
36 import java.util.concurrent.TimeUnit;
37 
38 /**
39  * Helper class used to wait for user-related intents broadcast by {@link BasicAdminReceiver}.
40  *
41  */
42 final class UserActionCallback {
43 
44     private static final String TAG = UserActionCallback.class.getSimpleName();
45 
46     private static final int BROADCAST_TIMEOUT = 300_000;
47 
48     private final int mExpectedSize;
49     private final List<String> mExpectedActions;
50     private final List<String> mPendingActions;
51 
52     private final List<UserHandle> mReceivedUsers = new ArrayList<>();
53     private final CountDownLatch mLatch;
54     private final Context mContext;
55 
56     private final BroadcastReceiver mReceiver = new BroadcastReceiver() {
57         @Override
58         public void onReceive(Context context, Intent intent) {
59             String action = intent.getAction();
60             if (intent.hasExtra(BasicAdminReceiver.EXTRA_USER_HANDLE)) {
61                 UserHandle userHandle = intent
62                         .getParcelableExtra(BasicAdminReceiver.EXTRA_USER_HANDLE);
63                 Log.d(TAG, "broadcast receiver received " + action + " with user " + userHandle);
64                 mReceivedUsers.add(userHandle);
65             } else {
66                 Log.e(TAG, "broadcast receiver received " + intent.getAction()
67                         + " WITHOUT " + BasicAdminReceiver.EXTRA_USER_HANDLE + " extra");
68             }
69             boolean removed = mPendingActions.remove(action);
70             if (!removed) {
71                 Log.e(TAG, "Unexpected action " + action + "; what's left is " + mPendingActions);
72                 return;
73             }
74             Log.d(TAG, "Counting down latch (id " + System.identityHashCode(mLatch)
75                     + ", current count " + mLatch.getCount() + ") on thread "
76                     + Thread.currentThread());
77             mLatch.countDown();
78         }
79     };
80 
UserActionCallback(Context context, String... actions)81     private UserActionCallback(Context context, String... actions) {
82         Log.d(TAG, "Constructed UserActionCallback for " + Arrays.toString(actions) + " on user "
83                 + context.getUserId());
84         mContext = context;
85         mExpectedSize = actions.length;
86         mExpectedActions = new ArrayList<>(mExpectedSize);
87         mPendingActions = new ArrayList<>(mExpectedSize);
88         for (String action : actions) {
89             mExpectedActions.add(action);
90             mPendingActions.add(action);
91         }
92         mLatch = new CountDownLatch(mExpectedSize);
93     }
94 
95     /**
96      * Creates a new {@link UserActionCallback} and registers it to receive user broadcasts in the
97      * given context.
98      *
99      * @param context context to register for.
100      * @param actions expected actions.
101      *
102      * @return a new {@link UserActionCallback}.
103      */
getCallbackForBroadcastActions(Context context, String... actions)104     public static UserActionCallback getCallbackForBroadcastActions(Context context,
105             String... actions) {
106         UserActionCallback callback = new UserActionCallback(context, actions);
107 
108         IntentFilter filter = new IntentFilter();
109         for (String action : actions) {
110             filter.addAction(action);
111         }
112 
113         LocalBroadcastManager.getInstance(context).registerReceiver(callback.mReceiver, filter);
114 
115         return callback;
116     }
117 
118     /**
119      * Runs the given operation, blocking until the broadcasts are received and automatically
120      * unregistering itself at the end.
121      *
122      * @param runnable operation to run.
123      *
124      * @return operation result.
125      */
callAndUnregisterSelf(Callable<V> callable)126     public <V> V callAndUnregisterSelf(Callable<V> callable)
127             throws Exception {
128         try {
129             return callable.call();
130         } finally {
131             unregisterSelf();
132         }
133     }
134 
135     /**
136      * Gets the list of {@link UserHandle} associated with the broadcasts received so far.
137      */
getUsersOnReceivedBroadcasts()138     public List<UserHandle> getUsersOnReceivedBroadcasts() {
139         return Collections.unmodifiableList(new ArrayList<>(mReceivedUsers));
140     }
141 
142     /**
143      * Runs the given operation, blocking until the broadcasts are received and automatically
144      * unregistering itself at the end.
145      *
146      * @param runnable operation to run.
147      *
148      * @return operation result.
149      */
runAndUnregisterSelf(ThrowingRunnable runnable)150     public void runAndUnregisterSelf(ThrowingRunnable runnable) throws Exception {
151         try {
152             runnable.run();
153             waitForBroadcasts();
154         } finally {
155             unregisterSelf();
156         }
157     }
158 
159     /**
160      * Unregister itself as a {@link BroadcastReceiver} for user events.
161      */
unregisterSelf()162     public void unregisterSelf() {
163         LocalBroadcastManager.getInstance(mContext).unregisterReceiver(mReceiver);
164     }
165 
166     /**
167      * Custom {@link Runnable} that throws an {@link Exception}.
168      */
169     interface ThrowingRunnable {
run()170         void run() throws Exception;
171     }
172 
waitForBroadcasts()173     private void waitForBroadcasts() throws Exception {
174         Log.d(TAG, "Waiting up to " + BROADCAST_TIMEOUT + " to receive " + mExpectedSize
175                 + " broadcasts");
176         boolean received = mLatch.await(BROADCAST_TIMEOUT, TimeUnit.MILLISECONDS);
177         try {
178             assertWithMessage("%s messages received in %s ms. Expected actions=%s, "
179                 + "pending=%s", mExpectedSize, BROADCAST_TIMEOUT, mExpectedActions,
180                 mPendingActions).that(received).isTrue();
181         } catch (Exception | Error e) {
182             Log.e(TAG, "waitForBroadcasts() failed: " + e);
183             throw e;
184         }
185         Log.d(TAG, "All broadcasts accounted for. Thank you and come again!");
186     }
187 }
188