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

Go to download
package com.undefinedlabs.scope.rules;

import com.undefinedlabs.scope.logger.ScopeLogger;
import com.undefinedlabs.scope.logger.ScopeLoggerResolver;
import com.undefinedlabs.scope.rules.transformer.ScopeAgentAdvicedTransformer;
import net.bytebuddy.agent.builder.AgentBuilder;
import net.bytebuddy.asm.Advice;

import java.lang.reflect.Method;
import java.net.HttpURLConnection;
import java.net.URLConnection;
import java.util.Collections;
import java.util.Map;
import java.util.UUID;

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

public class URLScopeAgentRule extends AbstractScopeAgentRule {

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

    @Override
    protected Iterable transformers() {
        return Collections.singleton(newAgentBuilder()
                .type(named("java.net.URL"), isBootstrapClassLoader()).transform(new ScopeAgentAdvicedTransformer(HttpURLOpenConnectionAdvice.class, named("openConnection")))
        );
    }

    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 Object urlConnectionObj) {
            final URLConnection urlConnection = (URLConnection) urlConnectionObj;
            if (!(urlConnection.getClass().getName().contains("HttpURLConnection")) || 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 Method logInfoMethod = ruleClass.getMethod("logInfo", String.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 (Exception e) {
                throw new RuntimeException(e);
            }
        }

    }
}