Skip to content

Instantly share code, notes, and snippets.

Show Gist options
  • Save bisignam/37c1465ba241e26f298e577d24a4c95e to your computer and use it in GitHub Desktop.
Save bisignam/37c1465ba241e26f298e577d24a4c95e to your computer and use it in GitHub Desktop.
<!-- Token Store -->
<hz:hazelcast id="tokenGrid">
<hz:config>
<hz:group name="${cluster.group.name}" password="${cluster.group.password}" />
<hz:map name="accessTokenStore" />
<hz:map name="authenticationToAccessTokenStore" />
<hz:map name="userNameToAccessTokenStore" />
<hz:map name="clientIdToAccessTokenStore" />
<hz:map name="refreshTokenStore" />
<hz:map name="accessTokenToRefreshTokenStore" />
<hz:map name="authenticationStore" />
<hz:map name="refreshTokenAuthenticationStore" />
<hz:map name="refreshTokenToAccessTokenStore" />
<hz:map name="expiryMap" />
<hz:map name="authorizationCodeStore" />
</hz:config>
</hz:hazelcast>
<bean id="tokenStore" class="org.zenithar.applications.portal.dao.oauth.HazelcastTokenStore">
<constructor-arg name="instance" ref="tokenGrid" />
</bean>
<bean id="authorizationCodeServices" class="org.zenithar.applications.portal.dao.oauth.HazelcastAuthorizationCodeServices">
<constructor-arg name="instance" ref="tokenGrid" />
</bean>
import com.hazelcast.core.HazelcastInstance;
import com.hazelcast.core.IMap;
import org.springframework.beans.factory.InitializingBean;
import org.springframework.security.oauth2.provider.code.AuthorizationRequestHolder;
import org.springframework.security.oauth2.provider.code.RandomValueAuthorizationCodeServices;
import org.springframework.util.Assert;
public class HazelcastAuthorizationCodeServices extends RandomValueAuthorizationCodeServices implements InitializingBean {
protected IMap<String, AuthorizationRequestHolder> authorizationCodeStore = null;
private HazelcastInstance instance;
public HazelcastAuthorizationCodeServices(HazelcastInstance instance) {
this.instance = instance;
}
@Override
protected void store(String code, AuthorizationRequestHolder authentication) {
this.authorizationCodeStore.put(code, authentication);
}
@Override
public AuthorizationRequestHolder remove(String code) {
AuthorizationRequestHolder auth = this.authorizationCodeStore.remove(code);
return auth;
}
@Override
public void afterPropertiesSet() throws Exception {
authorizationCodeStore = instance.getMap("authorizationCodeStore");
Assert.notNull(authorizationCodeStore, "AuthorizationCodeStore Map could not be found !");
}
}
import com.hazelcast.core.HazelcastInstance;
import com.hazelcast.core.IMap;
import org.springframework.beans.factory.InitializingBean;
import org.springframework.security.oauth2.common.OAuth2AccessToken;
import org.springframework.security.oauth2.common.OAuth2RefreshToken;
import org.springframework.security.oauth2.provider.OAuth2Authentication;
import org.springframework.security.oauth2.provider.token.AuthenticationKeyGenerator;
import org.springframework.security.oauth2.provider.token.DefaultAuthenticationKeyGenerator;
import org.springframework.security.oauth2.provider.token.TokenStore;
import org.springframework.util.Assert;
import java.util.Collection;
import java.util.Collections;
import java.util.HashSet;
import java.util.concurrent.TimeUnit;
public class HazelcastTokenStore implements TokenStore, InitializingBean {
private HazelcastInstance instance;
private IMap<String, OAuth2AccessToken> accessTokenStore;
private IMap<String, OAuth2AccessToken> authenticationToAccessTokenStore;
private IMap<String, Collection<OAuth2AccessToken>> userNameToAccessTokenStore;
private IMap<String, Collection<OAuth2AccessToken>> clientIdToAccessTokenStore;
private IMap<String, OAuth2RefreshToken> refreshTokenStore;
private IMap<String, String> accessTokenToRefreshTokenStore;
private IMap<String, OAuth2Authentication> authenticationStore;
private IMap<String, OAuth2Authentication> refreshTokenAuthenticationStore;
private IMap<String, String> refreshTokenToAcessTokenStore;
private IMap<String, String> expiryMap;
private AuthenticationKeyGenerator authenticationKeyGenerator = new DefaultAuthenticationKeyGenerator();
public HazelcastTokenStore(HazelcastInstance instance) {
this.instance = instance;
}
/**
* Convenience method for super admin users to remove all tokens (useful for testing, not really in production)
*/
public void clear() {
accessTokenStore.clear();
authenticationToAccessTokenStore.clear();
userNameToAccessTokenStore.clear();
clientIdToAccessTokenStore.clear();
refreshTokenStore.clear();
accessTokenToRefreshTokenStore.clear();
authenticationStore.clear();
refreshTokenAuthenticationStore.clear();
refreshTokenToAcessTokenStore.clear();
}
public void setAuthenticationKeyGenerator(AuthenticationKeyGenerator authenticationKeyGenerator) {
this.authenticationKeyGenerator = authenticationKeyGenerator;
}
public int getAccessTokenCount() {
Assert.state(accessTokenStore.isEmpty() || accessTokenStore.size() >= accessTokenToRefreshTokenStore.size(),
"Too many refresh tokens");
Assert.state(accessTokenStore.size() == authenticationToAccessTokenStore.size(),
"Inconsistent token store state");
Assert.state(accessTokenStore.size() <= authenticationStore.size(), "Inconsistent authentication store state");
return accessTokenStore.size();
}
public int getRefreshTokenCount() {
Assert.state(refreshTokenStore.size() == refreshTokenToAcessTokenStore.size(),
"Inconsistent refresh token store state");
return accessTokenStore.size();
}
public int getExpiryTokenCount() {
return expiryMap.size();
}
public OAuth2AccessToken getAccessToken(OAuth2Authentication authentication) {
String key = authenticationKeyGenerator.extractKey(authentication);
OAuth2AccessToken accessToken = authenticationToAccessTokenStore.get(key);
if (accessToken != null
&& !key.equals(authenticationKeyGenerator.extractKey(readAuthentication(accessToken.getValue())))) {
// Keep the stores consistent (maybe the same user is represented by this authentication but the details
// have changed)
storeAccessToken(accessToken, authentication);
}
return accessToken;
}
public OAuth2Authentication readAuthentication(OAuth2AccessToken token) {
return readAuthentication(token.getValue());
}
public OAuth2Authentication readAuthentication(String token) {
return this.authenticationStore.get(token);
}
public OAuth2Authentication readAuthenticationForRefreshToken(OAuth2RefreshToken token) {
return readAuthenticationForRefreshToken(token.getValue());
}
public OAuth2Authentication readAuthenticationForRefreshToken(String token) {
return this.refreshTokenAuthenticationStore.get(token);
}
public void storeAccessToken(OAuth2AccessToken token, OAuth2Authentication authentication) {
this.accessTokenStore.put(token.getValue(), token);
this.authenticationStore.put(token.getValue(), authentication);
this.authenticationToAccessTokenStore.put(authenticationKeyGenerator.extractKey(authentication), token);
if (!authentication.isClientOnly()) {
addToCollection(this.userNameToAccessTokenStore, authentication.getName(), token);
}
addToCollection(this.clientIdToAccessTokenStore, authentication.getAuthorizationRequest().getClientId(), token);
if (token.getExpiration() != null) {
// Remove existing expiry for this token if present
long ttl = System.currentTimeMillis() - token.getExpiration().getTime();
expiryMap.put(token.getValue(), token.getValue(), ttl, TimeUnit.MILLISECONDS);
}
if (token.getRefreshToken() != null && token.getRefreshToken().getValue() != null) {
this.refreshTokenToAcessTokenStore.put(token.getRefreshToken().getValue(), token.getValue());
this.accessTokenToRefreshTokenStore.put(token.getValue(), token.getRefreshToken().getValue());
}
}
private void addToCollection(IMap<String, Collection<OAuth2AccessToken>> store, String key,
OAuth2AccessToken token) {
if (!store.containsKey(key)) {
synchronized (store) {
if (!store.containsKey(key)) {
store.put(key, new HashSet<OAuth2AccessToken>());
}
}
}
store.get(key).add(token);
}
public void removeAccessToken(OAuth2AccessToken accessToken) {
removeAccessToken(accessToken.getValue());
}
public OAuth2AccessToken readAccessToken(String tokenValue) {
return this.accessTokenStore.get(tokenValue);
}
public void removeAccessToken(String tokenValue) {
OAuth2AccessToken removed = this.accessTokenStore.remove(tokenValue);
String refresh = this.accessTokenToRefreshTokenStore.remove(tokenValue);
if (refresh != null) {
// Don't remove the refresh token itself - it's up to the caller to do that
this.refreshTokenToAcessTokenStore.remove(tokenValue);
}
OAuth2Authentication authentication = this.authenticationStore.remove(tokenValue);
if (authentication != null) {
this.authenticationToAccessTokenStore.remove(authenticationKeyGenerator.extractKey(authentication));
Collection<OAuth2AccessToken> tokens;
tokens = this.userNameToAccessTokenStore.get(authentication.getName());
if (tokens != null) {
tokens.remove(removed);
}
String clientId = authentication.getAuthorizationRequest().getClientId();
tokens = this.clientIdToAccessTokenStore.get(clientId);
if (tokens != null) {
tokens.remove(removed);
}
this.authenticationToAccessTokenStore.remove(authenticationKeyGenerator.extractKey(authentication));
}
}
public void storeRefreshToken(OAuth2RefreshToken refreshToken, OAuth2Authentication authentication) {
this.refreshTokenStore.put(refreshToken.getValue(), refreshToken);
this.refreshTokenAuthenticationStore.put(refreshToken.getValue(), authentication);
}
public OAuth2RefreshToken readRefreshToken(String tokenValue) {
return this.refreshTokenStore.get(tokenValue);
}
public void removeRefreshToken(OAuth2RefreshToken refreshToken) {
removeRefreshToken(refreshToken.getValue());
}
public void removeRefreshToken(String tokenValue) {
this.refreshTokenStore.remove(tokenValue);
this.refreshTokenAuthenticationStore.remove(tokenValue);
this.refreshTokenToAcessTokenStore.remove(tokenValue);
}
public void removeAccessTokenUsingRefreshToken(OAuth2RefreshToken refreshToken) {
removeAccessTokenUsingRefreshToken(refreshToken.getValue());
}
private void removeAccessTokenUsingRefreshToken(String refreshToken) {
String accessToken = this.refreshTokenToAcessTokenStore.remove(refreshToken);
if (accessToken != null) {
removeAccessToken(accessToken);
}
}
public Collection<OAuth2AccessToken> findTokensByClientId(String clientId) {
Collection<OAuth2AccessToken> result = clientIdToAccessTokenStore.get(clientId);
return result != null ? Collections.<OAuth2AccessToken> unmodifiableCollection(result) : Collections
.<OAuth2AccessToken> emptySet();
}
public Collection<OAuth2AccessToken> findTokensByUserName(String userName) {
Collection<OAuth2AccessToken> result = userNameToAccessTokenStore.get(userName);
return result != null ? Collections.<OAuth2AccessToken> unmodifiableCollection(result) : Collections
.<OAuth2AccessToken> emptySet();
}
@Override
public void afterPropertiesSet() throws Exception {
accessTokenStore = instance.getMap("accessTokenStore");
Assert.notNull(accessTokenStore, "AccessTokenStore Map could not be found !");
authenticationToAccessTokenStore = instance.getMap("authenticationToAccessTokenStore");
Assert.notNull(authenticationToAccessTokenStore, "AuthenticationToAccessTokenStore Map could not be found !");
userNameToAccessTokenStore = instance.getMap("userNameToAccessTokenStore");
Assert.notNull(userNameToAccessTokenStore, "UserNameToAccessTokenStore Map could not be found !");
clientIdToAccessTokenStore = instance.getMap("clientIdToAccessTokenStore");
Assert.notNull(clientIdToAccessTokenStore, "ClientIdToAccessTokenStore Map could not be found !");
refreshTokenStore = instance.getMap("refreshTokenStore");
Assert.notNull(refreshTokenStore, "RefreshTokenStore Map could not be found !");
accessTokenToRefreshTokenStore = instance.getMap("accessTokenToRefreshTokenStore");
Assert.notNull(accessTokenToRefreshTokenStore, "AccessTokenToRefreshTokenStore Map could not be found !");
authenticationStore = instance.getMap("authenticationStore");
Assert.notNull(authenticationStore, "AuthenticationStore Map could not be found !");
refreshTokenAuthenticationStore = instance.getMap("refreshTokenAuthenticationStore");
Assert.notNull(refreshTokenAuthenticationStore, "RefreshTokenAuthenticationStore Map could not be found !");
refreshTokenToAcessTokenStore = instance.getMap("refreshTokenToAccessTokenStore");
Assert.notNull(refreshTokenToAcessTokenStore, "RefreshTokenToAccessTokenStore Map could not be found !");
expiryMap = instance.getMap("expiryMap");
Assert.notNull(expiryMap, "ExpiryMap Map could not be found !");
}
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment