From 7197820b156cfc052629f8b4bff2e2288555c1cb Mon Sep 17 00:00:00 2001
From: Motomu Utsumi <motomuman@google.com>
Date: Thu, 20 Feb 2025 10:59:33 +0900
Subject: [PATCH] Add method to destroy UDP socket

Upcoming CL adds feature to close quic connection which requires to
close UDP socket.

Bug: 311792075
Test: atest FrameworksNetTests CtsNetTestCases
Change-Id: I63b7eb9c31cca9d00297a7a67c0cb18d0aa97a59
---
 .../module/util/netlink/InetDiagMessage.java  | 37 ++++++++++++++++---
 1 file changed, 32 insertions(+), 5 deletions(-)

diff --git a/staticlibs/device/com/android/net/module/util/netlink/InetDiagMessage.java b/staticlibs/device/com/android/net/module/util/netlink/InetDiagMessage.java
index fecaa09289..c9a89ec1dd 100644
--- a/staticlibs/device/com/android/net/module/util/netlink/InetDiagMessage.java
+++ b/staticlibs/device/com/android/net/module/util/netlink/InetDiagMessage.java
@@ -309,16 +309,18 @@ public class InetDiagMessage extends NetlinkMessage {
     }
 
     private static void sendNetlinkDestroyRequest(FileDescriptor fd, int proto,
-            InetDiagMessage diagMsg) throws InterruptedIOException, ErrnoException {
+            StructInetDiagSockId id, short family, int state)
+            throws InterruptedIOException, ErrnoException {
+        // TODO: Investigate if it's fine to always set 0 to state and remove state from the arg
         final byte[] destroyMsg = InetDiagMessage.inetDiagReqV2(
                 proto,
-                diagMsg.inetDiagMsg.id,
-                diagMsg.inetDiagMsg.idiag_family,
+                id,
+                family,
                 SOCK_DESTROY,
                 (short) (StructNlMsgHdr.NLM_F_REQUEST | StructNlMsgHdr.NLM_F_ACK),
                 0 /* pad */,
                 0 /* idiagExt */,
-                1 << diagMsg.inetDiagMsg.idiag_state
+                state
         );
         NetlinkUtils.sendMessage(fd, destroyMsg, 0, destroyMsg.length, IO_TIMEOUT_MS);
         NetlinkUtils.receiveNetlinkAck(fd);
@@ -343,7 +345,8 @@ public class InetDiagMessage extends NetlinkMessage {
         Consumer<InetDiagMessage> handleNlDumpMsg = (diagMsg) -> {
             if (filter.test(diagMsg)) {
                 try {
-                    sendNetlinkDestroyRequest(destroyFd, proto, diagMsg);
+                    sendNetlinkDestroyRequest(destroyFd, proto, diagMsg.inetDiagMsg.id,
+                            diagMsg.inetDiagMsg.idiag_family, 1 << diagMsg.inetDiagMsg.idiag_state);
                     destroyedSockets.getAndIncrement();
                 } catch (InterruptedIOException | ErrnoException e) {
                     if (!(e instanceof ErrnoException
@@ -484,6 +487,30 @@ public class InetDiagMessage extends NetlinkMessage {
         Log.d(TAG, "Destroyed live tcp sockets for uids=" + ownerUids + " in " + durationMs + "ms");
     }
 
+    /**
+     * Close the udp socket which can be uniquely identified with the cookie and other information.
+     */
+    public static void destroyUdpSocket(final InetSocketAddress src, final InetSocketAddress dst,
+            final int ifIndex, final long cookie)
+            throws ErrnoException, SocketException, InterruptedIOException {
+        FileDescriptor fd = null;
+
+        try {
+            fd = NetlinkUtils.createNetLinkInetDiagSocket();
+            connectToKernel(fd);
+            final int family = (src.getAddress() instanceof Inet6Address) ? AF_INET6 : AF_INET;
+            final StructInetDiagSockId id = new StructInetDiagSockId(
+                    src,
+                    dst,
+                    ifIndex,
+                    cookie
+            );
+            sendNetlinkDestroyRequest(fd, IPPROTO_UDP, id, (short) family, 0 /* state */);
+        } finally {
+            closeSocketQuietly(fd);
+        }
+    }
+
     @Override
     public String toString() {
         return "InetDiagMessage{ "