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.ByteArrayOutputStream;
  import java.io.IOException;
  import java.io.InputStream;
  import java.io.Serializable;
  import java.io.UncheckedIOException;
  import java.lang.reflect.Method;
  import java.util.Collections;
  import java.util.EnumSet;
  import java.util.HashSet;
  import java.util.Map;
  import java.util.NavigableSet;
  import java.util.Set;
  import java.util.TreeSet;
  import java.util.concurrent.ConcurrentHashMap;
  import java.util.function.BiFunction;
  import java.util.function.Consumer;
  import java.util.function.Predicate;
  import java.util.regex.Pattern;
  
  import org.objectweb.asm.Opcodes;
  
  import static org.apache.cassandra.simulator.asm.InterceptClasses.Cached.Kind.MODIFIED;
  import static org.apache.cassandra.simulator.asm.InterceptClasses.Cached.Kind.UNMODIFIED;
  import static org.apache.cassandra.simulator.asm.InterceptClasses.Cached.Kind.UNSHAREABLE;
  
  // TODO (completeness): confirm that those classes we weave monitor-access for only extend other classes we also weave monitor access for
  // TODO (completeness): confirm that those classes we weave monitor access for only take monitors on types we also weave monitor access for (and vice versa)
  // WARNING: does not implement IClassTransformer directly as must be accessible to bootstrap class loader
  public class InterceptClasses implements BiFunction<String, byte[], byte[]>
  {
      public static final int BYTECODE_VERSION = Opcodes.ASM7;
  
      // TODO (cleanup): use annotations
      private static final Pattern MONITORS = Pattern.compile( "org[/.]apache[/.]cassandra[/.]utils[/.]concurrent[/.].*" +
                                                              "|org[/.]apache[/.]cassandra[/.]concurrent[/.].*" +
                                                              "|org[/.]apache[/.]cassandra[/.]simulator[/.]test.*" +
                                                              "|org[/.]apache[/.]cassandra[/.]db[/.]ColumnFamilyStore.*" +
                                                              "|org[/.]apache[/.]cassandra[/.]db[/.]Keyspace.*" +
                                                              "|org[/.]apache[/.]cassandra[/.]db[/.]SystemKeyspace.*" +
                                                              "|org[/.]apache[/.]cassandra[/.]index[/.].*" +
                                                              "|org[/.]apache[/.]cassandra[/.]streaming[/.].*" +
                                                              "|org[/.]apache[/.]cassandra[/.]db.streaming[/.].*" +
                                                              "|org[/.]apache[/.]cassandra[/.]distributed[/.]impl[/.]DirectStreamingConnectionFactory.*" +
                                                              "|org[/.]apache[/.]cassandra[/.]db[/.]commitlog[/.].*" +
                                                              "|org[/.]apache[/.]cassandra[/.]service[/.]paxos[/.].*" +
                                                              "|org[/.]apache[/.]cassandra[/.]service[/.]accord[/.].*" +
                                                              "|org[/.]apache[/.]cassandra[/.]journal[/.].*" +
                                                              "|accord[/.].*"
      );
  
      private static final Pattern GLOBAL_METHODS = Pattern.compile("org[/.]apache[/.]cassandra[/.](?!simulator[/.]).*" +
                                                                    "|org[/.]apache[/.]cassandra[/.]simulator[/.]test[/.].*" +
                                                                    "|org[/.]apache[/.]cassandra[/.]simulator[/.]cluster[/.].*" +
                                                                    "|io[/.]netty[/.]util[/.]concurrent[/.]FastThreadLocal" +
                                                                    "|accord[/.].*"
      ); // intercept IdentityHashMap for execution consistency
      private static final Pattern NEMESIS = GLOBAL_METHODS;
      private static final Set<String> WARNED = Collections.newSetFromMap(new ConcurrentHashMap<>());
  
      static final byte[] SENTINEL = new byte[0];
      static class Cached
      {
          enum Kind { MODIFIED, UNMODIFIED, UNSHAREABLE }
          final Kind kind;
          final byte[] bytes;
          final Set<String> uncacheablePeers;
          private Cached(Kind kind, byte[] bytes, Set<String> uncacheablePeers)
          {
              this.kind = kind;
              this.bytes = bytes;
              this.uncacheablePeers = uncacheablePeers;
          }
      }
  
      static class PeerGroup
      {
          final Set<String> uncacheablePeers = new TreeSet<>();
          final Cached unmodified = new Cached(UNMODIFIED, null, uncacheablePeers);
      }
  
      class SubTransformer implements BiFunction<String, byte[], byte[]>
      {
          private final Map<String, byte[]> isolatedCache = new ConcurrentHashMap<>();
          private final int id;
          SubTransformer(int id)
          {
              this.id = id;
          }
  
          @Override
          public byte[] apply(String name, byte[] bytes)
          {
              return transformTransitiveClosure(name, bytes, isolatedCache, id);
          }
      }
  
      private final Map<String, Cached> cache = new ConcurrentHashMap<>();
  
      private final int api;
      private final DeterministicChanceSupplier nemesisChance;
      private final DeterministicChanceSupplier monitorDelayChance;
      private final Hashcode insertHashcode;
      private final NemesisFieldKind.Selector nemesisFieldSelector;
      private final ClassLoader prewarmClassLoader;
      private final Predicate<String> prewarm;
      private final byte[] bufIn = new byte[4096];
      private final ByteArrayOutputStream bufOut = new ByteArrayOutputStream();
  
      public InterceptClasses(DeterministicChanceSupplier monitorDelayChance, DeterministicChanceSupplier nemesisChance, NemesisFieldKind.Selector nemesisFieldSelector, ClassLoader prewarmClassLoader, Predicate<String> prewarm)
      {
          this(BYTECODE_VERSION, monitorDelayChance, nemesisChance, nemesisFieldSelector, prewarmClassLoader, prewarm);
      }
  
      public InterceptClasses(int api, DeterministicChanceSupplier monitorDelayChance, DeterministicChanceSupplier nemesisChance, NemesisFieldKind.Selector nemesisFieldSelector, ClassLoader prewarmClassLoader, Predicate<String> prewarm)
      {
          this.api = api;
          this.nemesisChance = nemesisChance;
          this.monitorDelayChance = monitorDelayChance;
          this.insertHashcode = new Hashcode(api);
          this.nemesisFieldSelector = nemesisFieldSelector;
          this.prewarmClassLoader = prewarmClassLoader;
          this.prewarm = prewarm;
      }
  
      @Override
      public byte[] apply(String name, byte[] bytes)
      {
          return transformTransitiveClosure(name, bytes, null, 0);
      }
  
      private synchronized byte[] transformTransitiveClosure(String externalName, byte[] input, Map<String, byte[]> isolatedCache, int id)
      {
          if (input == null)
              return maybeSynthetic(externalName);
  
          String internalName = dotsToSlashes(externalName);
          if (isolatedCache != null)
          {
              byte[] isolatedCached = isolatedCache.get(internalName);
              if (isolatedCached != null)
                  return isolatedCached == SENTINEL ? input : isolatedCached;
          }
  
          Cached cached = cache.get(internalName);
          if (cached != null)
          {
              if (isolatedCache == null)
              {
                  switch (cached.kind)
                  {
                      default: throw new AssertionError();
                      case MODIFIED:
                          return cached.bytes;
                      case UNMODIFIED:
                          return input;
                      case UNSHAREABLE:
                          return transform(internalName, externalName, null, input, null, id, null);
                  }
              }
  
              for (String peer : cached.uncacheablePeers)
                  transform(peer, slashesToDots(peer), null, cache.get(peer).bytes, isolatedCache, id, null);
  
              switch (cached.kind)
              {
                  default: throw new AssertionError();
                  case MODIFIED:
                      return cached.bytes;
                  case UNMODIFIED:
                      return input;
                  case UNSHAREABLE:
                      return isolatedCache.get(internalName);
              }
          }
  
          Set<String> visited = new HashSet<>();
          visited.add(internalName);
          NavigableSet<String> load = new TreeSet<>();
          Consumer<String> dependentTypeConsumer = type -> {
              if (prewarm.test(type) && visited.add(type))
                  load.add(type);
          };
  
          final PeerGroup peerGroup = new PeerGroup();
          byte[] result = transform(internalName, externalName, peerGroup, input, isolatedCache, id, dependentTypeConsumer);
          for (String next = load.pollFirst(); next != null; next = load.pollFirst())
          {
              // TODO (now): otherwise merge peer groups
              Cached existing = cache.get(next);
              if (existing == null)
                  transform(next, slashesToDots(next), peerGroup, read(next), isolatedCache, id, dependentTypeConsumer);
          }
  
          return result;
      }
  
      private byte[] read(String name)
      {
          try (InputStream in = prewarmClassLoader.getResourceAsStream(dotsToSlashes(name) + ".class"))
          {
              if (in == null)
                  throw new NoClassDefFoundError(dotsToSlashes(name) + ".class");
  
              bufOut.reset();
              for (int c = in.read(bufIn) ; c >= 0 ; c = in.read(bufIn))
                  bufOut.write(bufIn, 0, c);
              return bufOut.toByteArray();
          }
          catch (IOException e)
          {
              throw new NoClassDefFoundError(name);
          }
      }
  
      private byte[] transform(String internalName, String externalName, PeerGroup peerGroup, byte[] input, Map<String, byte[]> isolatedCache, int id, Consumer<String> dependentTypes)
      {
          Hashcode hashcode = insertHashCode(externalName);
  
          EnumSet<Flag> flags = EnumSet.noneOf(Flag.class);
          if (MONITORS.matcher(internalName).matches())
          {
              flags.add(Flag.MONITORS);
          }
          if (GLOBAL_METHODS.matcher(internalName).matches())
          {
              flags.add(Flag.GLOBAL_METHODS);
              flags.add(Flag.LOCK_SUPPORT);
          }
          if (NEMESIS.matcher(internalName).matches())
          {
              flags.add(Flag.NEMESIS);
          }
  
          if (flags.isEmpty() && hashcode == null)
          {
              cache.put(internalName, peerGroup.unmodified);
              return input;
          }
  
          int chanceSeed = internalName.hashCode() * 31 + id;
          ClassTransformer transformer = new ClassTransformer(api, internalName, flags, monitorDelayChance.apply(chanceSeed), new NemesisGenerator(api, internalName, nemesisChance.apply(chanceSeed + 1)), nemesisFieldSelector, hashcode, dependentTypes);
          transformer.setUpdateVisibility(true);
          transformer.readAndTransform(input);
  
          if (!transformer.isTransformed())
          {
              cache.put(internalName, peerGroup.unmodified);
              return input;
          }
  
          byte[] output = transformer.toBytes();
          if (transformer.isCacheablyTransformed())
          {
              cache.put(internalName, new Cached(MODIFIED, output, peerGroup.uncacheablePeers));
          }
          else
          {
              if (peerGroup != null)
              {
                  cache.put(internalName, new Cached(UNSHAREABLE, input, peerGroup.uncacheablePeers));
                  peerGroup.uncacheablePeers.add(internalName);
              }
              if (isolatedCache != null)
                  isolatedCache.put(internalName, output);
          }
  
          return output;
      }
  
      static String dotsToSlashes(String className)
      {
          return className.replace('.', '/');
      }
  
      static String dotsToSlashes(Class<?> clazz)
      {
          return dotsToSlashes(clazz.getName());
      }
  
      static String slashesToDots(String className)
      {
          return className.replace('/', '.');
      }
  
      /**
       * Decide if we should insert our own hashCode() implementation that assigns deterministic hashes, i.e.
       *   - If it's one of our classes
       *   - If its parent is not one of our classes (else we'll assign it one anyway)
       *   - If it does not have its own hashCode() implementation that overrides Object's
       *   - If it is not Serializable OR it has a serialVersionUID
       *
       * Otherwise we either probably do not need it, or may break serialization between classloaders
       */
      private Hashcode insertHashCode(String externalName)
      {
          try
          {
              if (!externalName.startsWith("org.apache.cassandra"))
                  return null;
  
              Class<?> sharedClass = getClass().getClassLoader().loadClass(externalName);
              if (sharedClass.isInterface() || sharedClass.isEnum() || sharedClass.isArray() || sharedClass.isSynthetic())
                  return null;
  
              Class<?> parent = sharedClass.getSuperclass();
              if (parent.getName().startsWith("org.apache.cassandra"))
                  return null;
  
              try
              {
                  Method method = sharedClass.getMethod("hashCode");
                  if (method.getDeclaringClass() != Object.class)
                      return null;
              }
              catch (NoSuchMethodException ignore)
              {
              }
  
              if (!Serializable.class.isAssignableFrom(sharedClass))
                  return insertHashcode;
  
              try
              {
                  // if we haven't specified serialVersionUID we break ObjectInputStream transfers between class loaders
                  // (might be easiest to switch to serialization that doesn't require it)
                  sharedClass.getDeclaredField("serialVersionUID");
                  return insertHashcode;
              }
              catch (NoSuchFieldException e)
              {
                  if (!Throwable.class.isAssignableFrom(sharedClass) && WARNED.add(externalName))
                      System.err.println("No serialVersionUID on Serializable " + sharedClass);
                  return null;
              }
          }
          catch (ClassNotFoundException e)
          {
              System.err.println("Unable to determine if should insert hashCode() for " + externalName);
              e.printStackTrace();
          }
          return null;
      }
  
      static final String shadowRootExternalType = "org.apache.cassandra.simulator.systems.InterceptibleConcurrentHashMap";
      static final String shadowRootType = "org/apache/cassandra/simulator/systems/InterceptibleConcurrentHashMap";
      static final String originalRootType = Utils.toInternalName(ConcurrentHashMap.class);
      static final String shadowOuterTypePrefix = shadowRootType + '$';
      static final String originalOuterTypePrefix = originalRootType + '$';
  
      protected byte[] maybeSynthetic(String externalName)
      {
          if (!externalName.startsWith(shadowRootExternalType))
              return null;
  
          try
          {
              String originalType, shadowType = Utils.toInternalName(externalName);
              if (!shadowType.startsWith(shadowOuterTypePrefix))
                  originalType = originalRootType;
              else
                  originalType = originalOuterTypePrefix + externalName.substring(shadowOuterTypePrefix.length());
  
              EnumSet<Flag> flags = EnumSet.of(Flag.GLOBAL_METHODS, Flag.MONITORS, Flag.LOCK_SUPPORT);
              if (NEMESIS.matcher(externalName).matches()) flags.add(Flag.NEMESIS);
  
              int hashCode = externalName.hashCode();
              NemesisGenerator nemesis = new NemesisGenerator(api, externalName, nemesisChance.apply(hashCode));
  
              ShadowingTransformer transformer;
              transformer = new ShadowingTransformer(InterceptClasses.BYTECODE_VERSION,
                                                     originalType, shadowType, originalRootType, shadowRootType,
                                                     originalOuterTypePrefix, shadowOuterTypePrefix,
                                                     flags, monitorDelayChance.apply(hashCode), nemesis, nemesisFieldSelector, null);
              transformer.readAndTransform(Utils.readDefinition(originalType + ".class"));
              return transformer.toBytes();
          }
          catch (IOException e)
          {
              throw new UncheckedIOException(e);
          }
      }
  
  }