From 0ba2b637ae7ae725bdd7de24de58b517a0cce93f Mon Sep 17 00:00:00 2001
From: wuchangye <wuchangye@huawei.com>
Date: Sun, 24 Mar 2024 05:29:30 +0800
Subject: [PATCH] support querying udp multicast addresses

---
 src/api/sockets.c          | 43 ++++++++++++++++++++++++++++++++++++--
 src/include/lwip/sockets.h |  1 +
 2 files changed, 42 insertions(+), 2 deletions(-)

diff --git a/src/api/sockets.c b/src/api/sockets.c
index 66f6e2c..c544c22 100644
--- a/src/api/sockets.c
+++ b/src/api/sockets.c
@@ -4342,6 +4342,27 @@ lwip_inet_pton(int af, const char *src, void *dst)
 }
 
 #if LWIP_IGMP
+#if GAZELLE_ENABLE
+int lwip_socket_get_multi_memberships(int s, ip4_addr_t *if_addr, ip4_addr_t *multi_addr)
+{
+  struct lwip_sock *sock = get_socket(s);
+
+  if (!sock) {
+    return 0;
+  }
+
+  if (socket_ipv4_multicast_memberships[s].sock == sock) {
+    ip4_addr_copy(*if_addr, socket_ipv4_multicast_memberships[s].if_addr);
+    ip4_addr_copy(*multi_addr, socket_ipv4_multicast_memberships[s].multi_addr);
+    done_socket(sock);
+    return 1;
+  }
+
+  done_socket(sock);
+  return 0;
+}
+#endif
+
 /** Register a new IGMP membership. On socket close, the membership is dropped automatically.
  *
  * ATTENTION: this function is called from tcpip_thread (or under CORE_LOCK).
@@ -4357,8 +4378,11 @@ lwip_socket_register_membership(int s, const ip4_addr_t *if_addr, const ip4_addr
   if (!sock) {
     return 0;
   }
-
+#if GAZELLE_ENABLE
+  i = s;
+#else
   for (i = 0; i < LWIP_SOCKET_MAX_MEMBERSHIPS; i++) {
+#endif
     if (socket_ipv4_multicast_memberships[i].sock == NULL) {
       socket_ipv4_multicast_memberships[i].sock = sock;
       ip4_addr_copy(socket_ipv4_multicast_memberships[i].if_addr, *if_addr);
@@ -4366,7 +4390,9 @@ lwip_socket_register_membership(int s, const ip4_addr_t *if_addr, const ip4_addr
       done_socket(sock);
       return 1;
     }
+#if !GAZELLE_ENABLE
   }
+#endif
   done_socket(sock);
   return 0;
 }
@@ -4386,16 +4412,24 @@ lwip_socket_unregister_membership(int s, const ip4_addr_t *if_addr, const ip4_ad
     return;
   }
 
+#if GAZELLE_ENABLE
+  i = s;
+#else
   for (i = 0; i < LWIP_SOCKET_MAX_MEMBERSHIPS; i++) {
+#endif
     if ((socket_ipv4_multicast_memberships[i].sock == sock) &&
         ip4_addr_eq(&socket_ipv4_multicast_memberships[i].if_addr, if_addr) &&
         ip4_addr_eq(&socket_ipv4_multicast_memberships[i].multi_addr, multi_addr)) {
       socket_ipv4_multicast_memberships[i].sock = NULL;
       ip4_addr_set_zero(&socket_ipv4_multicast_memberships[i].if_addr);
       ip4_addr_set_zero(&socket_ipv4_multicast_memberships[i].multi_addr);
+#if !GAZELLE_ENABLE
       break;
+#endif
     }
+#if !GAZELLE_ENABLE
   }
+#endif
   done_socket(sock);
 }
 
@@ -4412,8 +4446,11 @@ lwip_socket_drop_registered_memberships(int s)
   if (!sock) {
     return;
   }
-
+#if GAZELLE_ENABLE
+  i = s;
+#else
   for (i = 0; i < LWIP_SOCKET_MAX_MEMBERSHIPS; i++) {
+#endif
     if (socket_ipv4_multicast_memberships[i].sock == sock) {
       ip_addr_t multi_addr, if_addr;
       ip_addr_copy_from_ip4(multi_addr, socket_ipv4_multicast_memberships[i].multi_addr);
@@ -4424,7 +4461,9 @@ lwip_socket_drop_registered_memberships(int s)
 
       netconn_join_leave_group(sock->conn, &multi_addr, &if_addr, NETCONN_LEAVE);
     }
+#if !GAZELLE_ENABLE
   }
+#endif
   done_socket(sock);
 }
 #endif /* LWIP_IGMP */
diff --git a/src/include/lwip/sockets.h b/src/include/lwip/sockets.h
index 4466f11..296d4eb 100644
--- a/src/include/lwip/sockets.h
+++ b/src/include/lwip/sockets.h
@@ -361,6 +361,7 @@ struct linger {
 #if GAZELLE_UDP_ENABLE
 #define IP_ADD_MEMBERSHIP  35
 #define IP_DROP_MEMBERSHIP 36
+int lwip_socket_get_multi_memberships(int s, ip4_addr_t *if_addr, ip4_addr_t *multi_addr);
 #else
 /*
  * Options and types related to multicast membership
-- 
2.33.0