import socket
from typing import Union


def receive_message(
    connection: socket.socket, header_len: int
) -> Union[dict[str, bytes], bool]:
    """
    Receives a message from a server or client.

    :param connection: The socket to listen to for messages.
        **MUST BE A SOCKET**
    :type connection: socket.socket
    :param header_len: The length of the header, so that
        it can successfully retrieve data without loss/gain of data
    :type header_len: int
    :return: A dictionary, with two key-value pairs;
        The first key-value pair refers to the header,
        while the second one refers to the actual data
    :rtype: Union[dict["header": bytes, "data": bytes], False
    """

    try:
        header_message = connection.recv(header_len)
        if header_message:
            message_len = int(header_message)
            data = connection.recv(message_len)
            return {"header": header_message, "data": data}
    except ConnectionResetError:
        # This is most likely where clients will disconnect
        pass
    return False