Add logic to handle EstablishedPeer messages

Author: Askaholic

server/game_connection_matrix.py

@@ -0,0 +1,27 @@
+from collections import defaultdict
+
+
+class ConnectionMatrix:
+    def __init__(self, established_peers: dict[int, set[int]]):
+        self.established_peers = established_peers
+
+    def get_unconnected_peer_ids(self) -> set[int]:
+        unconnected_peer_ids: set[int] = set()
+
+        # Group players by number of connected peers
+        players_by_num_peers = defaultdict(list)
+        for player_id, peer_ids in self.established_peers.items():
+            players_by_num_peers[len(peer_ids)].append((player_id, peer_ids))
+
+        # Mark players with least number of connections as unconnected if they
+        # don't meet the connection threshold. Each time a player is marked as
+        # 'unconnected', remaining players need 1 less connection to be
+        # considered connected.
+        connected_peers = dict(self.established_peers)
+        for num_connected, peers in sorted(players_by_num_peers.items()):
+            if num_connected < len(connected_peers) - 1:
+                for player_id, peer_ids in peers:
+                    unconnected_peer_ids.add(player_id)
+                    del connected_peers[player_id]
+
+        return unconnected_peer_ids

server/gameconnection.py

@@ -6,7 +6,7 @@
 import contextlib
 import json
 import logging
-from typing import Any
+from typing import Any, Optional
 
 from sqlalchemy import select
 
