From a26b9260a960302e3c053d31f3b4bad6f67c697a Mon Sep 17 00:00:00 2001
From: winlin <winlin@vip.126.com>
Date: Fri, 19 Mar 2021 18:43:08 +0800
Subject: [PATCH] Threads-SRTP: Use async encrypt SRTP packet

1. Async SRTP support protect RTP.
2. Send packet ignore when encrypt size is 0.
3. Callback to send packet if encrypt done.
---
 trunk/src/app/srs_app_rtc_conn.cpp | 26 ++++++++++++++++++++--
 trunk/src/app/srs_app_rtc_conn.hpp |  3 +++
 trunk/src/app/srs_app_threads.cpp  | 35 +++++++++++++++++++++++-------
 3 files changed, 54 insertions(+), 10 deletions(-)

diff --git a/trunk/src/app/srs_app_rtc_conn.cpp b/trunk/src/app/srs_app_rtc_conn.cpp
index fdf5cb8c66..fc00e20d82 100644
--- a/trunk/src/app/srs_app_rtc_conn.cpp
+++ b/trunk/src/app/srs_app_rtc_conn.cpp
@@ -209,6 +209,11 @@ srs_error_t SrsSecurityTransport::on_rtcp_plaintext(char* plaintext, int size)
     return session_->on_rtcp_plaintext(plaintext, size);
 }
 
+srs_error_t SrsSecurityTransport::on_rtp_cipher(char* cipher, int size)
+{
+    return session_->on_rtp_cipher(cipher, size);
+}
+
 srs_error_t SrsSecurityTransport::protect_rtp(void* packet, int* nb_cipher)
 {
     return srtp_->protect_rtp(packet, nb_cipher);
@@ -2053,6 +2058,17 @@ srs_error_t SrsRtcConnection::on_rtcp_plaintext(char* data, int nb_unprotected_b
     return err;
 }
 
+srs_error_t SrsRtcConnection::on_rtp_cipher(char* cipher, int size)
+{
+    srs_error_t err = srs_success;
+
+    if ((err = sendonly_skt->sendto(cipher, size, 0)) != srs_success) {
+        srs_error_reset(err); // Ignore any error.
+    }
+
+    return err;
+}
+
 srs_error_t SrsRtcConnection::dispatch_rtcp(SrsRtcpCommon* rtcp)
 {
     srs_error_t err = srs_success;
@@ -2592,6 +2608,11 @@ srs_error_t SrsRtcConnection::do_send_packet(SrsRtpPacket2* pkt)
         iov->iov_len = (size_t)nn_encrypt;
     }
 
+    // Async SRTP encrypt.
+    if (iov->iov_len <= 0) {
+        return err;
+    }
+
     // For NACK simulator, drop packet.
     if (nn_simulate_player_nack_drop) {
         simulate_player_drop_packet(&pkt->header, (int)iov->iov_len);
@@ -2601,8 +2622,9 @@ srs_error_t SrsRtcConnection::do_send_packet(SrsRtpPacket2* pkt)
 
     ++_srs_pps_srtps->sugar;
 
-    // TODO: FIXME: Handle error.
-    sendonly_skt->sendto(iov->iov_base, iov->iov_len, 0);
+    if ((err = sendonly_skt->sendto(iov->iov_base, iov->iov_len, 0)) != srs_success) {
+        srs_error_reset(err); // Ignore any error.
+    }
 
     // Detail log, should disable it in release version.
     srs_info("RTC: SEND PT=%u, SSRC=%#x, SEQ=%u, Time=%u, %u/%u bytes", pkt->header.get_payload_type(), pkt->header.get_ssrc(),
diff --git a/trunk/src/app/srs_app_rtc_conn.hpp b/trunk/src/app/srs_app_rtc_conn.hpp
index f49fb820fb..6fb6863194 100644
--- a/trunk/src/app/srs_app_rtc_conn.hpp
+++ b/trunk/src/app/srs_app_rtc_conn.hpp
@@ -145,6 +145,7 @@ class SrsSecurityTransport : public ISrsRtcTransport
 public:
     srs_error_t on_rtp_plaintext(char* plaintext, int size);
     srs_error_t on_rtcp_plaintext(char* plaintext, int size);
+    srs_error_t on_rtp_cipher(char* cipher, int size);
 };
 
 // Semi security transport, setup DTLS and SRTP, with SRTP decrypt, without SRTP encrypt.
@@ -496,6 +497,8 @@ class SrsRtcConnection : public ISrsResource
     srs_error_t on_rtcp(char* data, int nb_data);
 private:
     srs_error_t on_rtcp_plaintext(char* plaintext, int size);
+private:
+    srs_error_t on_rtp_cipher(char* cipher, int size);
 private:
     srs_error_t dispatch_rtcp(SrsRtcpCommon* rtcp);
 public:
diff --git a/trunk/src/app/srs_app_threads.cpp b/trunk/src/app/srs_app_threads.cpp
index 14a2cfa32a..d98f2fc767 100644
--- a/trunk/src/app/srs_app_threads.cpp
+++ b/trunk/src/app/srs_app_threads.cpp
@@ -785,8 +785,23 @@ srs_error_t SrsAsyncSRTP::protect_rtp(void* packet, int* nb_cipher)
         return srs_error_new(ERROR_RTC_SRTP_UNPROTECT, "not ready");
     }
 
-    // TODO: FIMXE: Remove it.
-    return SrsSRTP::protect_rtp(packet, nb_cipher);
+    int nb_plaintext = *nb_cipher;
+
+    // Note that we must allocate more bytes than the size of packet, because SRTP
+    // will write checksum at the end of buffer.
+    char* buf = new char[kRtcpPacketSize];
+    memcpy(buf, packet, nb_plaintext);
+
+    SrsAsyncSRTPPacket* pkt = new SrsAsyncSRTPPacket(task_);
+    pkt->msg_->wrap(buf, nb_plaintext);
+    pkt->is_rtp_ = true;
+    pkt->do_decrypt_ = false;
+    _srs_async_srtp->add_packet(pkt);
+
+    // Do the job asynchronously.
+    *nb_cipher = 0;
+
+    return srs_success;
 }
 
 srs_error_t SrsAsyncSRTP::protect_rtcp(void* packet, int* nb_cipher)
