use chrono::Utc;
use serde_json::Value;
use std::env::consts;
use std::io::Write;
use std::net::Shutdown;
use std::sync::{Arc, Mutex};
use std::time::{Duration as StdDuration, Instant};
use super::{ConnectionStage, GatewayError};
use websocket::client::Url;
use websocket::message::{CloseData, OwnedMessage};
use websocket::stream::sync::AsTcpStream;
use websocket::sync::client::{Client, ClientBuilder};
use websocket::sync::stream::{TcpStream, TlsStream};
use websocket::WebSocketError;
use ::constants::{self, OpCode, close_codes};
use ::internal::prelude::*;
use ::internal::ws_impl::SenderExt;
use ::model::event::{Event, GatewayEvent};
use ::model::{Game, GuildId, OnlineStatus};
#[cfg(feature="voice")]
use std::sync::mpsc::{self, Receiver as MpscReceiver};
#[cfg(feature="cache")]
use ::client::CACHE;
#[cfg(feature="voice")]
use ::ext::voice::Manager as VoiceManager;
#[cfg(feature="voice")]
use ::http;
#[cfg(feature="cache")]
use ::utils;
pub type WsClient = Client<TlsStream<TcpStream>>;
type CurrentPresence = (Option<Game>, OnlineStatus, bool);
pub struct Shard {
pub client: WsClient,
current_presence: CurrentPresence,
heartbeat_instants: (Option<Instant>, Option<Instant>),
heartbeat_interval: Option<u64>,
last_heartbeat_acknowledged: bool,
seq: u64,
session_id: Option<String>,
shard_info: [u64; 2],
stage: ConnectionStage,
token: Arc<Mutex<String>>,
ws_url: Arc<Mutex<String>>,
#[cfg(feature="voice")]
pub manager: VoiceManager,
#[cfg(feature="voice")]
manager_rx: MpscReceiver<Value>,
}
impl Shard {
pub fn new(ws_url: Arc<Mutex<String>>,
token: Arc<Mutex<String>>,
shard_info: [u64; 2])
-> Result<Shard> {
let client = connect(&*ws_url.lock().unwrap())?;
let current_presence = (None, OnlineStatus::Online, false);
let heartbeat_instants = (None, None);
let heartbeat_interval = None;
let last_heartbeat_acknowledged = true;
let seq = 0;
let stage = ConnectionStage::Handshake;
let session_id = None;
let mut shard = feature_voice! {{
let (tx, rx) = mpsc::channel();
let user = http::get_current_user()?;
Shard {
client,
current_presence,
heartbeat_instants,
heartbeat_interval,
last_heartbeat_acknowledged,
seq,
stage,
token,
session_id,
shard_info,
ws_url,
manager: VoiceManager::new(tx, user.id),
manager_rx: rx,
}
} else {
Shard {
client,
current_presence,
heartbeat_instants,
heartbeat_interval,
last_heartbeat_acknowledged,
seq,
stage,
token,
session_id,
shard_info,
ws_url,
}
}};
shard.identify()?;
Ok(shard)
}
pub fn shard_info(&self) -> [u64; 2] {
self.shard_info
}
pub fn set_afk(&mut self, afk: bool) {
self.current_presence.2 = afk;
self.update_presence();
}
pub fn set_game(&mut self, game: Option<Game>) {
self.current_presence.0 = game;
self.update_presence();
}
pub fn set_status(&mut self, online_status: OnlineStatus) {
self.current_presence.1 = match online_status {
OnlineStatus::Offline => OnlineStatus::Invisible,
other => other,
};
self.update_presence();
}
pub fn set_presence(&mut self,
game: Option<Game>,
mut status: OnlineStatus,
afk: bool) {
if status == OnlineStatus::Offline {
status = OnlineStatus::Invisible;
}
self.current_presence = (game, status, afk);
self.update_presence();
}
#[allow(cyclomatic_complexity)]
pub(crate) fn handle_event(&mut self, event: Result<GatewayEvent>) -> Result<Option<Event>> {
match event {
Ok(GatewayEvent::Dispatch(seq, event)) => {
match event {
Event::Ready(ref ready) => {
self.session_id = Some(ready.ready.session_id.clone());
self.stage = ConnectionStage::Connected;
set_client_timeout(&mut self.client)?;
},
Event::Resumed(_) => {
info!("[Shard {:?}] Resumed", self.shard_info);
self.stage = ConnectionStage::Connected;
},
ref _other => {
#[cfg(feature="voice")]
{
self.voice_dispatch(_other);
}
},
}
self.seq = seq;
Ok(Some(event))
},
Ok(GatewayEvent::Heartbeat(s)) => {
info!("[Shard {:?}] Received shard heartbeat", self.shard_info);
if s > self.seq + 1 {
info!("[Shard {:?}] Received off sequence (them: {}; us: {}); resuming",
self.shard_info,
s,
self.seq);
if self.stage == ConnectionStage::Handshake {
self.stage = ConnectionStage::Identifying;
self.identify()?;
} else {
warn!("[Shard {:?}] Heartbeat during non-Handshake; auto-reconnecting",
self.shard_info);
return self.autoreconnect().and(Ok(None));
}
}
let map = json!({
"d": Value::Null,
"op": OpCode::Heartbeat.num(),
});
self.client.send_json(&map)?;
Ok(None)
},
Ok(GatewayEvent::HeartbeatAck) => {
self.heartbeat_instants.1 = Some(Instant::now());
self.last_heartbeat_acknowledged = true;
Ok(None)
},
Ok(GatewayEvent::Hello(interval)) => {
if interval > 0 {
self.heartbeat_interval = Some(interval);
}
if self.stage == ConnectionStage::Handshake {
self.stage = ConnectionStage::Identifying;
Ok(None)
} else {
self.autoreconnect().and(Ok(None))
}
},
Ok(GatewayEvent::InvalidateSession) => {
info!("[Shard {:?}] Received session invalidation; re-identifying",
self.shard_info);
self.seq = 0;
self.session_id = None;
self.identify()?;
Ok(None)
},
Ok(GatewayEvent::Reconnect) => {
self.reconnect().and(Ok(None))
},
Err(Error::Gateway(GatewayError::Closed(data))) => {
let num = data.as_ref().map(|d| d.status_code);
let reason = data.map(|d| d.reason);
let clean = num == Some(1000);
{
let kind = if clean { "Cleanly" } else { "Uncleanly" };
info!("[Shard {:?}] {} closing with {:?}: {:?}",
self.shard_info,
kind,
num,
reason);
}
match num {
Some(close_codes::UNKNOWN_OPCODE) => warn!("Sent invalid opcode"),
Some(close_codes::DECODE_ERROR) => warn!("Sent invalid message"),
Some(close_codes::NOT_AUTHENTICATED) => {
warn!("Sent no authentication");
return Err(Error::Gateway(GatewayError::NoAuthentication));
},
Some(close_codes::AUTHENTICATION_FAILED) => {
warn!("Sent invalid authentication");
return Err(Error::Gateway(GatewayError::InvalidAuthentication));
},
Some(close_codes::ALREADY_AUTHENTICATED) => warn!("Already authenticated"),
Some(close_codes::INVALID_SEQUENCE) => {
warn!("[Shard {:?}] Sent invalid seq: {}", self.shard_info, self.seq);
self.seq = 0;
},
Some(close_codes::RATE_LIMITED) => warn!("Gateway ratelimited"),
Some(close_codes::INVALID_SHARD) => {
warn!("Sent invalid shard data");
return Err(Error::Gateway(GatewayError::InvalidShardData));
},
Some(close_codes::SHARDING_REQUIRED) => {
error!("Shard has too many guilds");
return Err(Error::Gateway(GatewayError::OverloadedShard));
},
Some(4006) | Some(close_codes::SESSION_TIMEOUT) => {
info!("[Shard {:?}] Invalid session", self.shard_info);
self.session_id = None;
},
Some(other) if !clean => {
warn!("[Shard {:?}] Unknown unclean close {}: {:?}",
self.shard_info,
other,
reason);
},
_ => {},
}
let resume = num.map(|x| {
x != 1000 && x != close_codes::AUTHENTICATION_FAILED && self.session_id.is_some()
}).unwrap_or(false);
if resume {
self.resume().or_else(|_| self.reconnect()).and(Ok(None))
} else {
self.reconnect().and(Ok(None))
}
},
Err(Error::WebSocket(why)) => {
if let WebSocketError::NoDataAvailable = why {
if self.heartbeat_instants.1.is_none() {
return Ok(None);
}
}
warn!("[Shard {:?}] Websocket error: {:?}", self.shard_info, why);
info!("[Shard {:?}] Will attempt to auto-reconnect", self.shard_info);
self.autoreconnect().and(Ok(None))
},
Err(error) => Err(error),
}
}
pub fn latency(&self) -> Option<StdDuration> {
if let (Some(received), Some(sent)) = self.heartbeat_instants {
Some(sent - received)
} else {
None
}
}
pub fn shutdown_clean(&mut self) -> Result<()> {
{
let message = OwnedMessage::Close(Some(CloseData {
status_code: 1000,
reason: String::new(),
}));
self.client.send_message(&message)?;
}
let mut stream = self.client.stream_ref().as_tcp();
stream.flush()?;
stream.shutdown(Shutdown::Both)?;
debug!("[Shard {:?}] Cleanly shutdown shard", self.shard_info);
Ok(())
}
pub fn shutdown(&mut self) -> Result<()> {
let mut stream = self.client.stream_ref().as_tcp();
stream.flush()?;
stream.shutdown(Shutdown::Both)?;
Ok(())
}
pub fn chunk_guilds(&mut self, guild_ids: &[GuildId], limit: Option<u16>, query: Option<&str>) {
let msg = json!({
"op": OpCode::GetGuildMembers.num(),
"d": {
"guild_id": guild_ids.iter().map(|x| x.0).collect::<Vec<u64>>(),
"limit": limit.unwrap_or(0),
"query": query.unwrap_or(""),
},
});
let _ = self.client.send_json(&msg);
}
#[cfg(feature="cache")]
pub fn guilds_handled(&self) -> u16 {
let cache = CACHE.read().unwrap();
let (shard_id, shard_count) = (self.shard_info[0], self.shard_info[1]);
cache.guilds
.keys()
.filter(|guild_id| utils::shard_id(guild_id.0, shard_count) == shard_id)
.count() as u16
}
#[cfg(feature="voice")]
fn voice_dispatch(&mut self, event: &Event) {
if let Event::VoiceStateUpdate(ref update) = *event {
if let Some(guild_id) = update.guild_id {
if let Some(handler) = self.manager.get(guild_id) {
handler.update_state(&update.voice_state);
}
}
}
if let Event::VoiceServerUpdate(ref update) = *event {
if let Some(guild_id) = update.guild_id {
if let Some(handler) = self.manager.get(guild_id) {
handler.update_server(&update.endpoint, &update.token);
}
}
}
}
#[cfg(feature="voice")]
pub(crate) fn cycle_voice_recv(&mut self) {
if let Ok(v) = self.manager_rx.try_recv() {
if let Err(why) = self.client.send_json(&v) {
warn!("[Shard {:?}] Err sending voice msg: {:?}", self.shard_info, why);
}
}
}
pub(crate) fn heartbeat(&mut self) -> Result<()> {
let map = json!({
"d": self.seq,
"op": OpCode::Heartbeat.num(),
});
trace!("[Shard {:?}] Sending heartbeat d: {}", self.shard_info, self.seq);
match self.client.send_json(&map) {
Ok(_) => {
self.heartbeat_instants.0 = Some(Instant::now());
self.last_heartbeat_acknowledged = false;
Ok(())
},
Err(why) => {
match why {
Error::WebSocket(WebSocketError::IoError(err)) => {
if err.raw_os_error() != Some(32) {
debug!("[Shard {:?}] Err w/ heartbeating: {:?}", self.shard_info, err);
}
},
other => {
warn!("[Shard {:?}] Other err w/ keepalive: {:?}", self.shard_info, other);
},
}
Err(Error::Gateway(GatewayError::HeartbeatFailed))
},
}
}
pub(crate) fn check_heartbeat(&mut self) -> Result<()> {
let heartbeat_interval = match self.heartbeat_interval {
Some(heartbeat_interval) => heartbeat_interval,
None => return Ok(()),
};
let wait = StdDuration::from_secs(heartbeat_interval / 1000);
if let Some(last_sent) = self.heartbeat_instants.0 {
if last_sent.elapsed() <= wait {
return Ok(());
}
}
if !self.last_heartbeat_acknowledged {
debug!("[Shard {:?}] Last heartbeat not acknowledged; re-connecting", self.shard_info);
return self.reconnect().map_err(|why| {
warn!("[Shard {:?}] Err auto-reconnecting from heartbeat check: {:?}",
self.shard_info,
why);
why
})
}
if let Err(why) = self.heartbeat() {
warn!("[Shard {:?}] Err heartbeating: {:?}", self.shard_info, why);
self.reconnect()
} else {
self.heartbeat_instants.0 = Some(Instant::now());
Ok(())
}
}
pub(crate) fn autoreconnect(&mut self) -> Result<()> {
if self.stage == ConnectionStage::Connecting {
return Ok(());
}
if self.session_id.is_some() {
debug!("[Shard {:?}] Autoreconnector choosing to resume", self.shard_info);
self.resume()
} else {
debug!("[Shard {:?}] Autoreconnector choosing to reconnect", self.shard_info);
self.reconnect()
}
}
#[inline]
pub(crate) fn heartbeat_interval(&self) -> Option<u64> {
self.heartbeat_interval
}
#[inline]
pub(crate) fn last_heartbeat_ack(&self) -> Option<Instant> {
self.heartbeat_instants.1
}
fn reconnect(&mut self) -> Result<()> {
info!("[Shard {:?}] Attempting to reconnect", self.shard_info);
self.reset();
self.initialize()
}
fn resume(&mut self) -> Result<()> {
self.send_resume().or_else(|why| {
warn!("Err sending resume: {:?}", why);
self.reconnect()
})
}
fn send_resume(&mut self) -> Result<()> {
let session_id = match self.session_id.clone() {
Some(session_id) => session_id,
None => return Err(Error::Gateway(GatewayError::NoSessionId)),
};
self.client.send_json(&json!({
"op": OpCode::Resume.num(),
"d": {
"session_id": session_id,
"seq": self.seq,
"token": &*self.token.lock().unwrap(),
},
}))
}
fn initialize(&mut self) -> Result<()> {
self.stage = ConnectionStage::Connecting;
self.client = connect(&self.ws_url.lock().unwrap())?;
self.identify()
}
fn identify(&mut self) -> Result<()> {
let identification = json!({
"op": OpCode::Identify.num(),
"d": {
"compression": true,
"large_threshold": constants::LARGE_THRESHOLD,
"shard": self.shard_info,
"token": &*self.token.lock().unwrap(),
"v": constants::GATEWAY_VERSION,
"properties": {
"$browser": "serenity",
"$device": "serenity",
"$os": consts::OS,
},
},
});
self.client.send_json(&identification)
}
fn reset(&mut self) {
self.heartbeat_instants = (Some(Instant::now()), None);
self.heartbeat_interval = None;
self.last_heartbeat_acknowledged = true;
self.stage = ConnectionStage::Disconnected;
self.seq = 0;
}
fn update_presence(&mut self) {
let (ref game, status, afk) = self.current_presence;
let now = Utc::now().timestamp() as u64;
let msg = json!({
"op": OpCode::StatusUpdate.num(),
"d": {
"afk": afk,
"since": now,
"status": status.name(),
"game": game.as_ref().map(|x| json!({
"name": x.name,
})),
},
});
if let Err(why) = self.client.send_json(&msg) {
warn!("[Shard {:?}] Err sending presence update: {:?}", self.shard_info, why);
}
#[cfg(feature="cache")]
{
let mut cache = CACHE.write().unwrap();
let current_user_id = cache.user.id;
cache.presences.get_mut(¤t_user_id).map(|presence| {
presence.game = game.clone();
presence.last_modified = Some(now);
});
}
}
}
fn connect(base_url: &str) -> Result<WsClient> {
let url = build_gateway_url(base_url)?;
let client = ClientBuilder::from_url(&url).connect_secure(None)?;
Ok(client)
}
fn set_client_timeout(client: &mut WsClient) -> Result<()> {
let stream = client.stream_ref().as_tcp();
stream.set_read_timeout(Some(StdDuration::from_millis(100)))?;
stream.set_write_timeout(Some(StdDuration::from_secs(5)))?;
Ok(())
}
fn build_gateway_url(base: &str) -> Result<Url> {
Url::parse(&format!("{}?v={}", base, constants::GATEWAY_VERSION))
.map_err(|_| Error::Gateway(GatewayError::BuildingUrl))
}