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 a8aef7b..eff1f25 100644
--- a/staticlibs/device/com/android/net/module/util/netlink/InetDiagMessage.java
+++ b/staticlibs/device/com/android/net/module/util/netlink/InetDiagMessage.java
@@ -106,21 +106,24 @@
         return bytes;
     }
 
-    public StructInetDiagMsg mStructInetDiagMsg;
+    public StructInetDiagMsg inetDiagMsg;
 
     private InetDiagMessage(@NonNull StructNlMsgHdr header) {
         super(header);
-        mStructInetDiagMsg = new StructInetDiagMsg();
+        inetDiagMsg = new StructInetDiagMsg();
     }
 
     /**
      * Parse an inet_diag_req_v2 message from buffer.
      */
-    @NonNull
+    @Nullable
     public static InetDiagMessage parse(@NonNull StructNlMsgHdr header,
             @NonNull ByteBuffer byteBuffer) {
         final InetDiagMessage msg = new InetDiagMessage(header);
-        msg.mStructInetDiagMsg = StructInetDiagMsg.parse(byteBuffer);
+        msg.inetDiagMsg = StructInetDiagMsg.parse(byteBuffer);
+        if (msg.inetDiagMsg == null) {
+            return null;
+        }
         return msg;
     }
 
@@ -133,12 +136,15 @@
         ByteBuffer response = NetlinkSocket.recvMessage(fd, DEFAULT_RECV_BUFSIZE, TIMEOUT_MS);
 
         final NetlinkMessage nlMsg = NetlinkMessage.parse(response, NETLINK_INET_DIAG);
+        if (nlMsg == null) {
+            return INVALID_UID;
+        }
         final StructNlMsgHdr hdr = nlMsg.getHeader();
         if (hdr.nlmsg_type == NetlinkConstants.NLMSG_DONE) {
             return INVALID_UID;
         }
         if (nlMsg instanceof InetDiagMessage) {
-            return ((InetDiagMessage) nlMsg).mStructInetDiagMsg.idiag_uid;
+            return ((InetDiagMessage) nlMsg).inetDiagMsg.idiag_uid;
         }
         return INVALID_UID;
     }
@@ -228,7 +234,7 @@
                 + "nlmsghdr{"
                 + (mHeader == null ? "" : mHeader.toString(NETLINK_INET_DIAG)) + "}, "
                 + "inet_diag_msg{"
-                + (mStructInetDiagMsg == null ? "" : mStructInetDiagMsg.toString()) + "} "
+                + (inetDiagMsg == null ? "" : inetDiagMsg.toString()) + "} "
                 + "}";
     }
 }
diff --git a/staticlibs/device/com/android/net/module/util/netlink/StructInetDiagMsg.java b/staticlibs/device/com/android/net/module/util/netlink/StructInetDiagMsg.java
index 205656e..ea018cf 100644
--- a/staticlibs/device/com/android/net/module/util/netlink/StructInetDiagMsg.java
+++ b/staticlibs/device/com/android/net/module/util/netlink/StructInetDiagMsg.java
@@ -16,6 +16,9 @@
 
 package com.android.net.module.util.netlink;
 
+import androidx.annotation.NonNull;
+import androidx.annotation.Nullable;
+
 import java.nio.ByteBuffer;
 
 /**
@@ -40,15 +43,28 @@
  */
 public class StructInetDiagMsg {
     public static final int STRUCT_SIZE = 4 + StructInetDiagSockId.STRUCT_SIZE + 20;
+    private static final int IDIAG_SOCK_ID_OFFSET = StructNlMsgHdr.STRUCT_SIZE + 4;
     private static final int IDIAG_UID_OFFSET = StructNlMsgHdr.STRUCT_SIZE + 4
             + StructInetDiagSockId.STRUCT_SIZE + 12;
     public int idiag_uid;
+    @NonNull
+    public StructInetDiagSockId id;
 
     /**
      * Parse inet diag netlink message from buffer.
      */
-    public static StructInetDiagMsg parse(ByteBuffer byteBuffer) {
+    @Nullable
+    public static StructInetDiagMsg parse(@NonNull ByteBuffer byteBuffer) {
+        if (byteBuffer.remaining() < STRUCT_SIZE) {
+            return null;
+        }
         StructInetDiagMsg struct = new StructInetDiagMsg();
+        final byte family = byteBuffer.get();
+        byteBuffer.position(IDIAG_SOCK_ID_OFFSET);
+        struct.id = StructInetDiagSockId.parse(byteBuffer, family);
+        if (struct.id == null) {
+            return null;
+        }
         struct.idiag_uid = byteBuffer.getInt(IDIAG_UID_OFFSET);
         return struct;
     }
@@ -57,6 +73,7 @@
     public String toString() {
         return "StructInetDiagMsg{ "
                 + "idiag_uid{" + idiag_uid + "}, "
+                + "id{" + id + "}, "
                 + "}";
     }
 }
diff --git a/staticlibs/device/com/android/net/module/util/netlink/StructInetDiagSockId.java b/staticlibs/device/com/android/net/module/util/netlink/StructInetDiagSockId.java
index 95723cb..648a020 100644
--- a/staticlibs/device/com/android/net/module/util/netlink/StructInetDiagSockId.java
+++ b/staticlibs/device/com/android/net/module/util/netlink/StructInetDiagSockId.java
@@ -106,7 +106,7 @@
             byteBuffer.get(srcAddrByte);
             byteBuffer.get(dstAddrByte);
         } else {
-            Log.e(TAG, "Invalid address family: " + family);
+            Log.wtf(TAG, "Invalid address family: " + family);
             return null;
         }
 