@@ -816,9 +831,7 @@ srs_error_t SrsAsyncSRTP::unprotect_rtp(void* packet, int* nb_plaintext)
     _srs_async_srtp->add_packet(pkt);
 
     // Do the job asynchronously.
-    if (nb_plaintext) {
-        *nb_plaintext = 0;
-    }
+    *nb_plaintext = 0;
 
     return srs_success;
 }
@@ -840,9 +853,7 @@ srs_error_t SrsAsyncSRTP::unprotect_rtcp(void* packet, int* nb_plaintext)
     _srs_async_srtp->add_packet(pkt);
 
     // Do the job asynchronously.
-    if (nb_plaintext) {
-        *nb_plaintext = 0;
-    }
+    *nb_plaintext = 0;
 
     return srs_success;
 }
@@ -896,6 +907,10 @@ srs_error_t SrsAsyncSRTPTask::cook(SrsAsyncSRTPPacket* pkt)
         } else {
             err = impl_->unprotect_rtcp(pkt->msg_->payload, &pkt->nb_consumed_);
         }
+    } else {
+        if (pkt->is_rtp_) {
+            err = impl_->protect_rtp(pkt->msg_->payload, &pkt->nb_consumed_);
+        }
     }
     if (err != srs_success) {
         return err;
@@ -921,6 +936,10 @@ srs_error_t SrsAsyncSRTPTask::consume(SrsAsyncSRTPPacket* pkt)
         } else {
             err = codec_->transport_->on_rtcp_plaintext(payload, pkt->nb_consumed_);
         }
+    } else {
+        if (pkt->is_rtp_) {
+            err = codec_->transport_->on_rtp_cipher(payload, pkt->nb_consumed_);
+        }
     }
 
     return err;