deploy/deployer/src/main/java/com/android/tools/deployer/JdiBasedClassRedefiner.java

/*
 * Copyright (C) 2018 The Android Open Source Project
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package com.android.tools.deployer;

import com.android.tools.deploy.proto.Deploy;
import com.google.common.collect.Lists;
import com.sun.jdi.ArrayReference;
import com.sun.jdi.ClassType;
import com.sun.jdi.Field;
import com.sun.jdi.Method;
import com.sun.jdi.ObjectReference;
import com.sun.jdi.ReferenceType;
import com.sun.jdi.ThreadReference;
import com.sun.jdi.VirtualMachine;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.concurrent.TimeUnit;

/**
 * An implementation of the {@link ClassRedefiner} that invoke the Android Virtual Machine's class
 * redefinition API by using JDWP's RedefineClasses command.
 */
public class JdiBasedClassRedefiner implements ClassRedefiner {
    private static final String AGENT_LOC_FORMAT =
            "/data/data/%s/code_cache/.studio/%s.so=irsocket";

    // Field has been present since Lollipop, so it's definitely present on Oreo.
    private static final String SUPPORTED_ABIS_FIELD = "SUPPORTED_64_BIT_ABIS";

    private final VirtualMachine vm;
    private final RedefineClassSupportState redefineSupportState;

    public JdiBasedClassRedefiner(
            VirtualMachine vm, RedefineClassSupportState redefineSupportState) {
        this.vm = vm;
        this.redefineSupportState = redefineSupportState;
    }

    @Override
    public Deploy.SwapResponse redefine(Deploy.SwapRequest request) throws DeployerException {
        Map<ReferenceType, byte[]> redefinitionRequest = new HashMap<>();

        for (Deploy.ClassDef redefinition : request.getModifiedClassesList()) {
            List<ReferenceType> classes = vm.classesByName(redefinition.getName());
            for (ReferenceType classRef : classes) {
                redefinitionRequest.put(classRef, redefinition.getDex().toByteArray());
            }
        }

        Deploy.SwapResponse.Builder response = Deploy.SwapResponse.newBuilder();
        switch (redefineSupportState.support) {
            case FULL:
                try {
                    vm.redefineClasses(redefinitionRequest);
                } catch (Throwable t) {
                    throw DeployerException.jdwpRedefineClassesException(t);
                }
                break;
            case MAIN_THREAD_RUNNING:
                // Nothing to do. The installer + agent will perform the swap for us.
                break;
            case NEEDS_AGENT_SERVER:
                List<ReferenceType> buildList = vm.classesByName("android.os.Build");
                ClassType build = (ClassType) buildList.get(0);

                Field abiField = build.fieldByName(SUPPORTED_ABIS_FIELD);
                if (abiField == null) {
                    throw DeployerException.abisFieldNotFound();
                }

                ArrayReference abis = (ArrayReference) build.getValue(abiField);
                String agentLoc =
                        String.format(
                                AGENT_LOC_FORMAT,
                                request.getPackageName(),
                                getAgentName(request.getArch(), abis.length() > 0));

                List<ReferenceType> debugList = vm.classesByName("dalvik.system.VMDebug");
                ClassType debug = (ClassType) debugList.get(0);
                List<ThreadReference> allThreads = vm.allThreads();
                for (ThreadReference thread : allThreads) {
                    if (thread.name().equals(redefineSupportState.targetThread)) {
                        Method attachAgentMethod =
                                debug.concreteMethodByName("attachAgent", "(Ljava/lang/String;)V");
                        if (attachAgentMethod == null) {
                            // This should not happen.
                            throw DeployerException.attachAgentNotFound();
                        }

                        try {
                            // The only time we are allowed to invoke a method like this is on a
                            // thread that is suspended by an event generated by itself. This is
                            // basically the same mechanism used for the debugger UI to print
                            // variable / expressions. The most likely case of this is when a
                            // thread is suspended because it arrived on a breakpoint. This will
                            // fail (with an exception that is caught below) should for cases such
                            // as a "suspend all" command.
                            debug.invokeMethod(
                                    thread,
                                    attachAgentMethod,
                                    Lists.newArrayList(vm.mirrorOf(agentLoc)),
                                    ObjectReference.INVOKE_SINGLE_THREADED);
                        } catch (Exception e) {
                            try {
                                // TODO: We can split up the agent extraction + server so we don't
                                //  need to do this. We can tell the installer only swap and wait
                                //  for us and don't install yet, we would no longer have to have
                                //  do a sleep and wait here.

                                // Re-try once more.
                                Thread.sleep(1000);
                                debug.invokeMethod(
                                        thread,
                                        attachAgentMethod,
                                        Lists.newArrayList(vm.mirrorOf(agentLoc)),
                                        ObjectReference.INVOKE_SINGLE_THREADED);
                            } catch (Exception e1) {
                                throw DeployerException.attachAgentException(e1);
                            }
                        }
                    }
                }
                break;
            default:
                // This should not happen.
                throw DeployerException.jdiInvalidState();
        }

        response.setStatus(Deploy.SwapResponse.Status.OK);
        return response.build();
    }

    @Override
    public Deploy.SwapResponse redefine(Deploy.OverlaySwapRequest request)
            throws DeployerException {
        Map<ReferenceType, byte[]> redefinitionRequest = new HashMap<>();

        for (Deploy.ClassDef redefinition : request.getModifiedClassesList()) {
            List<ReferenceType> classes = vm.classesByName(redefinition.getName());
            for (ReferenceType classRef : classes) {
                redefinitionRequest.put(classRef, redefinition.getDex().toByteArray());
            }
        }
        try {
            vm.redefineClasses(redefinitionRequest);
        } catch (Throwable t) {
            throw DeployerException.jdwpRedefineClassesException(t);
        }
        Deploy.SwapResponse.Builder response = Deploy.SwapResponse.newBuilder();
        response.setStatus(Deploy.SwapResponse.Status.OK);
        return response.build();
    }

    @Override
    public RedefineClassSupportState canRedefineClass() {
        return redefineSupportState;
    }

    private static String getAgentName(Deploy.Arch appArch, boolean deviceIs64Bit) {
        if (appArch == Deploy.Arch.ARCH_32_BIT && deviceIs64Bit) {
            return "agent-alt";
        }
        return "agent";
    }
}