Skip to content

Commit

Permalink
Merge pull request #5 from siyuniums/content_length
Browse files Browse the repository at this point in the history 
Content length
  • Loading branch information
@siyuniums
siyuniums committed Aug 5, 2022
2 parents 02f6ae3 + ac029ae commit a1b42a8
Show file tree
Hide file tree
Showing 10 changed files with 145 additions and 98 deletions.
32 changes: 23 additions & 9 deletions ...t/java/io/opentelemetry/javaagent/instrumentation/servlet/v3_0/snippet/InjectionTest.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -7,8 +7,10 @@

import static io.opentelemetry.javaagent.instrumentation.servlet.v3_0.snippet.TestUtil.readFile;
import static org.assertj.core.api.AssertionsForClassTypes.assertThat;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.when;

import io.opentelemetry.javaagent.bootstrap.servlet.SnippetHolder;
import io.opentelemetry.javaagent.bootstrap.servlet.ExperimentalSnippetHolder;
import java.io.IOException;
import java.io.StringWriter;
import java.nio.charset.StandardCharsets;
Expand All @@ -21,13 +23,16 @@ class InjectionTest {
@Test
void testInjectionForStringContainHeadTag() throws IOException {
String testSnippet = "\n <script type=\"text/javascript\"> Test </script>";
SnippetHolder.setSnippet(testSnippet);
ExperimentalSnippetHolder.setSnippet(testSnippet);
// read the originalFile
String original = readFile("staticHtmlOrigin.html");
// read the correct answer
String correct = readFile("staticHtmlAfter.html");
byte[] originalBytes = original.getBytes(StandardCharsets.UTF_8);
InjectionState obj = new InjectionState(StandardCharsets.UTF_8.name());
SnippetInjectingResponseWrapper response = mock(SnippetInjectingResponseWrapper.class);
when(response.isCommitted()).thenReturn(false);
when(response.getCharacterEncoding()).thenReturn(StandardCharsets.UTF_8.name());
InjectionState obj = new InjectionState(response);

StringWriter writer = new StringWriter();

Expand All @@ -54,13 +59,16 @@ public void write(int b) throws IOException {
@Disabled
void testInjectionForChinese() throws IOException {
String testSnippet = "\n <script type=\"text/javascript\"> Test </script>";
SnippetHolder.setSnippet(testSnippet);
ExperimentalSnippetHolder.setSnippet(testSnippet);
// read the originalFile
String original = readFile("staticHtmlChineseOrigin.html");
// read the correct answer
String correct = readFile("staticHtmlChineseAfter.html");
byte[] originalBytes = original.getBytes(StandardCharsets.UTF_8);
InjectionState obj = new InjectionState(StandardCharsets.UTF_8.name());
SnippetInjectingResponseWrapper response = mock(SnippetInjectingResponseWrapper.class);
when(response.isCommitted()).thenReturn(false);
when(response.getCharacterEncoding()).thenReturn(StandardCharsets.UTF_8.name());
InjectionState obj = new InjectionState(response);

StringWriter writer = new StringWriter();

Expand All @@ -86,12 +94,15 @@ public void write(int b) throws IOException {
@Test
void testInjectionForStringWithoutHeadTag() throws IOException {
String testSnippet = "\n <script type=\"text/javascript\"> Test </script>";
SnippetHolder.setSnippet(testSnippet);
ExperimentalSnippetHolder.setSnippet(testSnippet);
// read the originalFile
String original = readFile("htmlWithoutHeadTag.html");

byte[] originalBytes = original.getBytes(StandardCharsets.UTF_8);
InjectionState obj = new InjectionState(StandardCharsets.UTF_8.name());
SnippetInjectingResponseWrapper response = mock(SnippetInjectingResponseWrapper.class);
when(response.isCommitted()).thenReturn(false);
when(response.getCharacterEncoding()).thenReturn(StandardCharsets.UTF_8.name());
InjectionState obj = new InjectionState(response);
StringWriter writer = new StringWriter();

ServletOutputStream sp =
Expand All @@ -115,11 +126,14 @@ public void write(int b) throws IOException {
@Test
void testHalfHeadTag() throws IOException {
String testSnippet = "\n <script type=\"text/javascript\"> Test </script>";
SnippetHolder.setSnippet(testSnippet);
ExperimentalSnippetHolder.setSnippet(testSnippet);
// read the original string
String originalFirstPart = "<!DOCTYPE html>\n" + "<html lang=\"en\">\n" + "<he";
byte[] originalFirstPartBytes = originalFirstPart.getBytes(StandardCharsets.UTF_8);
InjectionState obj = new InjectionState(StandardCharsets.UTF_8.name());
SnippetInjectingResponseWrapper response = mock(SnippetInjectingResponseWrapper.class);
when(response.isCommitted()).thenReturn(false);
when(response.getCharacterEncoding()).thenReturn(StandardCharsets.UTF_8.name());
InjectionState obj = new InjectionState(response);
StringWriter writer = new StringWriter();

ServletOutputStream sp =
Expand Down
14 changes: 7 additions & 7 deletions ...y/javaagent/instrumentation/servlet/v3_0/snippet/SnippetInjectingResponseWrapperTest.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -10,7 +10,7 @@
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.when;

import io.opentelemetry.javaagent.bootstrap.servlet.SnippetHolder;
import io.opentelemetry.javaagent.bootstrap.servlet.ExperimentalSnippetHolder;
import java.io.IOException;
import java.io.PrintWriter;
import java.io.StringWriter;
Expand All @@ -33,7 +33,7 @@ void testInjectToTextHtml() throws IOException {
when(response.containsHeader("content-type")).thenReturn(true);
StringWriter writer = new StringWriter();
when(response.getWriter()).thenReturn(new PrintWriter(writer));
SnippetHolder.setSnippet("\n <script type=\"text/javascript\"> Test </script>");
ExperimentalSnippetHolder.setSnippet("\n <script type=\"text/javascript\"> Test </script>");
SnippetInjectingResponseWrapper responseWrapper = new SnippetInjectingResponseWrapper(response);
responseWrapper.getWriter().write(original);
responseWrapper.getWriter().flush();
Expand All @@ -58,7 +58,7 @@ void testInjectToChineseTextHtml() throws IOException {

StringWriter writer = new StringWriter();
when(response.getWriter()).thenReturn(new PrintWriter(writer));
SnippetHolder.setSnippet("\n <script type=\"text/javascript\"> Test </script>");
ExperimentalSnippetHolder.setSnippet("\n <script type=\"text/javascript\"> Test </script>");
SnippetInjectingResponseWrapper responseWrapper = new SnippetInjectingResponseWrapper(response);
responseWrapper.getWriter().write(original);
responseWrapper.getWriter().flush();
Expand All @@ -84,7 +84,7 @@ void shouldNotInjectToTextHtml() throws IOException {
when(response.containsHeader("content-type")).thenReturn(true);

when(response.getWriter()).thenReturn(new PrintWriter(writer, true));
SnippetHolder.setSnippet("\n <script type=\"text/javascript\"> Test </script>");
ExperimentalSnippetHolder.setSnippet("\n <script type=\"text/javascript\"> Test </script>");

SnippetInjectingResponseWrapper responseWrapper = new SnippetInjectingResponseWrapper(response);
responseWrapper.getWriter().write(original);
Expand All @@ -110,7 +110,7 @@ void testWriteInt() throws IOException {
StringWriter writer = new StringWriter();
// StringWriter correctWriter = new StringWriter();
when(response.getWriter()).thenReturn(new PrintWriter(writer));
SnippetHolder.setSnippet("\n <script type=\"text/javascript\"> Test </script>");
ExperimentalSnippetHolder.setSnippet("\n <script type=\"text/javascript\"> Test </script>");
SnippetInjectingResponseWrapper responseWrapper = new SnippetInjectingResponseWrapper(response);
byte[] originalBytes = original.getBytes(Charset.defaultCharset().name());
// byte[] correctBytes = correct.getBytes(UTF_8);
Expand Down Expand Up @@ -142,7 +142,7 @@ void testWriteCharArray() throws IOException {

StringWriter writer = new StringWriter();
when(response.getWriter()).thenReturn(new PrintWriter(writer));
SnippetHolder.setSnippet("\n <script type=\"text/javascript\"> Test </script>");
ExperimentalSnippetHolder.setSnippet("\n <script type=\"text/javascript\"> Test </script>");
SnippetInjectingResponseWrapper responseWrapper = new SnippetInjectingResponseWrapper(response);
char[] originalChars = original.toCharArray();
responseWrapper.getWriter().write(originalChars, 0, originalChars.length);
Expand Down Expand Up @@ -171,7 +171,7 @@ void testWriteWithOffset() throws IOException {

StringWriter writer = new StringWriter();
when(response.getWriter()).thenReturn(new PrintWriter(writer));
SnippetHolder.setSnippet("\n <script type=\"text/javascript\"> Test </script>");
ExperimentalSnippetHolder.setSnippet("\n <script type=\"text/javascript\"> Test </script>");
SnippetInjectingResponseWrapper responseWrapper = new SnippetInjectingResponseWrapper(response);

responseWrapper
Expand Down
4 changes: 2 additions & 2 deletions ...src/main/java/io/opentelemetry/javaagent/instrumentation/servlet/v3_0/Servlet3Advice.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -13,8 +13,8 @@
import io.opentelemetry.javaagent.bootstrap.CallDepth;
import io.opentelemetry.javaagent.bootstrap.Java8BytecodeBridge;
import io.opentelemetry.javaagent.bootstrap.servlet.AppServerBridge;
import io.opentelemetry.javaagent.bootstrap.servlet.ExperimentalSnippetHolder;
import io.opentelemetry.javaagent.bootstrap.servlet.MappingResolver;
import io.opentelemetry.javaagent.bootstrap.servlet.SnippetHolder;
import io.opentelemetry.javaagent.instrumentation.servlet.ServletRequestContext;
import io.opentelemetry.javaagent.instrumentation.servlet.v3_0.snippet.SnippetInjectingResponseWrapper;
import javax.servlet.Servlet;
Expand Down Expand Up @@ -43,7 +43,7 @@ public static void onEnter(
}
HttpServletRequest httpServletRequest = (HttpServletRequest) request;

if (!SnippetHolder.getSnippet().isEmpty()
if (!ExperimentalSnippetHolder.getSnippet().isEmpty()
&& !((HttpServletResponse) response).containsHeader(FAKE_SNIPPET_HEADER)) {
response = new SnippetInjectingResponseWrapper((HttpServletResponse) response);
}
Expand Down
26 changes: 8 additions & 18 deletions .../java/io/opentelemetry/javaagent/instrumentation/servlet/v3_0/snippet/InjectionState.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -5,48 +5,38 @@

package io.opentelemetry.javaagent.instrumentation.servlet.v3_0.snippet;

import javax.annotation.Nullable;

public class InjectionState {
private static final int ALREADY_INJECTED_FAKE_VALUE = -1;
private static final int HEAD_TAG_WRITTEN_FAKE_VALUE = -1;
private static final int HEAD_TAG_LENGTH = "<head>".length();

private final String characterEncoding;
@Nullable private final SnippetInjectingResponseWrapper wrapper;
private final SnippetInjectingResponseWrapper wrapper;
private int headTagBytesSeen = 0;

public InjectionState(SnippetInjectingResponseWrapper wrapper) {
this.wrapper = wrapper;
this.characterEncoding = wrapper.getCharacterEncoding();
}

public InjectionState(String characterEncoding) {
this.characterEncoding = characterEncoding;
this.wrapper = null;
}

public int getHeadTagBytesSeen() {
return headTagBytesSeen;
}

public String getCharacterEncoding() {
return characterEncoding;
return wrapper.getCharacterEncoding();
}

public void setAlreadyInjected() {
headTagBytesSeen = ALREADY_INJECTED_FAKE_VALUE;
public void setHeadTagWritten() {
headTagBytesSeen = HEAD_TAG_WRITTEN_FAKE_VALUE;
}

public boolean isAlreadyInjected() {
return headTagBytesSeen == ALREADY_INJECTED_FAKE_VALUE;
public boolean isHeadTagWritten() {
return headTagBytesSeen == HEAD_TAG_WRITTEN_FAKE_VALUE;
}

/**
* Returns true when the byte is the last character of "<head>" and now is the right time to
* inject. Otherwise, returns false.
*/
public boolean processByte(int b) {
if (isAlreadyInjected()) {
if (isHeadTagWritten()) {
return false;
}
if (inHeadTag(b)) {
Expand Down
43 changes: 29 additions & 14 deletions ...ry/javaagent/instrumentation/servlet/v3_0/snippet/ServletOutputStreamInjectionHelper.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -7,7 +7,7 @@

import static java.util.logging.Level.FINE;

import io.opentelemetry.javaagent.bootstrap.servlet.SnippetHolder;
import io.opentelemetry.javaagent.bootstrap.servlet.ExperimentalSnippetHolder;
import java.io.IOException;
import java.io.UnsupportedEncodingException;
import java.util.logging.Logger;
Expand All @@ -25,48 +25,63 @@ public class ServletOutputStreamInjectionHelper {
public static boolean handleWrite(
byte[] original, int off, int length, InjectionState state, ServletOutputStream out)
throws IOException {
if (state.isAlreadyInjected()) {
if (state.isHeadTagWritten()) {
return false;
}
int i;
boolean shouldInject = false;
boolean endOfHeadTagFound = false;
for (i = off; i < length && i - off < length; i++) {
if (state.processByte(original[i])) {
shouldInject = true;
endOfHeadTagFound = true;
break;
}
}
if (!shouldInject) {
if (!endOfHeadTagFound) {
return false;
}
state.setAlreadyInjected(); // set before write to avoid recursive loop
out.write(original, off, i + 1);
state.setHeadTagWritten(); // set before write to avoid recursive loop
if (state.getWrapper().isNotSafeToInject()) {
return false;
}
byte[] snippetBytes;
try {
byte[] snippetBytes = SnippetHolder.getSnippetBytes(state.getCharacterEncoding());
out.write(snippetBytes);
snippetBytes = ExperimentalSnippetHolder.getSnippetBytes(state.getCharacterEncoding());
} catch (UnsupportedEncodingException e) {
logger.log(FINE, "UnsupportedEncodingException", e);
return false;
}
// updating Content-Length before any further writing in case that writing triggers a flush
state.getWrapper().updateContentLengthIfPreviouslySet();
out.write(original, off, i + 1);
out.write(snippetBytes);
out.write(original, i + 1, length - i - 1);
return true;
}

public static boolean handleWrite(InjectionState state, ServletOutputStream out, int b)
throws IOException {
if (state.isAlreadyInjected()) {
if (state.isHeadTagWritten()) {
return false;
}
if (!state.processByte(b)) {
return false;
}
state.setAlreadyInjected(); // set before write to avoid recursive loop
out.write(b);
state.setHeadTagWritten(); // set before write to avoid recursive loop

if (state.getWrapper().isNotSafeToInject()) {
return false;
}
byte[] snippetBytes;
try {
byte[] snippetBytes = SnippetHolder.getSnippetBytes(state.getCharacterEncoding());
out.write(snippetBytes);
snippetBytes = ExperimentalSnippetHolder.getSnippetBytes(state.getCharacterEncoding());
} catch (UnsupportedEncodingException e) {
logger.log(FINE, "UnsupportedEncodingException", e);
return false;
}
state.getWrapper().updateContentLengthIfPreviouslySet();
out.write(b);

out.write(snippetBytes);
return true;
}
}
23 changes: 15 additions & 8 deletions ...telemetry/javaagent/instrumentation/servlet/v3_0/snippet/SnippetInjectingPrintWriter.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -11,9 +11,10 @@ public class SnippetInjectingPrintWriter extends PrintWriter {
private final String snippet;
private final InjectionState state;

public SnippetInjectingPrintWriter(PrintWriter writer, String snippet, String characterEncoding) {
public SnippetInjectingPrintWriter(
PrintWriter writer, String snippet, SnippetInjectingResponseWrapper wrapper) {
super(writer);
state = new InjectionState(characterEncoding);
state = new InjectionState(wrapper);
this.snippet = snippet;
}

Expand All @@ -26,14 +27,20 @@ public void write(String s, int off, int len) {

@Override
public void write(int b) {
boolean shouldInject = state.processByte(b);
super.write(b);
if (shouldInject) {
// set before write to avoid recursive loop since super.write(String) may delegate back to
// write(int)
state.setAlreadyInjected();
super.write(snippet);
if (state.isHeadTagWritten()) {
return;
}
boolean endOfHeadTagFound = state.processByte(b);
if (!endOfHeadTagFound) {
return;
}
state.setHeadTagWritten(); // set before write to avoid recursive loop
if (state.getWrapper().isNotSafeToInject()) {
return;
}
state.getWrapper().updateContentLengthIfPreviouslySet();
super.write(snippet);
}

@Override
Expand Down
Loading

0 comments on commit a1b42a8

Please sign in to comment.