Spring Data ACL permission filtering support for persistence layer

AclClass.java

import lombok.*;
import org.springframework.data.annotation.Immutable;

import javax.persistence.*;

@Entity
@Immutable
@Table(name = "acl_class")
@AllArgsConstructor
@NoArgsConstructor
@Getter
@EqualsAndHashCode
@ToString
public final class AclClass {

    @Id
    @GeneratedValue(strategy = GenerationType.IDENTITY)
    private Long id;

    @Column(name = "class", nullable = false)
    private String className;
}

AclEntry.java

import lombok.*;
import org.springframework.data.annotation.Immutable;

import javax.persistence.*;

@Entity
@Immutable
@Table(name = "acl_entry", uniqueConstraints = {
        @UniqueConstraint(name = "_ak_acl_object_identity_ace_order", columnNames = {"acl_object_identity", "ace_order"})
})
@AllArgsConstructor
@NoArgsConstructor
@Getter
@EqualsAndHashCode
@ToString
public final class AclEntry {

    @Id
    @GeneratedValue(strategy = GenerationType.IDENTITY)
    private Long id;

    @ManyToOne(optional = false)
    @JoinColumn(name = "acl_object_identity", referencedColumnName = "id", nullable = false)
    private AclObjectIdentity aclObjectIdentity;

    @Column(name = "ace_order", nullable = false)
    private int aceOrder;

    @ManyToOne(optional = false)
    @JoinColumn(name = "sid", referencedColumnName = "id", nullable = false)
    private AclSid aclSid;

    @Column(name = "mask", nullable = false)
    private int mask;

    @Column(name = "granting", nullable = false)
    private boolean granting;

    @Column(name = "audit_success", nullable = false)
    private boolean auditSuccess;

    @Column(name = "audit_failure", nullable = false)
    private boolean auditFailure;
}

AclJpaRepository.java

import org.springframework.data.domain.Page;
import org.springframework.data.domain.Pageable;
import org.springframework.data.jpa.domain.Specification;
import org.springframework.data.jpa.repository.JpaRepository;
import org.springframework.data.repository.NoRepositoryBean;
import org.springframework.security.acls.model.Permission;

import java.util.List;

/**
 * ACL specific extension of {@link JpaRepository}. Extends by supporting collection filtering
 * based on ACL {@link Permission permissions}.
 *
 * @param <T>  Entity domain type
 * @param <ID> Unique identifier's type
 * @author 0x1C1B
 */

@NoRepositoryBean
public interface AclJpaRepository<T, ID> extends JpaRepository<T, ID> {

    /**
     * Finds all available entities filtered by ACL permission.
     *
     * @param permission Permission filter criteria
     * @return Returns a list of all matching entities
     */

    List<T> findAll(Permission permission);

    /**
     * Fetches all available entities filtered by ACL permission as a {@link Page}.
     *
     * @param permission Permission filter criteria
     * @return Returns a Page of entities matching the permission criteria
     */

    Page<T> findAll(Pageable pageable, Permission permission);

    /**
     * Finds all available entities filtered by ACL permission and matching the given
     * {@link Specification}.
     *
     * @param spec       Given specification
     * @param permission Permission filter criteria
     * @return Returns a list of all matching entities
     */

    List<T> findAll(Specification<T> spec, Permission permission);

    /**
     * Fetches all available entities matching the given {@link Specification} and
     * filtered by ACL permission as a {@link Page}.
     *
     * @param spec       Given specification
     * @param permission Permission filter criteria
     * @return Returns a Page of entities matching the permission criteria
     */

    Page<T> findAll(Specification<T> spec, Pageable pageable, Permission permission);
}

AclObjectIdentity.java

import lombok.*;
import org.springframework.data.annotation.Immutable;

import javax.persistence.*;

@Entity
@Immutable
@Table(name = "acl_object_identity", uniqueConstraints = {
        @UniqueConstraint(name = "_ak_object_id_class_object_id_identity", columnNames = {"object_id_class", "object_id_identity"})
})
@AllArgsConstructor
@NoArgsConstructor
@Getter
@EqualsAndHashCode
@ToString
public final class AclObjectIdentity {

    @Id
    @GeneratedValue(strategy = GenerationType.IDENTITY)
    private Long id;

    @ManyToOne(optional = false)
    @JoinColumn(name = "object_id_class", referencedColumnName = "id", nullable = false)
    private AclClass objectIdClass;

