com.catchpoint.trace.instrument.integrations.spring.web.advice.SpringMockMvcAdviceHelper Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of cp-trace-instrument-integrations-spring-web Show documentation
There is a newer version: 0.0.20
Show newest version
package com.catchpoint.trace.instrument.integrations.spring.web.advice;

import com.catchpoint.trace.core.codec.impl.Map2CatchpointTextMapAdapter;
import com.catchpoint.trace.core.scope.CatchpointScope;
import com.catchpoint.trace.core.span.CommonTags;
import com.catchpoint.trace.core.span.CatchpointMutableSpan;
import com.catchpoint.trace.core.span.CatchpointSpan;
import com.catchpoint.trace.core.span.CatchpointFormat;
import com.catchpoint.trace.integrations.http.HttpTraceHelper;
import com.catchpoint.trace.common.logger.LoggerFactory;
import com.catchpoint.trace.integrations.http.HttpTags;
import org.slf4j.Logger;
import org.springframework.mock.web.MockHttpServletRequest;
import org.springframework.mock.web.MockHttpServletResponse;

import java.io.IOException;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

/**
 * @author serkan
 */
public final class SpringMockMvcAdviceHelper extends SpringHttpClientAdviceHelper {

    private static final Logger LOGGER = LoggerFactory.getLogger(SpringMockMvcAdviceHelper.class);

    private static final ThreadLocal requestScope = new ThreadLocal<>();

    private SpringMockMvcAdviceHelper() {
    }

    public static void onRequest(CatchpointSpan span, MockHttpServletRequest request) {
        if (!MASK_BODY) {
            try {
                if (request.getContentLength() <= MAX_BODY_SIZE) {
                    byte[] content = request.getContentAsByteArray();
                    if (content != null && content.length > 0 && content.length <= MAX_BODY_SIZE) {
                        span.setTag(HttpTags.BODY, new String(content, "UTF-8"));
                    }
                }
            } catch (IOException e) {
                LOGGER.error("Unable to get body of HTTP request!", e);
            }
        }

        if (!SpringMockMvcAdviceHelper.DISABLE_SPAN_CONTEXT_INJECTION) {
            Map headers = new HashMap<>();
            span.getTracer().inject(span.context(), CatchpointFormat.Builtin.TEXT_MAP, new Map2CatchpointTextMapAdapter(headers));
            headers.forEach((k, v) -> request.addHeader(k, v));
        }
    }

    public static void onResponse(CatchpointSpan span, MockHttpServletResponse response) {
        int statusCode = response.getStatus();

        span.setTag(HttpTags.STATUS_CODE, statusCode);
        HttpTraceHelper.setErrorTagsIfErroneous(span, statusCode);

        if (!MASK_RESPONSE_BODY) {
            try {
                if (response.getContentLength() <= MAX_RESPONSE_BODY_SIZE) {
                    byte[] content = response.getContentAsByteArray();
                    if (content != null && content.length > 0 && content.length <= MAX_RESPONSE_BODY_SIZE) {
                        span.setTag(HttpTags.RESPONSE_BODY, new String(content, "UTF-8"));
                    }
                }
            } catch (IOException e) {
                LOGGER.error("Unable to get body of HTTP response!", e);
            }
        }

        if (response.containsHeader(AWS_API_GW_ID_HEADER_NAME) && span instanceof CatchpointMutableSpan) {
            ((CatchpointMutableSpan) span).setClassName(AWS_API_GW_CLASS_NAME);
        }
        if (response.containsHeader(CATCHPOINT_RESOURCE_HEADER_NAME) && span instanceof CatchpointMutableSpan) {
            String resourceName = response.getHeaders(CATCHPOINT_RESOURCE_HEADER_NAME).get(0);
            span.setOperationName(resourceName);
        }

        if (response.containsHeader(AWS_TRACE_ID_HEADER_NAME)) {
            String traceHeader = response.getHeaders(AWS_TRACE_ID_HEADER_NAME).get(0);
            for (String traceHeaderPart : traceHeader.split(";")) {
                String[] traceInfo = traceHeaderPart.split("=");
                if (traceInfo.length == 2 && "Root".equals(traceInfo[0])) {
                    String traceId = traceInfo[1];
                    List traceLinks = span.getTag(CommonTags.TRACE_LINKS);
                    if (traceLinks == null) {
                        traceLinks = new ArrayList<>(1);
                        span.setTag(CommonTags.TRACE_LINKS, traceLinks);
                    }
                    traceLinks.add(traceId);
                }
            }
        }
    }

    public static CatchpointScope getCurrentRequestScope() {
        return requestScope.get();
    }

    public static void setCurrentRequestScope(CatchpointScope scope) {
        requestScope.set(scope);
    }

    public static void removeCurrentRequestScope() {
        requestScope.remove();
    }

}