Updated relationship between approved sites and access tokens, closes mitreid-connect#874

Author: jricher

openid-connect-common/src/main/java/org/mitre/oauth2/model/OAuth2AccessTokenEntity.java

@@ -47,6 +47,7 @@
 import javax.persistence.Transient;
 
 import org.mitre.oauth2.model.convert.JWTStringConverter;
+import org.mitre.openid.connect.model.ApprovedSite;
 import org.mitre.uma.model.Permission;
 import org.springframework.security.oauth2.common.OAuth2AccessToken;
 import org.springframework.security.oauth2.common.OAuth2AccessTokenJackson1Deserializer;
@@ -70,6 +71,7 @@
 @NamedQuery(name = OAuth2AccessTokenEntity.QUERY_BY_CLIENT, query = "select a from OAuth2AccessTokenEntity a where a.client = :" + OAuth2AccessTokenEntity.PARAM_CLIENT),
 @NamedQuery(name = OAuth2AccessTokenEntity.QUERY_BY_ID_TOKEN, query = "select a from OAuth2AccessTokenEntity a where a.idToken = :" + OAuth2AccessTokenEntity.PARAM_ID_TOKEN),
 @NamedQuery(name = OAuth2AccessTokenEntity.QUERY_BY_TOKEN_VALUE, query = "select a from OAuth2AccessTokenEntity a where a.jwt = :" + OAuth2AccessTokenEntity.PARAM_TOKEN_VALUE),
+@NamedQuery(name = OAuth2AccessTokenEntity.QUERY_BY_APPROVED_SITE, query = "select a from OAuth2AccessTokenEntity a where a.approvedSite = :" + OAuth2AccessTokenEntity.PARAM_APPROVED_SITE),
 @NamedQuery(name = OAuth2AccessTokenEntity.QUERY_BY_RESOURCE_SET, query = "select a from OAuth2AccessTokenEntity a join a.permissions p where p.resourceSet.id = :" + OAuth2AccessTokenEntity.PARAM_RESOURCE_SET_ID)
 })
 @org.codehaus.jackson.map.annotate.JsonSerialize(using = OAuth2AccessTokenJackson1Serializer.class)
@@ -78,6 +80,7 @@
 @com.fasterxml.jackson.databind.annotation.JsonDeserialize(using = OAuth2AccessTokenJackson2Deserializer.class)
 public class OAuth2AccessTokenEntity implements OAuth2AccessToken {
 
+public static final String QUERY_BY_APPROVED_SITE = "OAuth2AccessTokenEntity.getByApprovedSite";
 public static final String QUERY_BY_TOKEN_VALUE = "OAuth2AccessTokenEntity.getByTokenValue";
 public static final String QUERY_BY_ID_TOKEN = "OAuth2AccessTokenEntity.getByIdToken";
 public static final String QUERY_BY_CLIENT = "OAuth2AccessTokenEntity.getByClient";
@@ -92,6 +95,7 @@ public class OAuth2AccessTokenEntity implements OAuth2AccessToken {
 public static final String PARAM_REFERSH_TOKEN = "refreshToken";
 public static final String PARAM_DATE = "date";
 public static final String PARAM_RESOURCE_SET_ID = "rsid";
+public static final String PARAM_APPROVED_SITE = "approvedSite";
 
 public static final String ID_TOKEN_FIELD_NAME = "id_token";
 
@@ -114,6 +118,8 @@ public class OAuth2AccessTokenEntity implements OAuth2AccessToken {
 private Set<String> scope;
 
 private Set<Permission> permissions;
+
+private ApprovedSite approvedSite;
 
 /**
  * Create a new, blank access token
@@ -337,4 +343,13 @@ public void setPermissions(Set<Permission> permissions) {
 this.permissions = permissions;
 }
 
+@ManyToOne
+@JoinColumn(name="approved_site_id")
+public ApprovedSite getApprovedSite() {
+return approvedSite;
+}
+
+public void setApprovedSite(ApprovedSite approvedSite) {
+this.approvedSite = approvedSite;
+}
 }

openid-connect-common/src/main/java/org/mitre/oauth2/repository/OAuth2TokenRepository.java

@@ -22,6 +22,7 @@
 import org.mitre.oauth2.model.ClientDetailsEntity;
 import org.mitre.oauth2.model.OAuth2AccessTokenEntity;
 import org.mitre.oauth2.model.OAuth2RefreshTokenEntity;
+import org.mitre.openid.connect.model.ApprovedSite;
 import org.mitre.uma.model.ResourceSet;
 
 public interface OAuth2TokenRepository {
@@ -65,5 +66,7 @@ public interface OAuth2TokenRepository {
 public void clearDuplicateAccessTokens();
 
 public void clearDuplicateRefreshTokens();
+
+public List<OAuth2AccessTokenEntity> getAccessTokensForApprovedSite(ApprovedSite approvedSite);
 
 }

openid-connect-common/src/main/java/org/mitre/openid/connect/model/ApprovedSite.java

@@ -81,9 +81,6 @@ public class ApprovedSite {
 // this should include all information for what data to access
 private Set<String> allowedScopes;
 
-//Link to any access tokens approved through this stored decision
-private Set<OAuth2AccessTokenEntity> approvedAccessTokens = Sets.newHashSet();
-
 /**
  * Empty constructor
  */
@@ -229,16 +226,4 @@ public boolean isExpired() {
 }
 }
 
-@OneToMany(cascade=CascadeType.ALL, fetch=FetchType.LAZY)
-@JoinColumn(name="approved_site_id")
-public Set<OAuth2AccessTokenEntity> getApprovedAccessTokens() {
-return approvedAccessTokens;
-}
-
-/**
- * @param approvedAccessTokens the approvedAccessTokens to set
- */
-public void setApprovedAccessTokens(Set<OAuth2AccessTokenEntity> approvedAccessTokens) {
-this.approvedAccessTokens = approvedAccessTokens;
-}
 }

openid-connect-common/src/main/java/org/mitre/openid/connect/service/ApprovedSiteService.java

@@ -18,8 +18,10 @@
 
 import java.util.Collection;
 import java.util.Date;
+import java.util.List;
 import java.util.Set;
 
+import org.mitre.oauth2.model.OAuth2AccessTokenEntity;
 import org.mitre.openid.connect.model.ApprovedSite;
 import org.springframework.security.oauth2.provider.ClientDetails;
 
@@ -101,4 +103,11 @@ public interface ApprovedSiteService {
  * @return
  */
 public void clearExpiredSites();
+
+/**
+ * Return all approved access tokens for the site.
+ * @return
+ */
+public List<OAuth2AccessTokenEntity> getApprovedAccessTokens(ApprovedSite approvedSite);
+
 }

openid-connect-server/src/main/java/org/mitre/oauth2/repository/impl/JpaOAuth2TokenRepository.java

@@ -35,6 +35,7 @@
 import org.mitre.oauth2.model.OAuth2AccessTokenEntity;
 import org.mitre.oauth2.model.OAuth2RefreshTokenEntity;
 import org.mitre.oauth2.repository.OAuth2TokenRepository;
+import org.mitre.openid.connect.model.ApprovedSite;
 import org.mitre.uma.model.ResourceSet;
 import org.mitre.util.jpa.JpaUtil;
 import org.slf4j.Logger;
@@ -272,5 +273,13 @@ public void clearDuplicateRefreshTokens() {
 }
 
 }
