|  | 
|  | 1 | +/* | 
|  | 2 | + * Copyright The Hypertrace Authors | 
|  | 3 | + * | 
|  | 4 | + * Licensed under the Apache License, Version 2.0 (the "License"); | 
|  | 5 | + * you may not use this file except in compliance with the License. | 
|  | 6 | + * You may obtain a copy of the License at | 
|  | 7 | + * | 
|  | 8 | + *     http://www.apache.org/licenses/LICENSE-2.0 | 
|  | 9 | + * | 
|  | 10 | + * Unless required by applicable law or agreed to in writing, software | 
|  | 11 | + * distributed under the License is distributed on an "AS IS" BASIS, | 
|  | 12 | + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. | 
|  | 13 | + * See the License for the specific language governing permissions and | 
|  | 14 | + * limitations under the License. | 
|  | 15 | + */ | 
|  | 16 | + | 
|  | 17 | +package io.opentelemetry.javaagent.instrumentation.hypertrace.servlet.v5_0; | 
|  | 18 | + | 
|  | 19 | +import static io.opentelemetry.javaagent.extension.matcher.AgentElementMatchers.hasClassesNamed; | 
|  | 20 | +import static io.opentelemetry.javaagent.extension.matcher.AgentElementMatchers.hasSuperType; | 
|  | 21 | +import static net.bytebuddy.matcher.ElementMatchers.isPublic; | 
|  | 22 | +import static net.bytebuddy.matcher.ElementMatchers.named; | 
|  | 23 | +import static net.bytebuddy.matcher.ElementMatchers.namedOneOf; | 
|  | 24 | +import static net.bytebuddy.matcher.ElementMatchers.takesArgument; | 
|  | 25 | + | 
|  | 26 | +import io.opentelemetry.api.common.AttributeKey; | 
|  | 27 | +import io.opentelemetry.api.trace.Span; | 
|  | 28 | +import io.opentelemetry.instrumentation.api.util.VirtualField; | 
|  | 29 | +import io.opentelemetry.javaagent.bootstrap.Java8BytecodeBridge; | 
|  | 30 | +import io.opentelemetry.javaagent.extension.instrumentation.TypeInstrumentation; | 
|  | 31 | +import io.opentelemetry.javaagent.extension.instrumentation.TypeTransformer; | 
|  | 32 | +import jakarta.servlet.ServletInputStream; | 
|  | 33 | +import jakarta.servlet.ServletOutputStream; | 
|  | 34 | +import jakarta.servlet.ServletRequest; | 
|  | 35 | +import jakarta.servlet.ServletResponse; | 
|  | 36 | +import jakarta.servlet.http.HttpServletRequest; | 
|  | 37 | +import jakarta.servlet.http.HttpServletResponse; | 
|  | 38 | +import java.io.BufferedReader; | 
|  | 39 | +import java.io.IOException; | 
|  | 40 | +import java.io.PrintWriter; | 
|  | 41 | +import java.util.Collections; | 
|  | 42 | +import java.util.Enumeration; | 
|  | 43 | +import java.util.HashMap; | 
|  | 44 | +import java.util.Map; | 
|  | 45 | +import net.bytebuddy.asm.Advice; | 
|  | 46 | +import net.bytebuddy.description.type.TypeDescription; | 
|  | 47 | +import net.bytebuddy.matcher.ElementMatcher; | 
|  | 48 | +import org.hypertrace.agent.core.config.InstrumentationConfig; | 
|  | 49 | +import org.hypertrace.agent.core.filter.FilterResult; | 
|  | 50 | +import org.hypertrace.agent.core.instrumentation.HypertraceCallDepthThreadLocalMap; | 
|  | 51 | +import org.hypertrace.agent.core.instrumentation.HypertraceEvaluationException; | 
|  | 52 | +import org.hypertrace.agent.core.instrumentation.HypertraceSemanticAttributes; | 
|  | 53 | +import org.hypertrace.agent.core.instrumentation.SpanAndObjectPair; | 
|  | 54 | +import org.hypertrace.agent.core.instrumentation.buffer.BoundedByteArrayOutputStream; | 
|  | 55 | +import org.hypertrace.agent.core.instrumentation.buffer.BoundedCharArrayWriter; | 
|  | 56 | +import org.hypertrace.agent.core.instrumentation.buffer.ByteBufferSpanPair; | 
|  | 57 | +import org.hypertrace.agent.core.instrumentation.buffer.CharBufferSpanPair; | 
|  | 58 | +import org.hypertrace.agent.core.instrumentation.buffer.StringMapSpanPair; | 
|  | 59 | +import org.hypertrace.agent.core.instrumentation.utils.ContentTypeUtils; | 
|  | 60 | +import org.hypertrace.agent.filter.FilterRegistry; | 
|  | 61 | + | 
|  | 62 | +public class Servlet50AndFilterInstrumentation implements TypeInstrumentation { | 
|  | 63 | + | 
|  | 64 | +  @Override | 
|  | 65 | +  public ElementMatcher<ClassLoader> classLoaderOptimization() { | 
|  | 66 | +    return hasClassesNamed("jakarta.servlet.Filter"); | 
|  | 67 | +  } | 
|  | 68 | + | 
|  | 69 | +  @Override | 
|  | 70 | +  public ElementMatcher<TypeDescription> typeMatcher() { | 
|  | 71 | +    return hasSuperType(namedOneOf("jakarta.servlet.Filter", "jakarta.servlet.Servlet")); | 
|  | 72 | +  } | 
|  | 73 | + | 
|  | 74 | +  @Override | 
|  | 75 | +  public void transform(TypeTransformer transformer) { | 
|  | 76 | +    transformer.applyAdviceToMethod( | 
|  | 77 | +        namedOneOf("doFilter", "service") | 
|  | 78 | +            .and(takesArgument(0, named("jakarta.servlet.ServletRequest"))) | 
|  | 79 | +            .and(takesArgument(1, named("jakarta.servlet.ServletResponse"))) | 
|  | 80 | +            .and(isPublic()), | 
|  | 81 | +        Servlet50AndFilterInstrumentation.class.getName() + "$ServletAdvice"); | 
|  | 82 | +  } | 
|  | 83 | + | 
|  | 84 | +  public static class ServletAdvice { | 
|  | 85 | + | 
|  | 86 | +    @Advice.OnMethodEnter(suppress = Throwable.class, skipOn = Advice.OnNonDefaultValue.class) | 
|  | 87 | +    public static boolean start( | 
|  | 88 | +        @Advice.Argument(value = 0) ServletRequest request, | 
|  | 89 | +        @Advice.Argument(value = 1) ServletResponse response, | 
|  | 90 | +        @Advice.Local("currentSpan") Span currentSpan) { | 
|  | 91 | + | 
|  | 92 | +      int callDepth = | 
|  | 93 | +          HypertraceCallDepthThreadLocalMap.incrementCallDepth(Servlet50InstrumentationName.class); | 
|  | 94 | +      if (callDepth > 0) { | 
|  | 95 | +        return false; | 
|  | 96 | +      } | 
|  | 97 | +      if (!(request instanceof HttpServletRequest) || !(response instanceof HttpServletResponse)) { | 
|  | 98 | +        return false; | 
|  | 99 | +      } | 
|  | 100 | + | 
|  | 101 | +      HttpServletRequest httpRequest = (HttpServletRequest) request; | 
|  | 102 | +      HttpServletResponse httpResponse = (HttpServletResponse) response; | 
|  | 103 | +      currentSpan = Java8BytecodeBridge.currentSpan(); | 
|  | 104 | + | 
|  | 105 | +      InstrumentationConfig instrumentationConfig = InstrumentationConfig.ConfigProvider.get(); | 
|  | 106 | + | 
|  | 107 | +      Utils.addSessionId(currentSpan, httpRequest); | 
|  | 108 | + | 
|  | 109 | +      // set request headers | 
|  | 110 | +      Enumeration<String> headerNames = httpRequest.getHeaderNames(); | 
|  | 111 | +      Map<String, String> headers = new HashMap<>(); | 
|  | 112 | +      while (headerNames.hasMoreElements()) { | 
|  | 113 | +        String headerName = headerNames.nextElement(); | 
|  | 114 | +        String headerValue = httpRequest.getHeader(headerName); | 
|  | 115 | +        AttributeKey<String> attributeKey = | 
|  | 116 | +            HypertraceSemanticAttributes.httpRequestHeader(headerName); | 
|  | 117 | + | 
|  | 118 | +        if (instrumentationConfig.httpHeaders().request()) { | 
|  | 119 | +          currentSpan.setAttribute(attributeKey, headerValue); | 
|  | 120 | +        } | 
|  | 121 | +        headers.put(attributeKey.getKey(), headerValue); | 
|  | 122 | +      } | 
|  | 123 | + | 
|  | 124 | +      FilterResult filterResult = | 
|  | 125 | +          FilterRegistry.getFilter().evaluateRequestHeaders(currentSpan, headers); | 
|  | 126 | +      if (filterResult.shouldBlock()) { | 
|  | 127 | +        try { | 
|  | 128 | +          httpResponse.getWriter().write(filterResult.getBlockingMsg()); | 
|  | 129 | +        } catch (IOException ignored) { | 
|  | 130 | +        } | 
|  | 131 | +        httpResponse.setStatus(filterResult.getBlockingStatusCode()); | 
|  | 132 | +        // skip execution of the user code | 
|  | 133 | +        return true; | 
|  | 134 | +      } | 
|  | 135 | + | 
|  | 136 | +      if (instrumentationConfig.httpBody().request() | 
|  | 137 | +          && ContentTypeUtils.shouldCapture(httpRequest.getContentType())) { | 
|  | 138 | +        // The HttpServletRequest instrumentation uses this to | 
|  | 139 | +        // enable the instrumentation | 
|  | 140 | +        VirtualField.find(HttpServletRequest.class, SpanAndObjectPair.class) | 
|  | 141 | +            .set( | 
|  | 142 | +                httpRequest, | 
|  | 143 | +                new SpanAndObjectPair(currentSpan, Collections.unmodifiableMap(headers))); | 
|  | 144 | +      } | 
|  | 145 | +      return false; | 
|  | 146 | +    } | 
|  | 147 | + | 
|  | 148 | +    @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class) | 
|  | 149 | +    public static void exit( | 
|  | 150 | +        @Advice.Argument(0) ServletRequest request, | 
|  | 151 | +        @Advice.Argument(1) ServletResponse response, | 
|  | 152 | +        @Advice.Thrown(readOnly = false) Throwable throwable, | 
|  | 153 | +        @Advice.Local("currentSpan") Span currentSpan) { | 
|  | 154 | +      int callDepth = | 
|  | 155 | +          HypertraceCallDepthThreadLocalMap.decrementCallDepth(Servlet50InstrumentationName.class); | 
|  | 156 | +      if (callDepth > 0) { | 
|  | 157 | +        return; | 
|  | 158 | +      } | 
|  | 159 | +      // we are in the most outermost level of Servlet instrumentation | 
|  | 160 | + | 
|  | 161 | +      if (!(request instanceof HttpServletRequest) || !(response instanceof HttpServletResponse)) { | 
|  | 162 | +        return; | 
|  | 163 | +      } | 
|  | 164 | + | 
|  | 165 | +      HttpServletResponse httpResponse = (HttpServletResponse) response; | 
|  | 166 | +      HttpServletRequest httpRequest = (HttpServletRequest) request; | 
|  | 167 | +      InstrumentationConfig instrumentationConfig = InstrumentationConfig.ConfigProvider.get(); | 
|  | 168 | + | 
|  | 169 | +      try { | 
|  | 170 | +        // response context to capture body and clear the context | 
|  | 171 | +        VirtualField<HttpServletResponse, SpanAndObjectPair> responseContextStore = | 
|  | 172 | +            VirtualField.find(HttpServletResponse.class, SpanAndObjectPair.class); | 
|  | 173 | +        VirtualField<ServletOutputStream, BoundedByteArrayOutputStream> outputStreamContextStore = | 
|  | 174 | +            VirtualField.find(ServletOutputStream.class, BoundedByteArrayOutputStream.class); | 
|  | 175 | +        VirtualField<PrintWriter, BoundedCharArrayWriter> writerContextStore = | 
|  | 176 | +            VirtualField.find(PrintWriter.class, BoundedCharArrayWriter.class); | 
|  | 177 | + | 
|  | 178 | +        // request context to clear body buffer | 
|  | 179 | +        VirtualField<HttpServletRequest, SpanAndObjectPair> requestContextStore = | 
|  | 180 | +            VirtualField.find(HttpServletRequest.class, SpanAndObjectPair.class); | 
|  | 181 | +        VirtualField<ServletInputStream, ByteBufferSpanPair> inputStreamContextStore = | 
|  | 182 | +            VirtualField.find(ServletInputStream.class, ByteBufferSpanPair.class); | 
|  | 183 | +        VirtualField<BufferedReader, CharBufferSpanPair> readerContextStore = | 
|  | 184 | +            VirtualField.find(BufferedReader.class, CharBufferSpanPair.class); | 
|  | 185 | +        VirtualField<HttpServletRequest, StringMapSpanPair> urlEncodedMapContextStore = | 
|  | 186 | +            VirtualField.find(HttpServletRequest.class, StringMapSpanPair.class); | 
|  | 187 | + | 
|  | 188 | +        if (!request.isAsyncStarted()) { | 
|  | 189 | +          if (instrumentationConfig.httpHeaders().response()) { | 
|  | 190 | +            for (String headerName : httpResponse.getHeaderNames()) { | 
|  | 191 | +              String headerValue = httpResponse.getHeader(headerName); | 
|  | 192 | +              currentSpan.setAttribute( | 
|  | 193 | +                  HypertraceSemanticAttributes.httpResponseHeader(headerName), headerValue); | 
|  | 194 | +            } | 
|  | 195 | +          } | 
|  | 196 | + | 
|  | 197 | +          // capture response body | 
|  | 198 | +          if (instrumentationConfig.httpBody().response() | 
|  | 199 | +              && ContentTypeUtils.shouldCapture(httpResponse.getContentType())) { | 
|  | 200 | +            Utils.captureResponseBody( | 
|  | 201 | +                currentSpan, | 
|  | 202 | +                httpResponse, | 
|  | 203 | +                responseContextStore, | 
|  | 204 | +                outputStreamContextStore, | 
|  | 205 | +                writerContextStore); | 
|  | 206 | +          } | 
|  | 207 | + | 
|  | 208 | +          // remove request body buffers from context stores, otherwise they might get reused | 
|  | 209 | +          if (instrumentationConfig.httpBody().request() | 
|  | 210 | +              && ContentTypeUtils.shouldCapture(httpRequest.getContentType())) { | 
|  | 211 | +            Utils.resetRequestBodyBuffers( | 
|  | 212 | +                httpRequest, | 
|  | 213 | +                requestContextStore, | 
|  | 214 | +                inputStreamContextStore, | 
|  | 215 | +                readerContextStore, | 
|  | 216 | +                urlEncodedMapContextStore); | 
|  | 217 | +          } | 
|  | 218 | +        } | 
|  | 219 | +      } finally { | 
|  | 220 | +        Throwable tmp = throwable; | 
|  | 221 | +        while (tmp != null) { // loop in case our exception is nested (eg. springframework) | 
|  | 222 | +          if (tmp instanceof HypertraceEvaluationException) { | 
|  | 223 | +            FilterResult filterResult = ((HypertraceEvaluationException) tmp).getFilterResult(); | 
|  | 224 | +            try { | 
|  | 225 | +              httpResponse.getWriter().write(filterResult.getBlockingMsg()); | 
|  | 226 | +            } catch (IOException ignored) { | 
|  | 227 | +            } | 
|  | 228 | +            httpResponse.setStatus(filterResult.getBlockingStatusCode()); | 
|  | 229 | +            // bytebuddy treats the reassignment of this variable to null as an instruction to | 
|  | 230 | +            // suppress this exception, which is what we want | 
|  | 231 | +            throwable = null; | 
|  | 232 | +            break; | 
|  | 233 | +          } | 
|  | 234 | +          tmp = tmp.getCause(); | 
|  | 235 | +        } | 
|  | 236 | +      } | 
|  | 237 | +    } | 
|  | 238 | +  } | 
|  | 239 | +} | 
0 commit comments