Skip to content

Commit

Permalink
Additional EmbraceUrlConnectionDelegate unit tests and bug fixes
Browse files Browse the repository at this point in the history
  • Loading branch information
bidetofevil committed Nov 7, 2023
1 parent eded592 commit b16e8ba
Show file tree
Hide file tree
Showing 4 changed files with 517 additions and 144 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -73,8 +73,6 @@ interface EmbraceHttpUrlConnection {
@Nullable
InputStream getErrorStream();

boolean shouldInterceptHeaderRetrieval(@Nullable String key);

@Nullable
String getHeaderField(int n);

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -28,7 +28,7 @@ public String getHeaderByName(@NonNull String name) {
public String getOverriddenURL(@NonNull String pathOverride) {
try {
return new URL(connection.getURL().getProtocol(), connection.getURL().getHost(),
connection.getURL().getPort(), pathOverride).toString();
connection.getURL().getPort(), pathOverride + "?" + connection.getURL().getQuery()).toString();
} catch (MalformedURLException e) {
InternalStaticEmbraceLogger.logError("Failed to override path of " +
connection.getURL() + " with " + pathOverride);
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -63,12 +63,12 @@ class EmbraceUrlConnectionDelegate<T extends HttpURLConnection> implements Embra
/**
* The content encoding HTTP header.
*/
private static final String CONTENT_ENCODING = "Content-Encoding";
static final String CONTENT_ENCODING = "Content-Encoding";

/**
* The content length HTTP header.
*/
private static final String CONTENT_LENGTH = "Content-Length";
static final String CONTENT_LENGTH = "Content-Length";

/**
* Reference to the wrapped connection.
Expand Down Expand Up @@ -140,8 +140,7 @@ class EmbraceUrlConnectionDelegate<T extends HttpURLConnection> implements Embra
@Nullable
private volatile String traceparent = null;

@Nullable
private volatile byte[] responseBody = null;
private final boolean isSDKStarted;

/**
* Wraps an existing {@link HttpURLConnection} with the Embrace network logic.
Expand All @@ -160,6 +159,7 @@ public EmbraceUrlConnectionDelegate(@NonNull T connection, boolean enableWrapIoS
this.embrace = embrace;
this.createdTime = embrace.getInternalInterface().getSdkCurrentTime();
this.callId = UUID.randomUUID().toString();
this.isSDKStarted = embrace.isStarted();
}

@Override
Expand All @@ -169,21 +169,23 @@ public void addRequestProperty(@NonNull String key, @Nullable String value) {

@Override
public void connect() throws IOException {
identifyTraceId();
try {
if (embrace.getInternalInterface().isNetworkSpanForwardingEnabled()) {
traceparent = connection.getRequestProperty(TRACEPARENT_HEADER_NAME);
if (isSDKStarted) {
identifyTraceId();
try {
if (embrace.getInternalInterface().isNetworkSpanForwardingEnabled()) {
traceparent = connection.getRequestProperty(TRACEPARENT_HEADER_NAME);
}
} catch (Exception e) {
// Ignore traceparent if there was a problem obtaining it
}
} catch (Exception e) {
// Ignore traceparent if there was a problem obtaining it
}
this.connection.connect();
}

@Override
public void disconnect() {
// The network call must be logged before we close the transport
internalLogNetworkCall(this.createdTime);
internalLogNetworkCall(createdTime);
this.connection.disconnect();
}

Expand Down Expand Up @@ -287,8 +289,7 @@ public InputStream getErrorStream() {
return getWrappedInputStream(this.connection.getErrorStream());
}

@Override
public boolean shouldInterceptHeaderRetrieval(@Nullable String key) {
private boolean shouldInterceptHeaderRetrieval(@Nullable String key) {
return shouldUncompressGzip() && key != null && (key.equalsIgnoreCase(CONTENT_ENCODING) || key.equalsIgnoreCase(CONTENT_LENGTH));
}

Expand Down Expand Up @@ -358,7 +359,7 @@ public long getHeaderFieldLong(@NonNull String name, long defaultValue) {
@Nullable
public Map<String, List<String>> getHeaderFields() {
final long startTime = embrace.getInternalInterface().getSdkCurrentTime();
cacheResponseData();
cacheNetworkCallData();
internalLogNetworkCall(startTime);
return headerFields.get();
}
Expand All @@ -377,7 +378,7 @@ private <R> R retrieveHeaderField(@Nullable String name,
}

R result = action.invoke();
cacheResponseData();
cacheNetworkCallData();
internalLogNetworkCall(startTime);
return result;
}
Expand Down Expand Up @@ -473,7 +474,7 @@ public String getRequestProperty(@NonNull String key) {
public int getResponseCode() {
identifyTraceId();
long startTime = embrace.getInternalInterface().getSdkCurrentTime();
cacheResponseData();
cacheNetworkCallData();
internalLogNetworkCall(startTime);
return responseCode.get();
}
Expand All @@ -484,7 +485,7 @@ public String getResponseMessage() throws IOException {
identifyTraceId();
long startTime = embrace.getInternalInterface().getSdkCurrentTime();
String responseMsg = this.connection.getResponseMessage();
cacheResponseData();
cacheNetworkCallData();
internalLogNetworkCall(startTime);
return responseMsg;
}
Expand Down Expand Up @@ -547,7 +548,9 @@ public boolean usingProxy() {
* ignored.
*/
synchronized void internalLogNetworkCall(long startTime) {
internalLogNetworkCall(startTime, embrace.getInternalInterface().getSdkCurrentTime(), false, null);
if (isSDKStarted) {
internalLogNetworkCall(startTime, embrace.getInternalInterface().getSdkCurrentTime(), false, null);
}
}

/**
Expand Down Expand Up @@ -658,8 +661,7 @@ private CountingInputStreamWithCallback countingInputStream(InputStream inputStr
hasNetworkCaptureRules(),
(bytesCount, responseBody) -> {
if (this.startTime != null && this.endTime != null) {
this.responseBody = responseBody;
cacheResponseData();
cacheNetworkCallData(responseBody);
internalLogNetworkCall(
this.startTime,
this.endTime,
Expand Down Expand Up @@ -695,7 +697,7 @@ private boolean shouldUncompressGzip() {
}

private void identifyTraceId() {
if (traceId == null) {
if (isSDKStarted && traceId == null) {
try {
traceId = getRequestProperty(embrace.getTraceIdHeader());
} catch (Exception e) {
Expand Down Expand Up @@ -808,13 +810,13 @@ private InputStream getWrappedInputStream(InputStream connectionInputStream) {
countingInputStream(new BufferedInputStream(connectionInputStream)) : connectionInputStream;
}

cacheResponseData();
cacheNetworkCallData();
internalLogNetworkCall(startTime);
return in;
}

private boolean hasNetworkCaptureRules() {
if (this.connection.getURL() == null) {
if (!isSDKStarted || this.connection.getURL() == null) {
return false;
}
String url = this.connection.getURL().toString();
Expand All @@ -823,11 +825,17 @@ private boolean hasNetworkCaptureRules() {
return embrace.getInternalInterface().shouldCaptureNetworkBody(url, method);
}

private void cacheNetworkCallData() {
if (isSDKStarted) {
cacheNetworkCallData(null);
}
}

/**
* Cache values from response at the first point of availability so that we won't try to retrieve these values when the response
* is not available.
*/
private void cacheResponseData() {
private void cacheNetworkCallData(@Nullable byte[] responseBody) {
if (headerFields.get() == null) {
synchronized (headerFields) {
if (headerFields.get() == null) {
Expand Down Expand Up @@ -874,26 +882,42 @@ private void cacheResponseData() {
}
}

if (shouldCaptureNetworkData() && networkCaptureData.get() == null) {
if (shouldCaptureNetworkData()) {
// If we don't have network capture rules, it's unnecessary to save these values
synchronized (networkCaptureData) {
if (shouldCaptureNetworkData() && networkCaptureData.get() == null) {
if (shouldCaptureNetworkData()) {
try {
Map<String, String> requestHeaders = this.requestHeaders;
String requestQueryParams = connection.getURL().getQuery();
byte[] requestBody = this.outputStream != null ? this.outputStream.getRequestBody() : null;
Map<String, String> responseHeaders = getProcessedHeaders(headerFields.get());

networkCaptureData.set(
new NetworkCaptureData(
requestHeaders,
requestQueryParams,
requestBody,
responseHeaders,
responseBody,
null
)
);
NetworkCaptureData existingData = networkCaptureData.get();
if (existingData == null) {
Map<String, String> requestHeaders = this.requestHeaders;
String requestQueryParams = connection.getURL().getQuery();
byte[] requestBody = this.outputStream != null ? this.outputStream.getRequestBody() : null;
Map<String, String> responseHeaders = getProcessedHeaders(headerFields.get());

networkCaptureData.set(
new NetworkCaptureData(
requestHeaders,
requestQueryParams,
requestBody,
responseHeaders,
responseBody,
null
)
);
} else if (responseBody != null) {
// Update the response body field in the cached networkCaptureData object if a subsequent call
// is update to update the network logging with this data.
networkCaptureData.set(
new NetworkCaptureData(
existingData.getRequestHeaders(),
existingData.getRequestQueryParams(),
existingData.getCapturedRequestBody(),
existingData.getResponseHeaders(),
responseBody,
null
)
);
}
} catch (Exception e) {
lastConnectionAccessException = e;
}
Expand All @@ -903,6 +927,7 @@ private void cacheResponseData() {
}

private boolean shouldCaptureNetworkData() {
return hasNetworkCaptureRules() && (enableWrapIoStreams || inputStreamAccessException != null);
return (hasNetworkCaptureRules() && (enableWrapIoStreams || inputStreamAccessException != null)) &&
(networkCaptureData.get() == null || networkCaptureData.get().getCapturedResponseBody() == null);
}
}
Loading

0 comments on commit b16e8ba

Please sign in to comment.