xnu-4903.241.1.tar.gz

[apple/xnu.git] / bsd / netinet / flow_divert.c

diff --git a/bsd/netinet/flow_divert.c b/bsd/netinet/flow_divert.c
index cc2e2c8fe54a2c4c19d149d6368f12d40040d747..c3bc7930cd23061355fd8c48ad1a0f277e38df2f 100644 (file)
--- a/bsd/netinet/flow_divert.c
+++ b/bsd/netinet/flow_divert.c
@@ -1,5 +1,5 @@
 /*
 /*

- * Copyright (c) 2012-2016 Apple Inc. All rights reserved.
+ * Copyright (c) 2012-2017 Apple Inc. All rights reserved.

  *
  * @APPLE_OSREFERENCE_LICENSE_HEADER_START@
  * 
  *
  * @APPLE_OSREFERENCE_LICENSE_HEADER_START@
  * 


@@ -63,6 +63,7 @@
 #include <dev/random/randomdev.h>
 #include <libkern/crypto/sha1.h>
 #include <libkern/crypto/crypto_internal.h>
 #include <dev/random/randomdev.h>
 #include <libkern/crypto/sha1.h>
 #include <libkern/crypto/crypto_internal.h>

+#include <os/log.h>

 
 #define FLOW_DIVERT_CONNECT_STARTED            0x00000001
 #define FLOW_DIVERT_READ_CLOSED                        0x00000002
 
 #define FLOW_DIVERT_CONNECT_STARTED            0x00000001
 #define FLOW_DIVERT_READ_CLOSED                        0x00000002


@@ -72,17 +73,11 @@
 #define FLOW_DIVERT_TRANSFERRED                        0x00000020
 #define FLOW_DIVERT_HAS_HMAC            0x00000040
 
 #define FLOW_DIVERT_TRANSFERRED                        0x00000020
 #define FLOW_DIVERT_HAS_HMAC            0x00000040
 

-#define FDLOG(level, pcb, format, ...) do {                                                                                    \
-       if (level <= (pcb)->log_level) {                                                                                                \
-               log((level > LOG_NOTICE ? LOG_NOTICE : level), "%s (%u): " format "\n", __FUNCTION__, (pcb)->hash, __VA_ARGS__);        \
-       }                                                                                                                                                               \
-} while (0)
+#define FDLOG(level, pcb, format, ...) \
+       os_log_with_type(OS_LOG_DEFAULT, flow_divert_syslog_type_to_oslog_type(level), "(%u): " format "\n", (pcb)->hash, __VA_ARGS__)

 
 

-#define FDLOG0(level, pcb, msg) do {                                                                                           \
-       if (level <= (pcb)->log_level) {                                                                                                \
-               log((level > LOG_NOTICE ? LOG_NOTICE : level), "%s (%u): %s\n", __FUNCTION__, (pcb)->hash, msg);                                \
-       }                                                                                                                                                               \
-} while (0)
+#define FDLOG0(level, pcb, msg) \
+       os_log_with_type(OS_LOG_DEFAULT, flow_divert_syslog_type_to_oslog_type(level), "(%u): " msg "\n", (pcb)->hash)

 
 #define FDRETAIN(pcb)                  if ((pcb) != NULL) OSIncrementAtomic(&(pcb)->ref_count)
 #define FDRELEASE(pcb)                                                                                                         \
 
 #define FDRETAIN(pcb)                  if ((pcb) != NULL) OSIncrementAtomic(&(pcb)->ref_count)
 #define FDRELEASE(pcb)                                                                                                         \


@@ -95,7 +90,7 @@
 #define FDLOCK(pcb)                                            lck_mtx_lock(&(pcb)->mtx)
 #define FDUNLOCK(pcb)                                  lck_mtx_unlock(&(pcb)->mtx)
 
 #define FDLOCK(pcb)                                            lck_mtx_lock(&(pcb)->mtx)
 #define FDUNLOCK(pcb)                                  lck_mtx_unlock(&(pcb)->mtx)
 

-#define FD_CTL_SENDBUFF_SIZE                   (2 * FLOW_DIVERT_CHUNK_SIZE)
+#define FD_CTL_SENDBUFF_SIZE                   (128 * 1024)

 #define FD_CTL_RCVBUFF_SIZE                            (128 * 1024)
 
 #define GROUP_BIT_CTL_ENQUEUE_BLOCKED  0
 #define FD_CTL_RCVBUFF_SIZE                            (128 * 1024)
 
 #define GROUP_BIT_CTL_ENQUEUE_BLOCKED  0


@@ -103,30 +98,13 @@
 #define GROUP_COUNT_MAX                                        32
 #define FLOW_DIVERT_MAX_NAME_SIZE              4096
 #define FLOW_DIVERT_MAX_KEY_SIZE               1024
 #define GROUP_COUNT_MAX                                        32
 #define FLOW_DIVERT_MAX_NAME_SIZE              4096
 #define FLOW_DIVERT_MAX_KEY_SIZE               1024

-
-#define DNS_SERVICE_GROUP_UNIT                 (GROUP_COUNT_MAX + 1)
+#define FLOW_DIVERT_MAX_TRIE_MEMORY            (1024 * 1024)

 
 struct flow_divert_trie_node
 {
        uint16_t start;
        uint16_t length;
        uint16_t child_map;
 
 struct flow_divert_trie_node
 {
        uint16_t start;
        uint16_t length;
        uint16_t child_map;

-       uint32_t group_unit;
-};
-
-struct flow_divert_trie
-{
-       struct flow_divert_trie_node *nodes;
-       uint16_t *child_maps;
-       uint8_t *bytes;
-       void *memory;
-       size_t nodes_count;
-       size_t child_maps_count;
-       size_t bytes_count;
-       size_t nodes_free_next;
-       size_t child_maps_free_next;
-       size_t bytes_free_next;
-       uint16_t root;

 };
 
 #define CHILD_MAP_SIZE                 256
 };
 
 #define CHILD_MAP_SIZE                 256


@@ -140,7 +118,6 @@ static struct flow_divert_pcb               nil_pcb;
 decl_lck_rw_data(static, g_flow_divert_group_lck);
 static struct flow_divert_group                **g_flow_divert_groups                  = NULL;
 static uint32_t                                                g_active_group_count                    = 0;
 decl_lck_rw_data(static, g_flow_divert_group_lck);
 static struct flow_divert_group                **g_flow_divert_groups                  = NULL;
 static uint32_t                                                g_active_group_count                    = 0;

-static struct flow_divert_trie         g_signing_id_trie;

 
 static lck_grp_attr_t                          *flow_divert_grp_attr                   = NULL;
 static lck_attr_t                                      *flow_divert_mtx_attr                   = NULL;
 
 static lck_grp_attr_t                          *flow_divert_grp_attr                   = NULL;
 static lck_attr_t                                      *flow_divert_mtx_attr                   = NULL;


@@ -186,6 +163,17 @@ flow_divert_has_pcb_local_address(const struct inpcb *inp);
 static void
 flow_divert_disconnect_socket(struct socket *so);
 
 static void
 flow_divert_disconnect_socket(struct socket *so);
 

+static inline uint8_t
+flow_divert_syslog_type_to_oslog_type(int syslog_type)
+{
+       switch (syslog_type) {
+               case LOG_ERR: return OS_LOG_TYPE_ERROR;
+               case LOG_INFO: return OS_LOG_TYPE_INFO;
+               case LOG_DEBUG: return OS_LOG_TYPE_DEBUG;
+               default: return OS_LOG_TYPE_DEFAULT;
+       }
+}
+

 static inline int
 flow_divert_pcb_cmp(const struct flow_divert_pcb *pcb_a, const struct flow_divert_pcb *pcb_b)
 {
 static inline int
 flow_divert_pcb_cmp(const struct flow_divert_pcb *pcb_a, const struct flow_divert_pcb *pcb_b)
 {


@@ -211,8 +199,6 @@ flow_divert_packet_type2str(uint8_t packet_type)
                        return "read notification";
                case FLOW_DIVERT_PKT_PROPERTIES_UPDATE:
                        return "properties update";
                        return "read notification";
                case FLOW_DIVERT_PKT_PROPERTIES_UPDATE:
                        return "properties update";

-               case FLOW_DIVERT_PKT_APP_MAP_UPDATE:
-                       return "app map update";

                case FLOW_DIVERT_PKT_APP_MAP_CREATE:
                        return "app map create";
                default:
                case FLOW_DIVERT_PKT_APP_MAP_CREATE:
                        return "app map create";
                default:


@@ -418,7 +404,7 @@ flow_divert_packet_append_tlv(mbuf_t packet, uint8_t type, uint32_t length, cons
 
        error = mbuf_copyback(packet, mbuf_pkthdr_len(packet), sizeof(net_length), &net_length, MBUF_DONTWAIT);
        if (error) {
 
        error = mbuf_copyback(packet, mbuf_pkthdr_len(packet), sizeof(net_length), &net_length, MBUF_DONTWAIT);
        if (error) {

-               FDLOG(LOG_ERR, &nil_pcb, "failed to append the length (%lu)", length);
+               FDLOG(LOG_ERR, &nil_pcb, "failed to append the length (%u)", length);

                return error;
        }
 
                return error;
        }
 


@@ -615,6 +601,7 @@ flow_divert_add_data_statistics(struct flow_divert_pcb *fd_cb, int data_len, Boo
                INP_ADD_STAT(inp, cell, wifi, wired, rxpackets, 1);
                INP_ADD_STAT(inp, cell, wifi, wired, rxbytes, data_len);
        }
                INP_ADD_STAT(inp, cell, wifi, wired, rxpackets, 1);
                INP_ADD_STAT(inp, cell, wifi, wired, rxbytes, data_len);
        }

+       inp_set_activity_bitmap(inp);

 }
 
 static errno_t
 }
 
 static errno_t


@@ -1005,6 +992,8 @@ flow_divert_create_connect_packet(struct flow_divert_pcb *fd_cb, struct sockaddr
        char                    *signing_id = NULL;
        int                             free_signing_id = 0;
        mbuf_t                  connect_packet = NULL;
        char                    *signing_id = NULL;
        int                             free_signing_id = 0;
        mbuf_t                  connect_packet = NULL;

+       proc_t                  src_proc = p;
+       int                             release_proc = 0;

 
        error = flow_divert_packet_init(fd_cb, FLOW_DIVERT_PKT_CONNECT, &connect_packet);
        if (error) {
 
        error = flow_divert_packet_init(fd_cb, FLOW_DIVERT_PKT_CONNECT, &connect_packet);
        if (error) {


@@ -1027,69 +1016,70 @@ flow_divert_create_connect_packet(struct flow_divert_pcb *fd_cb, struct sockaddr
        }
 
        socket_unlock(so, 0);
        }
 
        socket_unlock(so, 0);

-       if (g_signing_id_trie.root != NULL_TRIE_IDX) {
-               proc_t src_proc = p;
-               int release_proc = 0;
-
-               if (signing_id == NULL) {
-                       release_proc = flow_divert_get_src_proc(so, &src_proc);
-                       if (src_proc != PROC_NULL) {
-                               proc_lock(src_proc);
-                               if (src_proc->p_csflags & CS_VALID) {
-                    const char * cs_id;
-                    cs_id = cs_identity_get(src_proc);
-                    signing_id = __DECONST(char *, cs_id);
-                               } else {
-                                       FDLOG0(LOG_WARNING, fd_cb, "Signature is invalid");
-                               }
+
+       if (signing_id == NULL) {
+               release_proc = flow_divert_get_src_proc(so, &src_proc);
+               if (src_proc != PROC_NULL) {
+                       proc_lock(src_proc);
+                       if (src_proc->p_csflags & (CS_VALID|CS_DEBUGGED)) {
+                               const char * cs_id;
+                               cs_id = cs_identity_get(src_proc);
+                               signing_id = __DECONST(char *, cs_id);

                        } else {
                        } else {

-                               FDLOG0(LOG_WARNING, fd_cb, "Failed to determine the current proc");
+                               FDLOG0(LOG_WARNING, fd_cb, "Signature is invalid");

                        }
                } else {
                        }
                } else {

-                       src_proc = PROC_NULL;
+                       FDLOG0(LOG_WARNING, fd_cb, "Failed to determine the current proc");

                }
                }

+       } else {
+               src_proc = PROC_NULL;
+       }

 
 

-               if (signing_id != NULL) {
-                       uint16_t result = NULL_TRIE_IDX;
-                       lck_rw_lock_shared(&g_flow_divert_group_lck);
-                       result = flow_divert_trie_search(&g_signing_id_trie, (uint8_t *)signing_id);
-                       lck_rw_done(&g_flow_divert_group_lck);
-                       if (result != NULL_TRIE_IDX) {
-                               error = 0;
-                               FDLOG(LOG_INFO, fd_cb, "%s matched", signing_id);
+       if (signing_id != NULL) {
+               uint16_t result = NULL_TRIE_IDX;
+               lck_rw_lock_shared(&fd_cb->group->lck);
+               if (fd_cb->group->flags & FLOW_DIVERT_GROUP_FLAG_NO_APP_MAP) {
+                       result = 1;
+               } else {
+                       result = flow_divert_trie_search(&fd_cb->group->signing_id_trie, (uint8_t *)signing_id);
+               }
+               lck_rw_done(&fd_cb->group->lck);
+               if (result != NULL_TRIE_IDX) {
+                       error = 0;
+                       FDLOG(LOG_INFO, fd_cb, "%s matched", signing_id);

 
 

-                               error = flow_divert_packet_append_tlv(connect_packet, FLOW_DIVERT_TLV_SIGNING_ID, strlen(signing_id), signing_id);
-                               if (error == 0) {
-                                       if (src_proc != PROC_NULL) {
-                                               unsigned char cdhash[SHA1_RESULTLEN];
-                                               error = proc_getcdhash(src_proc, cdhash);
-                                               if (error == 0) {
-                                                       error = flow_divert_packet_append_tlv(connect_packet, FLOW_DIVERT_TLV_CDHASH, sizeof(cdhash), cdhash);
-                                                       if (error) {
-                                                               FDLOG(LOG_ERR, fd_cb, "failed to append the cdhash: %d", error);
-                                                       }
-                                               } else {
-                                                       FDLOG(LOG_ERR, fd_cb, "failed to get the cdhash: %d", error);
+                       error = flow_divert_packet_append_tlv(connect_packet, FLOW_DIVERT_TLV_SIGNING_ID, strlen(signing_id), signing_id);
+                       if (error == 0) {
+                               if (src_proc != PROC_NULL) {
+                                       unsigned char cdhash[SHA1_RESULTLEN];
+                                       error = proc_getcdhash(src_proc, cdhash);
+                                       if (error == 0) {
+                                               error = flow_divert_packet_append_tlv(connect_packet, FLOW_DIVERT_TLV_CDHASH, sizeof(cdhash), cdhash);
+                                               if (error) {
+                                                       FDLOG(LOG_ERR, fd_cb, "failed to append the cdhash: %d", error);

                                                }
                                                }

+                                       } else {
+                                               FDLOG(LOG_ERR, fd_cb, "failed to get the cdhash: %d", error);

                                        }
                                        }

-                               } else {
-                                       FDLOG(LOG_ERR, fd_cb, "failed to append the signing ID: %d", error);

                                }
                        } else {
                                }
                        } else {

-                               FDLOG(LOG_WARNING, fd_cb, "%s did not match", signing_id);
+                               FDLOG(LOG_ERR, fd_cb, "failed to append the signing ID: %d", error);

                        }
                } else {
                        }
                } else {

-                       FDLOG0(LOG_WARNING, fd_cb, "Failed to get the code signing identity");
+                       FDLOG(LOG_WARNING, fd_cb, "%s did not match", signing_id);

                }
                }

+       } else {
+               FDLOG0(LOG_WARNING, fd_cb, "Failed to get the code signing identity");
+               if (fd_cb->group->flags & FLOW_DIVERT_GROUP_FLAG_NO_APP_MAP) {
+                       error = 0;
+               }
+       }

 
 

-               if (src_proc != PROC_NULL) {
-                       proc_unlock(src_proc);
-                       if (release_proc) {
-                               proc_rele(src_proc);
-                       }
+       if (src_proc != PROC_NULL) {
+               proc_unlock(src_proc);
+               if (release_proc) {
+                       proc_rele(src_proc);

                }
                }

-       } else {
-               FDLOG0(LOG_WARNING, fd_cb, "The signing ID trie is empty");

        }
        socket_lock(so, 0);
 
        }
        socket_lock(so, 0);
 


@@ -1196,7 +1186,7 @@ flow_divert_create_connect_packet(struct flow_divert_pcb *fd_cb, struct sockaddr
        if (fd_cb->local_address != NULL) {
                /* socket is bound. */
                error = flow_divert_packet_append_tlv(connect_packet, FLOW_DIVERT_TLV_LOCAL_ADDR,
        if (fd_cb->local_address != NULL) {
                /* socket is bound. */
                error = flow_divert_packet_append_tlv(connect_packet, FLOW_DIVERT_TLV_LOCAL_ADDR,

-                                                     sizeof(struct sockaddr_storage), fd_cb->local_address);
+                                                     fd_cb->local_address->sa_len, fd_cb->local_address);

                if (error) {
                        goto done;
                }
                if (error) {
                        goto done;
                }


@@ -1378,9 +1368,11 @@ flow_divert_send_data_packet(struct flow_divert_pcb *fd_cb, mbuf_t data, size_t
                }
        }
 
                }
        }
 

