@BeforeMethod public void init() { printerResult.getBuffer().setLength(0); respWrapper = new TagInjectionResponseWrapper(dummyRequest, dummyResponse, tracer, tagPrinter); respWrapper.proxyLinked(dummyProxy, linker); }
@Test public void testPlainTextNoInjection() throws IOException { respWrapper.getWriter().write(NON_HTML_TEST_CASE_A); assertThat(printerResult.toString(), equalTo(NON_HTML_TEST_CASE_A)); }
/** * Called when the headers are commited. At this point of time we have to decide whether we * force chunked encoding. */ private void commitHeaderData() { if (headerCommitted) { return; } headerCommitted = true; if (isNonHtmlContentTypeSet()) { // replicate the commands setting the content-length for (Runnable cmd : contentLengthHeaderModifications) { cmd.run(); } } else { setSessionIDCookie(); setTraceCorrelationInformation(); } }
/** * Proxy for {@link javax.servlet.ServletResponse#getWriter()}. * * @return the instrumented writer * @throws IOException * if an exception getting the original writer occurs. */ @ProxyMethod public PrintWriter getWriter() throws IOException { commitHeaderData(); if (wrappedWriter == null) { PrintWriter originalWriter = wrappedResponse.getWriter(); // avoid rewrapping or unnecessary wrapping if (isNonHtmlContentTypeSet() || (originalWriter instanceof TagInjectionPrintWriter)) { wrappedWriter = originalWriter; } else { wrappedWriter = new TagInjectionPrintWriter(originalWriter, tagToInject.printTags()); } } return wrappedWriter; }
/** * Lets InspectIT optionally instrument the given response by wrapping it. * * @param httpRequestObj * the request (an instance of javax.servlet.ServletRequest) * @param httpResponseObj * the response object (an instance of javax.servlet.ServletResponse) * @return the new response object to use, or the original one if it was not instrumented. */ public Object instrumentResponse(Object httpRequestObj, Object httpResponseObj) { try { if (configurationValid && WHttpServletResponse.isInstance(httpResponseObj)) { if (!linker.isProxyInstance(httpResponseObj, TagInjectionResponseWrapper.class)) { ClassLoader cl = httpResponseObj.getClass().getClassLoader(); TagInjectionResponseWrapper wrap = new TagInjectionResponseWrapper(httpRequestObj, httpResponseObj, tracer, scriptTags); Object proxy = linker.createProxy(TagInjectionResponseWrapper.class, wrap, cl); if (proxy == null) { return httpResponseObj; } else { return proxy; } } } } catch (Throwable e) { // NOPMD LOG.error("Error instrumenting response object.", e); } return httpResponseObj; // No instrumentation }
@SuppressWarnings("unchecked") @Test public void testBodyInjection() throws IOException { ArgumentCaptor<TagInjectionOutputStream> streamCaptor = ArgumentCaptor.forClass(TagInjectionOutputStream.class); respWrapper.getOutputStream(); verify(linker, times(1)).createProxy(any(Class.class), streamCaptor.capture(), any(ClassLoader.class)); TagInjectionOutputStream stream = streamCaptor.getValue(); byte[] bytes = HTML_TEST_CASE_B.getBytes(CHARACTER_ENCODING); int pos = 0; while (pos < bytes.length) { stream.write(bytes, pos, Math.min(3, bytes.length - pos)); pos += 3; // write 3 bytes at once } String result = new String(streamResult.toByteArray(), CHARACTER_ENCODING); assertThat(result, equalTo(HTML_TEST_CASE_B_REFERENCE)); }
/** * Proxy for {@link javax.servlet.ServletResponse#getOutputStream()}. * * @return the instrumented stream * @throws IOException * if an exception getting the original stream occurs. */ @ProxyMethod(returnType = "javax.servlet.ServletOutputStream") public OutputStream getOutputStream() throws IOException { commitHeaderData(); if (wrappedStream == null) { OutputStream originalStream = wrappedResponse.getOutputStream(); // avoid rewrapping or unnecessary wrapping if (isNonHtmlContentTypeSet() || linker.isProxyInstance(originalStream, TagInjectionOutputStream.class)) { wrappedStream = originalStream; } else { TagInjectionOutputStream resultStr = new TagInjectionOutputStream(originalStream, tagToInject.printTags()); resultStr.setEncoding(wrappedResponse.getCharacterEncoding()); ClassLoader cl = wrappedResponse.getWrappedElement().getClass().getClassLoader(); wrappedStream = (OutputStream) linker.createProxy(TagInjectionOutputStream.class, resultStr, cl); if (wrappedStream == null) { // fallback to the normal stream if it can not be linked wrappedStream = originalStream; } } } return wrappedStream; }
@SuppressWarnings("unchecked") @Test public void testPlainTextNoInjection() throws IOException { ArgumentCaptor<TagInjectionOutputStream> streamCaptor = ArgumentCaptor.forClass(TagInjectionOutputStream.class); respWrapper.getOutputStream(); verify(linker, times(1)).createProxy(any(Class.class), streamCaptor.capture(), any(ClassLoader.class)); TagInjectionOutputStream stream = streamCaptor.getValue(); byte[] bytes = NON_HTML_TEST_CASE_A.getBytes(CHARACTER_ENCODING); int pos = 0; while (pos < bytes.length) { stream.write(bytes, pos, Math.min(3, bytes.length - pos)); pos += 3; // write 3 bytes at once } String result = new String(streamResult.toByteArray(), CHARACTER_ENCODING); assertThat(result, equalTo(NON_HTML_TEST_CASE_A)); }
@SuppressWarnings("unchecked") @BeforeMethod public void init() { when(linker.isProxyInstance(any(Object.class), any(Class.class))).thenReturn(false); when(linker.createProxy(any(Class.class), any(TagInjectionOutputStream.class), any(ClassLoader.class))).then(new Answer<Object>() { @Override public Object answer(InvocationOnMock invocation) throws Throwable { TagInjectionOutputStream injectedStream = (TagInjectionOutputStream) invocation.getArguments()[1]; injectedStream.proxyLinked(dummyStreamProxy, linker); return dummyStreamProxy; } }); streamResult.reset(); respWrapper = new TagInjectionResponseWrapper(dummyRequest, dummyResponse, tracer, tagPrinter); respWrapper.proxyLinked(dummyProxy, linker); }
@Test public void testBodyInjection() throws IOException { respWrapper.getWriter().write(HTML_TEST_CASE_B); assertThat(printerResult.toString(), equalTo(HTML_TEST_CASE_B_REFERENCE)); }
@SuppressWarnings("unchecked") @Test public void testHeadInjection() throws IOException { ArgumentCaptor<TagInjectionOutputStream> streamCaptor = ArgumentCaptor.forClass(TagInjectionOutputStream.class); respWrapper.getOutputStream(); verify(linker, times(1)).createProxy(any(Class.class), streamCaptor.capture(), any(ClassLoader.class)); TagInjectionOutputStream stream = streamCaptor.getValue(); byte[] bytes = HTML_TEST_CASE_A.getBytes(CHARACTER_ENCODING); int pos = 0; while (pos < bytes.length) { stream.write(bytes, pos, Math.min(3, bytes.length - pos)); pos += 3; // write 3 bytes at once } String result = new String(streamResult.toByteArray(), CHARACTER_ENCODING); assertThat(result, equalTo(HTML_TEST_CASE_A_REFERENCE)); }
@Test public void testHeadInjection() throws IOException { respWrapper.getWriter().write(HTML_TEST_CASE_A); assertThat(printerResult.toString(), equalTo(HTML_TEST_CASE_A_REFERENCE)); }
@SuppressWarnings("unchecked") @Test public void testInvalidMarkupNoInjection() throws IOException { ArgumentCaptor<TagInjectionOutputStream> streamCaptor = ArgumentCaptor.forClass(TagInjectionOutputStream.class); respWrapper.getOutputStream(); verify(linker, times(1)).createProxy(any(Class.class), streamCaptor.capture(), any(ClassLoader.class)); TagInjectionOutputStream stream = streamCaptor.getValue(); byte[] bytes = NON_HTML_TEST_CASE_B.getBytes(CHARACTER_ENCODING); int pos = 0; while (pos < bytes.length) { stream.write(bytes, pos, Math.min(3, bytes.length - pos)); pos += 3; // write 3 bytes at once } String result = new String(streamResult.toByteArray(), CHARACTER_ENCODING); assertThat(result, equalTo(NON_HTML_TEST_CASE_B)); }
@Test public void testInvalidMarkupNoInjection() throws IOException { respWrapper.getWriter().write(NON_HTML_TEST_CASE_B); assertThat(printerResult.toString(), equalTo(NON_HTML_TEST_CASE_B)); }
@Test public void testPreventCookieOverwriting() throws IOException { hook = new EUMInstrumentationHook(linker, tracer, dataHandler, config, agentBuilder); String sessionId = "234587"; when(dummyRequest.getCookies()).thenReturn(new Cookie[] { new Cookie(JSAgentBuilder.SESSION_ID_COOKIE_NAME, sessionId) }); Object[] params = new Object[] { dummyRequest, dummyResponse }; boolean intercepted = null != hook.beforeBody(METHOD_ID, dummyServlet, params, ssc); respWrapper.getWriter(); assertThat(intercepted, equalTo(false)); verify(dummyResponse, never()).addCookie(any(Cookie.class)); }
@Test public void testCookieGeneration() throws IOException { hook = new EUMInstrumentationHook(linker, tracer, dataHandler, config, agentBuilder); Object[] params = new Object[] { dummyRequest, dummyResponse }; boolean intercepted = null != hook.beforeBody(METHOD_ID, dummyServlet, params, ssc); respWrapper.getWriter(); assertThat(intercepted, equalTo(false)); verify(dummyResponse, times(1)).addCookie(any(Cookie.class)); }
@Test public void testTraceCorrelationInfoSet() throws IOException { when(tagPrinter.clone()).thenReturn(tagPrinter); SpanContextImpl context = mock(SpanContextImpl.class); when(context.getTraceId()).thenReturn(1234L); String traceID = ConversionUtils.toHexString(1234L); when(tracer.getCurrentContext()).thenReturn(context); ArgumentCaptor<String> traceIdJSSetting = ArgumentCaptor.forClass(String.class); ArgumentCaptor<Cookie> cookieCapture = ArgumentCaptor.forClass(Cookie.class); respWrapper.getWriter(); verify(tagPrinter).setSetting(eq(JSAgentBuilder.TRACEID_CORRELATION_SETTING), traceIdJSSetting.capture()); verify(dummyResponse, atLeast(1)).addCookie(cookieCapture.capture()); boolean cookieFound = false; for (Cookie cookie : cookieCapture.getAllValues()) { if (cookie.getName().equals(JSAgentBuilder.TRACEID_CORRELATION_COOKIE_PREFIX + traceID)) { cookieFound = true; } } assertThat("Trace Correlation cookie not set", cookieFound); assertThat(traceIdJSSetting.getValue(), equalTo("\"" + traceID + "\"")); }