You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

688 lines
27 KiB

"""
The MIT License (MIT)
Copyright (c) 2015-present Rapptz
Permission is hereby granted, free of charge, to any person obtaining a
copy of this software and associated documentation files (the "Software"),
to deal in the Software without restriction, including without limitation
the rights to use, copy, modify, merge, publish, distribute, sublicense,
and/or sell copies of the Software, and to permit persons to whom the
Software is furnished to do so, subject to the following conditions:
The above copyright notice and this permission notice shall be included in
all copies or substantial portions of the Software.
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS
OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
DEALINGS IN THE SOFTWARE.
Some documentation to refer to:
- Our main web socket (mWS) sends opcode 4 with a guild ID and channel ID.
- The mWS receives VOICE_STATE_UPDATE and VOICE_SERVER_UPDATE.
- We pull the session_id from VOICE_STATE_UPDATE.
- We pull the token, endpoint and server_id from VOICE_SERVER_UPDATE.
- Then we initiate the voice web socket (vWS) pointing to the endpoint.
- We send opcode 0 with the user_id, server_id, session_id and token using the vWS.
- The vWS sends back opcode 2 with an ssrc, port, modes(array) and heartbeat_interval.
- We send a UDP discovery packet to endpoint:port and receive our IP and our port in LE.
- Then we send our IP and port via vWS with opcode 1.
- When that's all done, we receive opcode 4 from the vWS.
- Finally we can transmit data to endpoint:port.
"""
from __future__ import annotations
import select
import socket
import asyncio
import logging
import threading
from typing import TYPE_CHECKING, Optional, Dict, List, Callable, Coroutine, Any, Tuple
from .enums import Enum
from .utils import MISSING, sane_wait_for
from .errors import ConnectionClosed
from .backoff import ExponentialBackoff
from .gateway import DiscordVoiceWebSocket
if TYPE_CHECKING:
from . import abc
from .guild import Guild
from .user import ClientUser
from .member import VoiceState
from .voice_client import VoiceClient
from .types.voice import (
GuildVoiceState as GuildVoiceStatePayload,
VoiceServerUpdate as VoiceServerUpdatePayload,
SupportedModes,
)
WebsocketHook = Optional[Callable[[DiscordVoiceWebSocket, Dict[str, Any]], Coroutine[Any, Any, Any]]]
SocketReaderCallback = Callable[[bytes], Any]
__all__ = ('VoiceConnectionState',)
_log = logging.getLogger(__name__)
class SocketReader(threading.Thread):
def __init__(self, state: VoiceConnectionState, *, start_paused: bool = True) -> None:
super().__init__(daemon=True, name=f'voice-socket-reader:{id(self):#x}')
self.state: VoiceConnectionState = state
self.start_paused = start_paused
self._callbacks: List[SocketReaderCallback] = []
self._running = threading.Event()
self._end = threading.Event()
# If we have paused reading due to having no callbacks
self._idle_paused: bool = True
def register(self, callback: SocketReaderCallback) -> None:
self._callbacks.append(callback)
if self._idle_paused:
self._idle_paused = False
self._running.set()
def unregister(self, callback: SocketReaderCallback) -> None:
try:
self._callbacks.remove(callback)
except ValueError:
pass
else:
if not self._callbacks and self._running.is_set():
# If running is not set, we are either explicitly paused and
# should be explicitly resumed, or we are already idle paused
self._idle_paused = True
self._running.clear()
def pause(self) -> None:
self._idle_paused = False
self._running.clear()
def resume(self, *, force: bool = False) -> None:
if self._running.is_set():
return
# Don't resume if there are no callbacks registered
if not force and not self._callbacks:
# We tried to resume but there was nothing to do, so resume when ready
self._idle_paused = True
return
self._idle_paused = False
self._running.set()
def stop(self) -> None:
self._end.set()
self._running.set()
def run(self) -> None:
self._end.clear()
self._running.set()
if self.start_paused:
self.pause()
try:
self._do_run()
except Exception:
_log.exception('Error in %s', self)
finally:
self.stop()
self._running.clear()
self._callbacks.clear()
def _do_run(self) -> None:
while not self._end.is_set():
if not self._running.is_set():
self._running.wait()
continue
# Since this socket is a non blocking socket, select has to be used to wait on it for reading.
try:
readable, _, _ = select.select([self.state.socket], [], [], 30)
except (ValueError, TypeError, OSError) as e:
_log.debug(
"Select error handling socket in reader, this should be safe to ignore: %s: %s", e.__class__.__name__, e
)
# The socket is either closed or doesn't exist at the moment
continue
if not readable:
continue
try:
data = self.state.socket.recv(2048)
except OSError:
_log.debug('Error reading from socket in %s, this should be safe to ignore', self, exc_info=True)
else:
for cb in self._callbacks:
try:
cb(data)
except Exception:
_log.exception('Error calling %s in %s', cb, self)
class ConnectionFlowState(Enum):
"""Enum representing voice connection flow state."""
# fmt: off
disconnected = 0
set_guild_voice_state = 1
got_voice_state_update = 2
got_voice_server_update = 3
got_both_voice_updates = 4
websocket_connected = 5
got_websocket_ready = 6
got_ip_discovery = 7
connected = 8
# fmt: on
class VoiceConnectionState:
"""Represents the internal state of a voice connection."""
def __init__(self, voice_client: VoiceClient, *, hook: Optional[WebsocketHook] = None) -> None:
self.voice_client = voice_client
self.hook = hook
self.timeout: float = 30.0
self.reconnect: bool = True
self.self_deaf: bool = False
self.self_mute: bool = False
self.token: Optional[str] = None
self.session_id: Optional[str] = None
self.endpoint: Optional[str] = None
self.endpoint_ip: Optional[str] = None
self.server_id: Optional[int] = None
self.ip: Optional[str] = None
self.port: Optional[int] = None
self.voice_port: Optional[int] = None
self.secret_key: List[int] = MISSING
self.ssrc: int = MISSING
self.mode: SupportedModes = MISSING
self.socket: socket.socket = MISSING
self.ws: DiscordVoiceWebSocket = MISSING
self._state: ConnectionFlowState = ConnectionFlowState.disconnected
self._expecting_disconnect: bool = False
self._connected = threading.Event()
self._state_event = asyncio.Event()
self._disconnected = asyncio.Event()
self._runner: Optional[asyncio.Task] = None
self._connector: Optional[asyncio.Task] = None
self._socket_reader = SocketReader(self)
self._socket_reader.start()
@property
def state(self) -> ConnectionFlowState:
return self._state
@state.setter
def state(self, state: ConnectionFlowState) -> None:
if state is not self._state:
_log.debug('Connection state changed to %s', state.name)
self._state = state
self._state_event.set()
self._state_event.clear()
if state is ConnectionFlowState.connected:
self._connected.set()
else:
self._connected.clear()
@property
def guild(self) -> Guild:
return self.voice_client.guild
@property
def user(self) -> ClientUser:
return self.voice_client.user
@property
def supported_modes(self) -> Tuple[SupportedModes, ...]:
return self.voice_client.supported_modes
@property
def self_voice_state(self) -> Optional[VoiceState]:
return self.guild.me.voice
async def voice_state_update(self, data: GuildVoiceStatePayload) -> None:
channel_id = data['channel_id']
if channel_id is None:
self._disconnected.set()
# If we know we're going to get a voice_state_update where we have no channel due to
# being in the reconnect or disconnect flow, we ignore it. Otherwise, it probably wasn't from us.
if self._expecting_disconnect:
self._expecting_disconnect = False
else:
_log.debug('We were externally disconnected from voice.')
await self.disconnect()
return
channel_id = int(channel_id)
self.session_id = data['session_id']
# we got the event while connecting
if self.state in (ConnectionFlowState.set_guild_voice_state, ConnectionFlowState.got_voice_server_update):
if self.state is ConnectionFlowState.set_guild_voice_state:
self.state = ConnectionFlowState.got_voice_state_update
# we moved ourselves
if channel_id != self.voice_client.channel.id:
self._update_voice_channel(channel_id)
else:
self.state = ConnectionFlowState.got_both_voice_updates
return
if self.state is ConnectionFlowState.connected:
self._update_voice_channel(channel_id)
elif self.state is not ConnectionFlowState.disconnected:
if channel_id != self.voice_client.channel.id:
# For some unfortunate reason we were moved during the connection flow
_log.info('Handling channel move while connecting...')
self._update_voice_channel(channel_id)
await self.soft_disconnect(with_state=ConnectionFlowState.got_voice_state_update)
await self.connect(
reconnect=self.reconnect,
timeout=self.timeout,
self_deaf=(self.self_voice_state or self).self_deaf,
self_mute=(self.self_voice_state or self).self_mute,
resume=False,
wait=False,
)
else:
_log.debug('Ignoring unexpected voice_state_update event')
async def voice_server_update(self, data: VoiceServerUpdatePayload) -> None:
previous_token = self.token
previous_server_id = self.server_id
previous_endpoint = self.endpoint
self.token = data['token']
self.server_id = int(data['guild_id'])
endpoint = data.get('endpoint')
if self.token is None or endpoint is None:
_log.warning(
'Awaiting endpoint... This requires waiting. '
'If timeout occurred considering raising the timeout and reconnecting.'
)
return
self.endpoint, _, _ = endpoint.rpartition(':')
if self.endpoint.startswith('wss://'):
# Just in case, strip it off since we're going to add it later
self.endpoint = self.endpoint[6:]
# we got the event while connecting
if self.state in (ConnectionFlowState.set_guild_voice_state, ConnectionFlowState.got_voice_state_update):
# This gets set after READY is received
self.endpoint_ip = MISSING
self._create_socket()
if self.state is ConnectionFlowState.set_guild_voice_state:
self.state = ConnectionFlowState.got_voice_server_update
else:
self.state = ConnectionFlowState.got_both_voice_updates
elif self.state is ConnectionFlowState.connected:
_log.debug('Voice server update, closing old voice websocket')
await self.ws.close(4014)
self.state = ConnectionFlowState.got_voice_server_update
elif self.state is not ConnectionFlowState.disconnected:
# eventual consistency
if previous_token == self.token and previous_server_id == self.server_id and previous_endpoint == self.endpoint:
return
_log.debug('Unexpected server update event, attempting to handle')
await self.soft_disconnect(with_state=ConnectionFlowState.got_voice_server_update)
await self.connect(
reconnect=self.reconnect,
timeout=self.timeout,
self_deaf=(self.self_voice_state or self).self_deaf,
self_mute=(self.self_voice_state or self).self_mute,
resume=False,
wait=False,
)
self._create_socket()
async def connect(
self, *, reconnect: bool, timeout: float, self_deaf: bool, self_mute: bool, resume: bool, wait: bool = True
) -> None:
if self._connector:
self._connector.cancel()
self._connector = None
if self._runner:
self._runner.cancel()
self._runner = None
self.timeout = timeout
self.reconnect = reconnect
self._connector = self.voice_client.loop.create_task(
self._wrap_connect(reconnect, timeout, self_deaf, self_mute, resume), name='Voice connector'
)
if wait:
await self._connector
async def _wrap_connect(self, *args: Any) -> None:
try:
await self._connect(*args)
except asyncio.CancelledError:
_log.debug('Cancelling voice connection')
await self.soft_disconnect()
raise
except asyncio.TimeoutError:
_log.info('Timed out connecting to voice')
await self.disconnect()
raise
except Exception:
_log.exception('Error connecting to voice... disconnecting')
await self.disconnect()
raise
async def _inner_connect(self, reconnect: bool, self_deaf: bool, self_mute: bool, resume: bool) -> None:
for i in range(5):
_log.info('Starting voice handshake... (connection attempt %d)', i + 1)
await self._voice_connect(self_deaf=self_deaf, self_mute=self_mute)
# Setting this unnecessarily will break reconnecting
if self.state is ConnectionFlowState.disconnected:
self.state = ConnectionFlowState.set_guild_voice_state
await self._wait_for_state(ConnectionFlowState.got_both_voice_updates)
_log.info('Voice handshake complete. Endpoint found: %s', self.endpoint)
try:
self.ws = await self._connect_websocket(resume)
await self._handshake_websocket()
break
except ConnectionClosed:
if reconnect:
wait = 1 + i * 2.0
_log.exception('Failed to connect to voice... Retrying in %ss...', wait)
await self.disconnect(cleanup=False)
await asyncio.sleep(wait)
continue
else:
await self.disconnect()
raise
async def _connect(self, reconnect: bool, timeout: float, self_deaf: bool, self_mute: bool, resume: bool) -> None:
_log.info('Connecting to voice...')
await asyncio.wait_for(
self._inner_connect(reconnect=reconnect, self_deaf=self_deaf, self_mute=self_mute, resume=resume),
timeout=timeout,
)
_log.info('Voice connection complete.')
if not self._runner:
self._runner = self.voice_client.loop.create_task(self._poll_voice_ws(reconnect), name='Voice websocket poller')
async def disconnect(self, *, force: bool = True, cleanup: bool = True, wait: bool = False) -> None:
if not force and not self.is_connected():
return
try:
await self._voice_disconnect()
if self.ws:
await self.ws.close()
except Exception:
_log.debug('Ignoring exception disconnecting from voice', exc_info=True)
finally:
self.state = ConnectionFlowState.disconnected
self._socket_reader.pause()
# Stop threads before we unlock waiters so they end properly
if cleanup:
self._socket_reader.stop()
self.voice_client.stop()
# Flip the connected event to unlock any waiters
self._connected.set()
self._connected.clear()
if self.socket:
self.socket.close()
self.ip = MISSING
self.port = MISSING
# Skip this part if disconnect was called from the poll loop task
if wait and not self._inside_runner():
# Wait for the voice_state_update event confirming the bot left the voice channel.
# This prevents a race condition caused by disconnecting and immediately connecting again.
# The new VoiceConnectionState object receives the voice_state_update event containing channel=None while still
# connecting leaving it in a bad state. Since there's no nice way to transfer state to the new one, we have to do this.
try:
await asyncio.wait_for(self._disconnected.wait(), timeout=self.timeout)
except TimeoutError:
_log.debug('Timed out waiting for voice disconnection confirmation')
except asyncio.CancelledError:
pass
if cleanup:
self.voice_client.cleanup()
async def soft_disconnect(self, *, with_state: ConnectionFlowState = ConnectionFlowState.got_both_voice_updates) -> None:
_log.debug('Soft disconnecting from voice')
# Stop the websocket reader because closing the websocket will trigger an unwanted reconnect
if self._runner:
self._runner.cancel()
self._runner = None
try:
if self.ws:
await self.ws.close()
except Exception:
_log.debug('Ignoring exception soft disconnecting from voice', exc_info=True)
finally:
self.state = with_state
self._socket_reader.pause()
if self.socket:
self.socket.close()
self.ip = MISSING
self.port = MISSING
async def move_to(self, channel: Optional[abc.Snowflake], timeout: Optional[float]) -> None:
if channel is None:
# This function should only be called externally so its ok to wait for the disconnect.
await self.disconnect(wait=True)
return
if self.voice_client.channel and channel.id == self.voice_client.channel.id:
return
previous_state = self.state
# this is only an outgoing ws request
# if it fails, nothing happens and nothing changes (besides self.state)
await self._move_to(channel)
last_state = self.state
try:
await self.wait_async(timeout)
except asyncio.TimeoutError:
_log.warning('Timed out trying to move to channel %s in guild %s', channel.id, self.guild.id)
if self.state is last_state:
_log.debug('Reverting to previous state %s', previous_state.name)
self.state = previous_state
def wait(self, timeout: Optional[float] = None) -> bool:
return self._connected.wait(timeout)
async def wait_async(self, timeout: Optional[float] = None) -> None:
await self._wait_for_state(ConnectionFlowState.connected, timeout=timeout)
def is_connected(self) -> bool:
return self.state is ConnectionFlowState.connected
def send_packet(self, packet: bytes) -> None:
self.socket.sendall(packet)
def add_socket_listener(self, callback: SocketReaderCallback) -> None:
_log.debug('Registering socket listener callback %s', callback)
self._socket_reader.register(callback)
def remove_socket_listener(self, callback: SocketReaderCallback) -> None:
_log.debug('Unregistering socket listener callback %s', callback)
self._socket_reader.unregister(callback)
def _inside_runner(self) -> bool:
return self._runner is not None and asyncio.current_task() == self._runner
async def _wait_for_state(
self, state: ConnectionFlowState, *other_states: ConnectionFlowState, timeout: Optional[float] = None
) -> None:
states = (state, *other_states)
while True:
if self.state in states:
return
await sane_wait_for([self._state_event.wait()], timeout=timeout)
async def _voice_connect(self, *, self_deaf: bool = False, self_mute: bool = False) -> None:
channel = self.voice_client.channel
await channel.guild.change_voice_state(channel=channel, self_deaf=self_deaf, self_mute=self_mute)
async def _voice_disconnect(self) -> None:
_log.info(
'The voice handshake is being terminated for Channel ID %s (Guild ID %s)',
self.voice_client.channel.id,
self.voice_client.guild.id,
)
self.state = ConnectionFlowState.disconnected
await self.voice_client.channel.guild.change_voice_state(channel=None)
self._expecting_disconnect = True
self._disconnected.clear()
async def _connect_websocket(self, resume: bool) -> DiscordVoiceWebSocket:
ws = await DiscordVoiceWebSocket.from_connection_state(self, resume=resume, hook=self.hook)
self.state = ConnectionFlowState.websocket_connected
return ws
async def _handshake_websocket(self) -> None:
while not self.ip:
await self.ws.poll_event()
self.state = ConnectionFlowState.got_ip_discovery
while self.ws.secret_key is None:
await self.ws.poll_event()
self.state = ConnectionFlowState.connected
def _create_socket(self) -> None:
self.socket = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
self.socket.setblocking(False)
self._socket_reader.resume()
async def _poll_voice_ws(self, reconnect: bool) -> None:
backoff = ExponentialBackoff()
while True:
try:
await self.ws.poll_event()
except asyncio.CancelledError:
return
except (ConnectionClosed, asyncio.TimeoutError) as exc:
if isinstance(exc, ConnectionClosed):
# The following close codes are undocumented so I will document them here.
# 1000 - normal closure (obviously)
# 4014 - we were externally disconnected (voice channel deleted, we were moved, etc)
# 4015 - voice server has crashed
if exc.code in (1000, 4015):
# Don't call disconnect a second time if the websocket closed from a disconnect call
if not self._expecting_disconnect:
_log.info('Disconnecting from voice normally, close code %d.', exc.code)
await self.disconnect()
break
if exc.code == 4014:
# We were disconnected by discord
# This condition is a race between the main ws event and the voice ws closing
if self._disconnected.is_set():
_log.info('Disconnected from voice by discord, close code %d.', exc.code)
await self.disconnect()
break
# We may have been moved to a different channel
_log.info('Disconnected from voice by force... potentially reconnecting.')
successful = await self._potential_reconnect()
if not successful:
_log.info('Reconnect was unsuccessful, disconnecting from voice normally...')
# Don't bother to disconnect if already disconnected
if self.state is not ConnectionFlowState.disconnected:
await self.disconnect()
break
else:
continue
_log.debug('Not handling close code %s (%s)', exc.code, exc.reason or 'no reason')
if not reconnect:
await self.disconnect()
raise
retry = backoff.delay()
_log.exception('Disconnected from voice... Reconnecting in %.2fs.', retry)
await asyncio.sleep(retry)
await self.disconnect(cleanup=False)
try:
await self._connect(
reconnect=reconnect,
timeout=self.timeout,
self_deaf=(self.self_voice_state or self).self_deaf,
self_mute=(self.self_voice_state or self).self_mute,
resume=False,
)
except asyncio.TimeoutError:
# at this point we've retried 5 times... let's continue the loop.
_log.warning('Could not connect to voice... Retrying...')
continue
async def _potential_reconnect(self) -> bool:
try:
await self._wait_for_state(
ConnectionFlowState.got_voice_server_update,
ConnectionFlowState.got_both_voice_updates,
ConnectionFlowState.disconnected,
timeout=self.timeout,
)
except asyncio.TimeoutError:
return False
else:
if self.state is ConnectionFlowState.disconnected:
return False
previous_ws = self.ws
try:
self.ws = await self._connect_websocket(False)
await self._handshake_websocket()
except (ConnectionClosed, asyncio.TimeoutError):
return False
else:
return True
finally:
await previous_ws.close()
async def _move_to(self, channel: abc.Snowflake) -> None:
await self.voice_client.channel.guild.change_voice_state(channel=channel)
self.state = ConnectionFlowState.set_guild_voice_state
def _update_voice_channel(self, channel_id: Optional[int]) -> None:
self.voice_client.channel = channel_id and self.guild.get_channel(channel_id) # type: ignore