-       last = m_last(packet);
-       mbuf_setnext(last, data);
-       mbuf_pkthdr_adjustlen(packet, data_len);
+       if (data_len > 0 && data != NULL) {
+               last = m_last(packet);
+               mbuf_setnext(last, data);
+               mbuf_pkthdr_adjustlen(packet, data_len);
+       }

        error = flow_divert_send_packet(fd_cb, packet, force);
 
        if (error) {
        error = flow_divert_send_packet(fd_cb, packet, force);
 
        if (error) {


@@ -1465,11 +1457,15 @@ flow_divert_send_buffered_data(struct flow_divert_pcb *fd_cb, Boolean force)
                                }
                        }
                        data_len = mbuf_pkthdr_len(m);
                                }
                        }
                        data_len = mbuf_pkthdr_len(m);

-                       FDLOG(LOG_DEBUG, fd_cb, "mbuf_copym() data_len = %u", data_len);
-                       error = mbuf_copym(m, 0, data_len, MBUF_DONTWAIT, &data);
-                       if (error) {
-                               FDLOG(LOG_ERR, fd_cb, "mbuf_copym failed: %d", error);
-                               break;
+                       if (data_len > 0) {
+                               FDLOG(LOG_DEBUG, fd_cb, "mbuf_copym() data_len = %lu", data_len);
+                               error = mbuf_copym(m, 0, data_len, MBUF_DONTWAIT, &data);
+                               if (error) {
+                                       FDLOG(LOG_ERR, fd_cb, "mbuf_copym failed: %d", error);
+                                       break;
+                               }
+                       } else {
+                               data = NULL;

                        }
                        error = flow_divert_send_data_packet(fd_cb, data, data_len, toaddr, force);
                        if (error) {
                        }
                        error = flow_divert_send_data_packet(fd_cb, data, data_len, toaddr, force);
                        if (error) {


@@ -1569,10 +1565,10 @@ flow_divert_send_app_data(struct flow_divert_pcb *fd_cb, mbuf_t data, struct soc
                        }
                }
        } else if (SOCK_TYPE(fd_cb->so) == SOCK_DGRAM) {
                        }
                }
        } else if (SOCK_TYPE(fd_cb->so) == SOCK_DGRAM) {

-               if (to_send) {
+               if (to_send || mbuf_pkthdr_len(data) == 0) {

                        error = flow_divert_send_data_packet(fd_cb, data, to_send, toaddr, FALSE);
                        if (error) {
                        error = flow_divert_send_data_packet(fd_cb, data, to_send, toaddr, FALSE);
                        if (error) {

-                               FDLOG(LOG_ERR, fd_cb, "flow_divert_send_data_packet failed. send data size = %u", to_send);
+                               FDLOG(LOG_ERR, fd_cb, "flow_divert_send_data_packet failed. send data size = %lu", to_send);

                        } else {
                                fd_cb->send_window -= to_send;
                        }
                        } else {
                                fd_cb->send_window -= to_send;
                        }


@@ -1675,6 +1671,7 @@ flow_divert_handle_connect_result(struct flow_divert_pcb *fd_cb, mbuf_t packet,
        int                                                     out_if_index            = 0;
        struct sockaddr_storage         remote_address;
        uint32_t                                        send_window;
        int                                                     out_if_index            = 0;
        struct sockaddr_storage         remote_address;
        uint32_t                                        send_window;

+       uint32_t                                        app_data_length         = 0;

 
        memset(&local_address, 0, sizeof(local_address));
        memset(&remote_address, 0, sizeof(remote_address));
 
        memset(&local_address, 0, sizeof(local_address));
        memset(&remote_address, 0, sizeof(remote_address));


@@ -1695,32 +1692,37 @@ flow_divert_handle_connect_result(struct flow_divert_pcb *fd_cb, mbuf_t packet,
 
        error = flow_divert_packet_get_tlv(packet, offset, FLOW_DIVERT_TLV_CTL_UNIT, sizeof(ctl_unit), &ctl_unit, NULL);
        if (error) {
 
        error = flow_divert_packet_get_tlv(packet, offset, FLOW_DIVERT_TLV_CTL_UNIT, sizeof(ctl_unit), &ctl_unit, NULL);
        if (error) {

-               FDLOG(LOG_ERR, fd_cb, "failed to get the control unit: %d", error);
-               return;
+               FDLOG0(LOG_INFO, fd_cb, "No control unit provided in the connect result");

        }
 
        error = flow_divert_packet_get_tlv(packet, offset, FLOW_DIVERT_TLV_LOCAL_ADDR, sizeof(local_address), &local_address, NULL);
        if (error) {
        }
 
        error = flow_divert_packet_get_tlv(packet, offset, FLOW_DIVERT_TLV_LOCAL_ADDR, sizeof(local_address), &local_address, NULL);
        if (error) {

-               FDLOG0(LOG_NOTICE, fd_cb, "No local address provided");
+               FDLOG0(LOG_INFO, fd_cb, "No local address provided");

        }
 
        error = flow_divert_packet_get_tlv(packet, offset, FLOW_DIVERT_TLV_REMOTE_ADDR, sizeof(remote_address), &remote_address, NULL);
        if (error) {
        }
 
        error = flow_divert_packet_get_tlv(packet, offset, FLOW_DIVERT_TLV_REMOTE_ADDR, sizeof(remote_address), &remote_address, NULL);
        if (error) {

-               FDLOG0(LOG_NOTICE, fd_cb, "No remote address provided");
+               FDLOG0(LOG_INFO, fd_cb, "No remote address provided");

        }
 
        error = flow_divert_packet_get_tlv(packet, offset, FLOW_DIVERT_TLV_OUT_IF_INDEX, sizeof(out_if_index), &out_if_index, NULL);
        if (error) {
        }
 
        error = flow_divert_packet_get_tlv(packet, offset, FLOW_DIVERT_TLV_OUT_IF_INDEX, sizeof(out_if_index), &out_if_index, NULL);
        if (error) {

-               FDLOG0(LOG_NOTICE, fd_cb, "No output if index provided");
+               FDLOG0(LOG_INFO, fd_cb, "No output if index provided");

        }
 
        }
 

+       error = flow_divert_packet_get_tlv(packet, offset, FLOW_DIVERT_TLV_APP_DATA, 0, NULL, &app_data_length);
+       if (error) {
+               FDLOG0(LOG_INFO, fd_cb, "No application data provided in connect result");
+       }
+
+       error = 0;

        connect_error   = ntohl(connect_error);
        ctl_unit                = ntohl(ctl_unit);
 
        lck_rw_lock_shared(&g_flow_divert_group_lck);
 
        connect_error   = ntohl(connect_error);
        ctl_unit                = ntohl(ctl_unit);
 
        lck_rw_lock_shared(&g_flow_divert_group_lck);
 

-       if (connect_error == 0) {
-               if (ctl_unit == 0 || ctl_unit >= GROUP_COUNT_MAX) {
+       if (connect_error == 0 && ctl_unit > 0) {
+               if (ctl_unit >= GROUP_COUNT_MAX) {

                        FDLOG(LOG_ERR, fd_cb, "Connect result contains an invalid control unit: %u", ctl_unit);
                        error = EINVAL;
                } else if (g_flow_divert_groups == NULL || g_active_group_count == 0) {
                        FDLOG(LOG_ERR, fd_cb, "Connect result contains an invalid control unit: %u", ctl_unit);
                        error = EINVAL;
                } else if (g_flow_divert_groups == NULL || g_active_group_count == 0) {


@@ -1773,6 +1775,27 @@ flow_divert_handle_connect_result(struct flow_divert_pcb *fd_cb, mbuf_t packet,
                        goto set_socket_state;
                }
 
                        goto set_socket_state;
                }
 

+               if (app_data_length > 0) {
+                       uint8_t *app_data = NULL;
+                       MALLOC(app_data, uint8_t *, app_data_length, M_TEMP, M_WAITOK);
+                       if (app_data != NULL) {
+                               error = flow_divert_packet_get_tlv(packet, offset, FLOW_DIVERT_TLV_APP_DATA, app_data_length, app_data, NULL);
+                               if (error == 0) {
+                                       FDLOG(LOG_INFO, fd_cb, "Got %u bytes of app data from the connect result", app_data_length);
+                                       if (fd_cb->app_data != NULL) {
+                                               FREE(fd_cb->app_data, M_TEMP);
+                                       }
+                                       fd_cb->app_data = app_data;
+                                       fd_cb->app_data_length = app_data_length;
+                               } else {
+                                       FDLOG(LOG_ERR, fd_cb, "Failed to copy %u bytes of application data from the connect result packet", app_data_length);
+                                       FREE(app_data, M_TEMP);
+                               }
+                       } else {
+                               FDLOG(LOG_ERR, fd_cb, "Failed to allocate a buffer of size %u to hold the application data from the connect result", app_data_length);
+                       }
+               }
+

                ifnet_head_lock_shared();
                if (out_if_index > 0 && out_if_index <= if_index) {
                        ifp = ifindex2ifnet[out_if_index];
                ifnet_head_lock_shared();
                if (out_if_index > 0 && out_if_index <= if_index) {
                        ifp = ifindex2ifnet[out_if_index];


@@ -1794,20 +1817,22 @@ flow_divert_handle_connect_result(struct flow_divert_pcb *fd_cb, mbuf_t packet,
                        goto set_socket_state;
                }
 
                        goto set_socket_state;
                }
 

-               old_group = fd_cb->group;
+               if (grp != NULL) {
+                       old_group = fd_cb->group;

 
 

-               lck_rw_lock_exclusive(&old_group->lck);
-               lck_rw_lock_exclusive(&grp->lck);
+                       lck_rw_lock_exclusive(&old_group->lck);
+                       lck_rw_lock_exclusive(&grp->lck);

 
 

-               RB_REMOVE(fd_pcb_tree, &old_group->pcb_tree, fd_cb);
-               if (RB_INSERT(fd_pcb_tree, &grp->pcb_tree, fd_cb) != NULL) {
-                       panic("group with unit %u already contains a connection with hash %u", grp->ctl_unit, fd_cb->hash);
-               }
+                       RB_REMOVE(fd_pcb_tree, &old_group->pcb_tree, fd_cb);
+                       if (RB_INSERT(fd_pcb_tree, &grp->pcb_tree, fd_cb) != NULL) {
+                               panic("group with unit %u already contains a connection with hash %u", grp->ctl_unit, fd_cb->hash);
+                       }

 
 

-               fd_cb->group = grp;
+                       fd_cb->group = grp;

 
 

-               lck_rw_done(&grp->lck);
-               lck_rw_done(&old_group->lck);
+                       lck_rw_done(&grp->lck);
+                       lck_rw_done(&old_group->lck);
+               }

 
                fd_cb->send_window = ntohl(send_window);
 
 
                fd_cb->send_window = ntohl(send_window);
 


@@ -1888,18 +1913,20 @@ flow_divert_handle_close(struct flow_divert_pcb *fd_cb, mbuf_t packet, int offse
 static mbuf_t
 flow_divert_get_control_mbuf(struct flow_divert_pcb *fd_cb)
 {
 static mbuf_t
 flow_divert_get_control_mbuf(struct flow_divert_pcb *fd_cb)
 {

-       struct inpcb *inp = sotoinpcb(fd_cb->so);
-       if (inp->inp_vflag & INP_IPV4 && inp->inp_flags & INP_RECVDSTADDR) {
-               struct sockaddr_in *sin = (struct sockaddr_in *)(void *)fd_cb->local_address;
+       if (fd_cb->local_address != NULL) {
+               struct inpcb *inp = sotoinpcb(fd_cb->so);
+               if (inp->inp_vflag & INP_IPV4 && inp->inp_flags & INP_RECVDSTADDR) {
+                       struct sockaddr_in *sin = (struct sockaddr_in *)(void *)fd_cb->local_address;

 
 

-               return sbcreatecontrol((caddr_t) &sin->sin_addr, sizeof(struct in_addr), IP_RECVDSTADDR, IPPROTO_IP);
-       } else if (inp->inp_vflag & INP_IPV6 && (inp->inp_flags & IN6P_PKTINFO) != 0) {
-               struct sockaddr_in6 *sin6 = (struct sockaddr_in6 *)(void *)fd_cb->local_address;
-               struct in6_pktinfo pi6;
+                       return sbcreatecontrol((caddr_t) &sin->sin_addr, sizeof(struct in_addr), IP_RECVDSTADDR, IPPROTO_IP);
+               } else if (inp->inp_vflag & INP_IPV6 && (inp->inp_flags & IN6P_PKTINFO) != 0) {
+                       struct sockaddr_in6 *sin6 = (struct sockaddr_in6 *)(void *)fd_cb->local_address;
+                       struct in6_pktinfo pi6;

 
 

-               bcopy(&sin6->sin6_addr, &pi6.ipi6_addr, sizeof (struct in6_addr));
-               pi6.ipi6_ifindex = 0;
-               return sbcreatecontrol((caddr_t)&pi6, sizeof (struct in6_pktinfo), IPV6_PKTINFO, IPPROTO_IPV6);
+                       bcopy(&sin6->sin6_addr, &pi6.ipi6_addr, sizeof (struct in6_addr));
+                       pi6.ipi6_ifindex = 0;
+                       return sbcreatecontrol((caddr_t)&pi6, sizeof (struct in6_pktinfo), IPV6_PKTINFO, IPPROTO_IPV6);
+               }

        }
        return (NULL);
 }
        }
        return (NULL);
 }


@@ -1992,10 +2019,6 @@ flow_divert_handle_data(struct flow_divert_pcb *fd_cb, mbuf_t packet, size_t off
                        }
                }
                socket_unlock(fd_cb->so, 0);
                        }
                }
                socket_unlock(fd_cb->so, 0);

-
-               if (data != NULL) {
-                       mbuf_freem(data);
-               }

        }
        FDUNLOCK(fd_cb);
 }
        }
        FDUNLOCK(fd_cb);
 }


@@ -2030,6 +2053,7 @@ flow_divert_handle_group_init(struct flow_divert_group *group, mbuf_t packet, in
        int error = 0;
        uint32_t key_size = 0;
        int log_level;
        int error = 0;
        uint32_t key_size = 0;
        int log_level;

+       uint32_t flags = 0;

 
        error = flow_divert_packet_get_tlv(packet, offset, FLOW_DIVERT_TLV_TOKEN_KEY, 0, NULL, &key_size);
        if (error) {
 
        error = flow_divert_packet_get_tlv(packet, offset, FLOW_DIVERT_TLV_TOKEN_KEY, 0, NULL, &key_size);
        if (error) {


@@ -2038,7 +2062,7 @@ flow_divert_handle_group_init(struct flow_divert_group *group, mbuf_t packet, in
        }
 
        if (key_size == 0 || key_size > FLOW_DIVERT_MAX_KEY_SIZE) {
        }
 
        if (key_size == 0 || key_size > FLOW_DIVERT_MAX_KEY_SIZE) {

-               FDLOG(LOG_ERR, &nil_pcb, "Invalid key size: %lu", key_size);
+               FDLOG(LOG_ERR, &nil_pcb, "Invalid key size: %u", key_size);

                return;
        }
 
                return;
        }
 


@@ -2061,6 +2085,11 @@ flow_divert_handle_group_init(struct flow_divert_group *group, mbuf_t packet, in
 
        group->token_key_size = key_size;
 
 
        group->token_key_size = key_size;
 

+       error = flow_divert_packet_get_tlv(packet, offset, FLOW_DIVERT_TLV_FLAGS, sizeof(flags), &flags, NULL);
+       if (!error) {
+               group->flags = flags;
+       }
+

        lck_rw_done(&group->lck);
 }
 
        lck_rw_done(&group->lck);
 }
 


@@ -2167,7 +2196,7 @@ flow_divert_handle_properties_update(struct flow_divert_pcb *fd_cb, mbuf_t packe
 }
 
 static void
 }
 
 static void

-flow_divert_handle_app_map_create(mbuf_t packet, int offset)
+flow_divert_handle_app_map_create(struct flow_divert_group *group, mbuf_t packet, int offset)

 {
        size_t bytes_mem_size;
        size_t child_maps_mem_size;
 {
        size_t bytes_mem_size;
        size_t child_maps_mem_size;


@@ -2178,21 +2207,27 @@ flow_divert_handle_app_map_create(mbuf_t packet, int offset)
        size_t nodes_mem_size;
        int prefix_count = 0;
        int signing_id_count = 0;
        size_t nodes_mem_size;
        int prefix_count = 0;
        int signing_id_count = 0;

+       size_t trie_memory_size = 0;

 
 

-       lck_rw_lock_exclusive(&g_flow_divert_group_lck);
+       lck_rw_lock_exclusive(&group->lck);

 
        /* Re-set the current trie */
 
        /* Re-set the current trie */

-       if (g_signing_id_trie.memory != NULL) {
-               FREE(g_signing_id_trie.memory, M_TEMP);
+       if (group->signing_id_trie.memory != NULL) {
+               FREE(group->signing_id_trie.memory, M_TEMP);

        }
        }

-       memset(&g_signing_id_trie, 0, sizeof(g_signing_id_trie));
-       g_signing_id_trie.root = NULL_TRIE_IDX;
+       memset(&group->signing_id_trie, 0, sizeof(group->signing_id_trie));
+       group->signing_id_trie.root = NULL_TRIE_IDX;

 
        memset(&new_trie, 0, sizeof(new_trie));
 
        /* Get the number of shared prefixes in the new set of signing ID strings */
        flow_divert_packet_get_tlv(packet, offset, FLOW_DIVERT_TLV_PREFIX_COUNT, sizeof(prefix_count), &prefix_count, NULL);
 
 
        memset(&new_trie, 0, sizeof(new_trie));
 
        /* Get the number of shared prefixes in the new set of signing ID strings */
        flow_divert_packet_get_tlv(packet, offset, FLOW_DIVERT_TLV_PREFIX_COUNT, sizeof(prefix_count), &prefix_count, NULL);
 

+       if (prefix_count < 0) {
+               lck_rw_done(&group->lck);
+               return;
+       }
+

        /* Compute the number of signing IDs and the total amount of bytes needed to store them */
        for (cursor = flow_divert_packet_find_tlv(packet, offset, FLOW_DIVERT_TLV_SIGNING_ID, &error, 0);
             cursor >= 0;
        /* Compute the number of signing IDs and the total amount of bytes needed to store them */
        for (cursor = flow_divert_packet_find_tlv(packet, offset, FLOW_DIVERT_TLV_SIGNING_ID, &error, 0);
             cursor >= 0;


@@ -2205,7 +2240,7 @@ flow_divert_handle_app_map_create(mbuf_t packet, int offset)
        }
 
        if (signing_id_count == 0) {
        }
 
        if (signing_id_count == 0) {

-               lck_rw_done(&g_flow_divert_group_lck);
+               lck_rw_done(&group->lck);

                return;
        }
 
                return;
        }
 


@@ -2219,10 +2254,18 @@ flow_divert_handle_app_map_create(mbuf_t packet, int offset)
        child_maps_mem_size = (sizeof(*new_trie.child_maps) * CHILD_MAP_SIZE * new_trie.child_maps_count);
        bytes_mem_size = (sizeof(*new_trie.bytes) * new_trie.bytes_count);
 
        child_maps_mem_size = (sizeof(*new_trie.child_maps) * CHILD_MAP_SIZE * new_trie.child_maps_count);
        bytes_mem_size = (sizeof(*new_trie.bytes) * new_trie.bytes_count);
 

-       MALLOC(new_trie.memory, void *, nodes_mem_size + child_maps_mem_size + bytes_mem_size, M_TEMP, M_WAITOK);
+       trie_memory_size = nodes_mem_size + child_maps_mem_size + bytes_mem_size;
+       if (trie_memory_size > FLOW_DIVERT_MAX_TRIE_MEMORY) {
+               FDLOG(LOG_ERR, &nil_pcb, "Trie memory size (%lu) is too big (maximum is %u)", trie_memory_size, FLOW_DIVERT_MAX_TRIE_MEMORY);
+               lck_rw_done(&group->lck);
+               return;
+       }
+
+       MALLOC(new_trie.memory, void *, trie_memory_size, M_TEMP, M_WAITOK);

        if (new_trie.memory == NULL) {
                FDLOG(LOG_ERR, &nil_pcb, "Failed to allocate %lu bytes of memory for the signing ID trie",
                      nodes_mem_size + child_maps_mem_size + bytes_mem_size);
        if (new_trie.memory == NULL) {
                FDLOG(LOG_ERR, &nil_pcb, "Failed to allocate %lu bytes of memory for the signing ID trie",
                      nodes_mem_size + child_maps_mem_size + bytes_mem_size);

+               lck_rw_done(&group->lck);

                return;
        }
 
                return;
        }
 


@@ -2249,20 +2292,10 @@ flow_divert_handle_app_map_create(mbuf_t packet, int offset)
                uint32_t sid_size = 0;
                flow_divert_packet_get_tlv(packet, cursor, FLOW_DIVERT_TLV_SIGNING_ID, 0, NULL, &sid_size);
                if (new_trie.bytes_free_next + sid_size <= new_trie.bytes_count) {
                uint32_t sid_size = 0;
                flow_divert_packet_get_tlv(packet, cursor, FLOW_DIVERT_TLV_SIGNING_ID, 0, NULL, &sid_size);
                if (new_trie.bytes_free_next + sid_size <= new_trie.bytes_count) {

-                       boolean_t is_dns;

                        uint16_t new_node_idx;
                        flow_divert_packet_get_tlv(packet, cursor, FLOW_DIVERT_TLV_SIGNING_ID, sid_size, &TRIE_BYTE(&new_trie, new_trie.bytes_free_next), NULL);
                        uint16_t new_node_idx;
                        flow_divert_packet_get_tlv(packet, cursor, FLOW_DIVERT_TLV_SIGNING_ID, sid_size, &TRIE_BYTE(&new_trie, new_trie.bytes_free_next), NULL);

-                       is_dns = (sid_size == sizeof(FLOW_DIVERT_DNS_SERVICE_SIGNING_ID) - 1 && 
-                                 !memcmp(&TRIE_BYTE(&new_trie, new_trie.bytes_free_next),
-                                         FLOW_DIVERT_DNS_SERVICE_SIGNING_ID,
-                                         sid_size));

                        new_node_idx = flow_divert_trie_insert(&new_trie, new_trie.bytes_free_next, sid_size);
                        new_node_idx = flow_divert_trie_insert(&new_trie, new_trie.bytes_free_next, sid_size);

-                       if (new_node_idx != NULL_TRIE_IDX) {
-                               if (is_dns) {
-                                       FDLOG(LOG_INFO, &nil_pcb, "Setting group unit for %s to %d", FLOW_DIVERT_DNS_SERVICE_SIGNING_ID, DNS_SERVICE_GROUP_UNIT);
-                                       TRIE_NODE(&new_trie, new_node_idx).group_unit = DNS_SERVICE_GROUP_UNIT;
-                               }
-                       } else {
+                       if (new_node_idx == NULL_TRIE_IDX) {

                                insert_error = EINVAL;
                                break;
                        }
                                insert_error = EINVAL;
                                break;
                        }


@@ -2274,72 +2307,12 @@ flow_divert_handle_app_map_create(mbuf_t packet, int offset)
        }
 
        if (!insert_error) {
        }
 
        if (!insert_error) {

-               g_signing_id_trie = new_trie;
+               group->signing_id_trie = new_trie;

        } else {
                FREE(new_trie.memory, M_TEMP);
        }
 
        } else {
                FREE(new_trie.memory, M_TEMP);
        }
 

-       lck_rw_done(&g_flow_divert_group_lck);
-}
-
-static void
-flow_divert_handle_app_map_update(struct flow_divert_group *group, mbuf_t packet, int offset)
-{
-       int error = 0;
-       int cursor;
-       size_t max_size = 0;
-       uint8_t *signing_id;
-       uint32_t ctl_unit;
-
-       lck_rw_lock_shared(&group->lck);
-       ctl_unit = group->ctl_unit;

        lck_rw_done(&group->lck);
        lck_rw_done(&group->lck);

-
-       for (cursor = flow_divert_packet_find_tlv(packet, offset, FLOW_DIVERT_TLV_SIGNING_ID, &error, 0);
-            cursor >= 0;
-            cursor = flow_divert_packet_find_tlv(packet, cursor, FLOW_DIVERT_TLV_SIGNING_ID, &error, 1))
-       {
-               uint32_t sid_size = 0;
-               flow_divert_packet_get_tlv(packet, cursor, FLOW_DIVERT_TLV_SIGNING_ID, 0, NULL, &sid_size);
-               if (sid_size > max_size) {
-                       max_size = sid_size;
-               }
-       }
-
-       MALLOC(signing_id, uint8_t *, max_size + 1, M_TEMP, M_WAITOK);
-       if (signing_id == NULL) {
-               FDLOG(LOG_ERR, &nil_pcb, "Failed to allocate a string to hold the signing ID (size %lu)", max_size);
-               return;
-       }
-
-       for (cursor = flow_divert_packet_find_tlv(packet, offset, FLOW_DIVERT_TLV_SIGNING_ID, &error, 0);
-            cursor >= 0;
-            cursor = flow_divert_packet_find_tlv(packet, cursor, FLOW_DIVERT_TLV_SIGNING_ID, &error, 1))
-       {
-               uint32_t signing_id_len = 0;
-               uint16_t node;
-
-               flow_divert_packet_get_tlv(packet,
-                               cursor, FLOW_DIVERT_TLV_SIGNING_ID, max_size, signing_id, &signing_id_len);
-
-               signing_id[signing_id_len] = '\0';
-
-               lck_rw_lock_exclusive(&g_flow_divert_group_lck);
-
-               node = flow_divert_trie_search(&g_signing_id_trie, signing_id);
-               if (node != NULL_TRIE_IDX) {
-                       if (TRIE_NODE(&g_signing_id_trie, node).group_unit != DNS_SERVICE_GROUP_UNIT) {
-                               FDLOG(LOG_INFO, &nil_pcb, "Setting %s to ctl unit %u", signing_id, group->ctl_unit);
-                               TRIE_NODE(&g_signing_id_trie, node).group_unit = ctl_unit;
-                       }
-               } else {
-                       FDLOG(LOG_ERR, &nil_pcb, "Failed to find signing ID %s", signing_id);
-               }
-
-               lck_rw_done(&g_flow_divert_group_lck);
-       }
-
-       FREE(signing_id, M_TEMP);

 }
 
 static int
 }
 
 static int


@@ -2356,7 +2329,7 @@ flow_divert_input(mbuf_t packet, struct flow_divert_group *group)
        }
 
        if (mbuf_pkthdr_len(packet) > FD_CTL_RCVBUFF_SIZE) {
        }
 
        if (mbuf_pkthdr_len(packet) > FD_CTL_RCVBUFF_SIZE) {

-               FDLOG(LOG_ERR, &nil_pcb, "got a bad packet, length (%lu) > %lu", mbuf_pkthdr_len(packet), FD_CTL_RCVBUFF_SIZE);
+               FDLOG(LOG_ERR, &nil_pcb, "got a bad packet, length (%lu) > %d", mbuf_pkthdr_len(packet), FD_CTL_RCVBUFF_SIZE);

                error = EINVAL;
                goto done;
        }
                error = EINVAL;
                goto done;
        }


@@ -2376,10 +2349,7 @@ flow_divert_input(mbuf_t packet, struct flow_divert_group *group)
                                flow_divert_handle_group_init(group, packet, sizeof(hdr));
                                break;
                        case FLOW_DIVERT_PKT_APP_MAP_CREATE:
                                flow_divert_handle_group_init(group, packet, sizeof(hdr));
                                break;
                        case FLOW_DIVERT_PKT_APP_MAP_CREATE:

-                               flow_divert_handle_app_map_create(packet, sizeof(hdr));
-                               break;
-                       case FLOW_DIVERT_PKT_APP_MAP_UPDATE:
-                               flow_divert_handle_app_map_update(group, packet, sizeof(hdr));
+                               flow_divert_handle_app_map_create(group, packet, sizeof(hdr));

                                break;
                        default:
                                FDLOG(LOG_WARNING, &nil_pcb, "got an unknown message type: %d", hdr.packet_type);
                                break;
                        default:
                                FDLOG(LOG_WARNING, &nil_pcb, "got an unknown message type: %d", hdr.packet_type);


@@ -2452,6 +2422,7 @@ flow_divert_close_all(struct flow_divert_group *group)
                        flow_divert_pcb_remove(fd_cb);
                        flow_divert_update_closed_state(fd_cb, SHUT_RDWR, TRUE);
                        fd_cb->so->so_error = ECONNABORTED;
                        flow_divert_pcb_remove(fd_cb);
                        flow_divert_update_closed_state(fd_cb, SHUT_RDWR, TRUE);
                        fd_cb->so->so_error = ECONNABORTED;

+                       flow_divert_disconnect_socket(fd_cb->so);

                        socket_unlock(fd_cb->so, 0);
                }
                FDUNLOCK(fd_cb);
                        socket_unlock(fd_cb->so, 0);
                }
                FDUNLOCK(fd_cb);


