content: /*
   * Copyright The OpenTelemetry Authors
   * SPDX-License-Identifier: Apache-2.0
   */
  
  package io.opentelemetry.javaagent.instrumentation.liberty;
  
  import static io.opentelemetry.javaagent.instrumentation.liberty.LibertySingletons.helper;
  import static net.bytebuddy.matcher.ElementMatchers.named;
  import static net.bytebuddy.matcher.ElementMatchers.takesArgument;
  
  import io.opentelemetry.context.Context;
  import io.opentelemetry.context.Scope;
  import io.opentelemetry.javaagent.bootstrap.Java8BytecodeBridge;
  import io.opentelemetry.javaagent.bootstrap.http.HttpServerResponseCustomizerHolder;
  import io.opentelemetry.javaagent.extension.instrumentation.TypeInstrumentation;
  import io.opentelemetry.javaagent.extension.instrumentation.TypeTransformer;
  import io.opentelemetry.javaagent.instrumentation.servlet.ServletRequestContext;
  import io.opentelemetry.javaagent.instrumentation.servlet.v3_0.Servlet3Accessor;
  import javax.servlet.ServletRequest;
  import javax.servlet.ServletResponse;
  import javax.servlet.http.HttpServletRequest;
  import javax.servlet.http.HttpServletResponse;
  import net.bytebuddy.asm.Advice;
  import net.bytebuddy.description.type.TypeDescription;
  import net.bytebuddy.matcher.ElementMatcher;
  
  public class LibertyWebAppInstrumentation implements TypeInstrumentation {
  
    @Override
    public ElementMatcher<TypeDescription> typeMatcher() {
      return named("com.ibm.ws.webcontainer.webapp.WebApp");
    }
  
    @Override
    public void transform(TypeTransformer transformer) {
      // https://github.com/OpenLiberty/open-liberty/blob/master/dev/com.ibm.ws.webcontainer/src/com/ibm/ws/webcontainer/webapp/WebApp.java
      transformer.applyAdviceToMethod(
          named("handleRequest")
              .and(takesArgument(0, named("javax.servlet.ServletRequest")))
              .and(takesArgument(1, named("javax.servlet.ServletResponse"))),
          this.getClass().getName() + "$HandleRequestAdvice");
  
      // isForbidden is called from handleRequest
      transformer.applyAdviceToMethod(
          named("isForbidden").and(takesArgument(0, named(String.class.getName()))),
          this.getClass().getName() + "$IsForbiddenAdvice");
    }
  
    @SuppressWarnings("unused")
    public static class HandleRequestAdvice {
  
      @Advice.OnMethodEnter(suppress = Throwable.class)
      public static void onEnter(
          @Advice.Argument(value = 0) ServletRequest request,
          @Advice.Argument(value = 1) ServletResponse response,
          @Advice.Local("otelHandled") boolean handled) {
  
        // liberty has two handleRequest methods, skip processing when thread local context is already
        // set up
        handled = ThreadLocalContext.get() == null;
        if (!handled
            || !(request instanceof HttpServletRequest)
            || !(response instanceof HttpServletResponse)) {
          return;
        }
  
        HttpServletRequest httpServletRequest = (HttpServletRequest) request;
        // it is a bit too early to start span at this point because calling
        // some methods on HttpServletRequest will give a NPE
        // just remember the request and use it a bit later to start the span
        ThreadLocalContext.startRequest(httpServletRequest, (HttpServletResponse) response);
      }
  
      @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
      public static void stopSpan(
          @Advice.Argument(0) ServletRequest servletRequest,
          @Advice.Argument(1) ServletResponse servletResponse,
          @Advice.Thrown Throwable throwable,
          @Advice.Local("otelHandled") boolean handled) {
        if (!handled) {
          return;
        }
        ThreadLocalContext requestInfo = ThreadLocalContext.endRequest();
        if (requestInfo == null) {
          return;
        }
  
        HttpServletRequest request = (HttpServletRequest) servletRequest;
        HttpServletResponse response = (HttpServletResponse) servletResponse;
  
        helper()
            .end(
                requestInfo.getRequestContext(),
                request,
                response,
                throwable,
                requestInfo.getContext(),
                requestInfo.getScope());
      }
    }
  
    @SuppressWarnings("unused")
    public static class IsForbiddenAdvice {
  
      @Advice.OnMethodEnter(suppress = Throwable.class)
      public static void onEnter() {
        ThreadLocalContext requestInfo = ThreadLocalContext.get();
        if (requestInfo == null || !requestInfo.startSpan()) {
          return;
        }
  
        Context parentContext = Java8BytecodeBridge.currentContext();
        ServletRequestContext<HttpServletRequest> requestContext = requestInfo.getRequestContext();
  
        if (!helper().shouldStart(parentContext, requestContext)) {
          return;
        }
  
        Context context = helper().start(parentContext, requestContext);
        Scope scope = context.makeCurrent();
  
        requestInfo.setContext(context);
        requestInfo.setScope(scope);
  
        // Must be set here since Liberty RequestProcessors can use startAsync outside of servlet
        // scope.
        helper().setAsyncListenerResponse(context, requestInfo.getResponse());
  
        HttpServerResponseCustomizerHolder.getCustomizer()
            .customize(context, requestInfo.getResponse(), Servlet3Accessor.INSTANCE);
      }
    }
  }
  