From 4a6642df9fbfd9a9247e565ab3f4d8d1e6f6ed4f Mon Sep 17 00:00:00 2001
From: Craig Dazey <craigdazey@gmail.com>
Date: Tue, 23 Jul 2019 01:16:52 -0400
Subject: [PATCH] add support for changing voice state

---
 lib/nostrum/api.ex                       | 14 ++++++++++++++
 lib/nostrum/cache/mapping/guild_shard.ex |  2 --
 lib/nostrum/shard/payload.ex             | 11 +++++++++++
 lib/nostrum/shard/session.ex             |  9 +++++++++
 lib/nostrum/shard/supervisor.ex          | 21 +++++++++++++++++++++
 5 files changed, 55 insertions(+), 2 deletions(-)

diff --git a/lib/nostrum/api.ex b/lib/nostrum/api.ex
index 87ecfd335..a686bfb6f 100644
--- a/lib/nostrum/api.ex
+++ b/lib/nostrum/api.ex
@@ -127,6 +127,20 @@ defmodule Nostrum.Api do
     :ok
   end
 
+  @doc """
+  Joins, moves, or disconnects the bot from a voice channel.
+
+  The correct shard to send the update to will be inferred from the
+  `guild_id`. If a corresponding `guild_id` is not found a cache error will be
+  raised.
+
+  To disconnect from a channel, `channel_id` should be set to `nil`.
+  """
+  @spec update_voice_state(Guild.id(), Channel.id(), boolean, boolean) :: no_return | :ok
+  def update_voice_state(guild_id, channel_id, self_mute \\ false, self_deaf \\ false) do
+    Supervisor.update_voice_state(guild_id, channel_id, self_mute, self_deaf)
+  end
+
   @doc ~S"""
   Posts a message to a guild text or DM channel.
 
diff --git a/lib/nostrum/cache/mapping/guild_shard.ex b/lib/nostrum/cache/mapping/guild_shard.ex
index d9a2514e0..d8b714bf1 100644
--- a/lib/nostrum/cache/mapping/guild_shard.ex
+++ b/lib/nostrum/cache/mapping/guild_shard.ex
@@ -11,8 +11,6 @@ defmodule Nostrum.Cache.Mapping.GuildShard do
   """
   @spec get_shard(Guild.id()) :: {:ok, integer} | {:error, :id_not_found}
   def get_shard(guild_id) do
-    :ets.lookup_element(:guild_shard_map, guild_id, 2)
-
     case :ets.lookup(:guild_shard_map, guild_id) do
       [{_guild_id, shard_num}] -> {:ok, shard_num}
       [] -> {:error, :id_not_found}
diff --git a/lib/nostrum/shard/payload.ex b/lib/nostrum/shard/payload.ex
index 633480b5c..bc6ef9f5d 100644
--- a/lib/nostrum/shard/payload.ex
+++ b/lib/nostrum/shard/payload.ex
@@ -56,6 +56,17 @@ defmodule Nostrum.Shard.Payload do
     |> build_payload("STATUS_UPDATE")
   end
 
+  @doc false
+  def update_voice_state_payload(guild_id, channel_id, self_mute, self_deaf) do
+    %{
+      "guild_id" => guild_id,
+      "channel_id" => channel_id,
+      "self_mute" => self_mute,
+      "self_deaf" => self_deaf
+    }
+    |> build_payload("VOICE_STATUS_UPDATE")
+  end
+
   @doc false
   def request_members_payload(guild_id, limit) do
     %{
diff --git a/lib/nostrum/shard/session.ex b/lib/nostrum/shard/session.ex
index ef281294f..011e15eca 100644
--- a/lib/nostrum/shard/session.ex
+++ b/lib/nostrum/shard/session.ex
@@ -25,6 +25,11 @@ defmodule Nostrum.Shard.Session do
     WebSockex.cast(pid, {:status_update, payload})
   end
 
+  def update_voice_state(pid, guild_id, channel_id, self_mute, self_deaf) do
+    payload = Payload.update_voice_state_payload(guild_id, channel_id, self_mute, self_deaf)
+    WebSockex.cast(pid, {:update_voice_state, payload})
+  end
+
   def request_guild_members(pid, guild_id, limit \\ 0) do
     payload = Payload.request_members_payload(guild_id, limit)
     WebSockex.cast(pid, {:request_guild_members, payload})
@@ -92,6 +97,10 @@ defmodule Nostrum.Shard.Session do
     {:reply, {:binary, payload}, state}
   end
 
+  def handle_cast({:update_voice_state, payload}, state) do
+    {:reply, {:binary, payload}, state}
+  end
+
   def handle_cast({:request_guild_members, payload}, state) do
     {:reply, {:binary, payload}, state}
   end
diff --git a/lib/nostrum/shard/supervisor.ex b/lib/nostrum/shard/supervisor.ex
index 6db3cd73b..ece5ca086 100644
--- a/lib/nostrum/shard/supervisor.ex
+++ b/lib/nostrum/shard/supervisor.ex
@@ -4,6 +4,8 @@ defmodule Nostrum.Shard.Supervisor do
   use Supervisor
 
   alias Nostrum.{Shard, Util}
+  alias Nostrum.Cache.Mapping.GuildShard
+  alias Nostrum.Error.CacheError
   alias Nostrum.Shard.Session
   alias Nostrum.Shard.Stage.{Cache, Producer}
 
@@ -48,6 +50,25 @@ defmodule Nostrum.Shard.Supervisor do
     end)
   end
 
+  def update_voice_state(guild_id, channel_id, self_mute, self_deaf) do
+    case GuildShard.get_shard(guild_id) do
+      {:ok, shard_id} ->
+        ShardSupervisor
+        |> Supervisor.which_children()
+        |> Enum.filter(fn {_id, _pid, _type, [modules]} -> modules == Nostrum.Shard end)
+        |> Enum.filter(fn {id, _pid, _type, _modules} -> id == shard_id end)
+        |> Enum.map(fn {_id, pid, _type, _modules} -> Supervisor.which_children(pid) end)
+        |> List.flatten()
+        |> Enum.filter(fn {_id, _pid, _type, [modules]} -> modules == Nostrum.Shard.Session end)
+        |> List.first()
+        |> elem(1)
+        |> Session.update_voice_state(guild_id, channel_id, self_mute, self_deaf)
+
+      {:error, :id_not_found} ->
+        raise CacheError, key: guild_id, cache_name: GuildShardMapping
+    end
+  end
+
   @doc false
   def init([url, num_shards]) do
     children =