    @Column(name = "object_id_identity", nullable = false)
    private Long objectIdIdentity;

    @ManyToOne
    @JoinColumn(name = "parent_object", referencedColumnName = "id")
    private AclObjectIdentity parentObject;

    @ManyToOne(optional = false)
    @JoinColumn(name = "owner_sid", referencedColumnName = "id", nullable = false)
    private AclSid ownerSid;
}

AclSid.java

import lombok.*;
import org.springframework.data.annotation.Immutable;

import javax.persistence.*;

@Entity
@Immutable
@Table(name = "acl_sid", uniqueConstraints = {
        @UniqueConstraint(name = "_ak_sid_principal", columnNames = {"sid", "principal"})
})
@AllArgsConstructor
@NoArgsConstructor
@Getter
@EqualsAndHashCode
@ToString
public final class AclSid {

    @Id
    @GeneratedValue(strategy = GenerationType.IDENTITY)
    private Long id;

    @Column(name = "principal", nullable = false)
    private boolean principal;

    @Column(name = "sid", nullable = false, length = 100)
    private String sid;
}

SimpleAclJpaRepository.java

import org.springframework.data.domain.Page;
import org.springframework.data.domain.PageImpl;
import org.springframework.data.domain.Pageable;
import org.springframework.data.domain.Sort;
import org.springframework.data.jpa.domain.Specification;
import org.springframework.data.jpa.repository.query.QueryUtils;
import org.springframework.data.jpa.repository.support.JpaEntityInformation;
import org.springframework.data.jpa.repository.support.JpaEntityInformationSupport;
import org.springframework.data.jpa.repository.support.SimpleJpaRepository;
import org.springframework.data.repository.support.PageableExecutionUtils;
import org.springframework.lang.Nullable;
import org.springframework.security.acls.domain.PrincipalSid;
import org.springframework.security.acls.model.Permission;
import org.springframework.security.core.Authentication;
import org.springframework.security.core.context.SecurityContextHolder;
import org.springframework.security.core.userdetails.UserDetails;
import org.springframework.util.Assert;

import javax.persistence.EntityManager;
import javax.persistence.TypedQuery;
import javax.persistence.criteria.*;
import java.io.Serializable;
import java.util.Collections;
import java.util.List;

/**
 * Default implementation of the {@link AclJpaRepository} interface.
 * This will offer you a more sophisticated interface than the plain EntityManager.
 *
 * @param <T>  Entity domain type
 * @param <ID> Unique identifier's type
 */