+
+@Override
+public List<OAuth2AccessTokenEntity> getAccessTokensForApprovedSite(ApprovedSite approvedSite) {
+TypedQuery<OAuth2AccessTokenEntity> queryA = manager.createNamedQuery(OAuth2AccessTokenEntity.QUERY_BY_APPROVED_SITE, OAuth2AccessTokenEntity.class);
+queryA.setParameter(OAuth2AccessTokenEntity.PARAM_APPROVED_SITE, approvedSite);
+List<OAuth2AccessTokenEntity> accessTokens = queryA.getResultList();
+return accessTokens;
+}
 
 }

openid-connect-server/src/main/java/org/mitre/oauth2/service/impl/DefaultOAuth2ProviderTokenService.java

@@ -250,24 +250,21 @@ public OAuth2AccessTokenEntity createAccessToken(OAuth2Authentication authentica
 token.setRefreshToken(savedRefreshToken);
 }
 
-OAuth2AccessTokenEntity enhancedToken = (OAuth2AccessTokenEntity) tokenEnhancer.enhance(token, authentication);
-
-OAuth2AccessTokenEntity savedToken = tokenRepository.saveAccessToken(enhancedToken);
-
 //Add approved site reference, if any
 OAuth2Request originalAuthRequest = authHolder.getAuthentication().getOAuth2Request();
 
 if (originalAuthRequest.getExtensions() != null && originalAuthRequest.getExtensions().containsKey("approved_site")) {
 
 Long apId = Long.parseLong((String) originalAuthRequest.getExtensions().get("approved_site"));
 ApprovedSite ap = approvedSiteService.getById(apId);
-Set<OAuth2AccessTokenEntity> apTokens = ap.getApprovedAccessTokens();
-apTokens.add(savedToken);
-ap.setApprovedAccessTokens(apTokens);
-approvedSiteService.save(ap);
 
+token.setApprovedSite(ap);
 }
 
+OAuth2AccessTokenEntity enhancedToken = (OAuth2AccessTokenEntity) tokenEnhancer.enhance(token, authentication);
+
+OAuth2AccessTokenEntity savedToken = tokenRepository.saveAccessToken(enhancedToken);
+
 if (savedToken.getRefreshToken() != null) {
 tokenRepository.saveRefreshToken(savedToken.getRefreshToken()); // make sure we save any changes that might have been enhanced
 }

openid-connect-server/src/main/java/org/mitre/openid/connect/service/impl/DefaultApprovedSiteService.java

@@ -18,6 +18,7 @@
 
 import java.util.Collection;
 import java.util.Date;
+import java.util.List;
 import java.util.Set;
 
 import org.mitre.oauth2.model.OAuth2AccessTokenEntity;
