Skip to content

Commit

Permalink
Span changes in keypair generation, Fixes AB#3146656 (#2594)
Browse files Browse the repository at this point in the history
- Span.current() was used to capture an attribute on a span previously.
Changing it to just use the variable span.
- Captured span error status when legacy keypair generation fails.
- Catching Throwable instead of Exception.
- Made the code in the else block when flight is off to EXACTLY match
the old code.

Fixes
[AB#3146656](https://identitydivision.visualstudio.com/fac9d424-53d2-45c0-91b5-ef6ba7a6bf26/_workitems/edit/3146656)
  • Loading branch information
somalaya authored Feb 28, 2025
1 parent 42b9611 commit 53dba6f
Show file tree
Hide file tree
Showing 2 changed files with 26 additions and 17 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -259,25 +259,25 @@ private void saveSecretKeyToStorage(@NonNull final SecretKey unencryptedKey) thr
KeyPair keyPair = AndroidKeyStoreUtil.readKey(mAlias);
if (keyPair == null) {
Logger.info(methodTag, "No existing keypair. Generating a new one.");
final Span span = OTelUtility.createSpan(SpanName.KeyPairGeneration.name());
final long keypairGenStartTime = System.currentTimeMillis();
if (Build.VERSION.SDK_INT >= Build.VERSION_CODES.P && CommonFlightsManager.INSTANCE.getFlightsProvider().isFlightEnabled(CommonFlight.ENABLE_NEW_KEY_GEN_SPEC_FOR_WRAP)) {
final long keypairGenStartTime = System.currentTimeMillis();
final Span span = OTelUtility.createSpanFromParent(SpanName.KeyPairGeneration.name(), SpanExtension.current().getSpanContext());
try (final Scope scope = SpanExtension.makeCurrentSpan(span)) {
keyPair = attemptKeyPairGeneration(mAlias, true, keypairGenStartTime);
Logger.info(methodTag, "Successfully generated keypair with new KeyPairGeneratorSpec with wrap purpose.");
Span.current().setAttribute(AttributeName.key_pair_gen_successful_method.name(), "new_key_gen_spec_with_wrap");
span.setAttribute(AttributeName.key_pair_gen_successful_method.name(), "new_key_gen_spec_with_wrap");
span.setStatus(StatusCode.OK);
} catch (final ProviderException e) {
if ("SecureKeyImportUnavailableException".equals(e.getClass().getSimpleName())) {
Logger.warn(methodTag, "Wrap purpose may not be supported. Retrying without wrap.");
try {
keyPair = attemptKeyPairGeneration(mAlias, false, keypairGenStartTime);
Logger.info(methodTag, "Successfully generated keypair with new KeyPairGeneratorSpec without wrap purpose.");
Span.current().setAttribute(AttributeName.key_pair_gen_successful_method.name(), "new_key_gen_spec_without_wrap");
span.setAttribute(AttributeName.key_pair_gen_successful_method.name(), "new_key_gen_spec_without_wrap");
span.setStatus(StatusCode.OK);
} catch (Exception ex) {
} catch (final Exception ex) {
// 2nd fallback to legacy keygen spec
Logger.error(methodTag, "Second attempt without wrap also failed. Falling back to legacy spec.", ex);
Logger.warn(methodTag, "Second attempt without wrap also failed. Falling back to legacy spec."+ ex);
keyPair = generateKeyPairWithLegacySpec(mAlias, keypairGenStartTime);
if (e.getMessage() != null) {
span.setAttribute(AttributeName.keypair_gen_exception.name(), e.getMessage());
Expand All @@ -286,16 +286,16 @@ private void saveSecretKeyToStorage(@NonNull final SecretKey unencryptedKey) thr
span.setStatus(StatusCode.OK);
}
} else {
Logger.info(methodTag, "Some unknown exception occurred. Running legacy keygen spec logic.");
Logger.warn(methodTag, "Some unknown exception occurred. Running legacy keygen spec logic."+ e);
keyPair = generateKeyPairWithLegacySpec(mAlias, keypairGenStartTime);
Span.current().setAttribute(AttributeName.key_pair_gen_successful_method.name(), "legacy_key_gen_spec");
span.setAttribute(AttributeName.key_pair_gen_successful_method.name(), "legacy_key_gen_spec");
span.setStatus(StatusCode.OK);
}
} catch (final Exception e) {
Logger.warn(methodTag, "Unexpected error with new KeyPairGeneratorSpec. Falling back to legacy spec. "+ e);
} catch (final Throwable throwable) {
Logger.warn(methodTag, "Unexpected error with new KeyPairGeneratorSpec. Falling back to legacy spec. "+ throwable);
keyPair = generateKeyPairWithLegacySpec(mAlias, keypairGenStartTime);
if (e.getMessage() != null) {
span.setAttribute(AttributeName.keypair_gen_exception.name(), e.getMessage());
if (throwable.getMessage() != null) {
span.setAttribute(AttributeName.keypair_gen_exception.name(), throwable.getMessage());
}
span.setAttribute(AttributeName.key_pair_gen_successful_method.name(), "legacy_key_gen_spec");
span.setStatus(StatusCode.OK);
Expand All @@ -306,7 +306,8 @@ private void saveSecretKeyToStorage(@NonNull final SecretKey unencryptedKey) thr
else {
// If flight for using new keygen spec is not enabled, use the legacy spec.
Logger.info(methodTag, "Using legacy spec for keypair generation directly.");
keyPair = generateKeyPairWithLegacySpec(mAlias, keypairGenStartTime);
keyPair = AndroidKeyStoreUtil.generateKeyPair(
WRAP_KEY_ALGORITHM, getLegacySpecForKeyStoreKey(mContext, mAlias));
}
}
final byte[] keyWrapped = AndroidKeyStoreUtil.wrap(unencryptedKey, keyPair, WRAP_ALGORITHM);
Expand All @@ -322,10 +323,17 @@ private KeyPair attemptKeyPairGeneration(@NonNull final String alias, boolean us
}

private KeyPair generateKeyPairWithLegacySpec(@NonNull final String alias, long keypairGenStartTime) throws ClientException{
KeyPair keyPair = AndroidKeyStoreUtil.generateKeyPair(
WRAP_KEY_ALGORITHM, getLegacySpecForKeyStoreKey(mContext, alias));
recordKeyGenerationTime(keypairGenStartTime);
return keyPair;
try {
final KeyPair keyPair = AndroidKeyStoreUtil.generateKeyPair(
WRAP_KEY_ALGORITHM, getLegacySpecForKeyStoreKey(mContext, alias));
recordKeyGenerationTime(keypairGenStartTime);
return keyPair;
} catch (final ClientException e) {
SpanExtension.current().recordException(e);
SpanExtension.current().setStatus(StatusCode.ERROR);
Logger.error(TAG + ":generateKeyPairWithLegacySpec", "Error generating keypair with legacy spec.", e);
throw e;
}
}

private void recordKeyGenerationTime(long keypairGenStartTime) {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -24,6 +24,7 @@

/**
* Names of Open Telemetry Span Attributes we want to capture for broker's Spans.
* NOTE : Any changes to this enum should also be made in the corresponding enum in Broker.
*/
public enum AttributeName {
/**
Expand Down

0 comments on commit 53dba6f

Please sign in to comment.