@SuppressWarnings({"unchecked", "WeakerAccess"})
public class SimpleAclJpaRepository<T, ID extends Serializable>
        extends SimpleJpaRepository<T, ID> implements AclJpaRepository<T, ID> {

    private JpaEntityInformation<T, ?> entityInformation;
    private EntityManager entityManager;

    public SimpleAclJpaRepository(JpaEntityInformation<T, ?> entityInformation, EntityManager entityManager) {

        super(entityInformation, entityManager);

        this.entityInformation = entityInformation;
        this.entityManager = entityManager;
    }

    public SimpleAclJpaRepository(Class<T> domainClass, EntityManager entityManager) {
        this(JpaEntityInformationSupport.getEntityInformation(domainClass, entityManager), entityManager);
    }

    private static long executeCountQuery(TypedQuery<Long> query) {

        Assert.notNull(query, "TypedQuery must not be null!");
        List<Long> totals = query.getResultList();

        return totals.stream().mapToLong(total -> null == total ? 0 : total).sum();
    }

    @Override
    public List<T> findAll(Permission permission) {
        return findAll((Specification) null, permission);
    }

    @Override
    public Page<T> findAll(Pageable pageable, Permission permission) {
        return findAll(null, pageable, permission);
    }

    @Override
    public List<T> findAll(Specification<T> spec, Permission permission) {

        Authentication authentication = SecurityContextHolder.getContext().getAuthentication();

        if (null == authentication || !authentication.isAuthenticated()) {
            throw new IllegalStateException("Permission filtering not possible for anonymous user");
        }

        UserDetails userDetails = (UserDetails) authentication.getPrincipal();
        PrincipalSid sid = new PrincipalSid(userDetails.getUsername());

        return this.getQuery(spec, Sort.unsorted(), sid, permission).getResultList();
    }

    @Override
    public Page<T> findAll(Specification<T> spec, Pageable pageable, Permission permission) {

        Authentication authentication = SecurityContextHolder.getContext().getAuthentication();

        if (null == authentication || !authentication.isAuthenticated()) {
            throw new IllegalStateException("Permission filtering not possible for anonymous user");
        }

        UserDetails userDetails = (UserDetails) authentication.getPrincipal();
        PrincipalSid sid = new PrincipalSid(userDetails.getUsername());

        TypedQuery<T> query = getQuery(spec, pageable, sid, permission);

        return pageable.isUnpaged() ? new PageImpl<>(query.getResultList()) :
                readPage(query, this.getDomainClass(), pageable, spec, sid, permission);
    }

    /**
     * Reads the given {@link TypedQuery} into a {@link Page} applying the given {@link Pageable},
     * {@link Specification} and permission filter.
     *
     * @param query       Typed JPA query
     * @param domainClass Class of domain entity
     * @param pageable    Pageable configuration
     * @param spec        Additional specification
     * @param sid         ACL authorization principal
     * @param permission  Permission filter criteria
     * @param <S>         Domain type
     * @return Returns a Page of entities matching the permission criteria
     */

    protected <S extends T> Page<S> readPage(TypedQuery<S> query, Class<S> domainClass, Pageable pageable,
                                             @Nullable Specification<S> spec, PrincipalSid sid, Permission permission) {

        if (pageable.isPaged()) {

            query.setFirstResult((int) pageable.getOffset());
            query.setMaxResults(pageable.getPageSize());
        }

        return PageableExecutionUtils.getPage(query.getResultList(), pageable,
                () -> executeCountQuery(getCountQuery(spec, domainClass, sid, permission)));
    }

    /**
     * Creates a new {@link TypedQuery} based for the given {@link Specification specification} and
     * {@link Permission permission} filter criteria.
     *
     * @param spec       Additional specification
     * @param pageable   Pageable configuration
     * @param sid        ACL authorization principal
     * @param permission Permission filter criteria
     * @return Returns the related TypedQuery
     */

    protected TypedQuery<T> getQuery(@Nullable Specification<T> spec, Pageable pageable,
                                     PrincipalSid sid, Permission permission) {

        Sort sort = pageable.isPaged() ? pageable.getSort() : Sort.unsorted();
        return getQuery(spec, this.getDomainClass(), sort, sid, permission);
    }

    /**
     * Creates a new {@link TypedQuery} based for the given {@link Specification specification},
     * {@link Permission permission} filter criteria and {@link Sort}.
     *
     * @param spec       Additional specification
     * @param sort       Sorting configuration
     * @param sid        ACL authorization principal
     * @param permission Permission filter criteria
     * @return Returns the related TypedQuery
     */

    protected TypedQuery<T> getQuery(@Nullable Specification<T> spec, Sort sort,
                                     PrincipalSid sid, Permission permission) {

        return this.getQuery(spec, this.getDomainClass(), sort, sid, permission);
    }

    /**
     * Creates a new {@link TypedQuery} based for the given {@link Specification specification},
     * {@link Permission permission} filter criteria and {@link Sort}.
     *
     * @param spec        Additional specification
     * @param domainClass Class of domain entity
     * @param sort        Sorting configuration
     * @param sid         ACL authorization principal
     * @param permission  Permission filter criteria
     * @param <S>         Domain type
     * @return Returns the related TypedQuery
     */

    protected <S extends T> TypedQuery<S> getQuery(@Nullable Specification<S> spec, Class<S> domainClass, Sort sort,
                                                   PrincipalSid sid, Permission permission) {

        CriteriaBuilder criteriaBuilder = entityManager.getCriteriaBuilder();
        CriteriaQuery<S> criteriaQuery = criteriaBuilder.createQuery(domainClass);
        Root<S> root = applySpecificationToCriteria(spec, domainClass, criteriaQuery, sid, permission);

        criteriaQuery.select(root);

        if (sort.isSorted()) {
            criteriaQuery.orderBy(QueryUtils.toOrders(sort, root, criteriaBuilder));
        }

        return entityManager.createQuery(criteriaQuery);
    }

    protected <S extends T> TypedQuery<Long> getCountQuery(@Nullable Specification<S> spec, Class<S> domainClass,
                                                           PrincipalSid sid, Permission permission) {

        CriteriaBuilder criteriaBuilder = entityManager.getCriteriaBuilder();
        CriteriaQuery<Long> criteriaQuery = criteriaBuilder.createQuery(Long.class);
        Root<S> root = applySpecificationToCriteria(spec, domainClass, criteriaQuery, sid, permission);

        if (criteriaQuery.isDistinct()) {

            criteriaQuery.select(criteriaBuilder.countDistinct(root));

        } else {

            criteriaQuery.select(criteriaBuilder.count(root));
        }

        criteriaQuery.orderBy(Collections.emptyList());
        return entityManager.createQuery(criteriaQuery);
    }

    private <S, U extends T> Root<U> applySpecificationToCriteria(@Nullable Specification<U> spec,
                                                                  Class<U> domainClass, CriteriaQuery<S> query,
                                                                  PrincipalSid sid, Permission permission) {

        Assert.notNull(domainClass, "Domain class must not be null!");
        Assert.notNull(query, "CriteriaQuery must not be null!");

        Root<U> root = query.from(domainClass);

        if (null == spec) {

            query.where(filterPermitted(root, query, domainClass, sid, permission));
            return root;

        } else {

            CriteriaBuilder criteriaBuilder = entityManager.getCriteriaBuilder();
            Predicate predicate = spec.toPredicate(root, query, criteriaBuilder);

            if (null != predicate) {

                query.where(criteriaBuilder.and(predicate, filterPermitted(root, query, domainClass, sid, permission)));

            } else {

                query.where(filterPermitted(root, query, domainClass, sid, permission));
            }

            return root;
        }
    }

    private <S, U extends T> Predicate filterPermitted(Root<U> root, CriteriaQuery<S> query,
                                                       Class<U> domainClass, PrincipalSid sid, Permission permission) {

        return root.<Long>get(entityInformation.getRequiredIdAttribute().getName())
                .in(selectPermittedIds(query, domainClass, sid, permission));
    }

    private <S> Subquery<Long> selectPermittedIds(CriteriaQuery<S> query, Class<?> targetType, PrincipalSid sid,
                                                  Permission permission) {

        CriteriaBuilder criteriaBuilder = entityManager.getCriteriaBuilder();
        Subquery<Long> aclEntryQuery = query.subquery(Long.class);
        Root<AclEntry> root = aclEntryQuery.from(AclEntry.class);
        Join<AclEntry, AclObjectIdentity> aclObjectIdentityJoin = root.join("aclObjectIdentity");

        return aclEntryQuery.select(aclObjectIdentityJoin.get("objectIdIdentity"))
                .where(criteriaBuilder.and(
                        root.<Long>get("aclObjectIdentity").in(selectAclObjectIdentityId(aclEntryQuery, targetType)),
                        criteriaBuilder.equal(root.<Long>get("aclSid"), selectAclSidId(aclEntryQuery, sid)),
                        criteriaBuilder.equal(root.<Integer>get("mask"), permission.getMask())));
    }

    private <S> Subquery<Long> selectAclObjectIdentityId(Subquery<S> query, Class<?> targetType) {

        CriteriaBuilder criteriaBuilder = entityManager.getCriteriaBuilder();
        Subquery<Long> aclObjectIdentityQuery = query.subquery(Long.class);
        Root<AclObjectIdentity> root = aclObjectIdentityQuery.from(AclObjectIdentity.class);

        return aclObjectIdentityQuery.select(root.get("id"))
                .where(criteriaBuilder.equal(root.<Long>get("objectIdClass"),
                        selectAclClassId(aclObjectIdentityQuery, targetType)));
    }

    private <S> Subquery<Long> selectAclSidId(Subquery<S> query, PrincipalSid sid) {

        CriteriaBuilder criteriaBuilder = entityManager.getCriteriaBuilder();
        Subquery<Long> aclSidQuery = query.subquery(Long.class);
        Root<AclSid> root = aclSidQuery.from(AclSid.class);

        return aclSidQuery.select(root.get("id"))
                .where(criteriaBuilder.equal(root.<String>get("sid"), sid.getPrincipal()));
    }

    private <S> Subquery<Long> selectAclClassId(Subquery<S> query, Class<?> targetType) {

        CriteriaBuilder criteriaBuilder = entityManager.getCriteriaBuilder();
        Subquery<Long> aclClassQuery = query.subquery(Long.class);
        Root<AclClass> root = aclClassQuery.from(AclClass.class);

        return aclClassQuery.select(root.get("id"))
                .where(criteriaBuilder.equal(root.<String>get("className"), targetType.getSimpleName()));
    }
}

Hi, this is great! Do you have any usage examples that I can check? Thanks in advance!