From dcf9c6505738023ff821f38d6eb4e3695fc49d24 Mon Sep 17 00:00:00 2001
From: Riza Sulistyo <trengginas@users.noreply.github.com>
Date: Wed, 5 Feb 2025 23:21:13 +0700
Subject: [PATCH] Modification based on comments

- store the valid address to avoid checking with all of remote candidate
- add comment regarding the check source address in SRTP-DTLS is for non-ICE use
---
 pjmedia/src/pjmedia/transport_srtp_dtls.c | 18 +++++------
 pjnath/include/pjnath/ice_session.h       |  6 ++++
 pjnath/src/pjnath/ice_session.c           | 37 +++++++++++++++--------
 3 files changed, 40 insertions(+), 21 deletions(-)

diff --git a/pjmedia/src/pjmedia/transport_srtp_dtls.c b/pjmedia/src/pjmedia/transport_srtp_dtls.c
index 9643a92a96..4b9ba78e14 100644
--- a/pjmedia/src/pjmedia/transport_srtp_dtls.c
+++ b/pjmedia/src/pjmedia/transport_srtp_dtls.c
@@ -1374,16 +1374,15 @@ static pj_status_t dtls_on_recv(pjmedia_transport *tp, unsigned idx,
         (ds->setup == DTLS_SETUP_ACTPASS || ds->setup == DTLS_SETUP_PASSIVE))
     {
         pj_status_t status;
-        pj_bool_t check_hello_addr = PJ_FALSE;
 
 #if defined(PJMEDIA_SRTP_DTLS_CHECK_HELLO_ADDR) && \
             PJMEDIA_SRTP_DTLS_CHECK_HELLO_ADDR==1
 
-        if (!ds->use_ice)
-            check_hello_addr = PJ_TRUE;
-
-#endif
-        if (check_hello_addr) {
+        /* Check the souce address with the specified remote address from 
+         * the SDP. When ICE is used, the source address checking will be 
+         * done in ICE session.
+         */
+        if (!ds->use_ice) {
             pjmedia_transport_info info;
             pj_sockaddr src_addr;
             pj_bool_t src_addr_avail = PJ_TRUE;
@@ -1407,17 +1406,18 @@ static pj_status_t dtls_on_recv(pjmedia_transport *tp, unsigned idx,
                 char psrc_addr[PJ_INET6_ADDRSTRLEN] = "Unknown";
 
                 if (src_addr_avail) {
-                    pj_sockaddr_print(&src_addr, psrc_addr, 
+                    pj_sockaddr_print(&src_addr, psrc_addr,
                                       sizeof(psrc_addr), 3);
                 }
                 PJ_LOG(2, (ds->base.name, "DTLS-SRTP %s ignoring %lu bytes, "
-                    "from src addr [%s]", CHANNEL_TO_STRING(idx), 
-                    (unsigned long)size,   psrc_addr));
+                    "from unrecognized src addr [%s]", CHANNEL_TO_STRING(idx),
+                    (unsigned long)size, psrc_addr));
 
                 DTLS_UNLOCK(ds);
                 return PJ_SUCCESS;
             }
         }
+#endif
         ds->setup = DTLS_SETUP_PASSIVE;
         status = ssl_handshake_channel(ds, idx);
         if (status != PJ_SUCCESS) {
diff --git a/pjnath/include/pjnath/ice_session.h b/pjnath/include/pjnath/ice_session.h
index f224587e41..db0c89c9f3 100644
--- a/pjnath/include/pjnath/ice_session.h
+++ b/pjnath/include/pjnath/ice_session.h
@@ -201,6 +201,12 @@ typedef struct pj_ice_sess_comp
      */
     pj_stun_session     *stun_sess;
 
+    /**
+     * The remote candidate checked address. This is expected address that 
+     * the remote going to use.
+     */
+    pj_sockaddr         rcand_check_addr;
+
 } pj_ice_sess_comp;
 
 
diff --git a/pjnath/src/pjnath/ice_session.c b/pjnath/src/pjnath/ice_session.c
index 1e9ef4b2c4..309e73c22a 100644
--- a/pjnath/src/pjnath/ice_session.c
+++ b/pjnath/src/pjnath/ice_session.c
@@ -3747,21 +3747,34 @@ PJ_DEF(pj_status_t) pj_ice_sess_on_rx_pkt(pj_ice_sess *ice,
         PJ_RACE_ME(5);
 
         if (ice->opt.check_src_addr) {
-            for (i = 0; i < ice->rcand_cnt; ++i) {
-                if (ice->rcand[i].comp_id == comp_id && 
-                    ice->rcand[i].checked &&
-                    pj_sockaddr_cmp(src_addr, &ice->rcand[i].addr) == 0)
-                {
-                    break;
-                }
+            pj_bool_t check_addr = PJ_TRUE;
+            pj_sockaddr *raddr = &comp->rcand_check_addr;
+            char psrc_addr[PJ_INET6_ADDRSTRLEN] = {0};
+
+            if (pj_sockaddr_has_addr(src_addr)) {
+                pj_sockaddr_print(src_addr, psrc_addr, sizeof(psrc_addr), 3);
             }
-            if (i == ice->rcand_cnt) {
-                char psrc_addr[PJ_INET6_ADDRSTRLEN] = {0};
 
-                if (pj_sockaddr_has_addr(src_addr)) {
-                    pj_sockaddr_print(src_addr, psrc_addr,
-                                      sizeof(psrc_addr), 3);
+            if (!pj_sockaddr_has_addr(raddr)) {
+                for (i = 0; i < ice->rcand_cnt; ++i) {
+                    if (ice->rcand[i].comp_id == comp_id &&
+                        ice->rcand[i].checked &&
+                        pj_sockaddr_cmp(src_addr, &ice->rcand[i].addr) == 0)
+                    {
+                        pj_sockaddr_cp(raddr, src_addr);
+                        PJ_LOG(4, (ice->obj_name, "Using [%s] as valid address "
+                                   "for component [%d]",
+                                   psrc_addr, comp_id));
+
+                        check_addr = PJ_FALSE;
+                        break;
+                    }
                 }
+            }
+            if (check_addr && 
+                (!pj_sockaddr_has_addr(raddr) || 
+                 pj_sockaddr_cmp(src_addr, raddr) != 0))
+            {
                 PJ_LOG(4, (ice->obj_name, "Ignoring incoming message for "
                          "component [%d] because source addr [%s] unrecognized",
                          comp_id, psrc_addr));