com.undefinedlabs.scope.rules.UrlScopeAgentRule Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of scope-rule-javanet Show documentation
Scope is a APM for tests to give engineering teams unprecedented visibility into their CI process to quickly identify, troubleshoot and fix failed builds. This artifact contains the classes to instrument the Java NET package.
There is a newer version: 0.15.1-beta.2
Show newest version
package com.undefinedlabs.scope.rules;

import static net.bytebuddy.matcher.ElementMatchers.isBootstrapClassLoader;
import static net.bytebuddy.matcher.ElementMatchers.named;

import com.undefinedlabs.scope.logger.ScopeLogger;
import com.undefinedlabs.scope.logger.ScopeLoggerResolver;
import java.lang.reflect.Method;
import java.net.HttpURLConnection;
import java.net.URLConnection;
import java.util.HashMap;
import java.util.Map;
import java.util.UUID;
import net.bytebuddy.asm.Advice;
import net.bytebuddy.description.method.MethodDescription;
import net.bytebuddy.description.type.TypeDescription;
import net.bytebuddy.matcher.ElementMatcher;

public class UrlScopeAgentRule extends AbstractScopeAgentRule {

  private static final ScopeLogger LOGGER = ScopeLoggerResolver.INSTANCE.get();

  @Override
  protected ElementMatcher typeMatcher() {
    return named("java.net.URL");
  }

  @Override
  protected ElementMatcher classLoaderMatcher() {
    return isBootstrapClassLoader();
  }

  @Override
  protected Map, String> transformers() {
    final Map, String> transformers = new HashMap<>();
    transformers.put(
        named("openConnection"),
        UrlScopeAgentRule.class.getName() + "$HttpUrlOpenConnectionAdvice");
    return transformers;
  }

  public void putSpanByUuid(final String uuid, final Object span) {
    UrlConnectionScopeUtils.getSpanByUuid().put(uuid, span);
  }

  public void logInfo(final String message) {
    LOGGER.info(message);
  }

  public static class HttpUrlOpenConnectionAdvice {