@@ -2870,13 +2841,9 @@ done:
 }
 
 static int
 }
 
 static int

-flow_divert_connectx_out_common(struct socket *so, int af,
-    struct sockaddr_list **src_sl, struct sockaddr_list **dst_sl,
-    struct proc *p, uint32_t ifscope __unused, sae_associd_t aid __unused,
-    sae_connid_t *pcid, uint32_t flags __unused, void *arg __unused,
-    uint32_t arglen __unused, struct uio *auio, user_ssize_t *bytes_written)
+flow_divert_connectx_out_common(struct socket *so, struct sockaddr *dst,
+    struct proc *p, sae_connid_t *pcid, struct uio *auio, user_ssize_t *bytes_written)

 {
 {

-       struct sockaddr_entry *src_se = NULL, *dst_se = NULL;

        struct inpcb *inp = sotoinpcb(so);
        int error;
 
        struct inpcb *inp = sotoinpcb(so);
        int error;
 


@@ -2884,20 +2851,9 @@ flow_divert_connectx_out_common(struct socket *so, int af,
                return (EINVAL);
        }
 
                return (EINVAL);
        }
 

-       VERIFY(dst_sl != NULL);
-
-       /* select source (if specified) and destination addresses */
-       error = in_selectaddrs(af, src_sl, &src_se, dst_sl, &dst_se);
-       if (error != 0) {
-               return (error);
-       }
+       VERIFY(dst != NULL);

 
 

-       VERIFY(*dst_sl != NULL && dst_se != NULL);
-       VERIFY(src_se == NULL || *src_sl != NULL);
-       VERIFY(dst_se->se_addr->sa_family == af);
-       VERIFY(src_se == NULL || src_se->se_addr->sa_family == af);
-
-       error = flow_divert_connect_out(so, dst_se->se_addr, p);
+       error = flow_divert_connect_out(so, dst, p);

 
        if (error != 0) {
                return error;
 
        if (error != 0) {
                return error;


@@ -2940,26 +2896,22 @@ flow_divert_connectx_out_common(struct socket *so, int af,
 }
 
 static int
 }
 
 static int

-flow_divert_connectx_out(struct socket *so, struct sockaddr_list **src_sl,
-    struct sockaddr_list **dst_sl, struct proc *p, uint32_t ifscope,
-    sae_associd_t aid, sae_connid_t *pcid, uint32_t flags, void *arg,
-    uint32_t arglen, struct uio *uio, user_ssize_t *bytes_written)
+flow_divert_connectx_out(struct socket *so, struct sockaddr *src __unused,
+    struct sockaddr *dst, struct proc *p, uint32_t ifscope __unused,
+    sae_associd_t aid __unused, sae_connid_t *pcid, uint32_t flags __unused, void *arg __unused,
+    uint32_t arglen __unused, struct uio *uio, user_ssize_t *bytes_written)

 {
 {

-#pragma unused(uio, bytes_written)
-       return (flow_divert_connectx_out_common(so, AF_INET, src_sl, dst_sl,
-           p, ifscope, aid, pcid, flags, arg, arglen, uio, bytes_written));
+       return (flow_divert_connectx_out_common(so, dst, p, pcid, uio, bytes_written));

 }
 
 #if INET6
 static int
 }
 
 #if INET6
 static int

-flow_divert_connectx6_out(struct socket *so, struct sockaddr_list **src_sl,
-    struct sockaddr_list **dst_sl, struct proc *p, uint32_t ifscope,
-    sae_associd_t aid, sae_connid_t *pcid, uint32_t flags, void *arg,
-    uint32_t arglen, struct uio *uio, user_ssize_t *bytes_written)
+flow_divert_connectx6_out(struct socket *so, struct sockaddr *src __unused,
+    struct sockaddr *dst, struct proc *p, uint32_t ifscope __unused,
+    sae_associd_t aid __unused, sae_connid_t *pcid, uint32_t flags __unused, void *arg __unused,
+    uint32_t arglen __unused, struct uio *uio, user_ssize_t *bytes_written)

 {
 {

-#pragma unused(uio, bytes_written)
-       return (flow_divert_connectx_out_common(so, AF_INET6, src_sl, dst_sl,
-           p, ifscope, aid, pcid, flags, arg, arglen, uio, bytes_written));
+       return (flow_divert_connectx_out_common(so, dst, p, pcid, uio, bytes_written));

 }
 #endif /* INET6 */
 
 }
 #endif /* INET6 */
 


@@ -3199,7 +3151,7 @@ flow_divert_preconnect(struct socket *so)
                fd_cb->flags |= FLOW_DIVERT_CONNECT_STARTED;
        }
 
                fd_cb->flags |= FLOW_DIVERT_CONNECT_STARTED;
        }
 

