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

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

@@ -47,6 +47,7 @@ import javax.persistence.Temporal;
 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 @@ import com.nimbusds.jwt.JWT;
 	@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 @@ import com.nimbusds.jwt.JWT;
 @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 class OAuth2AccessTokenEntity implements OAuth2AccessToken {
 		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 java.util.Set;
 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 class ApprovedSite {
 		}
 	}
 
-	@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 @@ package org.mitre.openid.connect.service;
 
 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.ClientDetailsEntity;
 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 class JpaOAuth2TokenRepository implements OAuth2TokenRepository {
 		}
 
 	}
+	
+	@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,10 +250,6 @@ public class DefaultOAuth2ProviderTokenService implements OAuth2TokenEntityServi
 				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();
 
@@ -261,13 +257,14 @@ public class DefaultOAuth2ProviderTokenService implements OAuth2TokenEntityServi
 
 				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 @@ package org.mitre.openid.connect.service.impl;
 
 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 class DefaultApprovedSiteService implements ApprovedSiteService {
 	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 @@ public class DefaultApprovedSiteService implements ApprovedSiteService {
 		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 @@ public class MITREidDataService_1_0 extends MITREidDataServiceSupport implements
 		whitelistedSiteOldToNewIdMap.clear();
 		for (Long oldGrantId : grantToAccessTokensRefs.keySet()) {
 			Set<Long> oldAccessTokenIds = grantToAccessTokensRefs.get(oldGrantId);
-			Set<OAuth2AccessTokenEntity> tokens = new HashSet<>();
+
-			for(Long oldTokenId : oldAccessTokenIds) {
-				Long newTokenId = accessTokenOldToNewIdMap.get(oldTokenId);
-				tokens.add(tokenRepository.getAccessTokenById(newTokenId));
-			}
 			Long newGrantId = grantOldToNewIdMap.get(oldGrantId);
 			ApprovedSite site = approvedSiteRepository.getById(newGrantId);
-			site.setApprovedAccessTokens(tokens);
+
+			for(Long oldTokenId : oldAccessTokenIds) {
+				Long newTokenId = accessTokenOldToNewIdMap.get(oldTokenId);
+				OAuth2AccessTokenEntity token = tokenRepository.getAccessTokenById(newTokenId);
+				token.setApprovedSite(site);
+				tokenRepository.saveAccessToken(token);
+			}
+			
 			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 @@ public class MITREidDataService_1_1 extends MITREidDataServiceSupport implements
 		accessTokenToIdTokenRefs.clear();
 		for (Long oldGrantId : grantToAccessTokensRefs.keySet()) {
 			Set<Long> oldAccessTokenIds = grantToAccessTokensRefs.get(oldGrantId);
-			Set<OAuth2AccessTokenEntity> tokens = new HashSet<>();
+
-			for(Long oldTokenId : oldAccessTokenIds) {
-				Long newTokenId = accessTokenOldToNewIdMap.get(oldTokenId);
-				tokens.add(tokenRepository.getAccessTokenById(newTokenId));
-			}
 			Long newGrantId = grantOldToNewIdMap.get(oldGrantId);
 			ApprovedSite site = approvedSiteRepository.getById(newGrantId);
-			site.setApprovedAccessTokens(tokens);
+
+			for(Long oldTokenId : oldAccessTokenIds) {
+				Long newTokenId = accessTokenOldToNewIdMap.get(oldTokenId);
+				OAuth2AccessTokenEntity token = tokenRepository.getAccessTokenById(newTokenId);
+				token.setApprovedSite(site);
+				tokenRepository.saveAccessToken(token);
+			}
+			
 			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 @@ public class MITREidDataService_1_2 extends MITREidDataServiceSupport implements
 		accessTokenToIdTokenRefs.clear();
 		for (Long oldGrantId : grantToAccessTokensRefs.keySet()) {
 			Set<Long> oldAccessTokenIds = grantToAccessTokensRefs.get(oldGrantId);
-			Set<OAuth2AccessTokenEntity> tokens = new HashSet<OAuth2AccessTokenEntity>();
+
-			for(Long oldTokenId : oldAccessTokenIds) {
-				Long newTokenId = accessTokenOldToNewIdMap.get(oldTokenId);
-				tokens.add(tokenRepository.getAccessTokenById(newTokenId));
-			}
 			Long newGrantId = grantOldToNewIdMap.get(oldGrantId);
 			ApprovedSite site = approvedSiteRepository.getById(newGrantId);
-			site.setApprovedAccessTokens(tokens);
+
+			for(Long oldTokenId : oldAccessTokenIds) {
+				Long newTokenId = accessTokenOldToNewIdMap.get(oldTokenId);
+				OAuth2AccessTokenEntity token = tokenRepository.getAccessTokenById(newTokenId);
+				token.setApprovedSite(site);
+				tokenRepository.saveAccessToken(token);
+			}
+			
 			approvedSiteRepository.save(site);
 		}
 		accessTokenOldToNewIdMap.clear();

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

@@ -16,12 +16,17 @@
  *******************************************************************************/
 package org.mitre.openid.connect.service.impl;
 
+import static org.mitre.util.JsonUtils.readMap;
+import static org.mitre.util.JsonUtils.readSet;
+import static org.mitre.util.JsonUtils.writeNullSafeArray;
+
 import java.io.IOException;
 import java.io.Serializable;
 import java.text.ParseException;
 import java.util.Date;
 import java.util.HashMap;
 import java.util.HashSet;
+import java.util.List;
 import java.util.Map;
 import java.util.Map.Entry;
 import java.util.Set;
@@ -63,10 +68,6 @@ import com.nimbusds.jose.JWSAlgorithm;
 import com.nimbusds.jose.jwk.JWKSet;
 import com.nimbusds.jwt.JWTParser;
 
-import static org.mitre.util.JsonUtils.readMap;
-import static org.mitre.util.JsonUtils.readSet;
-import static org.mitre.util.JsonUtils.writeNullSafeArray;
-
 /**
  *
  * Data service to import and export MITREid 1.2 configuration.
@@ -367,7 +368,7 @@ public class MITREidDataService_1_3 extends MITREidDataServiceSupport implements
 			writer.name(USER_ID).value(site.getUserId());
 			writer.name(ALLOWED_SCOPES);
 			writeNullSafeArray(writer, site.getAllowedScopes());
-			Set<OAuth2AccessTokenEntity> tokens = site.getApprovedAccessTokens();
+			List<OAuth2AccessTokenEntity> tokens = tokenRepository.getAccessTokensForApprovedSite(site);
 			writer.name(APPROVED_ACCESS_TOKENS);
 			writer.beginArray();
 			for (OAuth2AccessTokenEntity token : tokens) {
@@ -1274,14 +1275,17 @@ public class MITREidDataService_1_3 extends MITREidDataServiceSupport implements
 		accessTokenToIdTokenRefs.clear();
 		for (Long oldGrantId : grantToAccessTokensRefs.keySet()) {
 			Set<Long> oldAccessTokenIds = grantToAccessTokensRefs.get(oldGrantId);
-			Set<OAuth2AccessTokenEntity> tokens = new HashSet<OAuth2AccessTokenEntity>();
+
-			for(Long oldTokenId : oldAccessTokenIds) {
-				Long newTokenId = accessTokenOldToNewIdMap.get(oldTokenId);
-				tokens.add(tokenRepository.getAccessTokenById(newTokenId));
-			}
 			Long newGrantId = grantOldToNewIdMap.get(oldGrantId);
 			ApprovedSite site = approvedSiteRepository.getById(newGrantId);
-			site.setApprovedAccessTokens(tokens);
+
+			for(Long oldTokenId : oldAccessTokenIds) {
+				Long newTokenId = accessTokenOldToNewIdMap.get(oldTokenId);
+				OAuth2AccessTokenEntity token = tokenRepository.getAccessTokenById(newTokenId);
+				token.setApprovedSite(site);
+				tokenRepository.saveAccessToken(token);
+			}
+			
 			approvedSiteRepository.save(site);
 		}
 		accessTokenOldToNewIdMap.clear();