    @Advice.OnMethodExit
    public static void exit(@Advice.Return final Object urlConnectionObj) {
      final URLConnection urlConnection = (URLConnection) urlConnectionObj;
      if (!(urlConnection.getClass().getName().contains("HttpURLConnection"))
          || urlConnection.getClass().getName().contains("Mockito")
          || urlConnection.getURL() == null
          || // Probably is a Mock
          urlConnection.getRequestProperty(UrlConnectionScopeUtils.SCOPE_OPEN_CONNECTION_CUUID_KEY)
              != null) {
        return;
      }

      try {

        final HttpURLConnection httpUrlConnection = (HttpURLConnection) urlConnection;
        final String uuid = UUID.randomUUID().toString();

        // It is necessary to load Scope classes using Reflection with SystemClassLoader because
        // java.net classes are loaded by BootstrapClassloader.
        final Class reflectionContextClass =
            Class.forName(
                "com.undefinedlabs.scope.jdk.reflection.ReflectionContext",
                false,
                ClassLoader.getSystemClassLoader());
        final Object reflCtxtInstance =
            reflectionContextClass.getDeclaredField("INSTANCE").get(null);
        final Method getScopeMethod =
            reflectionContextClass.getMethod(
                "getScopeMethod", String.class, String.class, Class[].class);
        final Method getScopeClass =
            reflectionContextClass.getMethod("getScopeClass", String.class);

        final Class ruleClass =
            Class.forName(
                "com.undefinedlabs.scope.rules.UrlScopeAgentRule",
                false,
                ClassLoader.getSystemClassLoader());
        final Method putSpanByUuidMethod =
            ruleClass.getMethod("putSpanByUuid", String.class, Object.class);
        final Object ruleInstance = ruleClass.newInstance();

        final Class globalTracerClass =
            (Class)
                getScopeClass.invoke(reflCtxtInstance, "com.undefinedlabs.scope.ScopeGlobalTracer");
        final Method getTracerMethod =
            (Method)
                getScopeMethod.invoke(
                    reflCtxtInstance, globalTracerClass.getName(), "get", new Class[] {});

        final Object tracer = getTracerMethod.invoke(null);
        final Class tracerClass = tracer.getClass();
        final Method activeSpanMethod =
            (Method)
                getScopeMethod.invoke(
                    reflCtxtInstance, tracerClass.getName(), "activeSpan", new Class[] {});
        final Object previousActiveSpan = activeSpanMethod.invoke(tracer);

        // If there is no previousActiveSpan, it is not created Span Client.
        if (previousActiveSpan == null) {
          return;
        }

        final Method buildSpanMethod =
            (Method)
                getScopeMethod.invoke(
                    reflCtxtInstance,
                    tracerClass.getName(),
                    "buildSpan",
                    new Class[] {String.class});

        final Object spanBuilder = buildSpanMethod.invoke(tracer, "temporal_op_name");
        final Class spanBuilderClass = spanBuilder.getClass();

        final Method withTagStringStringMethod =
            (Method)
                getScopeMethod.invoke(
                    reflCtxtInstance,
                    spanBuilderClass.getName(),
                    "withTag",
                    new Class[] {String.class, String.class});
        final Method withTagStringNumberMethod =
            (Method)
                getScopeMethod.invoke(
                    reflCtxtInstance,
                    spanBuilderClass.getName(),
                    "withTag",
                    new Class[] {String.class, Number.class});
        final Method startMethod =
            (Method)
                getScopeMethod.invoke(
                    reflCtxtInstance, spanBuilderClass.getName(), "start", new Class[] {});

        final Class opentracingSpanClass =
            (Class) getScopeClass.invoke(reflCtxtInstance, "io.opentracing.Span");
        final Class tagKeysClass =
            (Class)
                getScopeClass.invoke(reflCtxtInstance, "com.undefinedlabs.scope.utils.tag.TagKeys");
        final Class tagKeysNetworkClass =
            (Class)
                getScopeClass.invoke(
                    reflCtxtInstance, "com.undefinedlabs.scope.utils.tag.TagKeys$Network");
        final Class tagValuesNetworkClass =
            (Class)
                getScopeClass.invoke(
                    reflCtxtInstance, "com.undefinedlabs.scope.utils.tag.TagValues$Network");
        final Class tagValuesNetworkServiceClass =
            (Class)
                getScopeClass.invoke(
                    reflCtxtInstance,
                    "com.undefinedlabs.scope.utils.tag.TagValues$Network$Service");
        final Class tagValuesComponentClass =
            (Class)
                getScopeClass.invoke(
                    reflCtxtInstance, "com.undefinedlabs.scope.utils.tag.TagValues$Component");
        final Class baggageKeysClass =
            (Class)
                getScopeClass.invoke(
                    reflCtxtInstance, "com.undefinedlabs.scope.utils.baggage.BaggageKeys");
        final Class baggageValuesClass =
            (Class)
                getScopeClass.invoke(
                    reflCtxtInstance, "com.undefinedlabs.scope.utils.baggage.BaggageValues");

        withTagStringStringMethod.invoke(
            spanBuilder,
            tagKeysClass.getDeclaredField("SPAN_KIND").get(null),
            tagValuesNetworkClass.getDeclaredField("SPAN_KIND_CLIENT").get(null));
        withTagStringStringMethod.invoke(
            spanBuilder,
            tagKeysClass.getDeclaredField("COMPONENT").get(null),
            tagValuesComponentClass.getDeclaredField("HTTP").get(null));
        withTagStringStringMethod.invoke(
            spanBuilder,
            tagKeysNetworkClass.getDeclaredField("HTTP_URL").get(null),
            httpUrlConnection.getURL().toString());
        withTagStringStringMethod.invoke(
            spanBuilder,
            tagKeysNetworkClass.getDeclaredField("PEER_HOSTNAME").get(null),
            httpUrlConnection.getURL().getHost());
        withTagStringNumberMethod.invoke(
            spanBuilder,
            tagKeysNetworkClass.getDeclaredField("PEER_PORT").get(null),
            Integer.valueOf(httpUrlConnection.getURL().getPort()));
        withTagStringStringMethod.invoke(
            spanBuilder,
            tagKeysNetworkClass.getDeclaredField("PEER_SERVICE").get(null),
            tagValuesNetworkServiceClass.getDeclaredField("HTTP").get(null));
        final Object span = startMethod.invoke(spanBuilder);
        final Class spanClass = span.getClass();
        final Method setBaggageItemMethod =
            (Method)
                getScopeMethod.invoke(
                    reflCtxtInstance,
                    spanClass.getName(),
                    "setBaggageItem",
                    new Class[] {String.class, String.class});
        setBaggageItemMethod.invoke(
            span,
            baggageKeysClass.getDeclaredField("TRACE_KIND").get(null),
            baggageValuesClass.getDeclaredField("TRACE_KIND_TEST").get(null));

        // Request Headers
        final Class adapterClass =
            (Class)
                getScopeClass.invoke(
                    reflCtxtInstance,
                    "com.undefinedlabs.scope.rules.headers.UrlConnectionHeadersAdapter");
        final Class adapterInterfaceClass =
            (Class)
                getScopeClass.invoke(
                    reflCtxtInstance, "com.undefinedlabs.scope.rules.http.HttpHeadersAdapter");
        final Object adapterObj =
            adapterClass
                .getConstructor(Map.class)
                .newInstance(httpUrlConnection.getRequestProperties());
        final Class headersExtractorClass =
            (Class)
                getScopeClass.invoke(
                    reflCtxtInstance,
                    "com.undefinedlabs.scope.rules.http.HttpHeadersExtractorForSpan");
        final Method extractMethod =
            (Method)
                getScopeMethod.invoke(
                    reflCtxtInstance,
                    headersExtractorClass.getName(),
                    "extract",
                    new Class[] {adapterInterfaceClass});
        final Object headersExtractorObj =
            headersExtractorClass.getDeclaredField("DEFAULT").get(null);
        final Map httpResHeaders =
            (Map) extractMethod.invoke(headersExtractorObj, adapterObj);

        final Class spanUtilsClass =
            (Class)
                getScopeClass.invoke(reflCtxtInstance, "com.undefinedlabs.scope.utils.SpanUtils");
        final Method setTagObjectMethod =
            (Method)
                getScopeMethod.invoke(
                    reflCtxtInstance,
                    spanUtilsClass.getName(),
                    "setTagObject",
                    new Class[] {opentracingSpanClass, String.class, Object.class});
        final Object spanUtilsObj = spanUtilsClass.getDeclaredField("INSTANCE").get(null);
        setTagObjectMethod.invoke(
            spanUtilsObj,
            span,
            tagKeysNetworkClass.getDeclaredField("HTTP_REQUEST_HEADERS").get(null),
            httpResHeaders);

        final Method activateSpanMethod =
            (Method)
                getScopeMethod.invoke(
                    reflCtxtInstance,
                    tracerClass.getName(),
                    "activateSpan",
                    new Class[] {spanClass.getInterfaces()[0]});
        activateSpanMethod.invoke(tracer, span);

        final Class spanContextClass =
            (Class) getScopeClass.invoke(reflCtxtInstance, "io.opentracing.SpanContext");
        final Class formatClass =
            (Class) getScopeClass.invoke(reflCtxtInstance, "io.opentracing.propagation.Format");
        final Class formatBuiltinClass =
            (Class)
                getScopeClass.invoke(reflCtxtInstance, "io.opentracing.propagation.Format$Builtin");
        final Object httpHeadersConstant =
            formatBuiltinClass.getDeclaredField("HTTP_HEADERS").get(null);
        final Class urlConnInjectAdaptClass =
            (Class)
                getScopeClass.invoke(
                    reflCtxtInstance,
                    "com.undefinedlabs.scope.rules.propagation.carriers.UrlConnectionInjectAdapter");
        final Method getContextMethod =
            (Method)
                getScopeMethod.invoke(
                    reflCtxtInstance, spanClass.getName(), "context", new Class[] {});
        final Object spanContext = getContextMethod.invoke(span);
        final Object urlConnInjectAdaptInstance =
            urlConnInjectAdaptClass
                .getConstructor(URLConnection.class)
                .newInstance(httpUrlConnection);
        final Method injectMethod =
            (Method)
                getScopeMethod.invoke(
                    reflCtxtInstance,
                    tracerClass.getName(),
                    "inject",
                    new Class[] {spanContextClass, formatClass, Object.class});
        injectMethod.invoke(tracer, spanContext, httpHeadersConstant, urlConnInjectAdaptInstance);

        putSpanByUuidMethod.invoke(ruleInstance, uuid, span);

        httpUrlConnection.setRequestProperty(
            UrlConnectionScopeUtils.SCOPE_OPEN_CONNECTION_CUUID_KEY, uuid);
      } catch (final Exception e) {
        throw new RuntimeException(e);
      }
    }
  }
}