-       so->so_flags1 &= ~SOF1_PRECONNECT_DATA;
+       soclearfastopen(so);

 
        return error;
 }
 
        return error;
 }


@@ -3432,17 +3384,22 @@ flow_divert_token_set(struct socket *so, struct sockopt *sopt)
 
        error = soopt_getm(sopt, &token);
        if (error) {
 
        error = soopt_getm(sopt, &token);
        if (error) {

+               token = NULL;

                goto done;
        }
 
        error = soopt_mcopyin(sopt, token);
        if (error) {
                goto done;
        }
 
        error = soopt_mcopyin(sopt, token);
        if (error) {

+               token = NULL;

                goto done;
        }
 
        error = flow_divert_packet_get_tlv(token, 0, FLOW_DIVERT_TLV_KEY_UNIT, sizeof(key_unit), (void *)&key_unit, NULL);
        if (!error) {
                key_unit = ntohl(key_unit);
                goto done;
        }
 
        error = flow_divert_packet_get_tlv(token, 0, FLOW_DIVERT_TLV_KEY_UNIT, sizeof(key_unit), (void *)&key_unit, NULL);
        if (!error) {
                key_unit = ntohl(key_unit);

+               if (key_unit >= GROUP_COUNT_MAX) {
+                       key_unit = 0;
+               }

        } else if (error != ENOENT) {
                FDLOG(LOG_ERR, &nil_pcb, "Failed to get the key unit from the token: %d", error);
                goto done;
        } else if (error != ENOENT) {
                FDLOG(LOG_ERR, &nil_pcb, "Failed to get the key unit from the token: %d", error);
                goto done;


@@ -3595,6 +3552,12 @@ flow_divert_token_get(struct socket *so, struct sockopt *sopt)
                goto done;
        }
 
                goto done;
        }
 