@@ -82,7 +83,7 @@ public ApprovedSite getById(Long id) {
 public void remove(ApprovedSite approvedSite) {
 
 //Remove any associated access and refresh tokens
-Set<OAuth2AccessTokenEntity> accessTokens = approvedSite.getApprovedAccessTokens();
+List<OAuth2AccessTokenEntity> accessTokens = getApprovedAccessTokens(approvedSite);
 
 for (OAuth2AccessTokenEntity token : accessTokens) {
 if (token.getRefreshToken() != null) {
@@ -180,4 +181,11 @@ private Collection<ApprovedSite> getExpired() {
 return Collections2.filter(approvedSiteRepository.getAll(), isExpired);
 }
 
+@Override
+public List<OAuth2AccessTokenEntity> getApprovedAccessTokens(
+ApprovedSite approvedSite) {
+return tokenRepository.getAccessTokensForApprovedSite(approvedSite); 
+
+}
+
 }

openid-connect-server/src/main/java/org/mitre/openid/connect/service/impl/MITREidDataService_1_0.java

@@ -896,14 +896,17 @@ private void fixObjectReferences() {
 whitelistedSiteOldToNewIdMap.clear();
 for (Long oldGrantId : grantToAccessTokensRefs.keySet()) {
 Set<Long> oldAccessTokenIds = grantToAccessTokensRefs.get(oldGrantId);
-Set<OAuth2AccessTokenEntity> tokens = new HashSet<>();
+
+Long newGrantId = grantOldToNewIdMap.get(oldGrantId);
+ApprovedSite site = approvedSiteRepository.getById(newGrantId);
+
 for(Long oldTokenId : oldAccessTokenIds) {
 Long newTokenId = accessTokenOldToNewIdMap.get(oldTokenId);
-tokens.add(tokenRepository.getAccessTokenById(newTokenId));
+OAuth2AccessTokenEntity token = tokenRepository.getAccessTokenById(newTokenId);
+token.setApprovedSite(site);
+tokenRepository.saveAccessToken(token);
 }
-Long newGrantId = grantOldToNewIdMap.get(oldGrantId);
-ApprovedSite site = approvedSiteRepository.getById(newGrantId);
-site.setApprovedAccessTokens(tokens);
+
 approvedSiteRepository.save(site);
 }
 accessTokenOldToNewIdMap.clear();

openid-connect-server/src/main/java/org/mitre/openid/connect/service/impl/MITREidDataService_1_1.java

@@ -909,14 +909,17 @@ private void fixObjectReferences() {
 accessTokenToIdTokenRefs.clear();
 for (Long oldGrantId : grantToAccessTokensRefs.keySet()) {
 Set<Long> oldAccessTokenIds = grantToAccessTokensRefs.get(oldGrantId);
-Set<OAuth2AccessTokenEntity> tokens = new HashSet<>();
+
+Long newGrantId = grantOldToNewIdMap.get(oldGrantId);
+ApprovedSite site = approvedSiteRepository.getById(newGrantId);
+
 for(Long oldTokenId : oldAccessTokenIds) {
 Long newTokenId = accessTokenOldToNewIdMap.get(oldTokenId);
-tokens.add(tokenRepository.getAccessTokenById(newTokenId));
+OAuth2AccessTokenEntity token = tokenRepository.getAccessTokenById(newTokenId);
+token.setApprovedSite(site);
+tokenRepository.saveAccessToken(token);
 }
-Long newGrantId = grantOldToNewIdMap.get(oldGrantId);
-ApprovedSite site = approvedSiteRepository.getById(newGrantId);
-site.setApprovedAccessTokens(tokens);
+
 approvedSiteRepository.save(site);
 }
 accessTokenOldToNewIdMap.clear();

openid-connect-server/src/main/java/org/mitre/openid/connect/service/impl/MITREidDataService_1_2.java

@@ -900,14 +900,17 @@ private void fixObjectReferences() {
 accessTokenToIdTokenRefs.clear();
 for (Long oldGrantId : grantToAccessTokensRefs.keySet()) {
 Set<Long> oldAccessTokenIds = grantToAccessTokensRefs.get(oldGrantId);
-Set<OAuth2AccessTokenEntity> tokens = new HashSet<OAuth2AccessTokenEntity>();
+
+Long newGrantId = grantOldToNewIdMap.get(oldGrantId);
+ApprovedSite site = approvedSiteRepository.getById(newGrantId);
+
 for(Long oldTokenId : oldAccessTokenIds) {
 Long newTokenId = accessTokenOldToNewIdMap.get(oldTokenId);
-tokens.add(tokenRepository.getAccessTokenById(newTokenId));
+OAuth2AccessTokenEntity token = tokenRepository.getAccessTokenById(newTokenId);
+token.setApprovedSite(site);
+tokenRepository.saveAccessToken(token);
 }
-Long newGrantId = grantOldToNewIdMap.get(oldGrantId);
-ApprovedSite site = approvedSiteRepository.getById(newGrantId);
-site.setApprovedAccessTokens(tokens);
+
 approvedSiteRepository.save(site);
 }
 accessTokenOldToNewIdMap.clear();