@@ -118,7 +118,7 @@
         } catch (UnknownHostException e) {
             // Should not happen. UnknownHostException is thrown only if addr byte array is of
             // illegal length.
-            Log.e(TAG, "Failed to parse address: " + e);
+            Log.wtf(TAG, "Failed to parse address: " + e);
             return null;
         }
 
diff --git a/staticlibs/tests/unit/src/com/android/net/module/util/netlink/InetDiagSocketTest.java b/staticlibs/tests/unit/src/com/android/net/module/util/netlink/InetDiagSocketTest.java
index 6b1482d..c7e2a4d 100644
--- a/staticlibs/tests/unit/src/com/android/net/module/util/netlink/InetDiagSocketTest.java
+++ b/staticlibs/tests/unit/src/com/android/net/module/util/netlink/InetDiagSocketTest.java
@@ -31,6 +31,8 @@
 import static org.junit.Assert.assertTrue;
 import static org.junit.Assert.fail;
 
+import android.net.InetAddresses;
+
 import androidx.test.filters.SmallTest;
 import androidx.test.runner.AndroidJUnit4;
 
@@ -235,10 +237,10 @@
                 // inet_diag_sockid
                 "a817" +     // idiag_sport = 43031
                 "960f" +     // idiag_dport = 38415
-                "fe8000000000000086c9b2fffe6aed4b" + // idiag_src = fe80::86c9:b2ff:fe6a:ed4b
-                "00000000000000000000ffff08080808" + // idiag_dst = 8.8.8.8
-                "00000000" + // idiag_if
-                "ffffffffffffffff" + // idiag_cookie = INET_DIAG_NOCOOKIE
+                "20010db8000000000000000000000001" + // idiag_src = 2001:db8::1
+                "20010db8000000000000000000000002" + // idiag_dst = 2001:db8::2
+                "07000000" + // idiag_if = 7
+                "5800000000000000" + // idiag_cookie = 88
             "00000000" +     // idiag_expires
             "00000000" +     // idiag_rqueue
             "00000000" +     // idiag_wqueue
@@ -256,7 +258,16 @@
 
         assertTrue(msg instanceof InetDiagMessage);
         final InetDiagMessage inetDiagMsg = (InetDiagMessage) msg;
-        assertEquals(10147, inetDiagMsg.mStructInetDiagMsg.idiag_uid);
+        assertEquals(10147, inetDiagMsg.inetDiagMsg.idiag_uid);
+        final StructInetDiagSockId sockId = inetDiagMsg.inetDiagMsg.id;
+        assertEquals(43031, sockId.locSocketAddress.getPort());
+        assertEquals(InetAddresses.parseNumericAddress("2001:db8::1"),
+                sockId.locSocketAddress.getAddress());
+        assertEquals(38415, sockId.remSocketAddress.getPort());
+        assertEquals(InetAddresses.parseNumericAddress("2001:db8::2"),
+                sockId.remSocketAddress.getAddress());
+        assertEquals(7, sockId.ifIndex);
+        assertEquals(88, sockId.cookie);
 
         final StructNlMsgHdr hdr = inetDiagMsg.getHeader();
         assertNotNull(hdr);