Move DTLS utilities to DTLSTransport module (#10)

elixir-webrtc · Nov 6, 2023 · 1aeef5b · 1aeef5b
1 parent 8f3fac2
commit 1aeef5b
Show file tree

Hide file tree

Showing 3 changed files with 169 additions and 58 deletions.
diff --git a/examples/example.exs b/examples/example.exs
@@ -96,6 +96,7 @@ defmodule Peer do
     offer = %SessionDescription{type: :offer, sdp: sdp}
     :ok = PeerConnection.set_remote_description(state.peer_connection, offer)
     {:ok, answer} = PeerConnection.create_answer(state.peer_connection)
+    :ok = PeerConnection.set_local_description(state.peer_connection, answer)
     msg = %{"type" => "answer", "sdp" => answer.sdp}
     :gun.ws_send(state.conn, state.stream, {:text, Jason.encode!(msg)})
   end

diff --git a/lib/ex_webrtc/dtls_transport.ex b/lib/ex_webrtc/dtls_transport.ex
@@ -0,0 +1,123 @@
+defmodule ExWebRTC.DTLSTransport do
+  @moduledoc false
+
+  require Logger
+
+  alias ExICE.ICEAgent
+
+  defstruct [
+    :ice_agent,
+    :ice_state,
+    :client,
+    :buffered_packets,
+    :cert,
+    :pkey,
+    :fingerprint,
+    :mode,
+    finished: false
+  ]
+
+  def new(ice_agent) do
+    # temporary hack to generate certs
+    {:ok, cert_client} = ExDTLS.start_link(client_mode: true, dtls_srtp: true)
+    {:ok, cert} = ExDTLS.get_cert(cert_client)
+    {:ok, pkey} = ExDTLS.get_pkey(cert_client)
+    {:ok, fingerprint} = ExDTLS.get_cert_fingerprint(cert_client)
+    :ok = ExDTLS.stop(cert_client)
+
+    %__MODULE__{
+      ice_agent: ice_agent,
+      cert: cert,
+      pkey: pkey,
+      fingerprint: fingerprint
+    }
+  end
+
+  def start(dtls, :passive) do
+    {:ok, client} =
+      ExDTLS.start_link(
+        client_mode: false,
+        dtls_srtp: true,
+        pkey: dtls.pkey,
+        cert: dtls.cert
+      )
+
+    %__MODULE__{dtls | client: client, mode: :passive}
+  end
+
+  def start(dtls, :active) do
+    {:ok, client} =
+      ExDTLS.start_link(
+        client_mode: true,
+        dtls_srtp: true,
+        pkey: dtls.pkey,
+        cert: dtls.cert
+      )
+
+    # we assume that ICE in not in connected state yet
+    %__MODULE__{dtls | client: client, mode: :active}
+  end
+
+  def update_ice_state(dtls, :connected) do
+    if dtls.mode == :active do
+      {:ok, packets} = ExDTLS.do_handshake(dtls.client)
+      :ok = ICEAgent.send_data(dtls.ice_agent, packets)
+    end
+
+    dtls =
+      if dtls.buffered_packets do
+        Logger.debug("Sending buffered DTLS packets")
+        :ok = ICEAgent.send_data(dtls.ice_agent, dtls.buffered_packets)
+        %__MODULE__{dtls | buffered_packets: nil}
+      else
+        dtls
+      end
+
+    %__MODULE__{dtls | ice_state: :connected}
+  end
+
+  def update_ice_state(dtls, new_state) do
+    %__MODULE__{dtls | ice_state: new_state}
+  end
+
+  def handle_info(dtls, {:retransmit, packets})
+      when dtls.ice_state in [:connected, :completed] do
+    ICEAgent.send_data(dtls.ice_agent, packets)
+    dtls
+  end
+
+  def handle_info(%{buffered_packets: packets} = dtls, {:retransmit, packets}) do
+    # we got DTLS packets from the other side but
+    # we haven't established ICE connection yet so
+    # packets to retransmit have to be the same as dtls_buffered_packets
+    dtls
+  end
+
+  def process_data(dtls, data) do
+    case ExDTLS.process(dtls.client, data) do
+      {:handshake_packets, packets} when dtls.ice_state in [:connected, :completed] ->
+        :ok = ICEAgent.send_data(dtls.ice_agent, packets)
+        dtls
+
+      {:handshake_packets, packets} ->
+        Logger.debug("""
+        Generated local DTLS packets but ICE is not in the connected or completed state yet.
+        We will send those packets once ICE is ready.
+        """)
+
+        %__MODULE__{dtls | buffered_packets: packets}
+
+      {:handshake_finished, _keying_material, packets} ->
+        Logger.debug("DTLS handshake finished")
+        ICEAgent.send_data(dtls.ice_agent, packets)
+        %__MODULE__{dtls | finished: true}
+
+      {:handshake_finished, _keying_material} ->
+        Logger.debug("DTLS handshake finished")
+        %__MODULE__{dtls | finished: true}
+
+      :handshake_want_read ->
+        dtls
+    end
+  end
+end
diff --git a/lib/ex_webrtc/peer_connection.ex b/lib/ex_webrtc/peer_connection.ex
@@ -9,6 +9,7 @@ defmodule ExWebRTC.PeerConnection do
   alias ExICE.ICEAgent
 
   alias ExWebRTC.{
+    DTLSTransport,
     IceCandidate,
     MediaStreamTrack,
     RTPTransceiver,
@@ -37,9 +38,7 @@ defmodule ExWebRTC.PeerConnection do
                 :pending_remote_desc,
                 :ice_agent,
                 :ice_state,
-                :dtls_client,
-                :dtls_buffered_packets,
-                dtls_finished: false,
+                :dtls_transport,
                 transceivers: [],
                 signaling_state: :stable,
                 last_offer: nil,
@@ -114,13 +113,13 @@ defmodule ExWebRTC.PeerConnection do
       |> Enum.filter(&String.starts_with?(&1, "stun:"))
 
     {:ok, ice_agent} = ICEAgent.start_link(:controlled, stun_servers: stun_servers)
-    {:ok, dtls_client} = ExDTLS.start_link(client_mode: false, dtls_srtp: true)
+    dtls_transport = DTLSTransport.new(ice_agent)
 
     state = %__MODULE__{
       owner: owner,
       config: config,
       ice_agent: ice_agent,
-      dtls_client: dtls_client
+      dtls_transport: dtls_transport
     }
 
     {:ok, state}
@@ -144,7 +143,6 @@ defmodule ExWebRTC.PeerConnection do
     transceivers = assign_mids(state.transceivers, next_mid)
 
     {:ok, ice_ufrag, ice_pwd} = ICEAgent.get_local_credentials(state.ice_agent)
-    {:ok, dtls_fingerprint} = ExDTLS.get_cert_fingerprint(state.dtls_client)
 
     offer =
       %ExSDP{ExSDP.new() | timing: %ExSDP.Timing{start_time: 0, stop_time: 0}}
@@ -156,7 +154,7 @@ defmodule ExWebRTC.PeerConnection do
         ice_ufrag: ice_ufrag,
         ice_pwd: ice_pwd,
         ice_options: "trickle",
-        fingerprint: {:sha256, Utils.hex_dump(dtls_fingerprint)},
+        fingerprint: {:sha256, Utils.hex_dump(state.dtls_transport.fingerprint)},
         setup: :actpass,
         rtcp: true
       ]
@@ -200,7 +198,6 @@ defmodule ExWebRTC.PeerConnection do
     {:offer, remote_offer} = state.pending_remote_desc
 
     {:ok, ice_ufrag, ice_pwd} = ICEAgent.get_local_credentials(state.ice_agent)
-    {:ok, dtls_fingerprint} = ExDTLS.get_cert_fingerprint(state.dtls_client)
 
     answer =
       %ExSDP{ExSDP.new() | timing: %ExSDP.Timing{start_time: 0, stop_time: 0}}
@@ -212,7 +209,7 @@ defmodule ExWebRTC.PeerConnection do
         ice_ufrag: ice_ufrag,
         ice_pwd: ice_pwd,
         ice_options: "trickle",
-        fingerprint: {:sha256, Utils.hex_dump(dtls_fingerprint)},
+        fingerprint: {:sha256, Utils.hex_dump(state.dtls_transport.fingerprint)},
         setup: :active
       ]
 
