content: /*
   * Licensed to the Apache Software Foundation (ASF) under one
   * or more contributor license agreements.  See the NOTICE file
   * distributed with this work for additional information
   * regarding copyright ownership.  The ASF licenses this file
   * to you 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 org.apache.cassandra.simulator.asm;
  
  import java.io.IOException;
  import java.io.InputStream;
  import java.lang.instrument.ClassDefinition;
  import java.lang.instrument.ClassFileTransformer;
  import java.lang.instrument.IllegalClassFormatException;
  import java.lang.instrument.Instrumentation;
  import java.lang.instrument.UnmodifiableClassException;
  import java.security.ProtectionDomain;
  import java.util.ArrayList;
  import java.util.Arrays;
  import java.util.EnumSet;
  import java.util.List;
  import java.util.Objects;
  import java.util.function.BiFunction;
  import java.util.regex.Pattern;
  
  import org.objectweb.asm.ClassReader;
  import org.objectweb.asm.ClassVisitor;
  import org.objectweb.asm.ClassWriter;
  import org.objectweb.asm.FieldVisitor;
  import org.objectweb.asm.Label;
  import org.objectweb.asm.MethodVisitor;
  
  import static org.apache.cassandra.simulator.asm.Flag.DETERMINISTIC;
  import static org.apache.cassandra.simulator.asm.Flag.LOCK_SUPPORT;
  import static org.apache.cassandra.simulator.asm.Flag.NO_PROXY_METHODS;
  import static org.apache.cassandra.simulator.asm.Flag.SYSTEM_CLOCK;
  import static org.apache.cassandra.simulator.asm.InterceptClasses.BYTECODE_VERSION;
  import static org.objectweb.asm.Opcodes.ALOAD;
  import static org.objectweb.asm.Opcodes.GETFIELD;
  import static org.objectweb.asm.Opcodes.GETSTATIC;
  import static org.objectweb.asm.Opcodes.INVOKESPECIAL;
  import static org.objectweb.asm.Opcodes.INVOKESTATIC;
  import static org.objectweb.asm.Opcodes.INVOKEVIRTUAL;
  import static org.objectweb.asm.Opcodes.IRETURN;
  import static org.objectweb.asm.Opcodes.RETURN;
  
  // checkstyle: suppress below 'blockSystemPropertyUsage'
  
  /**
   * A mechanism for weaving classes loaded by the bootstrap classloader that we cannot override.
   * The design supports weaving of the internals of these classes, and in future we may want to
   * weave LockSupport or the internals of other blocking concurrency primitives.
   *
   * Ultimately this wasn't necessary for the initial functionality, but we have maintained
   * the layout so that it will be easier to enable such functionality in future should it be needed.
   *
   * To this end, the asm package and simulator-asm.jar is as self-contained set of classes for performing
   * simulator byteweaving, and simulator-bootstrap.jar contains a self-contained class and interface for
   * replacing important system methods.
   */
  public class InterceptAgent
  {
      public static void premain(final String agentArgs, final Instrumentation instrumentation) throws UnmodifiableClassException, ClassNotFoundException, IOException
      {
          setup(agentArgs, instrumentation);
      }
  
      public void agentmain(final String agentArgs, final Instrumentation instrumentation) throws UnmodifiableClassException, ClassNotFoundException, IOException
      {
          setup(agentArgs, instrumentation);
      }
  
      private static void setup(final String agentArgs, final Instrumentation instrumentation) throws UnmodifiableClassException, ClassNotFoundException, IOException
      {
          instrumentation.addTransformer(new ClassFileTransformer()
          {
              @Override
              public byte[] transform(ClassLoader loader, String className, Class<?> classBeingRedefined, ProtectionDomain protectionDomain, byte[] bytecode) throws IllegalClassFormatException
              {
                  if (className == null)
                      return null;
  
                  if (className.equals("java/lang/Object"))
                      return transformObject(bytecode);
  
                  if (className.equals("java/lang/Class"))
                      return transformClass(bytecode);
  
                  if (className.equals("java/lang/Enum"))
                      return transformEnum(bytecode);
  
                  if (className.equals("java/util/Random"))
                      return transformRandom(bytecode);
  
                  if (className.equals("java/util/concurrent/ThreadLocalRandom"))
                      return transformThreadLocalRandom(bytecode);
  
                  if (className.startsWith("java/util/concurrent/ConcurrentHashMap"))
                      return InterceptAgent.transform(className, bytecode, DETERMINISTIC, NO_PROXY_METHODS);
  
                  if (className.startsWith("java/util/concurrent/locks"))
                  {
                      if (className.equals("java/util/concurrent/locks/AbstractQueuedSynchronizer"))
                          return InterceptAgent.transformAbstractQueuedSynchronizer(className, bytecode, SYSTEM_CLOCK, LOCK_SUPPORT, NO_PROXY_METHODS);
                      return InterceptAgent.transform(className, bytecode, SYSTEM_CLOCK, LOCK_SUPPORT, NO_PROXY_METHODS);
                  }
  
                  return null;
              }
          });
  
          Pattern reloadPattern = Pattern.compile("java\\.(lang\\.Enum|util\\.concurrent\\.(locks\\..*|ConcurrentHashMap)|util\\.(concurrent\\.ThreadLocal)?Random|lang\\.Object)");
          List<ClassDefinition> redefine = new ArrayList<>();
          for (Class<?> loadedClass : instrumentation.getAllLoadedClasses())
          {
              if (reloadPattern.matcher(loadedClass.getName()).matches())
                  redefine.add(new ClassDefinition(loadedClass, readDefinition(loadedClass)));
          }
          if (!redefine.isEmpty())
              instrumentation.redefineClasses(redefine.toArray(new ClassDefinition[0]));
      }
  
      private static byte[] readDefinition(Class<?> clazz) throws IOException
      {
          return readDefinition(clazz.getName().replaceAll("\\.", "/"));
      }
  
      private static byte[] readDefinition(String className) throws IOException
      {
          byte[] bytes = new byte[1024];
          try (InputStream in = ClassLoader.getSystemResourceAsStream(className + ".class"))
          {
              int count = 0;
              while (true)
              {
                  int add = in.read(bytes, count, bytes.length - count);
                  if (add < 0)
                      break;
                  if (add == 0)
                      bytes = Arrays.copyOf(bytes, bytes.length * 2);
                  count += add;
              }
              return Arrays.copyOf(bytes, count);
          }
      }
  
      /**
       * We don't want Object.toString() to invoke our overridden identityHashCode by virtue of invoking some overridden hashCode()
       * So we overwrite Object.toString() to replace calls to Object.hashCode() with direct calls to System.identityHashCode()
       */
      private static byte[] transformObject(byte[] bytes)
      {
          class ObjectVisitor extends ClassVisitor
          {
              public ObjectVisitor(int api, ClassVisitor classVisitor)
              {
                  super(api, classVisitor);
              }
  
              @Override
              public MethodVisitor visitMethod(int access, String name, String descriptor, String signature, String[] exceptions)
              {
                  if (descriptor.equals("()Ljava/lang/String;") && name.equals("toString"))
                      return Utils.deterministicToString(super.visitMethod(access, name, descriptor, signature, exceptions));
                  else
                      return super.visitMethod(access, name, descriptor, signature, exceptions);
              }
          }
          return transform(bytes, ObjectVisitor::new);
      }
  
      /**
       * We don't want Object.toString() to invoke our overridden identityHashCode by virtue of invoking some overridden hashCode()
       * So we overwrite Object.toString() to replace calls to Object.hashCode() with direct calls to System.identityHashCode()
       */
      private static byte[] transformClass(byte[] bytes)
      {
          class ClazzVisitor extends ClassVisitor
          {
              public ClazzVisitor(int api, ClassVisitor classVisitor)
              {
                  super(api, classVisitor);
              }
  
              @Override
              public void visitEnd()
              {
                  new StringHashcode(api).accept(this);
                  super.visitEnd();
              }
          }
          return transform(bytes, ClazzVisitor::new);
      }
  
      /**
       * We want Enum to have a deterministic hashCode() so we simply forward calls to ordinal()
       */
      private static byte[] transformEnum(byte[] bytes)
      {
          class EnumVisitor extends ClassVisitor
          {
              public EnumVisitor(int api, ClassVisitor classVisitor)
              {
                  super(api, classVisitor);
              }
  
              @Override
              public MethodVisitor visitMethod(int access, String name, String descriptor, String signature, String[] exceptions)
              {
                  if (descriptor.equals("()I") && name.equals("hashCode"))
                  {
                      MethodVisitor visitor = super.visitMethod(access, name, descriptor, signature, exceptions);
                      visitor.visitLabel(new Label());
                      visitor.visitIntInsn(ALOAD, 0);
                      visitor.visitFieldInsn(GETFIELD, "java/lang/Enum", "ordinal", "I");
                      visitor.visitInsn(IRETURN);
                      visitor.visitLabel(new Label());
                      visitor.visitMaxs(1, 1);
                      visitor.visitEnd();
  
                      return new MethodVisitor(BYTECODE_VERSION) {};
                  }
                  else
                  {
                      return super.visitMethod(access, name, descriptor, signature, exceptions);
                  }
              }
          }
          return transform(bytes, EnumVisitor::new);
      }
  
      /**
       * We want Random to be initialised deterministically, so we modify the default constructor to fetch
       * some deterministically generated seed to pass to its seed constructor
       */
      private static byte[] transformRandom(byte[] bytes)
      {
          class RandomVisitor extends ClassVisitor
          {
              public RandomVisitor(int api, ClassVisitor classVisitor)
              {
                  super(api, classVisitor);
              }
  
              @Override
              public MethodVisitor visitMethod(int access, String name, String descriptor, String signature, String[] exceptions)
              {
                  if (descriptor.equals("()V") && name.equals("<init>"))
                  {
                      MethodVisitor visitor = super.visitMethod(access, name, descriptor, signature, exceptions);
                      visitor.visitLabel(new Label());
                      visitor.visitIntInsn(ALOAD, 0);
                      visitor.visitMethodInsn(INVOKESTATIC, "org/apache/cassandra/simulator/systems/InterceptorOfSystemMethods$Global", "randomSeed", "()J", false);
                      visitor.visitMethodInsn(INVOKESPECIAL, "java/util/Random", "<init>", "(J)V", false);
                      visitor.visitInsn(RETURN);
                      visitor.visitLabel(new Label());
                      visitor.visitMaxs(3, 1);
                      visitor.visitEnd();
  
                      return new MethodVisitor(BYTECODE_VERSION) {};
                  }
                  else
                  {
                      return super.visitMethod(access, name, descriptor, signature, exceptions);
                  }
              }
          }
          return transform(bytes, RandomVisitor::new);
      }
  
      /**
       * We require ThreadLocalRandom to be deterministic, so we modify its initialisation method to invoke a
       * global deterministic random value generator
       */
      private static byte[] transformThreadLocalRandom(byte[] bytes)
      {
          class ThreadLocalRandomVisitor extends ClassVisitor
          {
              // CassandraRelevantProperties is not available to us here
              final boolean determinismCheck = System.getProperty("cassandra.test.simulator.determinismcheck", "none").matches("relaxed|strict");
  
              public ThreadLocalRandomVisitor(int api, ClassVisitor classVisitor)
              {
                  super(api, classVisitor);
              }
  
              String unsafeDescriptor;
              String unsafeFieldName;
  
              @Override
              public FieldVisitor visitField(int access, String name, String descriptor, String signature, Object value)
              {
                  if (descriptor.equals("Lsun/misc/Unsafe;") || descriptor.equals("Ljdk/internal/misc/Unsafe;"))
                  {
                      unsafeFieldName = name;
                      unsafeDescriptor = descriptor;
                  }
                  return super.visitField(access, name, descriptor, signature, value);
              }
  
              @Override
              public MethodVisitor visitMethod(int access, String name, String descriptor, String signature, String[] exceptions)
              {
                  if (descriptor.equals("()V") && name.equals("localInit"))
                  {
                      if (unsafeFieldName == null)
                      {
                          String version = System.getProperty("java.version");
                          if (version.startsWith("11.")) { unsafeFieldName = "U"; unsafeDescriptor = "Ljdk/internal/misc/Unsafe;"; }
                          else if (version.startsWith("1.8")) { unsafeFieldName = "UNSAFE"; unsafeDescriptor = "Lsun/misc/Unsafe;"; }
                          else throw new AssertionError("Unsupported Java Version");
                      }
  
                      MethodVisitor visitor = super.visitMethod(access, name, descriptor, signature, exceptions);
                      visitor.visitLabel(new Label());
                      visitor.visitIntInsn(ALOAD, 0);
                      visitor.visitFieldInsn(GETSTATIC, "java/util/concurrent/ThreadLocalRandom", unsafeFieldName, unsafeDescriptor);
                      visitor.visitMethodInsn(INVOKESTATIC, "java/lang/Thread", "currentThread", "()Ljava/lang/Thread;", false);
                      visitor.visitFieldInsn(GETSTATIC, "java/util/concurrent/ThreadLocalRandom", "SEED", "J");
                      visitor.visitMethodInsn(INVOKESTATIC, "org/apache/cassandra/simulator/systems/InterceptorOfSystemMethods$Global", "randomSeed", "()J", false);
                      visitor.visitMethodInsn(INVOKEVIRTUAL, unsafeDescriptor, "putLong", "(Ljava/lang/Object;JJ)V", false);
                      visitor.visitFieldInsn(GETSTATIC, "java/util/concurrent/ThreadLocalRandom", unsafeFieldName, unsafeDescriptor);
                      visitor.visitMethodInsn(INVOKESTATIC, "java/lang/Thread", "currentThread", "()Ljava/lang/Thread;", false);
                      visitor.visitFieldInsn(GETSTATIC, "java/util/concurrent/ThreadLocalRandom", "PROBE", "J");
                      visitor.visitLdcInsn(0);
                      visitor.visitMethodInsn(INVOKEVIRTUAL, unsafeDescriptor, "putInt", "(Ljava/lang/Object;JI)V", false);
                      visitor.visitInsn(RETURN);
                      visitor.visitLabel(new Label());
                      visitor.visitMaxs(6, 1);
                      visitor.visitEnd();
  
                      return new MethodVisitor(BYTECODE_VERSION) {};
                  }
                  else
                  {
                      MethodVisitor mv = super.visitMethod(access, name, descriptor, signature, exceptions);
                      if (determinismCheck && (name.equals("nextSeed") || name.equals("nextSecondarySeed") || name.equals("advanceProbe")))
                          mv = new ThreadLocalRandomCheckTransformer(api, mv);
                      return mv;
                  }
              }
          }
          return transform(bytes, ThreadLocalRandomVisitor::new);
      }
  
      /**
       * We require ThreadLocalRandom to be deterministic, so we modify its initialisation method to invoke a
       * global deterministic random value generator
       */
      private static byte[] transformAbstractQueuedSynchronizer(String className, byte[] bytes, Flag flag, Flag ... flags)
      {
          class AbstractQueuedSynchronizerVisitor extends ClassVisitor
          {
              private long defaultSpinForTimeoutThreshold = 1000L;
  
              public AbstractQueuedSynchronizerVisitor(int api, ClassVisitor classVisitor)
              {
                  super(api, classVisitor);
              }
  
              @Override
              public FieldVisitor visitField(int access, String name, String descriptor, String signature, Object value)
              {
                  if (name.equals("SPIN_FOR_TIMEOUT_THRESHOLD"))
                  {
                      defaultSpinForTimeoutThreshold = (Long)value;
                      return super.visitField(access, name, descriptor, signature, 0L);
                  }
  
                  return super.visitField(access, name, descriptor, signature, value);
              }
  
              @Override
              public MethodVisitor visitMethod(int access, String name, String descriptor, String signature, String[] exceptions)
              {
                  /// !!!!! WARNING !!!!!
                  /// THIS IS SUPER BRITTLE BECAUSE rt.jar INLINES GETSTATIC AS LDC
                  // TODO (desired): visit constructor to fetch actual value of constant in case changes in future release -
                  //  but this is brittle enough changes upstream will likely need revisiting anyway
                  MethodVisitor mv = super.visitMethod(access, name, descriptor, signature, exceptions);
                  if (!name.equals("doAcquireNanos") && !name.equals("doAcquireSharedNanos"))
                      return mv;
  
                  return new MethodVisitor(api, mv)
                  {
                      @Override
                      public void visitLdcInsn(Object value)
                      {
                          if (Objects.equals(defaultSpinForTimeoutThreshold, value))
                              super.visitLdcInsn(0L);
                          else
                              super.visitLdcInsn(value);
                      }
                  };
              }
          }
          return transform(className, bytes, AbstractQueuedSynchronizerVisitor::new, flag, flags);
      }
  
      private static byte[] transform(byte[] bytes, BiFunction<Integer, ClassVisitor, ClassVisitor> constructor)
      {
          ClassWriter out = new ClassWriter(0);
          ClassReader in = new ClassReader(bytes);
          ClassVisitor transform = constructor.apply(BYTECODE_VERSION, out);
          in.accept(transform, 0);
          return out.toByteArray();
      }
  
      private static byte[] transform(String className, byte[] bytes, Flag flag, Flag ... flags)
      {
          ClassTransformer transformer = new ClassTransformer(BYTECODE_VERSION, className, EnumSet.of(flag, flags), null);
          transformer.readAndTransform(bytes);
          if (!transformer.isTransformed())
              return null;
          return transformer.toBytes();
      }
  
      private static byte[] transform(String className, byte[] bytes, BiFunction<Integer, ClassVisitor, ClassVisitor> constructor, Flag flag, Flag ... flags)
      {
          ClassReader in = new ClassReader(bytes);
          ClassTransformer transformer = new ClassTransformer(BYTECODE_VERSION, className, EnumSet.of(flag, flags), null);
          ClassVisitor extraTransformer = constructor.apply(BYTECODE_VERSION, transformer);
          in.accept(extraTransformer, 0);
          return transformer.toBytes();
      }
  }
  