JPAOAuthDataProvider.java
/**
* Licensed to the Apache Software Foundation (ASF) under one
* or more contributor license agreements. See the NOTICE file
* distributed with this work for additional information
* regarding copyright ownership. The ASF licenses this file
* to you under the Apache License, Version 2.0 (the
* "License"); you may not use this file except in compliance
* with the License. You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing,
* software distributed under the License is distributed on an
* "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
* KIND, either express or implied. See the License for the
* specific language governing permissions and limitations
* under the License.
*/
package org.apache.cxf.rs.security.oauth2.provider;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.LinkedList;
import java.util.List;
import java.util.Map;
import java.util.function.Function;
import jakarta.persistence.EntityManager;
import jakarta.persistence.EntityManagerFactory;
import jakarta.persistence.EntityTransaction;
import jakarta.persistence.TypedQuery;
import org.apache.cxf.helpers.CastUtils;
import org.apache.cxf.rs.security.oauth2.common.AccessTokenRegistration;
import org.apache.cxf.rs.security.oauth2.common.Client;
import org.apache.cxf.rs.security.oauth2.common.OAuthPermission;
import org.apache.cxf.rs.security.oauth2.common.ServerAccessToken;
import org.apache.cxf.rs.security.oauth2.common.UserSubject;
import org.apache.cxf.rs.security.oauth2.tokens.bearer.BearerAccessToken;
import org.apache.cxf.rs.security.oauth2.tokens.refresh.RefreshToken;
/**
* Provides a Jpa BMT implementation for OAuthDataProvider.
*
* If your application runs in a container and if you want to use
* container managed persistence, you'll have to override
* the following methods :
* <ul>
* <li> {@link #getEntityManager()}</li>
* <li> {@link #commitIfNeeded(EntityManager)}</li>
* <li> {@link #closeIfNeeded(EntityManager)}</li>
* </ul>
*/
public class JPAOAuthDataProvider extends AbstractOAuthDataProvider {
private static final String CLIENT_QUERY = "SELECT client FROM Client client"
+ " INNER JOIN client.resourceOwnerSubject ros";
private EntityManagerFactory entityManagerFactory;
public void setEntityManagerFactory(EntityManagerFactory emf) {
this.entityManagerFactory = emf;
}
@Override
public Client doGetClient(final String clientId) throws OAuthServiceException {
return execute(em -> {
return em.find(Client.class, clientId);
});
}
protected <T> T execute(EntityManagerOperation<T> operation) {
EntityManager em = getEntityManager();
T value;
try {
value = operation.apply(em);
} finally {
closeIfNeeded(em);
}
return value;
}
protected <T> T executeInTransaction(EntityManagerOperation<T> operation) {
EntityManager em = getEntityManager();
EntityTransaction transaction = null;
T value;
try {
transaction = beginIfNeeded(em);
value = operation.apply(em);
flushIfNeeded(em);
commitIfNeeded(em);
} catch (RuntimeException e) {
if (transaction != null) {
transaction.rollback();
}
throw e;
} finally {
closeIfNeeded(em);
}
return value;
}
public void setClient(final Client client) {
executeInTransaction(em -> {
if (client.getResourceOwnerSubject() != null) {
UserSubject sub =
em.find(UserSubject.class, client.getResourceOwnerSubject().getId());
if (sub == null) {
em.persist(client.getResourceOwnerSubject());
} else {
client.setResourceOwnerSubject(sub);
}
}
boolean clientExists = em.createQuery("SELECT count(client) from Client client "
+ "where client.clientId = :id", Long.class)
.setParameter("id", client.getClientId())
.getSingleResult() > 0;
if (clientExists) {
em.merge(client);
} else {
em.persist(client);
}
return null;
});
}
@Override
protected void doRemoveClient(final Client c) {
executeInTransaction(em -> {
Client clientToRemove = em.getReference(Client.class, c.getClientId());
em.remove(clientToRemove);
return null;
});
}
@Override
public List<Client> getClients(final UserSubject resourceOwner) {
return executeInTransaction(em -> {
return getClientsQuery(resourceOwner, em).getResultList();
});
}
@Override
public List<ServerAccessToken> getAccessTokens(final Client c, final UserSubject sub) {
return executeInTransaction(em -> {
return CastUtils.cast(getTokensQuery(c, sub, em).getResultList());
});
}
@Override
public List<RefreshToken> getRefreshTokens(final Client c, final UserSubject sub) {
return executeInTransaction(em -> {
return getRefreshTokensQuery(c, sub, em).getResultList();
});
}
@Override
public ServerAccessToken getAccessToken(final String accessToken) throws OAuthServiceException {
return executeInTransaction(em -> {
TypedQuery<BearerAccessToken> query = em.createQuery("SELECT t FROM BearerAccessToken t"
+ " WHERE t.tokenKey = :tokenKey", BearerAccessToken.class)
.setParameter("tokenKey", accessToken);
if (query.getResultList().isEmpty()) {
return null;
}
return query.getSingleResult();
});
}
@Override
protected void doRevokeAccessToken(final ServerAccessToken at) {
executeInTransaction(em -> {
ServerAccessToken tokenToRemove = em.getReference(at.getClass(), at.getTokenKey());
em.remove(tokenToRemove);
return null;
});
}
@Override
protected void linkRefreshTokenToAccessToken(final RefreshToken rt, final ServerAccessToken at) {
super.linkRefreshTokenToAccessToken(rt, at);
executeInTransaction(em -> {
em.merge(at);
return null;
});
}
@Override
protected RefreshToken getRefreshToken(final String refreshTokenKey) {
return executeInTransaction(em -> {
return em.find(RefreshToken.class, refreshTokenKey);
});
}
@Override
protected void doRevokeRefreshToken(final RefreshToken rt) {
executeInTransaction(em -> {
RefreshToken tokentoRemove = em.getReference(RefreshToken.class, rt.getTokenKey());
em.remove(tokentoRemove);
return null;
});
}
@Override
protected ServerAccessToken doCreateAccessToken(AccessTokenRegistration atReg) {
ServerAccessToken at = super.doCreateAccessToken(atReg);
// we override this in order to get rid of elementCollections directly injected
// from another entity
// this can be the case when using multiple cmt dataProvider operation in a single entityManager
// lifespan
if (at.getAudiences() != null) {
at.setAudiences(new ArrayList<>(at.getAudiences()));
}
if (at.getExtraProperties() != null) {
at.setExtraProperties(new HashMap<String, String>(at.getExtraProperties()));
}
if (at.getScopes() != null) {
at.setScopes(new ArrayList<>(at.getScopes()));
}
if (at.getParameters() != null) {
at.setParameters(new HashMap<String, String>(at.getParameters()));
}
return at;
}
protected void saveAccessToken(final ServerAccessToken serverToken) {
executeInTransaction(em -> {
List<OAuthPermission> perms = new LinkedList<>();
for (OAuthPermission perm : serverToken.getScopes()) {
OAuthPermission permSaved = em.find(OAuthPermission.class, perm.getPermission());
if (permSaved != null) {
perms.add(permSaved);
} else {
em.persist(perm);
perms.add(perm);
}
}
serverToken.setScopes(perms);
if (serverToken.getSubject() != null) {
UserSubject sub = em.find(UserSubject.class, serverToken.getSubject().getId());
if (sub == null) {
em.persist(serverToken.getSubject());
} else {
sub = em.merge(serverToken.getSubject());
serverToken.setSubject(sub);
}
}
// ensure we have a managed association
// (needed for OpenJPA : InvalidStateException: Encountered unmanaged object)
if (serverToken.getClient() != null) {
serverToken.setClient(em.find(Client.class, serverToken.getClient().getClientId()));
}
em.persist(serverToken);
return null;
});
}
protected void saveRefreshToken(RefreshToken refreshToken) {
persistEntity(refreshToken);
}
protected void persistEntity(final Object entity) {
executeInTransaction(em -> {
em.persist(entity);
return null;
});
}
protected void removeEntity(final Object entity) {
executeInTransaction(em -> {
em.remove(entity);
return null;
});
}
protected TypedQuery<Client> getClientsQuery(UserSubject resourceOwnerSubject, EntityManager entityManager) {
if (resourceOwnerSubject == null) {
return entityManager.createQuery(CLIENT_QUERY, Client.class);
}
return entityManager.createQuery(CLIENT_QUERY + " WHERE ros.login = :login", Client.class).
setParameter("login", resourceOwnerSubject.getLogin());
}
protected TypedQuery<BearerAccessToken> getTokensQuery(Client c, UserSubject resourceOwnerSubject,
EntityManager entityManager) {
return getQuery("BearerAccessToken", c, resourceOwnerSubject, entityManager, BearerAccessToken.class);
}
protected TypedQuery<RefreshToken> getRefreshTokensQuery(Client c, UserSubject resourceOwnerSubject,
EntityManager entityManager) {
return getQuery("RefreshToken", c, resourceOwnerSubject, entityManager, RefreshToken.class);
}
private static <T> TypedQuery<T> getQuery(String table, Client c, UserSubject resourceOwnerSubject,
EntityManager entityManager, Class<T> resultClass) {
StringBuilder query = new StringBuilder(64).append("SELECT t FROM ").append(table).append(" t");
Map<String, Object> parameterMap = new HashMap<>();
if (c != null || resourceOwnerSubject != null) {
query.append(" WHERE");
if (c != null) {
query.append(" t.client.clientId = :clientId");
parameterMap.put("clientId", c.getClientId());
}
if (resourceOwnerSubject != null) {
if (!parameterMap.isEmpty()) {
query.append(" AND");
}
query.append(" t.subject.login = :login");
parameterMap.put("login", resourceOwnerSubject.getLogin());
}
}
TypedQuery<T> typedQuery = entityManager.createQuery(query.toString(), resultClass);
for (Map.Entry<String, Object> entry : parameterMap.entrySet()) {
typedQuery.setParameter(entry.getKey(), entry.getValue());
}
return typedQuery;
}
/**
* Returns the entityManaged used for the current operation.
*/
protected EntityManager getEntityManager() {
return entityManagerFactory.createEntityManager();
}
/**
* Begins the current transaction.
*
* This method needs to be overridden in a CMT environment.
*/
protected EntityTransaction beginIfNeeded(EntityManager em) {
EntityTransaction tx = em.getTransaction();
tx.begin();
return tx;
}
/**
* Flush the current transaction.
*/
protected void flushIfNeeded(EntityManager em) {
em.flush();
}
/**
* Commits the current transaction.
*
* This method needs to be overridden in a CMT environment.
*/
protected void commitIfNeeded(EntityManager em) {
em.getTransaction().commit();
}
/**
* Closes the current em.
*
* This method needs to be overriden in a CMT environment.
*/
protected void closeIfNeeded(EntityManager em) {
em.close();
}
public interface EntityManagerOperation<T> extends Function<EntityManager, T> {
}
}