1 /*
2  * Copyright 2020 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 androidx.sqlite.inspection;
18 
19 import androidx.inspection.ArtTooling.EntryHook;
20 import androidx.inspection.ArtTooling.ExitHook;
21 import androidx.inspection.InspectorEnvironment;
22 
23 import org.jspecify.annotations.NonNull;
24 import org.jspecify.annotations.Nullable;
25 
26 import java.util.ArrayDeque;
27 import java.util.Deque;
28 import java.util.List;
29 
30 /**
31  * The class allows for observing method's EntryHook parameters in ExitHook.
32  * <p>
33  * It works by registering both (entry and exit) hooks and keeping its own method frame stack.
34  * On exit, it calls {@link OnExitCallback} provided by the user.
35  * <p>
36  * TODO: handle cases when frames could be dropped (e.g. because of an Exception) causing internal
37  * state to be corrupted.
38  * <p>
39  * Thread safe.
40  */
41 final class EntryExitMatchingHookRegistry {
42     private final InspectorEnvironment mEnvironment;
43     private final ThreadLocal<Deque<Frame>> mFrameStack;
44 
EntryExitMatchingHookRegistry(InspectorEnvironment environment)45     EntryExitMatchingHookRegistry(InspectorEnvironment environment) {
46         mEnvironment = environment;
47         mFrameStack = new ThreadLocal<Deque<Frame>>() {
48             @Override
49             protected @NonNull Deque<Frame> initialValue() {
50                 return new ArrayDeque<>();
51             }
52         };
53     }
54 
registerHook(Class<?> originClass, final String originMethod, final OnExitCallback onExitCallback)55     void registerHook(Class<?> originClass, final String originMethod,
56             final OnExitCallback onExitCallback) {
57         mEnvironment.artTooling().registerEntryHook(originClass, originMethod,
58                 new EntryHook() {
59                     @Override
60                     public void onEntry(@Nullable Object thisObject,
61                             @NonNull List<Object> args) {
62                         getFrameStack().addLast(new Frame(originMethod, thisObject, args, null));
63                     }
64                 });
65 
66         mEnvironment.artTooling().registerExitHook(originClass, originMethod,
67                 new ExitHook<Object>() {
68                     @Override
69                     public Object onExit(Object result) {
70                         Frame entryFrame = getFrameStack().pollLast();
71                         if (entryFrame == null || !originMethod.equals(entryFrame.mMethod)) {
72                             // TODO: make more specific and handle
73                             throw new IllegalStateException();
74                         }
75 
76                         onExitCallback.onExit(new Frame(entryFrame.mMethod, entryFrame.mThisObject,
77                                 entryFrame.mArgs, result));
78                         return result;
79                     }
80                 });
81     }
82 
getFrameStack()83     private @NonNull Deque<Frame> getFrameStack() {
84         /* It won't be null because of overridden {@link ThreadLocal#initialValue} */
85         //noinspection ConstantConditions
86         return mFrameStack.get();
87     }
88 
89     static final class Frame {
90         final String mMethod;
91         final Object mThisObject;
92         final List<Object> mArgs;
93         final Object mResult;
94 
Frame(String method, Object thisObject, List<Object> args, Object result)95         private Frame(String method, Object thisObject, List<Object> args, Object result) {
96             mMethod = method;
97             mThisObject = thisObject;
98             mArgs = args;
99             mResult = result;
100         }
101     }
102 
103     interface OnExitCallback {
onExit(Frame exitFrame)104         void onExit(Frame exitFrame);
105     }
106 }
107