AzureTokenProvider.java

  1. package no.nav.data.common.security.azure;

  3. import com.github.benmanes.caffeine.cache.Cache;
  4. import com.github.benmanes.caffeine.cache.Caffeine;
  5. import com.microsoft.aad.msal4j.*;
  6. import com.microsoft.graph.serviceclient.GraphServiceClient;
  7. import com.microsoft.kiota.authentication.AccessTokenProvider;
  8. import com.microsoft.kiota.authentication.AllowedHostsValidator;
  9. import com.microsoft.kiota.authentication.BaseBearerTokenAuthenticationProvider;
  10. import com.nimbusds.oauth2.sdk.pkce.CodeChallengeMethod;
  11. import io.prometheus.client.Summary;
  12. import lombok.extern.slf4j.Slf4j;
  13. import no.nav.data.common.exceptions.TechnicalException;
  14. import no.nav.data.common.security.AuthService;
  15. import no.nav.data.common.security.Encryptor;
  16. import no.nav.data.common.security.TokenProvider;
  17. import no.nav.data.common.security.azure.support.AuthResultExpiry;
  18. import no.nav.data.common.security.domain.Auth;
  19. import no.nav.data.common.security.dto.Credential;
  20. import no.nav.data.common.security.dto.OAuthState;
  21. import no.nav.data.common.utils.Constants;
  22. import no.nav.data.common.utils.MetricUtils;
  23. import org.apache.commons.codec.binary.Base64;
  24. import org.apache.commons.codec.digest.DigestUtils;
  25. import org.apache.commons.lang3.StringUtils;
  26. import org.springframework.stereotype.Service;
  27. import org.springframework.util.Assert;
  28. import org.springframework.util.ReflectionUtils;

  30. import java.lang.reflect.InvocationTargetException;
  31. import java.lang.reflect.Method;
  32. import java.net.URI;
  33. import java.net.URL;
  34. import java.time.Duration;
  35. import java.util.Map;
  36. import java.util.Set;

  38. import static java.util.Objects.requireNonNull;
  39. import static no.nav.data.common.security.SecurityConstants.SESS_ID_LEN;
  40. import static no.nav.data.common.security.SecurityConstants.TOKEN_TYPE;
  41. import static no.nav.data.common.security.azure.AzureConstants.MICROSOFT_GRAPH_SCOPES;
  42. import static no.nav.data.common.security.azure.AzureConstants.MICROSOFT_GRAPH_SCOPE_APP;

  44. @Slf4j
  45. @Service
  46. public class AzureTokenProvider implements TokenProvider {
  47.     private final Cache<String, IAuthenticationResult> accessTokenCache;

  49.     private final IConfidentialClientApplication msalClient;
  50.     private final AuthService authService;

  52.     private final AADAuthenticationProperties aadAuthProps;
  53.     private final Encryptor encryptor;

  55.     private final Summary tokenMetrics;
  56.     private CustomAccessTokenProvider customAccessTokenProvider;

  58.     public AzureTokenProvider(AADAuthenticationProperties aadAuthProps,
  59.             IConfidentialClientApplication msalClient,
  60.             AuthService authService, Encryptor encryptor) {
  61.         this.aadAuthProps = aadAuthProps;
  62.         this.msalClient = msalClient;
  63.         this.authService = authService;
  64.         this.encryptor = encryptor;
  65.         this.tokenMetrics = MetricUtils.summary()
  66.                 .labels("accessToken").labels("lookupGrantedAuthorities")
  67.                 .labelNames("action")
  68.                 .name(Constants.APP_ID.replace('-', '_') + "_token_summary")
  69.                 .help("Time taken for azure token lookups")
  70.                 .quantile(.5, .01).quantile(.9, .01).quantile(.99, .001)
  71.                 .maxAgeSeconds(Duration.ofHours(24).getSeconds())
  72.                 .ageBuckets(8)
  73.                 .register();

  75.         this.accessTokenCache = Caffeine.newBuilder().recordStats()
  76.                 .expireAfter(new AuthResultExpiry())
  77.                 .maximumSize(1000).build();
  78.         MetricUtils.register("accessTokenCache", accessTokenCache);

  80.         customAccessTokenProvider = new CustomAccessTokenProvider(this);
  81.     }

  83.     GraphServiceClient getGraphClient() {
  84.         return new GraphServiceClient(new BaseBearerTokenAuthenticationProvider(customAccessTokenProvider));
  85.     }

  87.     @Override
  88.     public String getConsumerToken(String resource) {
  89.         return Credential.getCredential()
  90.                 .filter(Credential::hasAuth)
  91.                 .map(cred -> TOKEN_TYPE + getAccessTokenForResource(cred.getAuth().decryptRefreshToken(), resource))
  92.                 .orElseGet(() -> TOKEN_TYPE + getApplicationTokenForResource(resource));
  93.     }

  95.     public Auth getAuth(String session) {
  96.         Assert.isTrue(session.length() > SESS_ID_LEN, "invalid session");
  97.         var sessionId = session.substring(0, SESS_ID_LEN);
  98.         var sessionKey = session.substring(SESS_ID_LEN);
  99.         var auth = authService.getAuth(sessionId, sessionKey);
  100.         try {
  101.             String accessToken = getAccessTokenForResource(auth.decryptRefreshToken(), resourceForAppId());
  102.             auth.addAccessToken(accessToken);
  103.             return auth;
  104.         } catch (RuntimeException e) {
  105.             throw new TechnicalException("Failed to get access token for userId=%s initiated=%s".formatted(auth.getUserId(), auth.getInitiated()));
  106.         }
  107.     }

  109.     @Override
  110.     public void destroySession() {
  111.         Credential.getCredential().map(Credential::getAuth).ifPresent(auth -> authService.endSession(auth.getId()));
  112.     }

  114.     @Override
  115.     public String createAuthRequestRedirectUrl(String postLoginRedirectUri, String postLoginErrorUri, String redirectUri) {
  116.         var auth = authService.createAuth();
  117.         var codeVerifier = auth.getCodeVerifier();
  118.         var s256 = DigestUtils.sha256(codeVerifier);
  119.         var codeChallenge = Base64.encodeBase64URLSafeString(s256);
  120.         URL url = msalClient.getAuthorizationRequestUrl(AuthorizationRequestUrlParameters
  121.                 .builder(redirectUri, MICROSOFT_GRAPH_SCOPES)
  122.                 .state(new OAuthState(auth.getId().toString(), postLoginRedirectUri, postLoginErrorUri).toJson(encryptor))
  123.                 .responseMode(ResponseMode.FORM_POST)
  124.                 .codeChallengeMethod(CodeChallengeMethod.S256.getValue())
  125.                 .codeChallenge(codeChallenge)
  126.                 .build());
  127.         return url.toString();
  128.     }

  130.     @Override
  131.     public String createSession(String sessionId, String code, String redirectUri) {
  132.         try {
  133.             log.debug("Looking up token for auth code");
  134.             var codeVerifier = authService.getCodeVerifier(sessionId);
  135.             var authResult = msalClient.acquireToken(AuthorizationCodeParameters
  136.                     .builder(code, new URI(redirectUri))
  137.                     .scopes(MICROSOFT_GRAPH_SCOPES)
  138.                     .codeVerifier(codeVerifier)
  139.                     .build()).get();
  140.             String userId = StringUtils.substringBefore(authResult.account().homeAccountId(), ".");
  141.             String refreshToken = getRefreshTokenFromAuthResult(authResult);
  142.             return authService.initAuth(userId, refreshToken, sessionId);
  143.         } catch (Exception e) {
  144.             log.error("Failed to get token for auth code", e);
  145.             throw new TechnicalException("Failed to get token for auth code", e);
  146.         }
  147.     }

  149.     private String getRefreshTokenFromAuthResult(IAuthenticationResult authResult) throws ClassNotFoundException, IllegalAccessException, InvocationTargetException {
  150.         // interface is missing refreshToken...
  151.         Method refreshTokenMethod = ReflectionUtils.findMethod(Class.forName("com.microsoft.aad.msal4j.AuthenticationResult"), "refreshToken");
  152.         Assert.notNull(refreshTokenMethod, "couldn't find refreshToken method");
  153.         refreshTokenMethod.setAccessible(true);
  154.         return (String) refreshTokenMethod.invoke(authResult);
  155.     }

  157.     private String resourceForAppId() {
  158.         return aadAuthProps.getClientId() + "/.default";
  159.     }

  161.     String getApplicationTokenForResource(String resource) {
  162.         log.trace("Getting application token for resource {}", resource);
  163.         return requireNonNull(accessTokenCache.get("credential" + resource, cacheKey -> acquireTokenByCredential(resource))).accessToken();
  164.     }

  166.     private String getAccessTokenForResource(String refreshToken, String resource) {
  167.         log.trace("Getting access token for resource {}", resource);
  168.         return requireNonNull(accessTokenCache.get("refresh" + refreshToken + resource, cacheKey -> acquireTokenByRefreshToken(refreshToken, resource))).accessToken();
  169.     }

  171.     private IAuthenticationResult acquireTokenByRefreshToken(String refreshToken, String resource) {
  172.         try (var ignored = tokenMetrics.labels("accessToken").startTimer()) {
  173.             log.debug("Looking up access token for resource {}", resource);
  174.             return msalClient.acquireToken(RefreshTokenParameters.builder(Set.of(resource), refreshToken).build()).get();
  175.         } catch (Exception e) {
  176.             throw new TechnicalException("Failed to get access token for refreshToken", e);
  177.         }
  178.     }

  180.     /**
  181.      * access token for app user
  182.      */
  183.     private IAuthenticationResult acquireTokenByCredential(String resource) {
  184.         try {
  185.             log.debug("Looking up application token for resource {}", resource);
  186.             return msalClient.acquireToken(ClientCredentialParameters.builder(Set.of(resource)).build()).get();
  187.         } catch (Exception e) {
  188.             throw new TechnicalException("Failed to get access token for credential", e);
  189.         }
  190.     }

  192.     class CustomAccessTokenProvider implements AccessTokenProvider {
  193.         AzureTokenProvider tokenProvider;

  195.         public CustomAccessTokenProvider(AzureTokenProvider tokenProvider) {
  196.             this.tokenProvider = tokenProvider;
  197.         }

  199.         @Override
  200.         public String getAuthorizationToken(URI uri, Map<String, Object> additionalAuthenticationContex) {
  201.             return tokenProvider.getApplicationTokenForResource(MICROSOFT_GRAPH_SCOPE_APP);
  202.         }

  204.         // Make sure to have the right set of hosts
  205.         private final AllowedHostsValidator validator = new AllowedHostsValidator("graph.microsoft.com");

  207.         @Override
  208.         public AllowedHostsValidator getAllowedHostsValidator() {
  209.             // Handle allowed hosts validation logic here
  210.             return validator;
  211.         }
  212.     }
  213. }
Created with JaCoCo 0.8.12.202403310830