@@ -62,6 +62,10 @@ def __init__(
         self.player = player
         player.game_connection = self  # Set up weak reference to self
         self.game = game
+        # None if the EstablishedPeers message is not implemented by the game
+        # version/mode used by the player. For instance, matchmaker might have
+        # it, but custom games might not.
+        self.established_peer_ids: Optional[set[int]] = None
 
         self.setup_timeout = setup_timeout
 
@@ -561,15 +565,21 @@ async def handle_established_peer(self, peer_id: str):
         - `peer_id`: The identifier of the peer that this connection received
             the message from
         """
-        pass
+        if self.established_peer_ids is None:
+            self.established_peer_ids = set()
+
+        self.established_peer_ids.add(int(peer_id))
 
     async def handle_disconnected_peer(self, peer_id: str):
         """
         Sent by the lobby when a player disconnects from a peer. This can happen
         when a peer is rejoining in which case that peer will have reported a
         "Rejoining" status, or if the peer has exited the game.
         """
-        pass
+        if self.established_peer_ids is None:
+            self.established_peer_ids = set()
+
+        self.established_peer_ids.discard(int(peer_id))
 
     def _mark_dirty(self):
         if self.game:

server/games/game.py

@@ -17,6 +17,7 @@
     game_stats,
     matchmaker_queue_game
 )
+from server.game_connection_matrix import ConnectionMatrix
 from server.games.game_results import (
     ArmyOutcome,
     ArmyReportedOutcome,
@@ -211,13 +212,41 @@ def players(self) -> list[Player]:
 
     def get_connected_players(self) -> list[Player]:
         """
-        Get a collection of all players currently connected to the game.
+        Get a collection of all players currently connected to the host.
         """
         return [
             player for player in self._connections.keys()
             if player.id in self._configured_player_ids
         ]
 
+    def get_unconnected_players_from_peer_matrix(
+        self,
+    ) -> Optional[list[Player]]:
+        """
+        Get a list of players who are not fully connected to the game based on
+        the established peers matrix if possible. The EstablishedPeers messages
+        might not be implemented by the game in which case this returns None.
+        """
+        if any(
+            conn.established_peer_ids is None
+            for conn in self._connections.values()
+        ):
+            return None
+
+        matrix = ConnectionMatrix(
+            established_peers={
+                player.id: conn.established_peer_ids
+                for player, conn in self._connections.items()
+            }
+        )
+        unconnected_peer_ids = matrix.get_unconnected_peer_ids()
+
+        return [
+            player
+            for player in self._connections.keys()
+            if player.id in unconnected_peer_ids
+        ]
+
     def _is_observer(self, player: Player) -> bool:
         army = self.get_player_option(player.id, "Army")
         return army is None or army < 0

server/ladder_service/ladder_service.py

@@ -667,6 +667,12 @@ async def launch_match(
         try:
             await game.wait_launched(60 + 10 * len(guests))
         except asyncio.TimeoutError:
+            unconnected_players = game.get_unconnected_players_from_peer_matrix()
+            if unconnected_players is not None:
+                raise NotConnectedError(unconnected_players)
+
+            # If the connection matrix was not available, fall back to looking
+            # at who was connected to the host only.
             connected_players = game.get_connected_players()
             raise NotConnectedError([
                 player for player in guests

tests/integration_tests/test_matchmaker_violations.py

@@ -1,11 +1,21 @@
 import asyncio
 from datetime import datetime, timezone
 
+import pytest
+
 from tests.utils import fast_forward
 
 from .conftest import connect_and_sign_in, read_until_command
-from .test_game import open_fa, queue_players_for_matchmaking, start_search
+from .test_game import (
+    client_response,
+    open_fa,
+    queue_players_for_matchmaking,
+    send_player_options,
+    start_search
+)
 from .test_parties import accept_party_invite, invite_to_party
+from .test_teammatchmaker import \
+    queue_players_for_matchmaking as queue_players_for_matchmaking_2v2
 
 
 @fast_forward(360)
@@ -18,8 +28,7 @@ async def test_violation_for_guest_timeout(mocker, lobby_server):
 
     # The player that queued last will be the host
     async def launch_game_and_timeout_guest():
-        await read_until_command(host, "game_launch")
-        await open_fa(host)
+        await client_response(host, timeout=60)
         await read_until_command(host, "game_info")
 
         await read_until_command(guest, "game_launch")
@@ -110,6 +119,126 @@ async def launch_game_and_timeout_guest():
     }
 
 
+@fast_forward(360)
+async def test_violation_established_peer(mocker, lobby_server):
+    mocker.patch(
+        "server.ladder_service.violation_service.datetime_now",
+        return_value=datetime(2022, 2, 5, tzinfo=timezone.utc)
+    )
+    protos, ids = await queue_players_for_matchmaking_2v2(lobby_server)
+    host, guest1, guest2, guest3 = protos
+    host_id, guest1_id, guest2_id, guest3_id = ids
+
+    # Connect all players to the host
+    await asyncio.gather(*[
+        client_response(proto, timeout=60)
+        for proto in protos
+    ])
+    await send_player_options(
+        host,
+        [host_id, "Color", 1],
+        [guest1_id, "Color", 2],
+        [guest2_id, "Color", 3],
+        [guest3_id, "Color", 4],
+    )
+
+    # Set up connection matrix
+    for id in (guest1_id, guest2_id, guest3_id):
+        await host.send_message({
+            "target": "game",
+            "command": "EstablishedPeer",
+            "args": [id],
+        })
+    for id in (host_id, guest2_id):
+        await guest1.send_message({
+            "target": "game",
+            "command": "EstablishedPeer",
+            "args": [id],
+        })
+    for id in (host_id, guest1_id):
+        await guest2.send_message({
+            "target": "game",
+            "command": "EstablishedPeer",
+            "args": [id],
+        })
+    # Guest3 only connects to the host
+    await guest3.send_message({
+        "target": "game",
+        "command": "EstablishedPeer",
+        "args": [host_id],
+    })
+
+    await read_until_command(host, "match_cancelled", timeout=120)
+    msg = await read_until_command(guest3, "search_violation", timeout=10)
+    assert msg == {
+        "command": "search_violation",
+        "count": 1,
+        "time": "2022-02-05T00:00:00+00:00",
+    }
+    for proto in (host, guest1, guest2):
+        with pytest.raises(asyncio.TimeoutError):
+            await read_until_command(proto, "search_violation", timeout=10)
+
+
+@fast_forward(360)
+async def test_violation_established_peer_multiple(mocker, lobby_server):
+    mocker.patch(
+        "server.ladder_service.violation_service.datetime_now",
+        return_value=datetime(2022, 2, 5, tzinfo=timezone.utc)
+    )
+    protos, ids = await queue_players_for_matchmaking_2v2(lobby_server)
+    host, guest1, guest2, guest3 = protos
+    host_id, guest1_id, guest2_id, guest3_id = ids
+
+    # Connect all players to the host
+    await asyncio.gather(*[
+        client_response(proto, timeout=60)
+        for proto in protos
+    ])
+    await send_player_options(
+        host,
+        [host_id, "Color", 1],
+        [guest1_id, "Color", 2],
+        [guest2_id, "Color", 3],
+        [guest3_id, "Color", 4],
+    )
+
+    # Set up connection matrix
+    for id in (guest1_id, guest2_id, guest3_id):
+        await host.send_message({
+            "target": "game",
+            "command": "EstablishedPeer",
+            "args": [id],
+        })
+    # Guests only connect to the host
+    await guest1.send_message({
+        "target": "game",
+        "command": "EstablishedPeer",
+        "args": [host_id],
+    })
+    await guest2.send_message({
+        "target": "game",
+        "command": "EstablishedPeer",
+        "args": [host_id],
+    })
+    await guest3.send_message({
+        "target": "game",
+        "command": "EstablishedPeer",
+        "args": [host_id],
+    })
+
+    await read_until_command(host, "match_cancelled", timeout=120)
+    for proto in (guest1, guest2, guest3):
+        msg = await read_until_command(proto, "search_violation", timeout=10)
+        assert msg == {
+            "command": "search_violation",
+            "count": 1,
+            "time": "2022-02-05T00:00:00+00:00",
+        }
+    with pytest.raises(asyncio.TimeoutError):
+        await read_until_command(host, "search_violation", timeout=10)
+
+
 @fast_forward(360)
 async def test_violation_persisted_across_logins(mocker, lobby_server):
     mocker.patch(