com.undefinedlabs.scope.rules.HttpURLConnectionScopeAgentRule 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.io.IOException;
import java.io.InputStream;
import java.lang.reflect.Method;
import java.net.HttpURLConnection;
import java.util.Collections;
import java.util.Map;

import static net.bytebuddy.matcher.ElementMatchers.*;

public class HttpURLConnectionScopeAgentRule extends AbstractScopeAgentRule {

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

    @Override
    protected Iterable transformers() {
        final Class instrumentedClass = lookUp("java.net.HttpURLConnection", getBootstrapClassLoader());
        return (instrumentedClass != null) ?
                Collections.singleton(newAgentBuilder().type(isSubTypeOf(instrumentedClass), isBootstrapClassLoader()).transform(new ScopeAgentAdvicedTransformer(HttpURLGetResponseCodeAdvice.class, named("getResponseCode"))))
                : emptyCollection();
    }

    public Object getSpanByUuid(final String uuid) {
        return URLConnectionScopeUtils.getSpanByUuid().get(uuid);
    }

    public void removeSpanByUuid(final String uuid) {
        URLConnectionScopeUtils.getSpanByUuid().remove(uuid);
    }

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

    public static class HttpURLGetResponseCodeAdvice {

        @Advice.OnMethodExit(onThrowable = IOException.class)
        public static void exit(@Advice.This HttpURLConnection httpURLConnection, @Advice.Return int responseCode, @Advice.Thrown Throwable throwable) {

            try {
                final String spanUuid = httpURLConnection.getRequestProperty(URLConnectionScopeUtils.SCOPE_OPEN_CONNECTION_CUUID_KEY);

                if (spanUuid == null) {
                    return;
                }

                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.HttpURLConnectionScopeAgentRule", false, ClassLoader.getSystemClassLoader());
                final Object ruleInstance = ruleClass.newInstance();

                //It is necessary to load Scope classes using Reflection with SystemClassLoader because java.net classes are loaded by BootstrapClassloader.
                final Method getSpanByUuidMethod = ruleClass.getMethod("getSpanByUuid", String.class);
                final Method removeSpanByUuidMethod = ruleClass.getMethod("removeSpanByUuid", String.class);
                final Method logInfoMethod = ruleClass.getMethod("logInfo", String.class);

                final Object span = getSpanByUuidMethod.invoke(ruleInstance, spanUuid);

                if (span == null) {
                    return;
                }


                final Class opentracingSpanClass = (Class) getScopeClass.invoke(reflCtxtInstance, "io.opentracing.Span");
                final Class spanClass = (Class) getScopeClass.invoke(reflCtxtInstance, span.getClass().getName());
                final Method setOperationNameMethod = (Method) getScopeMethod.invoke(reflCtxtInstance, spanClass.getName(), "setOperationName", new Class[]{String.class});
                final Method setTagStringStringMethod = (Method) getScopeMethod.invoke(reflCtxtInstance, spanClass.getName(), "setTag", new Class[]{String.class, String.class});
                final Method setTagStringBooleanMethod = (Method) getScopeMethod.invoke(reflCtxtInstance, spanClass.getName(), "setTag", new Class[]{String.class, boolean.class});
                final Method setTagStringNumberMethod = (Method) getScopeMethod.invoke(reflCtxtInstance, spanClass.getName(), "setTag", new Class[]{String.class, Number.class});
                final Method logFieldsSpanMethod = (Method) getScopeMethod.invoke(reflCtxtInstance, spanClass.getName(), "log", new Class[]{Map.class});
                final Method finishMethod = (Method) getScopeMethod.invoke(reflCtxtInstance, spanClass.getName(), "finish", new Class[]{});

                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 tagValuesClass = (Class) getScopeClass.invoke(reflCtxtInstance, "com.undefinedlabs.scope.utils.tag.TagValues");
                final Class tagValuesNetworkClass = (Class) getScopeClass.invoke(reflCtxtInstance, "com.undefinedlabs.scope.utils.tag.TagValues$Network");

                final String operationName = String.format("HTTP %s", httpURLConnection.getRequestMethod());
                setOperationNameMethod.invoke(span, operationName);
                setTagStringStringMethod.invoke(span, tagKeysNetworkClass.getDeclaredField("HTTP_METHOD").get(null), httpURLConnection.getRequestMethod());

                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.getHeaderFields());
                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_RESPONSE_HEADERS").get(null), httpResHeaders);


                final Class scopeSettingsResolverClass = (Class) getScopeClass.invoke(reflCtxtInstance, "com.undefinedlabs.scope.settings.ScopeSettingsResolver");
                final Object scopeSettingsResolverInstance = scopeSettingsResolverClass.getDeclaredField("INSTANCE").get(null);
                final Method scopeSettingsGetMethod = (Method) getScopeMethod.invoke(reflCtxtInstance, scopeSettingsResolverClass.getName(), "get", new Class[]{});
                final Object scopeSettingsObj = scopeSettingsGetMethod.invoke(scopeSettingsResolverInstance);
                final Method scopeSettingsGetSettingMethod = (Method) getScopeMethod.invoke(reflCtxtInstance, scopeSettingsObj.getClass().getName(), "getSetting", new Class[]{String.class});
                final Object scopeInstrumentationHttpPayloadsKey = scopeSettingsObj.getClass().getDeclaredField("SCOPE_INSTRUMENTATION_HTTP_PAYLOADS").get(null);
                final boolean httpPayloadsActive = (boolean) scopeSettingsGetSettingMethod.invoke(scopeSettingsObj, scopeInstrumentationHttpPayloadsKey);
                final Object scopeInstrumentationHttpStacktraceKey = scopeSettingsObj.getClass().getDeclaredField("SCOPE_INSTRUMENTATION_HTTP_STACKTRACE").get(null);
                final boolean httpStacktrace = (boolean) scopeSettingsGetSettingMethod.invoke(scopeSettingsObj, scopeInstrumentationHttpStacktraceKey);

                if (httpStacktrace) {
                    final Class tagKeysStacktraceClass = (Class) getScopeClass.invoke(reflCtxtInstance, "com.undefinedlabs.scope.utils.tag.TagKeys$Stacktrace");
                    final Class stacktraceUtilsClass = (Class) getScopeClass.invoke(reflCtxtInstance, "com.undefinedlabs.scope.utils.StacktraceUtils");
                    final Method createStackTraceFieldsMethod = (Method) getScopeMethod.invoke(reflCtxtInstance, stacktraceUtilsClass.getName(), "createStackTraceFields", new Class[]{StackTraceElement[].class});
                    final Object stacktraceUtilsObj = stacktraceUtilsClass.getDeclaredField("INSTANCE").get(null);
                    final Object stacktraceFields = createStackTraceFieldsMethod.invoke(stacktraceUtilsObj, (Object) Thread.currentThread().getStackTrace());
                    setTagObjectMethod.invoke(spanUtilsObj, span, tagKeysStacktraceClass.getDeclaredField("STACKTRACE").get(null), stacktraceFields);
                }

                // *
                // Request payload is not accessible without disturbing Streams.
                // *

                if(httpPayloadsActive) {
                    if("GET".equalsIgnoreCase(httpURLConnection.getRequestMethod())) {
                        setTagStringStringMethod.invoke(span, tagKeysNetworkClass.getDeclaredField("HTTP_REQUEST_PAYLOAD").get(null), "");
                    } else {
                        setTagStringStringMethod.invoke(span, tagKeysNetworkClass.getDeclaredField("HTTP_REQUEST_PAYLOAD_UNAVAILABLE").get(null),
                                tagValuesNetworkClass.getDeclaredField("HTTP_PAYLOAD_NOT_ACCESSIBLE").get(null));
                    }
                } else {
                    setTagStringStringMethod.invoke(span, tagKeysNetworkClass.getDeclaredField("HTTP_REQUEST_PAYLOAD_UNAVAILABLE").get(null),
                            tagValuesNetworkClass.getDeclaredField("HTTP_PAYLOAD_DISABLED").get(null));
                }

                if (throwable == null) {
                    setTagStringNumberMethod.invoke(span, tagKeysNetworkClass.getDeclaredField("HTTP_STATUS_CODE").get(null), responseCode);
                    setTagStringBooleanMethod.invoke(span, tagKeysClass.getDeclaredField("ERROR").get(null), (responseCode >= 400));

                    if (httpPayloadsActive) {
                        final Class scopeIOUtilsClass = (Class) getScopeClass.invoke(reflCtxtInstance, "com.undefinedlabs.scope.utils.ScopeIOUtils");
                        final Method getSubstringMethod = (Method) getScopeMethod.invoke(reflCtxtInstance, scopeIOUtilsClass.getName(), "getSubstring", new Class[]{InputStream.class});

                        final Object responseStr = httpURLConnection.getInputStream() != null && httpURLConnection.getInputStream().markSupported() ? getSubstringMethod.invoke(null, httpURLConnection.getInputStream()) : null;
                        setTagStringStringMethod.invoke(span, tagKeysNetworkClass.getDeclaredField("HTTP_RESPONSE_PAYLOAD_UNAVAILABLE").get(null), (responseStr != null && !"".equals(responseStr)) ? null : tagValuesNetworkClass.getDeclaredField("HTTP_PAYLOAD_NOT_ACCESSIBLE").get(null));
                        setTagStringStringMethod.invoke(span, tagKeysNetworkClass.getDeclaredField("HTTP_RESPONSE_PAYLOAD").get(null), (responseStr != null && !"".equals(responseStr)) ? responseStr : null);
                    } else {
                        setTagStringStringMethod.invoke(span, tagKeysNetworkClass.getDeclaredField("HTTP_RESPONSE_PAYLOAD_UNAVAILABLE").get(null),
                                tagValuesNetworkClass.getDeclaredField("HTTP_PAYLOAD_DISABLED").get(null));
                    }

                } else {

                    final Class throwableEventClass = (Class) getScopeClass.invoke(reflCtxtInstance, "com.undefinedlabs.scope.events.exception.ThrowableEvent");
                    final Method newBuilderMethod = (Method) getScopeMethod.invoke(reflCtxtInstance, throwableEventClass.getName(), "newBuilder", new Class[]{});

                    final Object throwableBuilder = newBuilderMethod.invoke(null);
                    final Class throwableBuilderClass = throwableBuilder.getClass();

                    final Method withEventTypeMethod = (Method) getScopeMethod.invoke(reflCtxtInstance, throwableBuilderClass.getName(), "withEventType", new Class[]{String.class});
                    final Method withThrowableMethod = (Method) getScopeMethod.invoke(reflCtxtInstance, throwableBuilderClass.getName(), "withThrowable", new Class[]{Throwable.class});
                    final Method withSourceMethod = (Method) getScopeMethod.invoke(reflCtxtInstance, throwableBuilderClass.getName(), "withSource", new Class[]{String.class});
                    final Method buildMethod = (Method) getScopeMethod.invoke(reflCtxtInstance, throwableBuilderClass.getName(), "build", new Class[]{});

                    final Class eventValuesGeneralClass = (Class) getScopeClass.invoke(reflCtxtInstance, "com.undefinedlabs.scope.utils.event.EventValues$General");
                    final Class sourceCodeFactoryClass = (Class) getScopeClass.invoke(reflCtxtInstance, "com.undefinedlabs.scope.utils.sourcecode.SourceCodeFrameFactory");

                    final Object sourceCodeFactoryInstance = sourceCodeFactoryClass.getDeclaredField("INSTANCE").get(null);
                    final Method createFrameMethod = (Method) getScopeMethod.invoke(reflCtxtInstance, sourceCodeFactoryClass.getName(), "createFrame", new Class[]{StackTraceElement[].class});

                    final Class eventFieldsFactoryClass = (Class) getScopeClass.invoke(reflCtxtInstance, "com.undefinedlabs.scope.events.EventFieldsFactory");

                    final Object eventFieldsFactoryInstance = eventFieldsFactoryClass.getDeclaredField("INSTANCE").get(null);
                    final Method createFieldsMethod = (Method) getScopeMethod.invoke(reflCtxtInstance, eventFieldsFactoryClass.getName(), "createFields", new Class[]{throwableEventClass});

                    setTagStringBooleanMethod.invoke(span, tagKeysClass.getDeclaredField("ERROR").get(null), true);

                    withEventTypeMethod.invoke(throwableBuilder, eventValuesGeneralClass.getDeclaredField("ERROR").get(null));
                    withThrowableMethod.invoke(throwableBuilder, throwable);

                    final Object sourceCodeFrameInstance = createFrameMethod.invoke(sourceCodeFactoryInstance, new Object[]{throwable.getStackTrace()});
                    final Method getLinkPathMethod = (Method) getScopeMethod.invoke(reflCtxtInstance, sourceCodeFrameInstance.getClass().getName(), "getLinkPathWithMethodLine", new Class[]{});

                    withSourceMethod.invoke(throwableBuilder, getLinkPathMethod.invoke(sourceCodeFrameInstance));

                    final Object throwableEventInstance = buildMethod.invoke(throwableBuilder);
                    final Object throwableEventFields = createFieldsMethod.invoke(eventFieldsFactoryInstance, throwableEventInstance);

                    logFieldsSpanMethod.invoke(span, throwableEventFields);
                }

                finishMethod.invoke(span);

                removeSpanByUuidMethod.invoke(ruleInstance, spanUuid);

            } catch (Exception e) {
                throw new RuntimeException(e);
            }
        }
    }
}