Merge pull request #218 from dmlloyd/identity

Make the security identity available to the association implementation
wildfly · Feb 24, 2017 · c32500b · c32500b
2 parents 396cb08 + 318fc69
commit c32500b
Show file tree

Hide file tree

Showing 2 changed files with 26 additions and 14 deletions.
diff --git a/src/main/java/org/jboss/ejb/protocol/remote/EJBServerChannel.java b/src/main/java/org/jboss/ejb/protocol/remote/EJBServerChannel.java
@@ -376,13 +376,11 @@ void handleSessionOpenRequest(final int invId, final MessageInputStream inputStr
             final Connection connection = channel.getConnection();
             final EJBIdentifier identifier = new EJBIdentifier(appName, moduleName, beanName, distName);
 
-            connection.getLocalIdentity(securityContext).runAs((Runnable) () ->
-                association.receiveSessionOpenRequest(new RemotingSessionOpenRequest(
-                    invId,
-                    identifier,
-                    transactionSupplier
-                ))
-            );
+            association.receiveSessionOpenRequest(new RemotingSessionOpenRequest(
+                invId,
+                identifier,
+                transactionSupplier,
+                connection.getLocalIdentity(securityContext)));
         }
 
         void handleInvocationRequest(final int invId, final InputStream input) throws IOException, ClassNotFoundException {
@@ -433,9 +431,9 @@ void handleInvocationRequest(final int invId, final InputStream input) throws IO
                 identity = connection.getLocalIdentity();
             }
             final RemotingInvocationRequest request = new RemotingInvocationRequest(
-                invId, connection, association, identifier, methodLocator, classResolver, unmarshaller
+                invId, connection, association, identifier, methodLocator, classResolver, unmarshaller, identity
             );
-            invocations.put(new InProgress(request, identity == null ? association.receiveInvocationRequest(request) : identity.runAsFunction(association::receiveInvocationRequest, request)));
+            invocations.put(new InProgress(request, association.receiveInvocationRequest(request)));
         }
     }
 
@@ -485,9 +483,11 @@ private void writeFailedResponse(final int invId, final Exception e) {
     abstract class RemotingRequest implements Request {
         final int invId;
         SessionID sessionId;
+        final SecurityIdentity identity;
 
-        protected RemotingRequest(final int invId) {
+        RemotingRequest(final int invId, final SecurityIdentity identity) {
             this.invId = invId;
+            this.identity = identity;
         }
 
         public Executor getRequestExecutor() {
@@ -510,6 +510,10 @@ public boolean isBlockingCaller() {
             return false;
         }
 
+        public SecurityIdentity getSecurityIdentity() {
+            return identity;
+        }
+
         public void writeNoSuchEJB() {
             final String message = Logs.REMOTING.remoteMessageNoSuchEJB(getEJBIdentifier());
             try (MessageOutputStream os = messageTracker.openMessageUninterruptibly()) {
@@ -617,8 +621,8 @@ final class RemotingSessionOpenRequest extends RemotingRequest implements Sessio
         final ExceptionSupplier<ImportResult<?>, SystemException> transactionSupplier;
         int txnCmd = 0; // assume nobody will ask about the transaction
 
-        RemotingSessionOpenRequest(final int invId, final EJBIdentifier identifier, final ExceptionSupplier<ImportResult<?>, SystemException> transactionSupplier) {
-            super(invId);
+        RemotingSessionOpenRequest(final int invId, final EJBIdentifier identifier, final ExceptionSupplier<ImportResult<?>, SystemException> transactionSupplier, final SecurityIdentity identity) {
+            super(invId, identity);
             this.transactionSupplier = transactionSupplier;
             this.identifier = identifier;
         }
@@ -692,8 +696,8 @@ final class RemotingInvocationRequest extends RemotingRequest implements Invocat
         final Unmarshaller remaining;
         int txnCmd = 0; // assume nobody will ask about the transaction
 
-        RemotingInvocationRequest(final int invId, final Connection connection, final Association association, final EJBIdentifier identifier, final EJBMethodLocator methodLocator, final ServerClassResolver classResolver, final Unmarshaller remaining) {
-            super(invId);
+        RemotingInvocationRequest(final int invId, final Connection connection, final Association association, final EJBIdentifier identifier, final EJBMethodLocator methodLocator, final ServerClassResolver classResolver, final Unmarshaller remaining, final SecurityIdentity identity) {
+            super(invId, identity);
             this.connection = connection;
             this.association = association;
             this.identifier = identifier;

diff --git a/src/main/java/org/jboss/ejb/server/Request.java b/src/main/java/org/jboss/ejb/server/Request.java
@@ -28,6 +28,7 @@
 import org.jboss.ejb.client.EJBIdentifier;
 import org.jboss.ejb.client.SessionID;
 import org.wildfly.common.annotation.NotNull;
+import org.wildfly.security.auth.server.SecurityIdentity;
 
 /**
  * The base type of any EJB server request.  This type is implemented by protocol implementations and consumed by
@@ -84,6 +85,13 @@ default SocketAddress getLocalAddress() {
     @NotNull
     EJBIdentifier getEJBIdentifier();
 
+    /**
+     * Get the security identity that is associated with this invocation.
+     *
+     * @return the security identity, or {@code null} if the connection is not bound to a security domain
+     */
+    SecurityIdentity getSecurityIdentity();
+
     /**
      * Write a message indicating that an exception was thrown by the operation.
      *