From 29d274d72b4967636ce9bdf44a057ec4f6f40fe9 Mon Sep 17 00:00:00 2001
From: Aaro Altonen <aaro.altonen@tuni.fi>
Date: Wed, 12 Aug 2020 16:19:31 +0300
Subject: [PATCH] Fix Media Stream initialization when ZRTP is used

---
 include/media_stream.hh |  3 +-
 src/media_stream.cc     | 77 ++++++++++++++++++++++++++++++-----------
 2 files changed, 59 insertions(+), 21 deletions(-)

diff --git a/include/media_stream.hh b/include/media_stream.hh
index 44b3775..07a64ad 100644
--- a/include/media_stream.hh
+++ b/include/media_stream.hh
@@ -199,8 +199,9 @@ namespace uvg_rtp {
             /* media stream type */
             enum mstream_type type_;
 
-            /* Primary handler's key for the RTP packet dispatcher */
+            /* Primary handler keys for the RTP packet dispatcher */
             uint32_t rtp_handler_key_;
+            uint32_t zrtp_handler_key_;
 
             /* RTP packet dispatcher for the receiver */
             uvg_rtp::pkt_dispatcher *pkt_dispatcher_;
diff --git a/src/media_stream.cc b/src/media_stream.cc
index b0fba9c..967f4da 100644
--- a/src/media_stream.cc
+++ b/src/media_stream.cc
@@ -160,28 +160,26 @@ rtp_error_t uvg_rtp::media_stream::init()
 #ifdef __RTP_CRYPTO__
 rtp_error_t uvg_rtp::media_stream::init(uvg_rtp::zrtp *zrtp)
 {
+    rtp_error_t ret;
+
     if (init_connection() != RTP_OK) {
         LOG_ERROR("Failed to initialize the underlying socket: %s!", strerror(errno));
         return RTP_GENERIC_ERROR;
     }
 
-    /* First initialize the RTP context for this media stream (SSRC, sequence number, etc.)
-     * Then initialize ZRTP and using ZRTP, initialize SRTP.
-     *
-     * When ZRTP and SRTP have been initialized, create sender and receiver for the media type
-     * before returning the media stream for user */
-    rtp_error_t ret = RTP_OK;
-
-    if ((rtp_ = new uvg_rtp::rtp(fmt_)) == nullptr)
+    if (!(pkt_dispatcher_ = new uvg_rtp::pkt_dispatcher()))
         return RTP_MEMORY_ERROR;
 
+    if (!(rtp_ = new uvg_rtp::rtp(fmt_))) {
+        delete pkt_dispatcher_;
+        return RTP_MEMORY_ERROR;
+    }
+
     if ((ret = zrtp->init(rtp_->get_ssrc(), socket_.get_raw_socket(), addr_out_)) != RTP_OK) {
         LOG_WARN("Failed to initialize ZRTP for media stream!");
         return ret;
     }
 
-    /* TODO: install zrtp packet handler */
-
     if ((srtp_ = new uvg_rtp::srtp()) == nullptr)
         return RTP_MEMORY_ERROR;
 
@@ -190,24 +188,62 @@ rtp_error_t uvg_rtp::media_stream::init(uvg_rtp::zrtp *zrtp)
         return ret;
     }
 
+    if (!(rtcp_ = new uvg_rtp::rtcp(rtp_))) {
+        delete rtp_;
+        delete pkt_dispatcher_;
+        return RTP_MEMORY_ERROR;
+    }
+
+    socket_.install_handler(rtcp_, rtcp_->send_packet_handler_buf);
+    socket_.install_handler(rtcp_, rtcp_->send_packet_handler_vec);
+
+    rtp_handler_key_  = pkt_dispatcher_->install_handler(rtp_->packet_handler);
+    zrtp_handler_key_ = pkt_dispatcher_->install_handler(zrtp->packet_handler);
+
+    pkt_dispatcher_->install_aux_handler(rtp_handler_key_, rtcp_, rtcp_->recv_packet_handler);
+    pkt_dispatcher_->install_aux_handler(rtp_handler_key_, srtp_, srtp_->recv_packet_handler);
+
+    switch (fmt_) {
+        case RTP_FORMAT_HEVC:
+            media_ = new uvg_rtp::formats::hevc(&socket_, rtp_, ctx_config_.flags);
+            pkt_dispatcher_->install_aux_handler(
+                rtp_handler_key_,
+                nullptr,
+                dynamic_cast<uvg_rtp::formats::hevc *>(media_)->packet_handler
+            );
+            break;
+
+        case RTP_FORMAT_OPUS:
+        case RTP_FORMAT_GENERIC:
+            media_ = new uvg_rtp::formats::media(&socket_, rtp_, ctx_config_.flags);
+            pkt_dispatcher_->install_aux_handler(rtp_handler_key_, nullptr, media_->packet_handler);
+            break;
+
+        default:
+            LOG_ERROR("Unknown payload format %u\n", fmt_);
+    }
+
+    if (!media_) {
+        delete rtp_;
+        delete pkt_dispatcher_;
+        return RTP_MEMORY_ERROR;
+    }
+
+    if (ctx_config_.flags & RCE_RTCP) {
+        rtcp_->add_participant(addr_, src_port_ + 1, dst_port_ + 1, rtp_->get_clock_rate());
+        rtcp_->start();
+    }
+
     if (ctx_config_.flags & RCE_SRTP_AUTHENTICATE_RTP)
         rtp_->set_payload_size(MAX_PAYLOAD - AUTH_TAG_LENGTH);
 
-    socket_.set_srtp(srtp_);
-
-    /* TODO: install srtp packet handler */
-
-    receiver_ = new uvg_rtp::receiver(socket_, ctx_config_, fmt_, rtp_);
-
-    receiver_->start();
-
     initialized_ = true;
-
-    return ret;
+    return pkt_dispatcher_->start(&socket_, ctx_config_.flags);
 }
 
 rtp_error_t uvg_rtp::media_stream::add_srtp_ctx(uint8_t *key, uint8_t *salt)
 {
+#if 0
     if (!key || !salt)
         return RTP_INVALID_VALUE;
 
@@ -245,6 +281,7 @@ rtp_error_t uvg_rtp::media_stream::add_srtp_ctx(uint8_t *key, uint8_t *salt)
     initialized_ = true;
 
     return ret;
+#endif
 }
 #endif