Skip to content

Commit

Permalink
Save reference
Browse files Browse the repository at this point in the history
  • Loading branch information
moarychan committed May 21, 2024
1 parent 79e0db9 commit 9b47ce5
Show file tree
Hide file tree
Showing 7 changed files with 345 additions and 36 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,8 @@
package com.azure.identity.extensions.implementation.credential.provider;

import com.azure.core.credential.TokenCredential;
import com.azure.core.http.policy.HttpLogOptions;
import com.azure.core.util.HttpClientOptions;
import com.azure.core.util.logging.ClientLogger;
import com.azure.identity.ClientCertificateCredentialBuilder;
import com.azure.identity.ClientSecretCredentialBuilder;
Expand Down Expand Up @@ -107,9 +109,7 @@ && hasText(options.getPassword())) {

LOGGER.verbose("Use the dac builder.");
ExecutorService executorService = Executors.newFixedThreadPool(1, new ThreadFactory() {

private int count = 0;

@Override
public Thread newThread(Runnable runnable) {
return new Thread(runnable, "az-id-test-" + count++);
Expand All @@ -120,7 +120,7 @@ public Thread newThread(Runnable runnable) {
.authorityHost(authorityHost)
.tenantId(tenantId)
.managedIdentityClientId(clientId)
.executorService(executorService)
.executorService(executorService).enableAccountIdentifierLogging().addPolicy(new HttpDebugLoggingPolicy())
.build();
}

Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,264 @@
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.

package com.azure.identity.extensions.implementation.credential.provider;

import com.azure.core.http.ContentType;
import com.azure.core.http.HttpHeader;
import com.azure.core.http.HttpHeaders;
import com.azure.core.http.HttpPipelineCallContext;
import com.azure.core.http.HttpPipelineNextPolicy;
import com.azure.core.http.HttpRequest;
import com.azure.core.http.HttpResponse;
import com.azure.core.http.policy.HttpLoggingPolicy;
import com.azure.core.http.policy.HttpPipelinePolicy;
import com.azure.core.util.CoreUtils;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.fasterxml.jackson.databind.SerializationFeature;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import reactor.core.publisher.Mono;

import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.io.UnsupportedEncodingException;
import java.nio.ByteBuffer;
import java.nio.channels.Channels;
import java.nio.channels.WritableByteChannel;
import java.nio.charset.StandardCharsets;
import java.util.HashSet;
import java.util.Locale;
import java.util.Optional;
import java.util.Set;
import java.util.concurrent.TimeUnit;

/**
* The pipeline policy that handles logging of HTTP requests and responses.
*/
public class HttpDebugLoggingPolicy implements HttpPipelinePolicy {

private static final ObjectMapper PRETTY_PRINTER = new ObjectMapper().enable(SerializationFeature.INDENT_OUTPUT);
private static final String REDACTED_PLACEHOLDER = "REDACTED";
private static final Set<String> DISALLOWED_HEADER_NAMES = new HashSet<>();
private static final boolean PRETTY_PRINT_BODY = true;

/**
* Creates an HttpDebugLoggingPolicy with the given log configurations.
*/
public HttpDebugLoggingPolicy() {
DISALLOWED_HEADER_NAMES.add("authorization");
}

@Override
public Mono<HttpResponse> process(HttpPipelineCallContext context, HttpPipelineNextPolicy next) {
final Logger logger = LoggerFactory.getLogger((String) context.getData("caller-method").orElse(""));
final long startNs = System.nanoTime();

return logRequest(logger, context.getHttpRequest(), context.getData(HttpLoggingPolicy.RETRY_COUNT_CONTEXT))
.then(next.process())
.flatMap(response -> logResponse(logger, response, startNs))
.doOnError(throwable -> logger.warn("<-- HTTP FAILED: ", throwable));
}

private Mono<Void> logRequest(final Logger logger, final HttpRequest request,
final Optional<Object> optionalRetryCount) {
if (!logger.isInfoEnabled()) {
return Mono.empty();
}

StringBuilder requestLogMessage = new StringBuilder();
requestLogMessage.append("--> ")
.append(request.getHttpMethod())
.append(" ")
.append(request.getUrl())
.append(System.lineSeparator());

optionalRetryCount.ifPresent(o -> requestLogMessage.append("Try count: ")
.append(o)
.append(System.lineSeparator()));

addHeadersToLogMessage(logger, request.getHeaders(), requestLogMessage);

if (request.getBody() == null) {
requestLogMessage.append("(empty body)")
.append(System.lineSeparator())
.append("--> END ")
.append(request.getHttpMethod())
.append(System.lineSeparator());

return logAndReturn(logger, requestLogMessage, null);
}

String contentType = request.getHeaders().getValue("Content-Type");
long contentLength = getContentLength(logger, request.getHeaders());

if (shouldBodyBeLogged(contentType, contentLength)) {
ByteArrayOutputStream outputStream = new ByteArrayOutputStream((int) contentLength);
WritableByteChannel bodyContentChannel = Channels.newChannel(outputStream);

// Add non-mutating operators to the data stream.
request.setBody(
request.getBody()
.flatMap(byteBuffer -> writeBufferToBodyStream(bodyContentChannel, byteBuffer))
.doFinally(ignored -> {
requestLogMessage.append(contentLength)
.append("-byte body:")
.append(System.lineSeparator())
.append(prettyPrintIfNeeded(logger, contentType,
convertStreamToString(outputStream, logger)))
.append(System.lineSeparator())
.append("--> END ")
.append(request.getHttpMethod())
.append(System.lineSeparator());

logger.info(requestLogMessage.toString());
}));

return Mono.empty();
} else {
requestLogMessage.append(contentLength)
.append("-byte body: (content not logged)")
.append(System.lineSeparator())
.append("--> END ")
.append(request.getHttpMethod())
.append(System.lineSeparator());

return logAndReturn(logger, requestLogMessage, null);
}
}

/*
* Logs thr HTTP response.
*
* @param logger Logger used to log the response.
* @param response HTTP response returned from Azure.
* @param startNs Nanosecond representation of when the request was sent.
* @return A Mono containing the HTTP response.
*/
private Mono<HttpResponse> logResponse(final Logger logger, final HttpResponse response, long startNs) {
if (!logger.isInfoEnabled()) {
return Mono.just(response);
}

long tookMs = TimeUnit.NANOSECONDS.toMillis(System.nanoTime() - startNs);

String contentLengthString = response.getHeaderValue("Content-Length");
String bodySize = (CoreUtils.isNullOrEmpty(contentLengthString))
? "unknown-length body"
: contentLengthString + "-byte body";

StringBuilder responseLogMessage = new StringBuilder();
responseLogMessage.append("<-- ")
.append(response.getStatusCode())
.append(" ")
.append(response.getRequest().getUrl())
.append(" (")
.append(tookMs)
.append(" ms, ")
.append(bodySize)
.append(")")
.append(System.lineSeparator());

addHeadersToLogMessage(logger, response.getHeaders(), responseLogMessage);

String contentTypeHeader = response.getHeaderValue("Content-Type");
long contentLength = getContentLength(logger, response.getHeaders());

if (shouldBodyBeLogged(contentTypeHeader, contentLength)) {
HttpResponse bufferedResponse = response.buffer();
ByteArrayOutputStream outputStream = new ByteArrayOutputStream((int) contentLength);
WritableByteChannel bodyContentChannel = Channels.newChannel(outputStream);
return bufferedResponse.getBody()
.flatMap(byteBuffer -> writeBufferToBodyStream(bodyContentChannel, byteBuffer))
.doFinally(ignored -> {
responseLogMessage.append("Response body:")
.append(System.lineSeparator())
.append(prettyPrintIfNeeded(logger, contentTypeHeader,
convertStreamToString(outputStream, logger)))
.append(System.lineSeparator())
.append("<-- END HTTP");

logger.info(responseLogMessage.toString());
}).then(Mono.just(bufferedResponse));
} else {
responseLogMessage.append("(body content not logged)")
.append(System.lineSeparator())
.append("<-- END HTTP");

return logAndReturn(logger, responseLogMessage, response);
}
}

private <T> Mono<T> logAndReturn(Logger logger, StringBuilder logMessageBuilder, T data) {
logger.info(logMessageBuilder.toString());
return Mono.justOrEmpty(data);
}

private void addHeadersToLogMessage(Logger logger, HttpHeaders headers, StringBuilder sb) {
for (HttpHeader header : headers) {
String headerName = header.getName();
sb.append(headerName).append(":");
if (!DISALLOWED_HEADER_NAMES.contains(headerName.toLowerCase(Locale.ROOT))) {
sb.append(header.getValue());
} else {
sb.append(REDACTED_PLACEHOLDER);
}
sb.append(System.lineSeparator());
}
}

private String prettyPrintIfNeeded(Logger logger, String contentType, String body) {
String result = body;
if (PRETTY_PRINT_BODY && contentType != null
&& (contentType.startsWith(ContentType.APPLICATION_JSON) || contentType.startsWith("text/json"))) {
try {
final Object deserialized = PRETTY_PRINTER.readTree(body);
result = PRETTY_PRINTER.writeValueAsString(deserialized);
} catch (Exception e) {
logger.warn("Failed to pretty print JSON: {}", e.getMessage());
}
}
return result;
}

private long getContentLength(Logger logger, HttpHeaders headers) {
long contentLength = 0;

String contentLengthString = headers.getValue("Content-Length");
if (CoreUtils.isNullOrEmpty(contentLengthString)) {
return contentLength;
}

try {
contentLength = Long.parseLong(contentLengthString);
} catch (NumberFormatException | NullPointerException e) {
logger.warn("Could not parse the HTTP header content-length: '{}'.",
headers.getValue("content-length"), e);
}

return contentLength;
}

private boolean shouldBodyBeLogged(String contentTypeHeader, long contentLength) {
return !ContentType.APPLICATION_OCTET_STREAM.equalsIgnoreCase(contentTypeHeader)
&& contentLength != 0;
}

private static String convertStreamToString(ByteArrayOutputStream stream, Logger logger) {
try {
return stream.toString(StandardCharsets.UTF_8.name());
} catch (UnsupportedEncodingException ex) {
logger.error(ex.toString());
throw new RuntimeException(ex);
}
}

private static Mono<ByteBuffer> writeBufferToBodyStream(WritableByteChannel channel, ByteBuffer byteBuffer) {
try {
channel.write(byteBuffer.duplicate());
return Mono.just(byteBuffer);
} catch (IOException ex) {
return Mono.error(ex);
}
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -126,7 +126,7 @@ TokenCredentialProvider getTokenCredentialProvider() {
}

Duration getBlockTimeout() {
return Duration.ofSeconds(30 * 4);
return Duration.ofSeconds(30 * 3);
}

AtomicBoolean getIsInitialized() {
Expand Down
2 changes: 1 addition & 1 deletion sdk/identity/azure-identity/pom.xml
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,7 @@

<groupId>com.azure</groupId>
<artifactId>azure-identity</artifactId>
<version>1.11.0-beta.1</version><!-- {x-version-update;com.azure:azure-identity;current} -->
<version>1.10.0</version><!-- {x-version-update;com.azure:azure-identity;current} -->

<name>Microsoft Azure client library for Identity</name>
<description>This module contains client library for Microsoft Azure Identity.</description>
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -953,9 +953,12 @@ private Mono<AccessToken> authenticateToArcManagedIdentityEndpoint(String identi
* @return a Publisher that emits an AccessToken
*/
public Mono<AccessToken> authenticateWithExchangeToken(TokenRequestContext request) {

LOGGER.verbose("wi check exec authenticateWithExchangeToken");
return clientAssertionAccessor.getValue()
.flatMap(assertionToken -> Mono.fromCallable(() -> authenticateWithExchangeTokenHelper(request, assertionToken)));
.flatMap(assertionToken -> Mono.fromCallable(() -> {
LOGGER.verbose("wi check exec from callable, assertionToken: " + assertionToken);
return authenticateWithExchangeTokenHelper(request, assertionToken);
}));
}

/**
Expand Down Expand Up @@ -1289,14 +1292,17 @@ private boolean isADFSTenant() {

Function<AppTokenProviderParameters, CompletableFuture<TokenProviderResult>> getWorkloadIdentityTokenProvider() {
return appTokenProviderParameters -> {
LOGGER.verbose("wi check exec appTokenProvider");
TokenRequestContext trc = new TokenRequestContext()
.setScopes(new ArrayList<>(appTokenProviderParameters.scopes))
.setClaims(appTokenProviderParameters.claims)
.setTenantId(appTokenProviderParameters.tenantId);

LOGGER.verbose("wi check authenticateWithExchangeTokenSync");
Mono<AccessToken> accessTokenAsync = authenticateWithExchangeToken(trc);

LOGGER.verbose("wi check return mono token");
return accessTokenAsync.map(accessToken -> {
LOGGER.verbose("wi check access token returned: " + accessToken.getToken());
TokenProviderResult result = new TokenProviderResult();
result.setAccessToken(accessToken.getToken());
result.setTenantId(trc.getTenantId());
Expand Down
Loading

0 comments on commit 9b47ce5

Please sign in to comment.