@@ -286,6 +283,11 @@ defmodule ExWebRTC.PeerConnection do
     end
   end
 
+  @impl true
+  def handle_call({:add_ice_candidate, _}, _from, %{current_remote_desc: nil} = state) do
+    {:reply, {:error, :no_remote_description}, state}
+  end
+
   @impl true
   def handle_call({:add_ice_candidate, candidate}, _from, state) do
     with "candidate:" <> attr <- candidate.candidate do
@@ -322,12 +324,8 @@ defmodule ExWebRTC.PeerConnection do
 
   @impl true
   def handle_info({:ex_ice, _from, :connected}, state) do
-    if state.dtls_buffered_packets do
-      Logger.debug("Sending buffered DTLS packets")
-      ICEAgent.send_data(state.ice_agent, state.dtls_buffered_packets)
-    end
-
-    {:noreply, %__MODULE__{state | ice_state: :connected, dtls_buffered_packets: nil}}
+    dtls = DTLSTransport.update_ice_state(state.dtls_transport, :connected)
+    {:noreply, %__MODULE__{state | dtls_transport: dtls, ice_state: :connected}}
   end
 
   @impl true
@@ -345,50 +343,16 @@ defmodule ExWebRTC.PeerConnection do
   end
 
   @impl true
-  def handle_info({:ex_ice, _from, {:data, data}}, %{dtls_finished: false} = state) do
-    case ExDTLS.process(state.dtls_client, data) do
-      {:handshake_packets, packets} when state.ice_state in [:connected, :completed] ->
-        :ok = ICEAgent.send_data(state.ice_agent, packets)
-        {:noreply, state}
-
-      {:handshake_packets, packets} ->
-        Logger.debug("""
-        Generated local DTLS packets but ICE is not in the connected or completed state yet.
-        We will send those packets once ICE is ready.
-        """)
-
-        {:noreply, %__MODULE__{state | dtls_buffered_packets: packets}}
-
-      {:handshake_finished, _keying_material, packets} ->
-        Logger.debug("DTLS handshake finished")
-        ICEAgent.send_data(state.ice_agent, packets)
-        {:noreply, %__MODULE__{state | dtls_finished: true}}
-
-      {:handshake_finished, _keying_material} ->
-        Logger.debug("DTLS handshake finished")
-        {:noreply, %__MODULE__{state | dtls_finished: true}}
-
-      :handshake_want_read ->
-        {:noreply, state}
-    end
+  def handle_info({:ex_ice, _from, {:data, data}}, state)
+      when not state.dtls_transport.finished do
+    dtls = DTLSTransport.process_data(state.dtls_transport, data)
+    {:noreply, %__MODULE__{state | dtls_transport: dtls}}
   end
 
   @impl true
