Make permissions checks for ongoing connections more granular (#1563)

Signed-off-by: Adrian Sutton <adrian.sutton@consensys.net>

ethereum/p2p/src/main/java/tech/pegasys/pantheon/ethereum/p2p/network/NodePermissioningAdapter.java

@@ -71,10 +71,10 @@ class NodePermissioningAdapter extends PeerPermissions {
         return inboundIsPermitted(localNode, remotePeer);
       case RLPX_ALLOW_NEW_OUTBOUND_CONNECTION:
         return outboundIsPermitted(localNode, remotePeer);
-      case RLPX_ALLOW_ONGOING_CONNECTION:
+      case RLPX_ALLOW_ONGOING_LOCALLY_INITIATED_CONNECTION:
-        // TODO: This should probably check outbound || inbound, or the check should be more
-        // granular
         return outboundIsPermitted(localNode, remotePeer);
+      case RLPX_ALLOW_ONGOING_REMOTELY_INITIATED_CONNECTION:
+        return inboundIsPermitted(localNode, remotePeer);
       default:
         // Return false for unknown / unhandled permissions
         LOG.error(

ethereum/p2p/src/main/java/tech/pegasys/pantheon/ethereum/p2p/permissions/PeerPermissions.java

@@ -48,7 +48,8 @@ public abstract class PeerPermissions implements AutoCloseable {
     DISCOVERY_ALLOW_OUTBOUND_NEIGHBORS_REQUEST,
     RLPX_ALLOW_NEW_INBOUND_CONNECTION,
     RLPX_ALLOW_NEW_OUTBOUND_CONNECTION,
-    RLPX_ALLOW_ONGOING_CONNECTION
+    RLPX_ALLOW_ONGOING_REMOTELY_INITIATED_CONNECTION,
+    RLPX_ALLOW_ONGOING_LOCALLY_INITIATED_CONNECTION
   }
 
   /**

ethereum/p2p/src/main/java/tech/pegasys/pantheon/ethereum/p2p/rlpx/RlpxAgent.java

@@ -164,7 +164,9 @@ public class RlpxAgent {
 
   public Optional<CompletableFuture<PeerConnection>> getPeerConnection(final Peer peer) {
     final RlpxConnection connection = connectionsById.get(peer.getId());
-    return isNull(connection) ? Optional.empty() : Optional.of(connection.getFuture());
+    return Optional.ofNullable(connection)
+        .filter(conn -> !conn.isFailedOrDisconnected())
+        .map(RlpxConnection::getFuture);
   }
 
   /**
@@ -278,7 +280,8 @@ public class RlpxAgent {
 
     connectionsToCheck.forEach(
         connection -> {
-          if (!peerPermissions.allowOngoingConnection(connection.getPeer())) {
+          if (!peerPermissions.allowOngoingConnection(
+              connection.getPeer(), connection.initiatedRemotely())) {
             connection.disconnect(DisconnectReason.REQUESTED);
           }
         });

ethereum/p2p/src/main/java/tech/pegasys/pantheon/ethereum/p2p/rlpx/connections/PeerRlpxPermissions.java

@@ -48,12 +48,15 @@ public class PeerRlpxPermissions {
         localNode.getPeer(), peer, Action.RLPX_ALLOW_NEW_INBOUND_CONNECTION);
   }
 
-  public boolean allowOngoingConnection(final Peer peer) {
+  public boolean allowOngoingConnection(final Peer peer, final boolean remotelyInitiated) {
     if (!localNode.isReady()) {
       return false;
     }
-    return peerPermissions.isPermitted(
+    final Action action =
-        localNode.getPeer(), peer, Action.RLPX_ALLOW_ONGOING_CONNECTION);
+        remotelyInitiated
+            ? Action.RLPX_ALLOW_ONGOING_REMOTELY_INITIATED_CONNECTION
+            : Action.RLPX_ALLOW_ONGOING_LOCALLY_INITIATED_CONNECTION;
+    return peerPermissions.isPermitted(localNode.getPeer(), peer, action);
   }
 
   public void subscribeUpdate(final PermissionsUpdateCallback callback) {

ethereum/p2p/src/test/java/tech/pegasys/pantheon/ethereum/p2p/network/NodePermissioningAdapterTest.java

@@ -322,8 +322,8 @@ public class NodePermissioningAdapterTest {
   }
 
   @Test
-  public void allowOngoingConnection() {
+  public void allowOngoingLocallyInitiatedConnection() {
-    final Action action = Action.RLPX_ALLOW_ONGOING_CONNECTION;
+    final Action action = Action.RLPX_ALLOW_ONGOING_LOCALLY_INITIATED_CONNECTION;
 
     mockControllerPermissions(true, false);
     assertThat(adapter.isPermitted(localNode, remoteNode, action)).isTrue();
@@ -338,6 +338,23 @@ public class NodePermissioningAdapterTest {
     assertThat(adapter.isPermitted(localNode, remoteNode, action)).isFalse();
   }
 
+  @Test
+  public void allowOngoingRemotelyInitiatedConnection() {
+    final Action action = Action.RLPX_ALLOW_ONGOING_REMOTELY_INITIATED_CONNECTION;
+
+    mockControllerPermissions(true, false);
+    assertThat(adapter.isPermitted(localNode, remoteNode, action)).isFalse();
+
+    mockControllerPermissions(false, true);
+    assertThat(adapter.isPermitted(localNode, remoteNode, action)).isTrue();
+
+    mockControllerPermissions(true, true);
+    assertThat(adapter.isPermitted(localNode, remoteNode, action)).isTrue();
+
+    mockControllerPermissions(false, false);
+    assertThat(adapter.isPermitted(localNode, remoteNode, action)).isFalse();
+  }
+
   @Test
   public void subscribeUpdate_firesWhenBlockAdded() {
     final AtomicBoolean updateDispatched = new AtomicBoolean(false);

ethereum/p2p/src/test/java/tech/pegasys/pantheon/ethereum/p2p/rlpx/RlpxAgentTest.java

@@ -557,7 +557,7 @@ public class RlpxAgentTest {
         .isPermitted(
             eq(localNode.getPeer()),
             eq(nonPermittedPeer),
-            eq(Action.RLPX_ALLOW_ONGOING_CONNECTION));
+            eq(Action.RLPX_ALLOW_ONGOING_LOCALLY_INITIATED_CONNECTION));
     peerPermissions.testDispatchUpdate(true, Optional.empty());
 
     assertThat(agent.getConnectionCount()).isEqualTo(1);
@@ -584,7 +584,7 @@ public class RlpxAgentTest {
         .isPermitted(
             eq(localNode.getPeer()),
             eq(nonPermittedPeer),
-            eq(Action.RLPX_ALLOW_ONGOING_CONNECTION));
+            eq(Action.RLPX_ALLOW_ONGOING_LOCALLY_INITIATED_CONNECTION));
     peerPermissions.testDispatchUpdate(true, Optional.of(Arrays.asList(nonPermittedPeer)));
 
     assertThat(agent.getConnectionCount()).isEqualTo(1);
@@ -594,6 +594,64 @@ public class RlpxAgentTest {
     assertThat(nonPermittedConnection.isDisconnected()).isTrue();
   }
 
+  @Test
+  public void permissionsUpdate_permissionsRestrictedForRemotelyInitiatedConnections()
+      throws ExecutionException, InterruptedException {
+    final Peer locallyConnectedPeer = createPeer();
+    final Peer remotelyConnectedPeer = createPeer();
+    startAgent();
+    final PeerConnection locallyInitiatedConnection = agent.connect(locallyConnectedPeer).get();
+    final PeerConnection remotelyInitiatedConnection =
+        MockPeerConnection.create(remotelyConnectedPeer);
+    connectionInitializer.simulateIncomingConnection(remotelyInitiatedConnection);
+
+    // Sanity check
+    assertThat(agent.getConnectionCount()).isEqualTo(2);
+
+    doReturn(false)
+        .when(peerPermissions)
+        .isPermitted(
+            eq(localNode.getPeer()),
+            any(),
+            eq(Action.RLPX_ALLOW_ONGOING_REMOTELY_INITIATED_CONNECTION));
+    peerPermissions.testDispatchUpdate(true, Optional.empty());
+
+    assertThat(agent.getConnectionCount()).isEqualTo(1);
+    assertThat(agent.getPeerConnection(locallyConnectedPeer)).isNotEmpty();
+    assertThat(agent.getPeerConnection(remotelyConnectedPeer)).isEmpty();
+    assertThat(locallyInitiatedConnection.isDisconnected()).isFalse();
+    assertThat(remotelyInitiatedConnection.isDisconnected()).isTrue();
+  }
+
+  @Test
+  public void permissionsUpdate_permissionsRestrictedForLocallyInitiatedConnections()
+      throws ExecutionException, InterruptedException {
+    final Peer locallyConnectedPeer = createPeer();
+    final Peer remotelyConnectedPeer = createPeer();
+    startAgent();
+    final PeerConnection locallyInitiatedConnection = agent.connect(locallyConnectedPeer).get();
+    final PeerConnection remotelyInitiatedConnection =
+        MockPeerConnection.create(remotelyConnectedPeer);
+    connectionInitializer.simulateIncomingConnection(remotelyInitiatedConnection);
+
+    // Sanity check
+    assertThat(agent.getConnectionCount()).isEqualTo(2);
+
+    doReturn(false)
+        .when(peerPermissions)
+        .isPermitted(
+            eq(localNode.getPeer()),
+            any(),
+            eq(Action.RLPX_ALLOW_ONGOING_LOCALLY_INITIATED_CONNECTION));
+    peerPermissions.testDispatchUpdate(true, Optional.empty());
+
+    assertThat(agent.getConnectionCount()).isEqualTo(1);
+    assertThat(agent.getPeerConnection(locallyConnectedPeer)).isEmpty();
+    assertThat(agent.getPeerConnection(remotelyConnectedPeer)).isNotEmpty();
+    assertThat(locallyInitiatedConnection.isDisconnected()).isTrue();
+    assertThat(remotelyInitiatedConnection.isDisconnected()).isFalse();
+  }
+
   @Test
   public void permissionsUpdate_permissionsRelaxedWithNoListOfPeers()
       throws ExecutionException, InterruptedException {