+       if (sopt->sopt_val == USER_ADDR_NULL) {
+               /* If the caller passed NULL to getsockopt, just set the size of the token and return */
+               sopt->sopt_valsize = mbuf_pkthdr_len(token);
+               goto done;
+       }
+

        error = soopt_mcopyout(sopt, token);
        if (error) {
                token = NULL;   /* For some reason, soopt_mcopyout() frees the mbuf if it fails */
        error = soopt_mcopyout(sopt, token);
        if (error) {
                token = NULL;   /* For some reason, soopt_mcopyout() frees the mbuf if it fails */


@@ -3634,6 +3597,7 @@ flow_divert_kctl_connect(kern_ctl_ref kctlref __unused, struct sockaddr_ctl *sac
        RB_INIT(&new_group->pcb_tree);
        new_group->ctl_unit = sac->sc_unit;
        MBUFQ_INIT(&new_group->send_queue);
        RB_INIT(&new_group->pcb_tree);
        new_group->ctl_unit = sac->sc_unit;
        MBUFQ_INIT(&new_group->send_queue);

+       new_group->signing_id_trie.root = NULL_TRIE_IDX;

 
        lck_rw_lock_exclusive(&g_flow_divert_group_lck);
 
 
        lck_rw_lock_exclusive(&g_flow_divert_group_lck);
 


@@ -3670,7 +3634,6 @@ flow_divert_kctl_disconnect(kern_ctl_ref kctlref __unused, uint32_t unit, void *
 {
        struct flow_divert_group        *group  = NULL;
        errno_t                                         error   = 0;
 {
        struct flow_divert_group        *group  = NULL;
        errno_t                                         error   = 0;

-       uint16_t                                        node    = 0;

 
        if (unit >= GROUP_COUNT_MAX) {
                return EINVAL;
 
        if (unit >= GROUP_COUNT_MAX) {
                return EINVAL;


@@ -3699,6 +3662,14 @@ flow_divert_kctl_disconnect(kern_ctl_ref kctlref __unused, uint32_t unit, void *
                        group->token_key = NULL;
                        group->token_key_size = 0;
                }
                        group->token_key = NULL;
                        group->token_key_size = 0;
                }

+
+               /* Re-set the current trie */
+               if (group->signing_id_trie.memory != NULL) {
+                       FREE(group->signing_id_trie.memory, M_TEMP);
+               }
+               memset(&group->signing_id_trie, 0, sizeof(group->signing_id_trie));
+               group->signing_id_trie.root = NULL_TRIE_IDX;
+

                FREE_ZONE(group, sizeof(*group), M_FLOW_DIVERT_GROUP);
                g_flow_divert_groups[unit] = NULL;
                g_active_group_count--;
                FREE_ZONE(group, sizeof(*group), M_FLOW_DIVERT_GROUP);
                g_flow_divert_groups[unit] = NULL;
                g_active_group_count--;


@@ -3711,13 +3682,6 @@ flow_divert_kctl_disconnect(kern_ctl_ref kctlref __unused, uint32_t unit, void *
                g_flow_divert_groups = NULL;
        }
 
                g_flow_divert_groups = NULL;
        }
 

-       /* Remove all signing IDs that point to this unit */
-       for (node = 0; node < g_signing_id_trie.nodes_count; node++) {
-               if (TRIE_NODE(&g_signing_id_trie, node).group_unit == unit) {
-                       TRIE_NODE(&g_signing_id_trie, node).group_unit = 0;
-               }
-       }
-

        lck_rw_done(&g_flow_divert_group_lck);
 
        return error;
        lck_rw_done(&g_flow_divert_group_lck);
 
        return error;


@@ -3978,9 +3942,6 @@ flow_divert_init(void)
 
        lck_rw_init(&g_flow_divert_group_lck, flow_divert_mtx_grp, flow_divert_mtx_attr);
 
 
        lck_rw_init(&g_flow_divert_group_lck, flow_divert_mtx_grp, flow_divert_mtx_attr);
 

-       memset(&g_signing_id_trie, 0, sizeof(g_signing_id_trie));
-       g_signing_id_trie.root = NULL_TRIE_IDX;
-

 done:
        if (g_init_result != 0) {
                if (flow_divert_mtx_attr != NULL) {
 done:
        if (g_init_result != 0) {
                if (flow_divert_mtx_attr != NULL) {