-  def handle_info({:ex_dtls, _from, {:retransmit, packets}}, state)
-      when state.ice_state in [:connected, :completed] do
-    ICEAgent.send_data(state.ice_agent, packets)
-    {:noreply, state}
-  end
-
-  @impl true
-  def handle_info(
-        {:ex_dtls, _from, {:retransmit, packets}},
-        %{dtls_buffered_packets: packets} = state
-      ) do
-    # we got DTLS packets from the other side but
-    # we haven't established ICE connection yet so
-    # packets to retransmit have to be the same as dtls_buffered_packets
-    {:noreply, state}
+  def handle_info({:ex_dtls, _from, msg}, state) do
+    dtls = DTLSTransport.handle_info(state.dtls_transport, msg)
+    {:noreply, %__MODULE__{state | dtls_transport: dtls}}
   end
 
   @impl true
@@ -401,7 +365,15 @@ defmodule ExWebRTC.PeerConnection do
     new_transceivers = update_local_transceivers(type, state.transceivers, sdp)
     state = set_description(:local, type, sdp, state)
 
-    {:ok, %{state | transceivers: new_transceivers}}
+    dtls =
+      if type == :answer do
+        {:setup, setup} = ExSDP.Media.get_attribute(hd(sdp.media), :setup)
+        DTLSTransport.start(state.dtls_transport, setup)
+      else
+        state.dtls_transport
+      end
+
+    {:ok, %{state | transceivers: new_transceivers, dtls_transport: dtls}}
   end
 
   defp update_local_transceivers(:offer, transceivers, _sdp) do
@@ -437,7 +409,22 @@ defmodule ExWebRTC.PeerConnection do
 
       state = set_description(:remote, type, sdp, state)
 
-      {:ok, %{state | transceivers: new_transceivers}}
+      dtls =
+        if type == :answer do
+          {:setup, setup} = ExSDP.Media.get_attribute(hd(sdp.media), :setup)
+
+          setup =
+            case setup do
+              :active -> :passive
+              :passive -> :active
+            end
+
+          DTLSTransport.start(state.dtls_transport, setup)
+        else
+          state.dtls_transport
+        end
+
+      {:ok, %{state | transceivers: new_transceivers, dtls_transport: dtls}}
     else
       error -> error
     end