]> git.saurik.com Git - apple/xnu.git/blobdiff - bsd/netinet/flow_divert.c
xnu-4903.270.47.tar.gz
[apple/xnu.git] / bsd / netinet / flow_divert.c
index 8fbe88a80aa84f9b85346889ceca36e7541c7db9..13b9cad5d7ff1637fb7d40f064de93a33f40cc6f 100644 (file)
@@ -1,8 +1,8 @@
 /*
- * Copyright (c) 2012-2013 Apple Inc. All rights reserved.
+ * Copyright (c) 2012-2017 Apple Inc. All rights reserved.
  *
  * @APPLE_OSREFERENCE_LICENSE_HEADER_START@
- * 
+ *
  * This file contains Original Code and/or Modifications of Original Code
  * as defined in and that are subject to the Apple Public Source License
  * Version 2.0 (the 'License'). You may not use this file except in
  * unlawful or unlicensed copies of an Apple operating system, or to
  * circumvent, violate, or enable the circumvention or violation of, any
  * terms of an Apple operating system software license agreement.
- * 
+ *
  * Please obtain a copy of the License at
  * http://www.opensource.apple.com/apsl/ and read it before using this file.
- * 
+ *
  * The Original Code and all software distributed under the License are
  * distributed on an 'AS IS' basis, WITHOUT WARRANTY OF ANY KIND, EITHER
  * EXPRESS OR IMPLIED, AND APPLE HEREBY DISCLAIMS ALL SUCH WARRANTIES,
@@ -22,7 +22,7 @@
  * FITNESS FOR A PARTICULAR PURPOSE, QUIET ENJOYMENT OR NON-INFRINGEMENT.
  * Please see the License for the specific language governing rights and
  * limitations under the License.
- * 
+ *
  * @APPLE_OSREFERENCE_LICENSE_HEADER_END@
  */
 
 #include <netinet/flow_divert.h>
 #include <netinet/flow_divert_proto.h>
 #if INET6
+#include <netinet6/in6_pcb.h>
 #include <netinet6/ip6protosw.h>
-#endif /* INET6 */
+#endif  /* INET6 */
 #include <dev/random/randomdev.h>
 #include <libkern/crypto/sha1.h>
 #include <libkern/crypto/crypto_internal.h>
-
-#define FLOW_DIVERT_CONNECT_STARTED            0x00000001
-#define FLOW_DIVERT_READ_CLOSED                        0x00000002
-#define FLOW_DIVERT_WRITE_CLOSED               0x00000004
-#define FLOW_DIVERT_TUNNEL_RD_CLOSED   0x00000008
-#define FLOW_DIVERT_TUNNEL_WR_CLOSED   0x00000010
-#define FLOW_DIVERT_TRANSFERRED                        0x00000020
-
-#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 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 FDRETAIN(pcb)                  if ((pcb) != NULL) OSIncrementAtomic(&(pcb)->ref_count)
-#define FDRELEASE(pcb)                                                                                                         \
-       do {                                                                                                                                    \
-               if ((pcb) != NULL && 1 == OSDecrementAtomic(&(pcb)->ref_count)) {       \
-                       flow_divert_pcb_destroy(pcb);                                                                   \
-               }                                                                                                                                       \
+#include <os/log.h>
+
+#define FLOW_DIVERT_CONNECT_STARTED             0x00000001
+#define FLOW_DIVERT_READ_CLOSED                 0x00000002
+#define FLOW_DIVERT_WRITE_CLOSED                0x00000004
+#define FLOW_DIVERT_TUNNEL_RD_CLOSED    0x00000008
+#define FLOW_DIVERT_TUNNEL_WR_CLOSED    0x00000010
+#define FLOW_DIVERT_TRANSFERRED                 0x00000020
+#define FLOW_DIVERT_HAS_HMAC            0x00000040
+
+#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) \
+       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)                                                                                                          \
+       do {                                                                                                                                    \
+               if ((pcb) != NULL && 1 == OSDecrementAtomic(&(pcb)->ref_count)) {       \
+                       flow_divert_pcb_destroy(pcb);                                                                   \
+               }                                                                                                                                       \
        } while (0)
 
-#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_RCVBUFF_SIZE                            (128 * 1024)
+#define FDLOCK(pcb)                                             lck_mtx_lock(&(pcb)->mtx)
+#define FDUNLOCK(pcb)                                   lck_mtx_unlock(&(pcb)->mtx)
 
-#define GROUP_BIT_CTL_ENQUEUE_BLOCKED  0
+#define FD_CTL_SENDBUFF_SIZE                    (128 * 1024)
+#define FD_CTL_RCVBUFF_SIZE                             (128 * 1024)
 
-#define GROUP_COUNT_MAX                                        32
-#define FLOW_DIVERT_MAX_NAME_SIZE              4096
-#define FLOW_DIVERT_MAX_KEY_SIZE               1024
+#define GROUP_BIT_CTL_ENQUEUE_BLOCKED   0
 
-#define DNS_SERVICE_GROUP_UNIT                 (GROUP_COUNT_MAX + 1)
+#define GROUP_COUNT_MAX                                 32
+#define FLOW_DIVERT_MAX_NAME_SIZE               4096
+#define FLOW_DIVERT_MAX_KEY_SIZE                1024
+#define FLOW_DIVERT_MAX_TRIE_MEMORY             (1024 * 1024)
 
-struct flow_divert_trie_node
-{
+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 NULL_TRIE_IDX                  0xffff
-#define TRIE_NODE(t, i)                        ((t)->nodes[(i)])
-#define TRIE_CHILD(t, i, b)            (((t)->child_maps + (CHILD_MAP_SIZE * TRIE_NODE(t, i).child_map))[(b)])
-#define TRIE_BYTE(t, i)                        ((t)->bytes[(i)])
+#define CHILD_MAP_SIZE                  256
+#define NULL_TRIE_IDX                   0xffff
+#define TRIE_NODE(t, i)                 ((t)->nodes[(i)])
+#define TRIE_CHILD(t, i, b)             (((t)->child_maps + (CHILD_MAP_SIZE * TRIE_NODE(t, i).child_map))[(b)])
+#define TRIE_BYTE(t, i)                 ((t)->bytes[(i)])
 
-static struct flow_divert_pcb          nil_pcb;
+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;
-static struct flow_divert_trie         g_signing_id_trie;
+static struct flow_divert_group         **g_flow_divert_groups                  = NULL;
+static uint32_t                                         g_active_group_count                    = 0;
 
-static lck_grp_attr_t                          *flow_divert_grp_attr                   = NULL;
-static lck_attr_t                                      *flow_divert_mtx_attr                   = NULL;
-static lck_grp_t                                       *flow_divert_mtx_grp                    = NULL;
-static errno_t                                         g_init_result                                   = 0;
+static  lck_grp_attr_t                          *flow_divert_grp_attr                   = NULL;
+static  lck_attr_t                                      *flow_divert_mtx_attr                   = NULL;
+static  lck_grp_t                                       *flow_divert_mtx_grp                    = NULL;
+static  errno_t                                         g_init_result                                   = 0;
 
-static kern_ctl_ref                            g_flow_divert_kctl_ref                  = NULL;
+static  kern_ctl_ref                            g_flow_divert_kctl_ref                  = NULL;
 
-static struct protosw                          g_flow_divert_in_protosw;
-static struct pr_usrreqs                       g_flow_divert_in_usrreqs;
+static struct protosw                           g_flow_divert_in_protosw;
+static struct pr_usrreqs                        g_flow_divert_in_usrreqs;
+static struct protosw                           g_flow_divert_in_udp_protosw;
+static struct pr_usrreqs                        g_flow_divert_in_udp_usrreqs;
 #if INET6
-static struct ip6protosw                       g_flow_divert_in6_protosw;
-static struct pr_usrreqs                       g_flow_divert_in6_usrreqs;
-#endif /* INET6 */
+static struct ip6protosw                        g_flow_divert_in6_protosw;
+static struct pr_usrreqs                        g_flow_divert_in6_usrreqs;
+static struct ip6protosw                        g_flow_divert_in6_udp_protosw;
+static struct pr_usrreqs                        g_flow_divert_in6_udp_usrreqs;
+#endif  /* INET6 */
+
+static struct protosw                           *g_tcp_protosw                                  = NULL;
+static struct ip6protosw                        *g_tcp6_protosw                                 = NULL;
+static struct protosw                           *g_udp_protosw                                  = NULL;
+static struct ip6protosw                        *g_udp6_protosw                                 = NULL;
+
+static errno_t
+flow_divert_dup_addr(sa_family_t family, struct sockaddr *addr, struct sockaddr **dup);
+
+static errno_t
+flow_divert_inp_to_sockaddr(const struct inpcb *inp, struct sockaddr **local_socket);
+
+static boolean_t
+flow_divert_is_sockaddr_valid(struct sockaddr *addr);
+
+static int
+flow_divert_append_target_endpoint_tlv(mbuf_t connect_packet, struct sockaddr *toaddr);
+
+struct sockaddr *
+flow_divert_get_buffered_target_address(mbuf_t buffer);
+
+static boolean_t
+flow_divert_has_pcb_local_address(const struct inpcb *inp);
+
+static void
+flow_divert_disconnect_socket(struct socket *so);
 
-static struct protosw                          *g_tcp_protosw                                  = NULL;
-static struct ip6protosw                       *g_tcp6_protosw                                 = NULL;
+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)
@@ -170,32 +186,30 @@ static const char *
 flow_divert_packet_type2str(uint8_t packet_type)
 {
        switch (packet_type) {
-               case FLOW_DIVERT_PKT_CONNECT:
-                       return "connect";
-               case FLOW_DIVERT_PKT_CONNECT_RESULT:
-                       return "connect result";
-               case FLOW_DIVERT_PKT_DATA:
-                       return "data";
-               case FLOW_DIVERT_PKT_CLOSE:
-                       return "close";
-               case FLOW_DIVERT_PKT_READ_NOTIFY:
-                       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:
-                       return "unknown";
+       case FLOW_DIVERT_PKT_CONNECT:
+               return "connect";
+       case FLOW_DIVERT_PKT_CONNECT_RESULT:
+               return "connect result";
+       case FLOW_DIVERT_PKT_DATA:
+               return "data";
+       case FLOW_DIVERT_PKT_CLOSE:
+               return "close";
+       case FLOW_DIVERT_PKT_READ_NOTIFY:
+               return "read notification";
+       case FLOW_DIVERT_PKT_PROPERTIES_UPDATE:
+               return "properties update";
+       case FLOW_DIVERT_PKT_APP_MAP_CREATE:
+               return "app map create";
+       default:
+               return "unknown";
        }
 }
 
 static struct flow_divert_pcb *
 flow_divert_pcb_lookup(uint32_t hash, struct flow_divert_group *group)
 {
-       struct flow_divert_pcb  key_item;
-       struct flow_divert_pcb  *fd_cb          = NULL;
+       struct flow_divert_pcb  key_item;
+       struct flow_divert_pcb  *fd_cb          = NULL;
 
        key_item.hash = hash;
 
@@ -210,13 +224,12 @@ flow_divert_pcb_lookup(uint32_t hash, struct flow_divert_group *group)
 static errno_t
 flow_divert_pcb_insert(struct flow_divert_pcb *fd_cb, uint32_t ctl_unit)
 {
-       int                                                     error                                           = 0;
-       struct                                          flow_divert_pcb *exist          = NULL;
-       struct flow_divert_group        *group;
-       static uint32_t                         g_nextkey                                       = 1;
-       static uint32_t                         g_hash_seed                                     = 0;
-       errno_t                                         result                                          = 0;
-       int                                                     try_count                                       = 0;
+       errno_t                                                 error                                           = 0;
+       struct                                          flow_divert_pcb *exist          = NULL;
+       struct flow_divert_group        *group;
+       static uint32_t                         g_nextkey                                       = 1;
+       static uint32_t                         g_hash_seed                                     = 0;
+       int                                                     try_count                                       = 0;
 
        if (ctl_unit == 0 || ctl_unit >= GROUP_COUNT_MAX) {
                return EINVAL;
@@ -241,8 +254,8 @@ flow_divert_pcb_insert(struct flow_divert_pcb *fd_cb, uint32_t ctl_unit)
        socket_lock(fd_cb->so, 0);
 
        do {
-               uint32_t        key[2];
-               uint32_t        idx;
+               uint32_t        key[2];
+               uint32_t        idx;
 
                key[0] = g_nextkey++;
                key[1] = RandomULong();
@@ -274,10 +287,10 @@ flow_divert_pcb_insert(struct flow_divert_pcb *fd_cb, uint32_t ctl_unit)
 
        if (exist == NULL) {
                fd_cb->group = group;
-               FDRETAIN(fd_cb);                /* The group now has a reference */
+               FDRETAIN(fd_cb);                /* The group now has a reference */
        } else {
                fd_cb->hash = 0;
-               result = EEXIST;
+               error = EEXIST;
        }
 
        socket_unlock(fd_cb->so, 0);
@@ -286,13 +299,13 @@ done:
        lck_rw_done(&g_flow_divert_group_lck);
        socket_lock(fd_cb->so, 0);
 
-       return result;
+       return error;
 }
 
 static struct flow_divert_pcb *
 flow_divert_pcb_create(socket_t so)
 {
-       struct flow_divert_pcb  *new_pcb        = NULL;
+       struct flow_divert_pcb  *new_pcb        = NULL;
 
        MALLOC_ZONE(new_pcb, struct flow_divert_pcb *, sizeof(*new_pcb), M_FLOW_DIVERT_PCB, M_WAITOK);
        if (new_pcb == NULL) {
@@ -306,7 +319,7 @@ flow_divert_pcb_create(socket_t so)
        new_pcb->so = so;
        new_pcb->log_level = nil_pcb.log_level;
 
-       FDRETAIN(new_pcb);      /* Represents the socket's reference */
+       FDRETAIN(new_pcb);      /* Represents the socket's reference */
 
        return new_pcb;
 }
@@ -315,7 +328,7 @@ static void
 flow_divert_pcb_destroy(struct flow_divert_pcb *fd_cb)
 {
        FDLOG(LOG_INFO, fd_cb, "Destroying, app tx %u, app rx %u, tunnel tx %u, tunnel rx %u",
-                       fd_cb->bytes_written_by_app, fd_cb->bytes_read_by_app, fd_cb->bytes_sent, fd_cb->bytes_received);
+           fd_cb->bytes_written_by_app, fd_cb->bytes_read_by_app, fd_cb->bytes_sent, fd_cb->bytes_received);
 
        if (fd_cb->local_address != NULL) {
                FREE(fd_cb->local_address, M_SONAME);
@@ -326,6 +339,12 @@ flow_divert_pcb_destroy(struct flow_divert_pcb *fd_cb)
        if (fd_cb->connect_token != NULL) {
                mbuf_freem(fd_cb->connect_token);
        }
+       if (fd_cb->connect_packet != NULL) {
+               mbuf_freem(fd_cb->connect_packet);
+       }
+       if (fd_cb->app_data != NULL) {
+               FREE(fd_cb->app_data, M_TEMP);
+       }
        FREE_ZONE(fd_cb, sizeof(*fd_cb), M_FLOW_DIVERT_PCB);
 }
 
@@ -338,7 +357,7 @@ flow_divert_pcb_remove(struct flow_divert_pcb *fd_cb)
                FDLOG(LOG_INFO, fd_cb, "Removing from group %d, ref count = %d", group->ctl_unit, fd_cb->ref_count);
                RB_REMOVE(fd_pcb_tree, &group->pcb_tree, fd_cb);
                fd_cb->group = NULL;
-               FDRELEASE(fd_cb);                               /* Release the group's reference */
+               FDRELEASE(fd_cb);                               /* Release the group's reference */
                lck_rw_done(&group->lck);
        }
 }
@@ -346,8 +365,8 @@ flow_divert_pcb_remove(struct flow_divert_pcb *fd_cb)
 static int
 flow_divert_packet_init(struct flow_divert_pcb *fd_cb, uint8_t packet_type, mbuf_t *packet)
 {
-       struct flow_divert_packet_header        hdr;
-       int                                     error           = 0;
+       struct flow_divert_packet_header        hdr;
+       int                                     error           = 0;
 
        error = mbuf_gethdr(MBUF_DONTWAIT, MBUF_TYPE_HEADER, packet);
        if (error) {
@@ -371,10 +390,10 @@ flow_divert_packet_init(struct flow_divert_pcb *fd_cb, uint8_t packet_type, mbuf
 }
 
 static int
-flow_divert_packet_append_tlv(mbuf_t packet, uint8_t type, size_t length, const void *value)
+flow_divert_packet_append_tlv(mbuf_t packet, uint8_t type, uint32_t length, const void *value)
 {
-       size_t  net_length      = htonl(length);
-       int             error           = 0;
+       uint32_t        net_length      = htonl(length);
+       int                     error           = 0;
 
        error = mbuf_copyback(packet, mbuf_pkthdr_len(packet), sizeof(type), &type, MBUF_DONTWAIT);
        if (error) {
@@ -384,7 +403,7 @@ flow_divert_packet_append_tlv(mbuf_t packet, uint8_t type, size_t length, const
 
        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;
        }
 
@@ -400,10 +419,10 @@ flow_divert_packet_append_tlv(mbuf_t packet, uint8_t type, size_t length, const
 static int
 flow_divert_packet_find_tlv(mbuf_t packet, int offset, uint8_t type, int *err, int next)
 {
-       size_t  cursor                  = offset;
-       int             error                   = 0;
-       size_t  curr_length;
-       uint8_t curr_type;
+       size_t          cursor                  = offset;
+       int                     error                   = 0;
+       uint32_t        curr_length;
+       uint8_t         curr_type;
 
        *err = 0;
 
@@ -435,11 +454,11 @@ flow_divert_packet_find_tlv(mbuf_t packet, int offset, uint8_t type, int *err, i
 }
 
 static int
-flow_divert_packet_get_tlv(mbuf_t packet, int offset, uint8_t type, size_t buff_len, void *buff, size_t *val_size)
+flow_divert_packet_get_tlv(mbuf_t packet, int offset, uint8_t type, size_t buff_len, void *buff, uint32_t *val_size)
 {
-       int             error           = 0;
-       size_t  length;
-       int             tlv_offset;
+       int                     error           = 0;
+       uint32_t        length;
+       int                     tlv_offset;
 
        tlv_offset = flow_divert_packet_find_tlv(packet, offset, type, &error, 0);
        if (tlv_offset < 0) {
@@ -471,7 +490,7 @@ flow_divert_packet_get_tlv(mbuf_t packet, int offset, uint8_t type, size_t buff_
 static int
 flow_divert_packet_compute_hmac(mbuf_t packet, struct flow_divert_group *group, uint8_t *hmac)
 {
-       mbuf_t  curr_mbuf       = packet;
+       mbuf_t  curr_mbuf       = packet;
 
        if (g_crypto_funcs == NULL || group->token_key == NULL) {
                return ENOPROTOOPT;
@@ -493,12 +512,12 @@ flow_divert_packet_compute_hmac(mbuf_t packet, struct flow_divert_group *group,
 static int
 flow_divert_packet_verify_hmac(mbuf_t packet, uint32_t ctl_unit)
 {
-       int                                                     error = 0;
-       struct flow_divert_group        *group = NULL;
-       int                                                     hmac_offset;
-       uint8_t                                         packet_hmac[SHA_DIGEST_LENGTH];
-       uint8_t                                         computed_hmac[SHA_DIGEST_LENGTH];
-       mbuf_t                                          tail;
+       int                                                     error = 0;
+       struct flow_divert_group        *group = NULL;
+       int                                                     hmac_offset;
+       uint8_t                                         packet_hmac[SHA_DIGEST_LENGTH];
+       uint8_t                                         computed_hmac[SHA_DIGEST_LENGTH];
+       mbuf_t                                          tail;
 
        lck_rw_lock_shared(&g_flow_divert_group_lck);
 
@@ -560,7 +579,8 @@ flow_divert_add_data_statistics(struct flow_divert_pcb *fd_cb, int data_len, Boo
        struct ifnet *ifp = NULL;
        Boolean cell = FALSE;
        Boolean wifi = FALSE;
-       
+       Boolean wired = FALSE;
+
        inp = sotoinpcb(fd_cb->so);
        if (inp == NULL) {
                return;
@@ -570,33 +590,44 @@ flow_divert_add_data_statistics(struct flow_divert_pcb *fd_cb, int data_len, Boo
        if (ifp != NULL) {
                cell = IFNET_IS_CELLULAR(ifp);
                wifi = (!cell && IFNET_IS_WIFI(ifp));
+               wired = (!wifi && IFNET_IS_WIRED(ifp));
        }
-       
+
        if (send) {
-               INP_ADD_STAT(inp, cell, wifi, txpackets, 1);
-               INP_ADD_STAT(inp, cell, wifi, txbytes, data_len);
+               INP_ADD_STAT(inp, cell, wifi, wired, txpackets, 1);
+               INP_ADD_STAT(inp, cell, wifi, wired, txbytes, data_len);
        } else {
-               INP_ADD_STAT(inp, cell, wifi, rxpackets, 1);
-               INP_ADD_STAT(inp, cell, wifi, 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
 flow_divert_check_no_cellular(struct flow_divert_pcb *fd_cb)
 {
        struct inpcb *inp = NULL;
-       struct ifnet *ifp = NULL;
 
        inp = sotoinpcb(fd_cb->so);
-       if ((inp != NULL) && (inp->inp_flags & INP_NO_IFT_CELLULAR)) {
-               ifp = inp->inp_last_outifp;
-               if (ifp != NULL) {
-                       if (IFNET_IS_CELLULAR(ifp)) {
-                               return EHOSTUNREACH;
-                       }
-               }
+       if (inp && INP_NO_CELLULAR(inp) && inp->inp_last_outifp &&
+           IFNET_IS_CELLULAR(inp->inp_last_outifp)) {
+               return EHOSTUNREACH;
+       }
+
+       return 0;
+}
+
+static errno_t
+flow_divert_check_no_expensive(struct flow_divert_pcb *fd_cb)
+{
+       struct inpcb *inp = NULL;
+
+       inp = sotoinpcb(fd_cb->so);
+       if (inp && INP_NO_EXPENSIVE(inp) && inp->inp_last_outifp &&
+           IFNET_IS_EXPENSIVE(inp->inp_last_outifp)) {
+               return EHOSTUNREACH;
        }
-       
+
        return 0;
 }
 
@@ -676,10 +707,12 @@ flow_divert_trie_insert(struct flow_divert_trie *trie, uint16_t string_start, si
                current_end = TRIE_NODE(trie, current).start + TRIE_NODE(trie, current).length;
 
                for (node_idx = TRIE_NODE(trie, current).start;
-                    node_idx < current_end &&
-                    string_idx < string_end &&
-                    TRIE_BYTE(trie, node_idx) == TRIE_BYTE(trie, string_idx);
-                    node_idx++, string_idx++);
+                   node_idx < current_end &&
+                   string_idx < string_end &&
+                   TRIE_BYTE(trie, node_idx) == TRIE_BYTE(trie, string_idx);
+                   node_idx++, string_idx++) {
+                       ;
+               }
 
                string_remainder = string_end - string_idx;
 
@@ -769,8 +802,9 @@ flow_divert_trie_insert(struct flow_divert_trie *trie, uint16_t string_start, si
        return current;
 }
 
+#define APPLE_WEBCLIP_ID_PREFIX "com.apple.webapp"
 static uint16_t
-flow_divert_trie_search(struct flow_divert_trie *trie, const uint8_t *string_bytes)
+flow_divert_trie_search(struct flow_divert_trie *trie, uint8_t *string_bytes)
 {
        uint16_t current = trie->root;
        uint16_t string_idx = 0;
@@ -781,12 +815,18 @@ flow_divert_trie_search(struct flow_divert_trie *trie, const uint8_t *string_byt
                uint16_t node_idx;
 
                for (node_idx = TRIE_NODE(trie, current).start;
-                    node_idx < node_end && string_bytes[string_idx] != '\0' && string_bytes[string_idx] == TRIE_BYTE(trie, node_idx);
-                    node_idx++, string_idx++);
+                   node_idx < node_end && string_bytes[string_idx] != '\0' && string_bytes[string_idx] == TRIE_BYTE(trie, node_idx);
+                   node_idx++, string_idx++) {
+                       ;
+               }
 
                if (node_idx == node_end) {
                        if (string_bytes[string_idx] == '\0') {
                                return current; /* Got an exact match */
+                       } else if (string_idx == strlen(APPLE_WEBCLIP_ID_PREFIX) &&
+                           0 == strncmp((const char *)string_bytes, APPLE_WEBCLIP_ID_PREFIX, string_idx)) {
+                               string_bytes[string_idx] = '\0';
+                               return current; /* Got an apple webclip id prefix match */
                        } else if (TRIE_NODE(trie, current).child_map != NULL_TRIE_IDX) {
                                next = TRIE_CHILD(trie, current, string_bytes[string_idx]);
                        }
@@ -797,17 +837,106 @@ flow_divert_trie_search(struct flow_divert_trie *trie, const uint8_t *string_byt
        return NULL_TRIE_IDX;
 }
 
+struct uuid_search_info {
+       uuid_t target_uuid;
+       char *found_signing_id;
+       boolean_t found_multiple_signing_ids;
+       proc_t found_proc;
+};
+
+static int
+flow_divert_find_proc_by_uuid_callout(proc_t p, void *arg)
+{
+       struct uuid_search_info *info = (struct uuid_search_info *)arg;
+       int result = PROC_RETURNED_DONE; /* By default, we didn't find the process */
+
+       if (info->found_signing_id != NULL) {
+               if (!info->found_multiple_signing_ids) {
+                       /* All processes that were found had the same signing identifier, so just claim this first one and be done. */
+                       info->found_proc = p;
+                       result = PROC_CLAIMED_DONE;
+               } else {
+                       uuid_string_t uuid_str;
+                       uuid_unparse(info->target_uuid, uuid_str);
+                       FDLOG(LOG_WARNING, &nil_pcb, "Found multiple processes with UUID %s with different signing identifiers", uuid_str);
+               }
+               FREE(info->found_signing_id, M_TEMP);
+               info->found_signing_id = NULL;
+       }
+
+       if (result == PROC_RETURNED_DONE) {
+               uuid_string_t uuid_str;
+               uuid_unparse(info->target_uuid, uuid_str);
+               FDLOG(LOG_WARNING, &nil_pcb, "Failed to find a process with UUID %s", uuid_str);
+       }
+
+       return result;
+}
+
+static int
+flow_divert_find_proc_by_uuid_filter(proc_t p, void *arg)
+{
+       struct uuid_search_info *info = (struct uuid_search_info *)arg;
+       int include = 0;
+
+       if (info->found_multiple_signing_ids) {
+               return include;
+       }
+
+       include = (uuid_compare(p->p_uuid, info->target_uuid) == 0);
+       if (include) {
+               const char *signing_id = cs_identity_get(p);
+               if (signing_id != NULL) {
+                       FDLOG(LOG_INFO, &nil_pcb, "Found process %d with signing identifier %s", p->p_pid, signing_id);
+                       size_t signing_id_size = strlen(signing_id) + 1;
+                       if (info->found_signing_id == NULL) {
+                               MALLOC(info->found_signing_id, char *, signing_id_size, M_TEMP, M_WAITOK);
+                               memcpy(info->found_signing_id, signing_id, signing_id_size);
+                       } else if (memcmp(signing_id, info->found_signing_id, signing_id_size)) {
+                               info->found_multiple_signing_ids = TRUE;
+                       }
+               } else {
+                       info->found_multiple_signing_ids = TRUE;
+               }
+               include = !info->found_multiple_signing_ids;
+       }
+
+       return include;
+}
+
+static proc_t
+flow_divert_find_proc_by_uuid(uuid_t uuid)
+{
+       struct uuid_search_info info;
+
+       if (LOG_INFO <= nil_pcb.log_level) {
+               uuid_string_t uuid_str;
+               uuid_unparse(uuid, uuid_str);
+               FDLOG(LOG_INFO, &nil_pcb, "Looking for process with UUID %s", uuid_str);
+       }
+
+       memset(&info, 0, sizeof(info));
+       info.found_proc = PROC_NULL;
+       uuid_copy(info.target_uuid, uuid);
+
+       proc_iterate(PROC_ALLPROCLIST, flow_divert_find_proc_by_uuid_callout, &info, flow_divert_find_proc_by_uuid_filter, &info);
+
+       return info.found_proc;
+}
+
 static int
-flow_divert_get_src_proc(struct socket *so, proc_t *proc, boolean_t match_delegate)
+flow_divert_get_src_proc(struct socket *so, proc_t *proc)
 {
        int release = 0;
 
-       if (!match_delegate && 
-           (so->so_flags & SOF_DELEGATED) &&
-           (*proc == PROC_NULL || (*proc)->p_pid != so->e_pid))
-       {
-               *proc = proc_find(so->e_pid);
-               release = 1;
+       if (so->so_flags & SOF_DELEGATED) {
+               if ((*proc)->p_pid != so->e_pid) {
+                       *proc = proc_find(so->e_pid);
+                       release = 1;
+               } else if (uuid_compare((*proc)->p_uuid, so->e_uuid)) {
+                       *proc = flow_divert_find_proc_by_uuid(so->e_uuid);
+                       release = 1;
+               }
        } else if (*proc == PROC_NULL) {
                *proc = current_proc();
        }
@@ -828,11 +957,11 @@ flow_divert_get_src_proc(struct socket *so, proc_t *proc, boolean_t match_delega
 static int
 flow_divert_send_packet(struct flow_divert_pcb *fd_cb, mbuf_t packet, Boolean enqueue)
 {
-       int             error;
+       int             error;
 
        if (fd_cb->group == NULL) {
                fd_cb->so->so_error = ECONNABORTED;
-               soisdisconnected(fd_cb->so);
+               flow_divert_disconnect_socket(fd_cb->so);
                return ECONNABORTED;
        }
 
@@ -861,53 +990,166 @@ flow_divert_send_packet(struct flow_divert_pcb *fd_cb, mbuf_t packet, Boolean en
 }
 
 static int
-flow_divert_send_connect(struct flow_divert_pcb *fd_cb, struct sockaddr *to, proc_t p)
+flow_divert_create_connect_packet(struct flow_divert_pcb *fd_cb, struct sockaddr *to, struct socket *so, proc_t p, mbuf_t *out_connect_packet)
 {
-       mbuf_t                  connect_packet  = NULL;
-       int                             error                   = 0;
+       int                             error                   = 0;
+       int                             flow_type               = 0;
+       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) {
                goto done;
        }
 
+       error = EPERM;
+
+       if (fd_cb->connect_token != NULL && (fd_cb->flags & FLOW_DIVERT_HAS_HMAC)) {
+               uint32_t sid_size = 0;
+               int find_error = flow_divert_packet_get_tlv(fd_cb->connect_token, 0, FLOW_DIVERT_TLV_SIGNING_ID, 0, NULL, &sid_size);
+               if (find_error == 0 && sid_size > 0) {
+                       MALLOC(signing_id, char *, sid_size + 1, M_TEMP, M_WAITOK | M_ZERO);
+                       if (signing_id != NULL) {
+                               flow_divert_packet_get_tlv(fd_cb->connect_token, 0, FLOW_DIVERT_TLV_SIGNING_ID, sid_size, signing_id, NULL);
+                               FDLOG(LOG_INFO, fd_cb, "Got %s from token", signing_id);
+                               free_signing_id = 1;
+                       }
+               }
+       }
+
+       socket_unlock(so, 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 | CS_DEBUGGED)) {
+                               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");
+                       }
+               } else {
+                       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(&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);
+                                       }
+                               }
+                       } else {
+                               FDLOG(LOG_ERR, fd_cb, "failed to append the signing ID: %d", error);
+                       }
+               } else {
+                       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);
+               }
+       }
+       socket_lock(so, 0);
+
+       if (free_signing_id) {
+               FREE(signing_id, M_TEMP);
+       }
+
+       if (error) {
+               goto done;
+       }
+
+       error = flow_divert_packet_append_tlv(connect_packet,
+           FLOW_DIVERT_TLV_TRAFFIC_CLASS,
+           sizeof(fd_cb->so->so_traffic_class),
+           &fd_cb->so->so_traffic_class);
+       if (error) {
+               goto done;
+       }
+
+       if (SOCK_TYPE(fd_cb->so) == SOCK_STREAM) {
+               flow_type = FLOW_DIVERT_FLOW_TYPE_TCP;
+       } else if (SOCK_TYPE(fd_cb->so) == SOCK_DGRAM) {
+               flow_type = FLOW_DIVERT_FLOW_TYPE_UDP;
+       } else {
+               error = EINVAL;
+               goto done;
+       }
        error = flow_divert_packet_append_tlv(connect_packet,
-                                             FLOW_DIVERT_TLV_TRAFFIC_CLASS,
-                                             sizeof(fd_cb->so->so_traffic_class),
-                                             &fd_cb->so->so_traffic_class);
+           FLOW_DIVERT_TLV_FLOW_TYPE,
+           sizeof(flow_type),
+           &flow_type);
+
        if (error) {
                goto done;
        }
 
        if (fd_cb->so->so_flags & SOF_DELEGATED) {
                error = flow_divert_packet_append_tlv(connect_packet,
-                                                     FLOW_DIVERT_TLV_PID,
-                                                     sizeof(fd_cb->so->e_pid),
-                                                     &fd_cb->so->e_pid);
+                   FLOW_DIVERT_TLV_PID,
+                   sizeof(fd_cb->so->e_pid),
+                   &fd_cb->so->e_pid);
                if (error) {
                        goto done;
                }
 
                error = flow_divert_packet_append_tlv(connect_packet,
-                                                     FLOW_DIVERT_TLV_UUID,
-                                                     sizeof(fd_cb->so->e_uuid),
-                                                     &fd_cb->so->e_uuid);
+                   FLOW_DIVERT_TLV_UUID,
+                   sizeof(fd_cb->so->e_uuid),
+                   &fd_cb->so->e_uuid);
                if (error) {
                        goto done;
                }
        } else {
                error = flow_divert_packet_append_tlv(connect_packet,
-                                                     FLOW_DIVERT_TLV_PID,
-                                                     sizeof(fd_cb->so->e_pid),
-                                                     &fd_cb->so->last_pid);
+                   FLOW_DIVERT_TLV_PID,
+                   sizeof(fd_cb->so->e_pid),
+                   &fd_cb->so->last_pid);
                if (error) {
                        goto done;
                }
 
                error = flow_divert_packet_append_tlv(connect_packet,
-                                                     FLOW_DIVERT_TLV_UUID,
-                                                     sizeof(fd_cb->so->e_uuid),
-                                                     &fd_cb->so->last_uuid);
+                   FLOW_DIVERT_TLV_UUID,
+                   sizeof(fd_cb->so->e_uuid),
+                   &fd_cb->so->last_uuid);
                if (error) {
                        goto done;
                }
@@ -920,66 +1162,54 @@ flow_divert_send_connect(struct flow_divert_pcb *fd_cb, struct sockaddr *to, pro
                fd_cb->connect_token = NULL;
        } else {
                uint32_t ctl_unit = htonl(fd_cb->control_group_unit);
-               int port;
-               int release_proc;
 
                error = flow_divert_packet_append_tlv(connect_packet, FLOW_DIVERT_TLV_CTL_UNIT, sizeof(ctl_unit), &ctl_unit);
                if (error) {
                        goto done;
                }
 
-               error = flow_divert_packet_append_tlv(connect_packet, FLOW_DIVERT_TLV_TARGET_ADDRESS, to->sa_len, to);
+               error = flow_divert_append_target_endpoint_tlv(connect_packet, to);
                if (error) {
                        goto done;
                }
+       }
 
-               if (to->sa_family == AF_INET) {
-                       port = ntohs((satosin(to))->sin_port);
-               }
-#if INET6
-               else {
-                       port = ntohs((satosin6(to))->sin6_port);
+       if (fd_cb->local_address != NULL) {
+               error = EALREADY;
+               goto done;
+       } else {
+               struct inpcb *inp = sotoinpcb(so);
+               if (flow_divert_has_pcb_local_address(inp)) {
+                       error = flow_divert_inp_to_sockaddr(inp, &fd_cb->local_address);
+                       if (error) {
+                               FDLOG0(LOG_ERR, fd_cb, "failed to get the local socket address.");
+                               goto done;
+                       }
                }
-#endif
+       }
 
-               error = flow_divert_packet_append_tlv(connect_packet, FLOW_DIVERT_TLV_TARGET_PORT, sizeof(port), &port);
+       if (fd_cb->local_address != NULL) {
+               /* socket is bound. */
+               error = flow_divert_packet_append_tlv(connect_packet, FLOW_DIVERT_TLV_LOCAL_ADDR,
+                   fd_cb->local_address->sa_len, fd_cb->local_address);
                if (error) {
                        goto done;
                }
-
-               release_proc = flow_divert_get_src_proc(fd_cb->so, &p, FALSE);
-               if (p != PROC_NULL) {
-                       proc_lock(p);
-                       if (p->p_csflags & CS_VALID) {
-                               const char *signing_id = cs_identity_get(p);
-                               if (signing_id != NULL) {
-                                       error = flow_divert_packet_append_tlv(connect_packet, FLOW_DIVERT_TLV_SIGNING_ID, strlen(signing_id), signing_id);
-                               }
-
-                               if (error == 0) {
-                                       unsigned char cdhash[SHA1_RESULTLEN];
-                                       error = proc_getcdhash(p, cdhash);
-                                       if (error == 0) {
-                                               error = flow_divert_packet_append_tlv(connect_packet, FLOW_DIVERT_TLV_CDHASH, sizeof(cdhash), cdhash);
-                                       }
-                               }
-                       }
-                       proc_unlock(p);
-
-                       if (release_proc) {
-                               proc_rele(p);
-                       }
-               }
        }
 
-       error = flow_divert_send_packet(fd_cb, connect_packet, TRUE);
-       if (error) {
-               goto done;
+       if (so->so_flags1 & SOF1_DATA_IDEMPOTENT) {
+               uint32_t flags = FLOW_DIVERT_TOKEN_FLAG_TFO;
+               error = flow_divert_packet_append_tlv(connect_packet, FLOW_DIVERT_TLV_FLAGS, sizeof(flags), &flags);
+               if (error) {
+                       goto done;
+               }
        }
 
 done:
-       if (error && connect_packet != NULL) {
-               mbuf_free(connect_packet);
+       if (!error) {
+               *out_connect_packet = connect_packet;
+       } else if (connect_packet != NULL) {
+               mbuf_freem(connect_packet);
        }
 
        return error;
@@ -988,9 +1218,9 @@ done:
 static int
 flow_divert_send_connect_result(struct flow_divert_pcb *fd_cb)
 {
-       int             error                   = 0;
-       mbuf_t  packet                  = NULL;
-       int             rbuff_space             = 0;
+       int             error                   = 0;
+       mbuf_t  packet                  = NULL;
+       int             rbuff_space             = 0;
 
        error = flow_divert_packet_init(fd_cb, FLOW_DIVERT_PKT_CONNECT_RESULT, &packet);
        if (error) {
@@ -998,15 +1228,15 @@ flow_divert_send_connect_result(struct flow_divert_pcb *fd_cb)
                goto done;
        }
 
-       rbuff_space = sbspace(&fd_cb->so->so_rcv);
+       rbuff_space = fd_cb->so->so_rcv.sb_hiwat;
        if (rbuff_space < 0) {
                rbuff_space = 0;
        }
        rbuff_space = htonl(rbuff_space);
        error = flow_divert_packet_append_tlv(packet,
-                                             FLOW_DIVERT_TLV_SPACE_AVAILABLE,
-                                             sizeof(rbuff_space),
-                                             &rbuff_space);
+           FLOW_DIVERT_TLV_SPACE_AVAILABLE,
+           sizeof(rbuff_space),
+           &rbuff_space);
        if (error) {
                goto done;
        }
@@ -1018,7 +1248,7 @@ flow_divert_send_connect_result(struct flow_divert_pcb *fd_cb)
 
 done:
        if (error && packet != NULL) {
-               mbuf_free(packet);
+               mbuf_freem(packet);
        }
 
        return error;
@@ -1027,9 +1257,9 @@ done:
 static int
 flow_divert_send_close(struct flow_divert_pcb *fd_cb, int how)
 {
-       int             error   = 0;
-       mbuf_t  packet  = NULL;
-       uint32_t        zero    = 0;
+       int             error   = 0;
+       mbuf_t  packet  = NULL;
+       uint32_t        zero    = 0;
 
        error = flow_divert_packet_init(fd_cb, FLOW_DIVERT_PKT_CLOSE, &packet);
        if (error) {
@@ -1066,9 +1296,8 @@ done:
 static int
 flow_divert_tunnel_how_closed(struct flow_divert_pcb *fd_cb)
 {
-       if ((fd_cb->flags & (FLOW_DIVERT_TUNNEL_RD_CLOSED|FLOW_DIVERT_TUNNEL_WR_CLOSED)) == 
-                       (FLOW_DIVERT_TUNNEL_RD_CLOSED|FLOW_DIVERT_TUNNEL_WR_CLOSED))
-       {
+       if ((fd_cb->flags & (FLOW_DIVERT_TUNNEL_RD_CLOSED | FLOW_DIVERT_TUNNEL_WR_CLOSED)) ==
+           (FLOW_DIVERT_TUNNEL_RD_CLOSED | FLOW_DIVERT_TUNNEL_WR_CLOSED)) {
                return SHUT_RDWR;
        } else if (fd_cb->flags & FLOW_DIVERT_TUNNEL_RD_CLOSED) {
                return SHUT_RD;
@@ -1086,15 +1315,15 @@ flow_divert_tunnel_how_closed(struct flow_divert_pcb *fd_cb)
 static void
 flow_divert_send_close_if_needed(struct flow_divert_pcb *fd_cb)
 {
-       int             how             = -1;
+       int             how             = -1;
 
        /* Do not send any close messages if there is still data in the send buffer */
        if (fd_cb->so->so_snd.sb_cc == 0) {
-               if ((fd_cb->flags & (FLOW_DIVERT_READ_CLOSED|FLOW_DIVERT_TUNNEL_RD_CLOSED)) == FLOW_DIVERT_READ_CLOSED) {
+               if ((fd_cb->flags & (FLOW_DIVERT_READ_CLOSED | FLOW_DIVERT_TUNNEL_RD_CLOSED)) == FLOW_DIVERT_READ_CLOSED) {
                        /* Socket closed reads, but tunnel did not. Tell tunnel to close reads */
                        how = SHUT_RD;
                }
-               if ((fd_cb->flags & (FLOW_DIVERT_WRITE_CLOSED|FLOW_DIVERT_TUNNEL_WR_CLOSED)) == FLOW_DIVERT_WRITE_CLOSED) {
+               if ((fd_cb->flags & (FLOW_DIVERT_WRITE_CLOSED | FLOW_DIVERT_TUNNEL_WR_CLOSED)) == FLOW_DIVERT_WRITE_CLOSED) {
                        /* Socket closed writes, but tunnel did not. Tell tunnel to close writes */
                        if (how == SHUT_RD) {
                                how = SHUT_RDWR;
@@ -1118,16 +1347,16 @@ flow_divert_send_close_if_needed(struct flow_divert_pcb *fd_cb)
        }
 
        if (flow_divert_tunnel_how_closed(fd_cb) == SHUT_RDWR) {
-               soisdisconnected(fd_cb->so);
+               flow_divert_disconnect_socket(fd_cb->so);
        }
 }
 
 static errno_t
-flow_divert_send_data_packet(struct flow_divert_pcb *fd_cb, mbuf_t data, size_t data_len, Boolean force)
+flow_divert_send_data_packet(struct flow_divert_pcb *fd_cb, mbuf_t data, size_t data_len, struct sockaddr *toaddr, Boolean force)
 {
-       mbuf_t  packet;
-       mbuf_t  last;
-       int             error   = 0;
+       mbuf_t  packet;
+       mbuf_t  last;
+       int             error   = 0;
 
        error = flow_divert_packet_init(fd_cb, FLOW_DIVERT_PKT_DATA, &packet);
        if (error) {
@@ -1135,15 +1364,24 @@ flow_divert_send_data_packet(struct flow_divert_pcb *fd_cb, mbuf_t data, size_t
                return error;
        }
 
-       last = m_last(packet);
-       mbuf_setnext(last, data);
-       mbuf_pkthdr_adjustlen(packet, data_len);
+       if (toaddr != NULL) {
+               error = flow_divert_append_target_endpoint_tlv(packet, toaddr);
+               if (error) {
+                       FDLOG(LOG_ERR, fd_cb, "flow_divert_append_target_endpoint_tlv() failed: %d", error);
+                       return error;
+               }
+       }
 
+       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) {
                mbuf_setnext(last, NULL);
-               mbuf_free(packet);
+               mbuf_freem(packet);
        } else {
                fd_cb->bytes_sent += data_len;
                flow_divert_add_data_statistics(fd_cb, data_len, TRUE);
@@ -1155,10 +1393,10 @@ flow_divert_send_data_packet(struct flow_divert_pcb *fd_cb, mbuf_t data, size_t
 static void
 flow_divert_send_buffered_data(struct flow_divert_pcb *fd_cb, Boolean force)
 {
-       size_t  to_send;
-       size_t  sent    = 0;
-       int             error   = 0;
-       mbuf_t  buffer;
+       size_t  to_send;
+       size_t  sent    = 0;
+       int             error   = 0;
+       mbuf_t  buffer;
 
        to_send = fd_cb->so->so_snd.sb_cc;
        buffer = fd_cb->so->so_snd.sb_mb;
@@ -1173,28 +1411,76 @@ flow_divert_send_buffered_data(struct flow_divert_pcb *fd_cb, Boolean force)
                to_send = fd_cb->send_window;
        }
 
-       while (sent < to_send) {
-               mbuf_t  data;
-               size_t  data_len;
+       if (SOCK_TYPE(fd_cb->so) == SOCK_STREAM) {
+               while (sent < to_send) {
+                       mbuf_t  data;
+                       size_t  data_len;
 
-               data_len = to_send - sent;
-               if (data_len > FLOW_DIVERT_CHUNK_SIZE) {
-                       data_len = FLOW_DIVERT_CHUNK_SIZE;
-               }
+                       data_len = to_send - sent;
+                       if (data_len > FLOW_DIVERT_CHUNK_SIZE) {
+                               data_len = FLOW_DIVERT_CHUNK_SIZE;
+                       }
 
-               error = mbuf_copym(buffer, sent, data_len, MBUF_DONTWAIT, &data);
-               if (error) {
-                       FDLOG(LOG_ERR, fd_cb, "mbuf_copym failed: %d", error);
-                       break;
-               }
+                       error = mbuf_copym(buffer, sent, data_len, MBUF_DONTWAIT, &data);
+                       if (error) {
+                               FDLOG(LOG_ERR, fd_cb, "mbuf_copym failed: %d", error);
+                               break;
+                       }
 
-               error = flow_divert_send_data_packet(fd_cb, data, data_len, force);
-               if (error) {
-                       mbuf_free(data);
-                       break;
-               }
+                       error = flow_divert_send_data_packet(fd_cb, data, data_len, NULL, force);
+                       if (error) {
+                               mbuf_freem(data);
+                               break;
+                       }
 
-               sent += data_len;
+                       sent += data_len;
+               }
+               sbdrop(&fd_cb->so->so_snd, sent);
+               sowwakeup(fd_cb->so);
+       } else if (SOCK_TYPE(fd_cb->so) == SOCK_DGRAM) {
+               mbuf_t data;
+               mbuf_t m;
+               size_t data_len;
+
+               while (buffer) {
+                       struct sockaddr *toaddr = flow_divert_get_buffered_target_address(buffer);
+
+                       m = buffer;
+                       if (toaddr != NULL) {
+                               /* look for data in the chain */
+                               do {
+                                       m = m->m_next;
+                                       if (m != NULL && m->m_type == MT_DATA) {
+                                               break;
+                                       }
+                               } while (m);
+                               if (m == NULL) {
+                                       /* unexpected */
+                                       FDLOG0(LOG_ERR, fd_cb, "failed to find type MT_DATA in the mbuf chain.");
+                                       goto move_on;
+                               }
+                       }
+                       data_len = mbuf_pkthdr_len(m);
+                       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) {
+                               mbuf_freem(data);
+                               break;
+                       }
+                       sent += data_len;
+move_on:
+                       buffer = buffer->m_nextpkt;
+                       (void) sbdroprecord(&(fd_cb->so->so_snd));
+               }
        }
 
        if (sent > 0) {
@@ -1204,79 +1490,111 @@ flow_divert_send_buffered_data(struct flow_divert_pcb *fd_cb, Boolean force)
                } else {
                        fd_cb->send_window = 0;
                }
-               sbdrop(&fd_cb->so->so_snd, sent);
-               sowwakeup(fd_cb->so);
        }
 }
 
 static int
-flow_divert_send_app_data(struct flow_divert_pcb *fd_cb, mbuf_t data)
+flow_divert_send_app_data(struct flow_divert_pcb *fd_cb, mbuf_t data, struct sockaddr *toaddr)
 {
-       size_t  to_send         = mbuf_pkthdr_len(data);
-       size_t  sent            = 0;
-       int             error           = 0;
-       mbuf_t  remaining_data  = data;
-       mbuf_t  pkt_data        = NULL;
+       size_t  to_send         = mbuf_pkthdr_len(data);
+       int     error           = 0;
 
        if (to_send > fd_cb->send_window) {
                to_send = fd_cb->send_window;
        }
 
        if (fd_cb->so->so_snd.sb_cc > 0) {
-               to_send = 0;    /* If the send buffer is non-empty, then we can't send anything */
+               to_send = 0;    /* If the send buffer is non-empty, then we can't send anything */
        }
 
-       while (sent < to_send) {
-               size_t  pkt_data_len;
+       if (SOCK_TYPE(fd_cb->so) == SOCK_STREAM) {
+               size_t  sent            = 0;
+               mbuf_t  remaining_data  = data;
+               mbuf_t  pkt_data        = NULL;
+               while (sent < to_send && remaining_data != NULL) {
+                       size_t  pkt_data_len;
+
+                       pkt_data = remaining_data;
+
+                       if ((to_send - sent) > FLOW_DIVERT_CHUNK_SIZE) {
+                               pkt_data_len = FLOW_DIVERT_CHUNK_SIZE;
+                       } else {
+                               pkt_data_len = to_send - sent;
+                       }
+
+                       if (pkt_data_len < mbuf_pkthdr_len(pkt_data)) {
+                               error = mbuf_split(pkt_data, pkt_data_len, MBUF_DONTWAIT, &remaining_data);
+                               if (error) {
+                                       FDLOG(LOG_ERR, fd_cb, "mbuf_split failed: %d", error);
+                                       pkt_data = NULL;
+                                       break;
+                               }
+                       } else {
+                               remaining_data = NULL;
+                       }
 
-               pkt_data = remaining_data;
+                       error = flow_divert_send_data_packet(fd_cb, pkt_data, pkt_data_len, NULL, FALSE);
 
-               if ((to_send - sent) > FLOW_DIVERT_CHUNK_SIZE) {
-                       pkt_data_len = FLOW_DIVERT_CHUNK_SIZE;
-                       error = mbuf_split(pkt_data, pkt_data_len, MBUF_DONTWAIT, &remaining_data);
                        if (error) {
-                               FDLOG(LOG_ERR, fd_cb, "mbuf_split failed: %d", error);
-                               pkt_data = NULL;
                                break;
                        }
-               } else {
-                       pkt_data_len = to_send - sent;
-                       remaining_data = NULL;
-               }
 
-               error = flow_divert_send_data_packet(fd_cb, pkt_data, pkt_data_len, FALSE);
-
-               if (error) {
-                       break;
+                       pkt_data = NULL;
+                       sent += pkt_data_len;
                }
 
-               pkt_data = NULL;
-               sent += pkt_data_len;
-       }
-
-       fd_cb->send_window -= sent;
+               fd_cb->send_window -= sent;
 
-       error = 0;
+               error = 0;
 
-       if (pkt_data != NULL) {
-               if (sbspace(&fd_cb->so->so_snd) > 0) {
-                       if (!sbappendstream(&fd_cb->so->so_snd, pkt_data)) {
-                               FDLOG(LOG_ERR, fd_cb, "sbappendstream failed with pkt_data, send buffer size = %u, send_window = %u\n",
-                                               fd_cb->so->so_snd.sb_cc, fd_cb->send_window);
+               if (pkt_data != NULL) {
+                       if (sbspace(&fd_cb->so->so_snd) > 0) {
+                               if (!sbappendstream(&fd_cb->so->so_snd, pkt_data)) {
+                                       FDLOG(LOG_ERR, fd_cb, "sbappendstream failed with pkt_data, send buffer size = %u, send_window = %u\n",
+                                           fd_cb->so->so_snd.sb_cc, fd_cb->send_window);
+                               }
+                       } else {
+                               error = ENOBUFS;
                        }
-               } else {
-                       error = ENOBUFS;
                }
-       }
 
-       if (remaining_data != NULL) {
-               if (sbspace(&fd_cb->so->so_snd) > 0) {
-                       if (!sbappendstream(&fd_cb->so->so_snd, remaining_data)) {
-                               FDLOG(LOG_ERR, fd_cb, "sbappendstream failed with remaining_data, send buffer size = %u, send_window = %u\n",
-                                               fd_cb->so->so_snd.sb_cc, fd_cb->send_window);
+               if (remaining_data != NULL) {
+                       if (sbspace(&fd_cb->so->so_snd) > 0) {
+                               if (!sbappendstream(&fd_cb->so->so_snd, remaining_data)) {
+                                       FDLOG(LOG_ERR, fd_cb, "sbappendstream failed with remaining_data, send buffer size = %u, send_window = %u\n",
+                                           fd_cb->so->so_snd.sb_cc, fd_cb->send_window);
+                               }
+                       } else {
+                               error = ENOBUFS;
+                       }
+               }
+       } else if (SOCK_TYPE(fd_cb->so) == SOCK_DGRAM) {
+               if (to_send || mbuf_pkthdr_len(data) == 0) {
+                       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 = %lu", to_send);
+                       } else {
+                               fd_cb->send_window -= to_send;
                        }
                } else {
-                       error = ENOBUFS;
+                       /* buffer it */
+                       if (sbspace(&fd_cb->so->so_snd) >= (int)mbuf_pkthdr_len(data)) {
+                               if (toaddr != NULL) {
+                                       if (!sbappendaddr(&fd_cb->so->so_snd, toaddr, data, NULL, &error)) {
+                                               FDLOG(LOG_ERR, fd_cb,
+                                                   "sbappendaddr failed. send buffer size = %u, send_window = %u, error = %d\n",
+                                                   fd_cb->so->so_snd.sb_cc, fd_cb->send_window, error);
+                                       }
+                               } else {
+                                       if (!sbappendrecord(&fd_cb->so->so_snd, data)) {
+                                               FDLOG(LOG_ERR, fd_cb,
+                                                   "sbappendrecord failed. send buffer size = %u, send_window = %u, error = %d\n",
+                                                   fd_cb->so->so_snd.sb_cc, fd_cb->send_window, error);
+                                       }
+                               }
+                       } else {
+                               error = ENOBUFS;
+                       }
                }
        }
 
@@ -1286,9 +1604,9 @@ flow_divert_send_app_data(struct flow_divert_pcb *fd_cb, mbuf_t data)
 static int
 flow_divert_send_read_notification(struct flow_divert_pcb *fd_cb, uint32_t read_count)
 {
-       int             error           = 0;
-       mbuf_t  packet          = NULL;
-       uint32_t        net_read_count  = htonl(read_count);
+       int             error           = 0;
+       mbuf_t  packet          = NULL;
+       uint32_t        net_read_count  = htonl(read_count);
 
        error = flow_divert_packet_init(fd_cb, FLOW_DIVERT_PKT_READ_NOTIFY, &packet);
        if (error) {
@@ -1318,8 +1636,8 @@ done:
 static int
 flow_divert_send_traffic_class_update(struct flow_divert_pcb *fd_cb, int traffic_class)
 {
-       int             error           = 0;
-       mbuf_t  packet          = NULL;
+       int             error           = 0;
+       mbuf_t  packet          = NULL;
 
        error = flow_divert_packet_init(fd_cb, FLOW_DIVERT_PKT_PROPERTIES_UPDATE, &packet);
        if (error) {
@@ -1349,14 +1667,15 @@ done:
 static void
 flow_divert_handle_connect_result(struct flow_divert_pcb *fd_cb, mbuf_t packet, int offset)
 {
-       uint32_t                                        connect_error;
-       uint32_t                                        ctl_unit                        = 0;
-       int                                                     error                           = 0;
-       struct flow_divert_group        *grp                            = NULL;
-       struct sockaddr_storage         local_address;
-       int                                                     out_if_index            = 0;
-       struct sockaddr_storage         remote_address;
-       uint32_t                                        send_window;
+       uint32_t                                        connect_error;
+       uint32_t                                        ctl_unit                        = 0;
+       int                                                     error                           = 0;
+       struct flow_divert_group        *grp                            = NULL;
+       struct sockaddr_storage         local_address;
+       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));
@@ -1377,32 +1696,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) {
-               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) {
-               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) {
-               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) {
-               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");
        }
 
-       connect_error   = ntohl(connect_error);
-       ctl_unit                = ntohl(ctl_unit);
+       error = 0;
+       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) {
@@ -1418,9 +1742,9 @@ flow_divert_handle_connect_result(struct flow_divert_pcb *fd_cb, mbuf_t packet,
 
        FDLOCK(fd_cb);
        if (fd_cb->so != NULL) {
-               struct inpcb                            *inp = NULL;
-               struct ifnet                            *ifp = NULL;
-               struct flow_divert_group        *old_group;
+               struct inpcb                            *inp = NULL;
+               struct ifnet                            *ifp = NULL;
+               struct flow_divert_group        *old_group;
 
                socket_lock(fd_cb->so, 0);
 
@@ -1434,14 +1758,15 @@ flow_divert_handle_connect_result(struct flow_divert_pcb *fd_cb, mbuf_t packet,
                        goto set_socket_state;
                }
 
-               if (local_address.ss_family != 0) {
+               if (local_address.ss_family == 0 && fd_cb->local_address == NULL) {
+                       error = EINVAL;
+                       goto set_socket_state;
+               }
+               if (local_address.ss_family != 0 && fd_cb->local_address == NULL) {
                        if (local_address.ss_len > sizeof(local_address)) {
                                local_address.ss_len = sizeof(local_address);
                        }
                        fd_cb->local_address = dup_sockaddr((struct sockaddr *)&local_address, 1);
-               } else {
-                       error = EINVAL;
-                       goto set_socket_state;
                }
 
                if (remote_address.ss_family != 0) {
@@ -1454,8 +1779,29 @@ flow_divert_handle_connect_result(struct flow_divert_pcb *fd_cb, mbuf_t packet,
                        goto set_socket_state;
                }
 
-               ifnet_head_lock_shared();
-               if (out_if_index > 0 && out_if_index <= if_index) {
+               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];
                }
 
@@ -1475,23 +1821,24 @@ flow_divert_handle_connect_result(struct flow_divert_pcb *fd_cb, mbuf_t packet,
                        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);
-               flow_divert_send_buffered_data(fd_cb, FALSE);
 
 set_socket_state:
                if (!connect_error && !error) {
@@ -1508,8 +1855,9 @@ set_socket_state:
                                flow_divert_update_closed_state(fd_cb, SHUT_RDWR, TRUE);
                                fd_cb->so->so_error = connect_error;
                        }
-                       soisdisconnected(fd_cb->so);
+                       flow_divert_disconnect_socket(fd_cb->so);
                } else {
+                       flow_divert_send_buffered_data(fd_cb, FALSE);
                        soisconnected(fd_cb->so);
                }
 
@@ -1524,9 +1872,9 @@ done:
 static void
 flow_divert_handle_close(struct flow_divert_pcb *fd_cb, mbuf_t packet, int offset)
 {
-       uint32_t        close_error;
-       int                     error                   = 0;
-       int                     how;
+       uint32_t        close_error;
+       int                     error                   = 0;
+       int                     how;
 
        error = flow_divert_packet_get_tlv(packet, offset, FLOW_DIVERT_TLV_ERROR_CODE, sizeof(close_error), &close_error, NULL);
        if (error) {
@@ -1551,10 +1899,10 @@ flow_divert_handle_close(struct flow_divert_pcb *fd_cb, mbuf_t packet, int offse
                fd_cb->so->so_error = ntohl(close_error);
 
                flow_divert_update_closed_state(fd_cb, how, TRUE);
-               
+
                how = flow_divert_tunnel_how_closed(fd_cb);
                if (how == SHUT_RDWR) {
-                       soisdisconnected(fd_cb->so);
+                       flow_divert_disconnect_socket(fd_cb->so);
                } else if (how == SHUT_RD) {
                        socantrcvmore(fd_cb->so);
                } else if (how == SHUT_WR) {
@@ -1566,55 +1914,129 @@ flow_divert_handle_close(struct flow_divert_pcb *fd_cb, mbuf_t packet, int offse
        FDUNLOCK(fd_cb);
 }
 
-static void
-flow_divert_handle_data(struct flow_divert_pcb *fd_cb, mbuf_t packet, size_t offset)
+static mbuf_t
+flow_divert_get_control_mbuf(struct flow_divert_pcb *fd_cb)
 {
-       int             error           = 0;
-       mbuf_t  data            = NULL;
-       size_t  data_size;
-
-       data_size = (mbuf_pkthdr_len(packet) - offset);
+       if (fd_cb->local_address != NULL) {
+               struct inpcb *inp = sotoinpcb(fd_cb->so);
+               if ((inp->inp_vflag & INP_IPV4) &&
+                   (inp->inp_flags & INP_RECVDSTADDR) &&
+                   fd_cb->local_address->sa_family == AF_INET &&
+                   fd_cb->local_address->sa_len >= sizeof(struct sockaddr_in)) {
+                       struct sockaddr_in *sin = (struct sockaddr_in *)(void *)fd_cb->local_address;
 
-       FDLOG(LOG_DEBUG, fd_cb, "received %lu bytes of data", data_size);
+                       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) &&
+                   fd_cb->local_address->sa_family == AF_INET6 &&
+                   fd_cb->local_address->sa_len >= sizeof(struct sockaddr_in6)) {
+                       struct sockaddr_in6 *sin6 = (struct sockaddr_in6 *)(void *)fd_cb->local_address;
+                       struct in6_pktinfo pi6;
 
-       error = mbuf_split(packet, offset, MBUF_DONTWAIT, &data);
-       if (error || data == NULL) {
-               FDLOG(LOG_ERR, fd_cb, "mbuf_split failed: %d", error);
-               return;
+                       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;
+}
 
+static void
+flow_divert_handle_data(struct flow_divert_pcb *fd_cb, mbuf_t packet, size_t offset)
+{
        FDLOCK(fd_cb);
        if (fd_cb->so != NULL) {
+               int             error           = 0;
+               mbuf_t  data            = NULL;
+               size_t  data_size;
+               struct sockaddr_storage remote_address;
+               boolean_t got_remote_sa = FALSE;
+
                socket_lock(fd_cb->so, 0);
-               if (flow_divert_check_no_cellular(fd_cb)) {
-                       flow_divert_update_closed_state(fd_cb, SHUT_RDWR, TRUE);
-                       flow_divert_send_close(fd_cb, SHUT_RDWR);
-                       soisdisconnected(fd_cb->so);
-               } else if (!(fd_cb->so->so_state & SS_CANTRCVMORE)) {
-                       if (sbappendstream(&fd_cb->so->so_rcv, data)) {
-                               fd_cb->bytes_received += data_size;
-                               flow_divert_add_data_statistics(fd_cb, data_size, FALSE);
-                               fd_cb->sb_size = fd_cb->so->so_rcv.sb_cc;
-                               sorwakeup(fd_cb->so);
-                               data = NULL;
+
+               if (SOCK_TYPE(fd_cb->so) == SOCK_DGRAM) {
+                       uint32_t val_size = 0;
+
+                       /* check if we got remote address with data */
+                       memset(&remote_address, 0, sizeof(remote_address));
+                       error = flow_divert_packet_get_tlv(packet, offset, FLOW_DIVERT_TLV_REMOTE_ADDR, sizeof(remote_address), &remote_address, &val_size);
+                       if (error || val_size > sizeof(remote_address)) {
+                               FDLOG0(LOG_INFO, fd_cb, "No remote address provided");
+                               error = 0;
                        } else {
-                               FDLOG0(LOG_ERR, fd_cb, "received data, but appendstream failed");
+                               /* validate the address */
+                               if (flow_divert_is_sockaddr_valid((struct sockaddr *)&remote_address)) {
+                                       got_remote_sa = TRUE;
+                               }
+                               offset += (sizeof(uint8_t) + sizeof(uint32_t) + val_size);
+                       }
+               }
+
+               data_size = (mbuf_pkthdr_len(packet) - offset);
+
+               FDLOG(LOG_DEBUG, fd_cb, "received %lu bytes of data", data_size);
+
+               error = mbuf_split(packet, offset, MBUF_DONTWAIT, &data);
+               if (error || data == NULL) {
+                       FDLOG(LOG_ERR, fd_cb, "mbuf_split failed: %d", error);
+               } else {
+                       if (flow_divert_check_no_cellular(fd_cb) ||
+                           flow_divert_check_no_expensive(fd_cb)) {
+                               flow_divert_update_closed_state(fd_cb, SHUT_RDWR, TRUE);
+                               flow_divert_send_close(fd_cb, SHUT_RDWR);
+                               flow_divert_disconnect_socket(fd_cb->so);
+                       } else if (!(fd_cb->so->so_state & SS_CANTRCVMORE)) {
+                               if (SOCK_TYPE(fd_cb->so) == SOCK_STREAM) {
+                                       if (sbappendstream(&fd_cb->so->so_rcv, data)) {
+                                               fd_cb->bytes_received += data_size;
+                                               flow_divert_add_data_statistics(fd_cb, data_size, FALSE);
+                                               fd_cb->sb_size = fd_cb->so->so_rcv.sb_cc;
+                                               sorwakeup(fd_cb->so);
+                                               data = NULL;
+                                       } else {
+                                               FDLOG0(LOG_ERR, fd_cb, "received data, but appendstream failed");
+                                       }
+                               } else if (SOCK_TYPE(fd_cb->so) == SOCK_DGRAM) {
+                                       struct sockaddr *append_sa;
+                                       mbuf_t mctl;
+
+                                       if (got_remote_sa == TRUE) {
+                                               error = flow_divert_dup_addr(fd_cb->so->so_proto->pr_domain->dom_family,
+                                                   (struct sockaddr *)&remote_address, &append_sa);
+                                       } else {
+                                               error = flow_divert_dup_addr(fd_cb->so->so_proto->pr_domain->dom_family,
+                                                   fd_cb->remote_address, &append_sa);
+                                       }
+                                       if (error) {
+                                               FDLOG0(LOG_ERR, fd_cb, "failed to dup the socket address.");
+                                       }
+
+                                       mctl = flow_divert_get_control_mbuf(fd_cb);
+                                       if (sbappendaddr(&fd_cb->so->so_rcv, append_sa, data, mctl, NULL)) {
+                                               fd_cb->bytes_received += data_size;
+                                               flow_divert_add_data_statistics(fd_cb, data_size, FALSE);
+                                               fd_cb->sb_size = fd_cb->so->so_rcv.sb_cc;
+                                               sorwakeup(fd_cb->so);
+                                               data = NULL;
+                                       } else {
+                                               FDLOG0(LOG_ERR, fd_cb, "received data, but sbappendaddr failed");
+                                       }
+                                       if (!error) {
+                                               FREE(append_sa, M_TEMP);
+                                       }
+                               }
                        }
                }
                socket_unlock(fd_cb->so, 0);
        }
        FDUNLOCK(fd_cb);
-
-       if (data != NULL) {
-               mbuf_free(data);
-       }
 }
 
 static void
 flow_divert_handle_read_notification(struct flow_divert_pcb *fd_cb, mbuf_t packet, int offset)
 {
-       uint32_t        read_count;
-       int             error                   = 0;
+       uint32_t        read_count;
+       int             error                   = 0;
 
        error = flow_divert_packet_get_tlv(packet, offset, FLOW_DIVERT_TLV_READ_COUNT, sizeof(read_count), &read_count, NULL);
        if (error) {
@@ -1622,7 +2044,7 @@ flow_divert_handle_read_notification(struct flow_divert_pcb *fd_cb, mbuf_t packe
                return;
        }
 
-       FDLOG(LOG_DEBUG, fd_cb, "received a read notification for %u bytes", read_count);
+       FDLOG(LOG_DEBUG, fd_cb, "received a read notification for %u bytes", ntohl(read_count));
 
        FDLOCK(fd_cb);
        if (fd_cb->so != NULL) {
@@ -1638,8 +2060,9 @@ static void
 flow_divert_handle_group_init(struct flow_divert_group *group, mbuf_t packet, int offset)
 {
        int error = 0;
-       size_t key_size = 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) {
@@ -1648,7 +2071,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) {
-               FDLOG(LOG_ERR, &nil_pcb, "Invalid key size: %lu", key_size);
+               FDLOG(LOG_ERR, &nil_pcb, "Invalid key size: %u", key_size);
                return;
        }
 
@@ -1671,16 +2094,22 @@ flow_divert_handle_group_init(struct flow_divert_group *group, mbuf_t packet, in
 
        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);
 }
 
 static void
 flow_divert_handle_properties_update(struct flow_divert_pcb *fd_cb, mbuf_t packet, int offset)
 {
-       int                                                     error                           = 0;
-       struct sockaddr_storage         local_address;
-       int                                                     out_if_index            = 0;
-       struct sockaddr_storage         remote_address;
+       int                                                     error                           = 0;
+       struct sockaddr_storage         local_address;
+       int                                                     out_if_index            = 0;
+       struct sockaddr_storage         remote_address;
+       uint32_t                                        app_data_length         = 0;
 
        FDLOG0(LOG_INFO, fd_cb, "received a properties update");
 
@@ -1689,32 +2118,36 @@ flow_divert_handle_properties_update(struct flow_divert_pcb *fd_cb, mbuf_t packe
 
        error = flow_divert_packet_get_tlv(packet, offset, FLOW_DIVERT_TLV_LOCAL_ADDR, sizeof(local_address), &local_address, NULL);
        if (error) {
-               FDLOG0(LOG_INFO, fd_cb, "No local address provided");
+               FDLOG0(LOG_INFO, fd_cb, "No local address provided in properties update");
        }
 
        error = flow_divert_packet_get_tlv(packet, offset, FLOW_DIVERT_TLV_REMOTE_ADDR, sizeof(remote_address), &remote_address, NULL);
        if (error) {
-               FDLOG0(LOG_INFO, fd_cb, "No remote address provided");
+               FDLOG0(LOG_INFO, fd_cb, "No remote address provided in properties update");
        }
 
        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_INFO, fd_cb, "No output if index provided");
+               FDLOG0(LOG_INFO, fd_cb, "No output if index provided in properties update");
+       }
+
+       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 properties update");
        }
 
        FDLOCK(fd_cb);
        if (fd_cb->so != NULL) {
-               struct inpcb                            *inp = NULL;
-               struct ifnet                            *ifp = NULL;
-
                socket_lock(fd_cb->so, 0);
 
-               inp = sotoinpcb(fd_cb->so);
-
                if (local_address.ss_family != 0) {
                        if (local_address.ss_len > sizeof(local_address)) {
                                local_address.ss_len = sizeof(local_address);
                        }
+                       if (fd_cb->local_address != NULL) {
+                               FREE(fd_cb->local_address, M_SONAME);
+                               fd_cb->local_address = NULL;
+                       }
                        fd_cb->local_address = dup_sockaddr((struct sockaddr *)&local_address, 1);
                }
 
@@ -1722,18 +2155,49 @@ flow_divert_handle_properties_update(struct flow_divert_pcb *fd_cb, mbuf_t packe
                        if (remote_address.ss_len > sizeof(remote_address)) {
                                remote_address.ss_len = sizeof(remote_address);
                        }
+                       if (fd_cb->remote_address != NULL) {
+                               FREE(fd_cb->remote_address, M_SONAME);
+                               fd_cb->remote_address = NULL;
+                       }
                        fd_cb->remote_address = dup_sockaddr((struct sockaddr *)&remote_address, 1);
                }
 
-               ifnet_head_lock_shared();
-               if (out_if_index > 0 && out_if_index <= if_index) {
-                       ifp = ifindex2ifnet[out_if_index];
+               if (out_if_index > 0) {
+                       struct inpcb *inp = NULL;
+                       struct ifnet *ifp = NULL;
+
+                       inp = sotoinpcb(fd_cb->so);
+
+                       ifnet_head_lock_shared();
+                       if (out_if_index <= if_index) {
+                               ifp = ifindex2ifnet[out_if_index];
+                       }
+
+                       if (ifp != NULL) {
+                               inp->inp_last_outifp = ifp;
+                       }
+                       ifnet_head_done();
                }
 
-               if (ifp != NULL) {
-                       inp->inp_last_outifp = ifp;
+               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) {
+                                       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 properties update 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 properties update", app_data_length);
+                       }
                }
-               ifnet_head_done();
 
                socket_unlock(fd_cb->so, 0);
        }
@@ -1741,7 +2205,7 @@ flow_divert_handle_properties_update(struct flow_divert_pcb *fd_cb, mbuf_t packe
 }
 
 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;
@@ -1752,34 +2216,39 @@ 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 trie_memory_size = 0;
 
-       lck_rw_lock_exclusive(&g_flow_divert_group_lck);
+       lck_rw_lock_exclusive(&group->lck);
 
        /* 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);
 
+       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;
-            cursor = flow_divert_packet_find_tlv(packet, cursor, FLOW_DIVERT_TLV_SIGNING_ID, &error, 1))
-       {
-               size_t sid_size = 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);
                new_trie.bytes_count += sid_size;
                signing_id_count++;
        }
 
        if (signing_id_count == 0) {
-               lck_rw_done(&g_flow_divert_group_lck);
+               lck_rw_done(&group->lck);
                return;
        }
 
@@ -1787,16 +2256,24 @@ flow_divert_handle_app_map_create(mbuf_t packet, int offset)
        new_trie.child_maps_count = (prefix_count + 1); /* + 1 for the root node */
 
        FDLOG(LOG_INFO, &nil_pcb, "Nodes count = %lu, child maps count = %lu, bytes_count = %lu",
-                       new_trie.nodes_count, new_trie.child_maps_count, new_trie.bytes_count);
+           new_trie.nodes_count, new_trie.child_maps_count, new_trie.bytes_count);
 
        nodes_mem_size = (sizeof(*new_trie.nodes) * new_trie.nodes_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);
+                   nodes_mem_size + child_maps_mem_size + bytes_mem_size);
+               lck_rw_done(&group->lck);
                return;
        }
 
@@ -1817,26 +2294,15 @@ flow_divert_handle_app_map_create(mbuf_t packet, int offset)
 
        /* Add each signing ID to the trie */
        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))
-       {
-               size_t sid_size = 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 (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);
-                       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);
-                       if (new_node_idx != NULL_TRIE_IDX) {
-                               if (is_dns) {
-                                       FDLOG(LOG_NOTICE, &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;
                        }
@@ -1848,80 +2314,20 @@ flow_divert_handle_app_map_create(mbuf_t packet, int offset)
        }
 
        if (!insert_error) {
-               g_signing_id_trie = new_trie;
+               group->signing_id_trie = new_trie;
        } 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);
-
-       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))
-       {
-               size_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))
-       {
-               size_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
 flow_divert_input(mbuf_t packet, struct flow_divert_group *group)
 {
-       struct flow_divert_packet_header        hdr;
-       int                                                                     error           = 0;
-       struct flow_divert_pcb                          *fd_cb;
+       struct flow_divert_packet_header        hdr;
+       int                                                                     error           = 0;
+       struct flow_divert_pcb                          *fd_cb;
 
        if (mbuf_pkthdr_len(packet) < sizeof(hdr)) {
                FDLOG(LOG_ERR, &nil_pcb, "got a bad packet, length (%lu) < sizeof hdr (%lu)", mbuf_pkthdr_len(packet), sizeof(hdr));
@@ -1929,6 +2335,12 @@ flow_divert_input(mbuf_t packet, struct flow_divert_group *group)
                goto done;
        }
 
+       if (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 = mbuf_copydata(packet, 0, sizeof(hdr), &hdr);
        if (error) {
                FDLOG(LOG_ERR, &nil_pcb, "mbuf_copydata failed for the header: %d", error);
@@ -1940,23 +2352,20 @@ flow_divert_input(mbuf_t packet, struct flow_divert_group *group)
 
        if (hdr.conn_id == 0) {
                switch (hdr.packet_type) {
-                       case FLOW_DIVERT_PKT_GROUP_INIT:
-                               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));
-                               break;
-                       default:
-                               FDLOG(LOG_WARNING, &nil_pcb, "got an unknown message type: %d", hdr.packet_type);
-                               break;
+               case FLOW_DIVERT_PKT_GROUP_INIT:
+                       flow_divert_handle_group_init(group, packet, sizeof(hdr));
+                       break;
+               case FLOW_DIVERT_PKT_APP_MAP_CREATE:
+                       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;
                }
                goto done;
        }
 
-       fd_cb = flow_divert_pcb_lookup(hdr.conn_id, group);             /* This retains the PCB */
+       fd_cb = flow_divert_pcb_lookup(hdr.conn_id, group);             /* This retains the PCB */
        if (fd_cb == NULL) {
                if (hdr.packet_type != FLOW_DIVERT_PKT_CLOSE && hdr.packet_type != FLOW_DIVERT_PKT_READ_NOTIFY) {
                        FDLOG(LOG_NOTICE, &nil_pcb, "got a %s message from group %d for an unknown pcb: %u", flow_divert_packet_type2str(hdr.packet_type), group->ctl_unit, hdr.conn_id);
@@ -1965,38 +2374,38 @@ flow_divert_input(mbuf_t packet, struct flow_divert_group *group)
        }
 
        switch (hdr.packet_type) {
-               case FLOW_DIVERT_PKT_CONNECT_RESULT:
-                       flow_divert_handle_connect_result(fd_cb, packet, sizeof(hdr));
-                       break;
-               case FLOW_DIVERT_PKT_CLOSE:
-                       flow_divert_handle_close(fd_cb, packet, sizeof(hdr));
-                       break;
-               case FLOW_DIVERT_PKT_DATA:
-                       flow_divert_handle_data(fd_cb, packet, sizeof(hdr));
-                       break;
-               case FLOW_DIVERT_PKT_READ_NOTIFY:
-                       flow_divert_handle_read_notification(fd_cb, packet, sizeof(hdr));
-                       break;
-               case FLOW_DIVERT_PKT_PROPERTIES_UPDATE:
-                       flow_divert_handle_properties_update(fd_cb, packet, sizeof(hdr));
-                       break;
-               default:
-                       FDLOG(LOG_WARNING, fd_cb, "got an unknown message type: %d", hdr.packet_type);
-                       break;
+       case FLOW_DIVERT_PKT_CONNECT_RESULT:
+               flow_divert_handle_connect_result(fd_cb, packet, sizeof(hdr));
+               break;
+       case FLOW_DIVERT_PKT_CLOSE:
+               flow_divert_handle_close(fd_cb, packet, sizeof(hdr));
+               break;
+       case FLOW_DIVERT_PKT_DATA:
+               flow_divert_handle_data(fd_cb, packet, sizeof(hdr));
+               break;
+       case FLOW_DIVERT_PKT_READ_NOTIFY:
+               flow_divert_handle_read_notification(fd_cb, packet, sizeof(hdr));
+               break;
+       case FLOW_DIVERT_PKT_PROPERTIES_UPDATE:
+               flow_divert_handle_properties_update(fd_cb, packet, sizeof(hdr));
+               break;
+       default:
+               FDLOG(LOG_WARNING, fd_cb, "got an unknown message type: %d", hdr.packet_type);
+               break;
        }
 
        FDRELEASE(fd_cb);
 
 done:
-       mbuf_free(packet);
+       mbuf_freem(packet);
        return error;
 }
 
 static void
 flow_divert_close_all(struct flow_divert_group *group)
 {
-       struct flow_divert_pcb                  *fd_cb;
-       SLIST_HEAD(, flow_divert_pcb)   tmp_list;
+       struct flow_divert_pcb                  *fd_cb;
+       SLIST_HEAD(, flow_divert_pcb)   tmp_list;
 
        SLIST_INIT(&tmp_list);
 
@@ -2020,6 +2429,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_disconnect_socket(fd_cb->so);
                        socket_unlock(fd_cb->so, 0);
                }
                FDUNLOCK(fd_cb);
@@ -2030,7 +2440,7 @@ flow_divert_close_all(struct flow_divert_group *group)
 void
 flow_divert_detach(struct socket *so)
 {
-       struct flow_divert_pcb  *fd_cb          = so->so_fd_pcb;
+       struct flow_divert_pcb  *fd_cb          = so->so_fd_pcb;
 
        VERIFY((so->so_flags & SOF_FLOW_DIVERT) && so->so_fd_pcb != NULL);
 
@@ -2043,6 +2453,8 @@ flow_divert_detach(struct socket *so)
                /* Last-ditch effort to send any buffered data */
                flow_divert_send_buffered_data(fd_cb, TRUE);
 
+               flow_divert_update_closed_state(fd_cb, SHUT_RDWR, FALSE);
+               flow_divert_send_close_if_needed(fd_cb);
                /* Remove from the group */
                flow_divert_pcb_remove(fd_cb);
        }
@@ -2053,20 +2465,22 @@ flow_divert_detach(struct socket *so)
        FDUNLOCK(fd_cb);
        socket_lock(so, 0);
 
-       FDRELEASE(fd_cb);       /* Release the socket's reference */
+       FDRELEASE(fd_cb);       /* Release the socket's reference */
 }
 
 static int
 flow_divert_close(struct socket *so)
 {
-       struct flow_divert_pcb  *fd_cb          = so->so_fd_pcb;
+       struct flow_divert_pcb  *fd_cb          = so->so_fd_pcb;
 
        VERIFY((so->so_flags & SOF_FLOW_DIVERT) && so->so_fd_pcb != NULL);
 
        FDLOG0(LOG_INFO, fd_cb, "Closing");
 
-       soisdisconnecting(so);
-       sbflush(&so->so_rcv);
+       if (SOCK_TYPE(so) == SOCK_STREAM) {
+               soisdisconnecting(so);
+               sbflush(&so->so_rcv);
+       }
 
        flow_divert_send_buffered_data(fd_cb, TRUE);
        flow_divert_update_closed_state(fd_cb, SHUT_RDWR, FALSE);
@@ -2079,19 +2493,20 @@ flow_divert_close(struct socket *so)
 }
 
 static int
-flow_divert_disconnectx(struct socket *so, associd_t aid, connid_t cid __unused)
+flow_divert_disconnectx(struct socket *so, sae_associd_t aid,
+    sae_connid_t cid __unused)
 {
-       if (aid != ASSOCID_ANY && aid != ASSOCID_ALL) {
-               return (EINVAL);
+       if (aid != SAE_ASSOCID_ANY && aid != SAE_ASSOCID_ALL) {
+               return EINVAL;
        }
 
-       return (flow_divert_close(so));
+       return flow_divert_close(so);
 }
 
 static int
 flow_divert_shutdown(struct socket *so)
 {
-       struct flow_divert_pcb  *fd_cb          = so->so_fd_pcb;
+       struct flow_divert_pcb  *fd_cb          = so->so_fd_pcb;
 
        VERIFY((so->so_flags & SOF_FLOW_DIVERT) && so->so_fd_pcb != NULL);
 
@@ -2108,9 +2523,9 @@ flow_divert_shutdown(struct socket *so)
 static int
 flow_divert_rcvd(struct socket *so, int flags __unused)
 {
-       struct flow_divert_pcb  *fd_cb                  = so->so_fd_pcb;
-       uint32_t                                latest_sb_size;
-       uint32_t                                read_count;
+       struct flow_divert_pcb  *fd_cb                  = so->so_fd_pcb;
+       uint32_t                                latest_sb_size;
+       uint32_t                                read_count;
 
        VERIFY((so->so_flags & SOF_FLOW_DIVERT) && so->so_fd_pcb != NULL);
 
@@ -2118,7 +2533,7 @@ flow_divert_rcvd(struct socket *so, int flags __unused)
 
        if (fd_cb->sb_size < latest_sb_size) {
                panic("flow divert rcvd event handler (%u): saved rcv buffer size (%u) is less than latest rcv buffer size (%u)",
-                               fd_cb->hash, fd_cb->sb_size, latest_sb_size);
+                   fd_cb->hash, fd_cb->sb_size, latest_sb_size);
        }
 
        read_count = fd_cb->sb_size - latest_sb_size;
@@ -2133,13 +2548,112 @@ flow_divert_rcvd(struct socket *so, int flags __unused)
        return 0;
 }
 
+static int
+flow_divert_append_target_endpoint_tlv(mbuf_t connect_packet, struct sockaddr *toaddr)
+{
+       int error = 0;
+       int port  = 0;
+
+       error = flow_divert_packet_append_tlv(connect_packet, FLOW_DIVERT_TLV_TARGET_ADDRESS, toaddr->sa_len, toaddr);
+       if (error) {
+               goto done;
+       }
+
+       if (toaddr->sa_family == AF_INET) {
+               port = ntohs((satosin(toaddr))->sin_port);
+       }
+#if INET6
+       else {
+               port = ntohs((satosin6(toaddr))->sin6_port);
+       }
+#endif
+
+       error = flow_divert_packet_append_tlv(connect_packet, FLOW_DIVERT_TLV_TARGET_PORT, sizeof(port), &port);
+       if (error) {
+               goto done;
+       }
+
+done:
+       return error;
+}
+
+struct sockaddr *
+flow_divert_get_buffered_target_address(mbuf_t buffer)
+{
+       if (buffer != NULL && buffer->m_type == MT_SONAME) {
+               struct sockaddr *toaddr = mtod(buffer, struct sockaddr *);
+               if (toaddr != NULL && flow_divert_is_sockaddr_valid(toaddr)) {
+                       return toaddr;
+               }
+       }
+       return NULL;
+}
+
+static boolean_t
+flow_divert_is_sockaddr_valid(struct sockaddr *addr)
+{
+       switch (addr->sa_family) {
+       case AF_INET:
+               if (addr->sa_len != sizeof(struct sockaddr_in)) {
+                       return FALSE;
+               }
+               break;
+#if INET6
+       case AF_INET6:
+               if (addr->sa_len != sizeof(struct sockaddr_in6)) {
+                       return FALSE;
+               }
+               break;
+#endif  /* INET6 */
+       default:
+               return FALSE;
+       }
+       return TRUE;
+}
+
+static errno_t
+flow_divert_inp_to_sockaddr(const struct inpcb *inp, struct sockaddr **local_socket)
+{
+       int error = 0;
+       union sockaddr_in_4_6 sin46;
+
+       bzero(&sin46, sizeof(sin46));
+       if (inp->inp_vflag & INP_IPV4) {
+               struct sockaddr_in  *sin = &sin46.sin;
+
+               sin->sin_family = AF_INET;
+               sin->sin_len = sizeof(*sin);
+               sin->sin_port = inp->inp_lport;
+               sin->sin_addr = inp->inp_laddr;
+       } else if (inp->inp_vflag & INP_IPV6) {
+               struct sockaddr_in6 *sin6 = &sin46.sin6;
+
+               sin6->sin6_len = sizeof(*sin6);
+               sin6->sin6_family = AF_INET6;
+               sin6->sin6_port = inp->inp_lport;
+               sin6->sin6_addr = inp->in6p_laddr;
+       }
+       *local_socket = dup_sockaddr((struct sockaddr *)&sin46, 1);
+       if (*local_socket == NULL) {
+               error = ENOBUFS;
+       }
+       return error;
+}
+
+static boolean_t
+flow_divert_has_pcb_local_address(const struct inpcb *inp)
+{
+       return inp->inp_lport != 0
+              && (inp->inp_laddr.s_addr != INADDR_ANY || !IN6_IS_ADDR_UNSPECIFIED(&inp->in6p_laddr));
+}
+
 static errno_t
 flow_divert_dup_addr(sa_family_t family, struct sockaddr *addr,
-                     struct sockaddr **dup)
+    struct sockaddr **dup)
 {
-       int                                             error           = 0;
-       struct sockaddr                 *result;
-       struct sockaddr_storage ss;
+       int                                             error           = 0;
+       struct sockaddr                 *result;
+       struct sockaddr_storage ss;
 
        if (addr != NULL) {
                result = addr;
@@ -2153,7 +2667,7 @@ flow_divert_dup_addr(sa_family_t family, struct sockaddr *addr,
                else if (ss.ss_family == AF_INET6) {
                        ss.ss_len = sizeof(struct sockaddr_in6);
                }
-#endif /* INET6 */
+#endif  /* INET6 */
                else {
                        error = EINVAL;
                }
@@ -2170,34 +2684,53 @@ flow_divert_dup_addr(sa_family_t family, struct sockaddr *addr,
        return error;
 }
 
+static void
+flow_divert_disconnect_socket(struct socket *so)
+{
+       soisdisconnected(so);
+       if (SOCK_TYPE(so) == SOCK_DGRAM) {
+               struct inpcb *inp = NULL;
+
+               inp = sotoinpcb(so);
+               if (inp != NULL) {
+#if INET6
+                       if (SOCK_CHECK_DOM(so, PF_INET6)) {
+                               in6_pcbdetach(inp);
+                       } else
+#endif /* INET6 */
+                       in_pcbdetach(inp);
+               }
+       }
+}
+
 static errno_t
 flow_divert_getpeername(struct socket *so, struct sockaddr **sa)
 {
-       struct flow_divert_pcb  *fd_cb  = so->so_fd_pcb;
+       struct flow_divert_pcb  *fd_cb  = so->so_fd_pcb;
 
        VERIFY((so->so_flags & SOF_FLOW_DIVERT) && so->so_fd_pcb != NULL);
 
-       return flow_divert_dup_addr(so->so_proto->pr_domain->dom_family, 
-                                   fd_cb->remote_address,
-                                   sa);
+       return flow_divert_dup_addr(so->so_proto->pr_domain->dom_family,
+                  fd_cb->remote_address,
+                  sa);
 }
 
 static errno_t
 flow_divert_getsockaddr(struct socket *so, struct sockaddr **sa)
 {
-       struct flow_divert_pcb  *fd_cb  = so->so_fd_pcb;
+       struct flow_divert_pcb  *fd_cb  = so->so_fd_pcb;
 
        VERIFY((so->so_flags & SOF_FLOW_DIVERT) && so->so_fd_pcb != NULL);
 
-       return flow_divert_dup_addr(so->so_proto->pr_domain->dom_family, 
-                                   fd_cb->local_address,
-                                   sa);
+       return flow_divert_dup_addr(so->so_proto->pr_domain->dom_family,
+                  fd_cb->local_address,
+                  sa);
 }
 
 static errno_t
 flow_divert_ctloutput(struct socket *so, struct sockopt *sopt)
 {
-       struct flow_divert_pcb  *fd_cb  = so->so_fd_pcb;
+       struct flow_divert_pcb  *fd_cb  = so->so_fd_pcb;
 
        VERIFY((so->so_flags & SOF_FLOW_DIVERT) && so->so_fd_pcb != NULL);
 
@@ -2221,10 +2754,12 @@ flow_divert_ctloutput(struct socket *so, struct sockopt *sopt)
 errno_t
 flow_divert_connect_out(struct socket *so, struct sockaddr *to, proc_t p)
 {
-       struct flow_divert_pcb  *fd_cb  = so->so_fd_pcb;
-       int                                             error   = 0;
-       struct inpcb                    *inp    = sotoinpcb(so);
-       struct sockaddr_in              *sinp;
+       struct flow_divert_pcb  *fd_cb  = so->so_fd_pcb;
+       int                                             error   = 0;
+       struct inpcb                    *inp    = sotoinpcb(so);
+       struct sockaddr_in              *sinp;
+       mbuf_t                                  connect_packet = NULL;
+       int                                             do_send = 1;
 
        VERIFY((so->so_flags & SOF_FLOW_DIVERT) && so->so_fd_pcb != NULL);
 
@@ -2246,12 +2781,6 @@ flow_divert_connect_out(struct socket *so, struct sockaddr *to, proc_t p)
                goto done;
        }
 
-       sinp = (struct sockaddr_in *)(void *)to;
-       if (sinp->sin_family == AF_INET && IN_MULTICAST(ntohl(sinp->sin_addr.s_addr))) {
-               error = EAFNOSUPPORT;
-               goto done;
-       }
-
        if ((fd_cb->flags & FLOW_DIVERT_CONNECT_STARTED) && !(fd_cb->flags & FLOW_DIVERT_TRANSFERRED)) {
                error = EALREADY;
                goto done;
@@ -2268,88 +2797,140 @@ flow_divert_connect_out(struct socket *so, struct sockaddr *to, proc_t p)
 
        FDLOG0(LOG_INFO, fd_cb, "Connecting");
 
-       error = flow_divert_send_connect(fd_cb, to, p);
-       if (error) {
-               goto done;
+       if (fd_cb->connect_packet == NULL) {
+               if (to == NULL) {
+                       FDLOG0(LOG_ERR, fd_cb, "No destination address available when creating connect packet");
+                       error = EINVAL;
+                       goto done;
+               }
+
+               sinp = (struct sockaddr_in *)(void *)to;
+               if (sinp->sin_family == AF_INET && IN_MULTICAST(ntohl(sinp->sin_addr.s_addr))) {
+                       error = EAFNOSUPPORT;
+                       goto done;
+               }
+
+               error = flow_divert_create_connect_packet(fd_cb, to, so, p, &connect_packet);
+               if (error) {
+                       goto done;
+               }
+
+               if (so->so_flags1 & SOF1_PRECONNECT_DATA) {
+                       FDLOG0(LOG_INFO, fd_cb, "Delaying sending the connect packet until send or receive");
+                       do_send = 0;
+               }
+       } else {
+               FDLOG0(LOG_INFO, fd_cb, "Sending saved connect packet");
+               connect_packet = fd_cb->connect_packet;
+               fd_cb->connect_packet = NULL;
        }
 
-       fd_cb->flags |= FLOW_DIVERT_CONNECT_STARTED;
+       if (do_send) {
+               error = flow_divert_send_packet(fd_cb, connect_packet, TRUE);
+               if (error) {
+                       goto done;
+               }
+
+               fd_cb->flags |= FLOW_DIVERT_CONNECT_STARTED;
+       } else {
+               fd_cb->connect_packet = connect_packet;
+               connect_packet = NULL;
+       }
 
        soisconnecting(so);
 
 done:
+       if (error && connect_packet != NULL) {
+               mbuf_freem(connect_packet);
+       }
        return error;
 }
 
 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, associd_t aid __unused,
-    connid_t *pcid, uint32_t flags __unused, void *arg __unused,
-    uint32_t arglen __unused)
+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;
 
        if (inp == NULL) {
-               return (EINVAL);
+               return EINVAL;
        }
 
-       VERIFY(dst_sl != NULL);
+       VERIFY(dst != NULL);
+
+       error = flow_divert_connect_out(so, dst, p);
 
-       /* select source (if specified) and destination addresses */
-       error = in_selectaddrs(af, src_sl, &src_se, dst_sl, &dst_se);
        if (error != 0) {
-               return (error);
+               return error;
        }
 
-       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);
+       /* if there is data, send it */
+       if (auio != NULL) {
+               user_ssize_t datalen = 0;
+
+               socket_unlock(so, 0);
+
+               VERIFY(bytes_written != NULL);
+
+               datalen = uio_resid(auio);
+               error = so->so_proto->pr_usrreqs->pru_sosend(so, NULL, (uio_t)auio, NULL, NULL, 0);
+               socket_lock(so, 0);
+
+               if (error == 0 || error == EWOULDBLOCK) {
+                       *bytes_written = datalen - uio_resid(auio);
+               }
 
-       error = flow_divert_connect_out(so, dst_se->se_addr, p);
+               /*
+                * sosend returns EWOULDBLOCK if it's a non-blocking
+                * socket or a timeout occured (this allows to return
+                * the amount of queued data through sendit()).
+                *
+                * However, connectx() returns EINPROGRESS in case of a
+                * blocking socket. So we change the return value here.
+                */
+               if (error == EWOULDBLOCK) {
+                       error = EINPROGRESS;
+               }
+       }
 
        if (error == 0 && pcid != NULL) {
-               *pcid = 1;      /* there is only 1 connection for a TCP */
+               *pcid = 1;      /* there is only 1 connection for a TCP */
        }
 
-       return (error);
+       return error;
 }
 
 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,
-    associd_t aid, connid_t *pcid, uint32_t flags, void *arg,
-    uint32_t arglen)
+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)
 {
-       return (flow_divert_connectx_out_common(so, AF_INET, src_sl, dst_sl,
-           p, ifscope, aid, pcid, flags, arg, arglen));
+       return flow_divert_connectx_out_common(so, dst, p, pcid, uio, bytes_written);
 }
 
 #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,
-    associd_t aid, connid_t *pcid, uint32_t flags, void *arg,
-    uint32_t arglen)
+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)
 {
-       return (flow_divert_connectx_out_common(so, AF_INET6, src_sl, dst_sl,
-           p, ifscope, aid, pcid, flags, arg, arglen));
+       return flow_divert_connectx_out_common(so, dst, p, pcid, uio, bytes_written);
 }
 #endif /* INET6 */
 
 static int
-flow_divert_getconninfo(struct socket *so, connid_t cid, uint32_t *flags,
-                        uint32_t *ifindex, int32_t *soerror, user_addr_t src, socklen_t *src_len,
-                        user_addr_t dst, socklen_t *dst_len, uint32_t *aux_type,
-                        user_addr_t aux_data __unused, uint32_t *aux_len)
+flow_divert_getconninfo(struct socket *so, sae_connid_t cid, uint32_t *flags,
+    uint32_t *ifindex, int32_t *soerror, user_addr_t src, socklen_t *src_len,
+    user_addr_t dst, socklen_t *dst_len, uint32_t *aux_type,
+    user_addr_t aux_data __unused, uint32_t *aux_len)
 {
-       int                                             error   = 0;
-       struct flow_divert_pcb  *fd_cb  = so->so_fd_pcb;
-       struct ifnet                    *ifp    = NULL;
-       struct inpcb                    *inp    = sotoinpcb(so);
+       int                                             error   = 0;
+       struct flow_divert_pcb  *fd_cb  = so->so_fd_pcb;
+       struct ifnet                    *ifp    = NULL;
+       struct inpcb                    *inp    = sotoinpcb(so);
 
        VERIFY((so->so_flags & SOF_FLOW_DIVERT));
 
@@ -2358,7 +2939,7 @@ flow_divert_getconninfo(struct socket *so, connid_t cid, uint32_t *flags,
                goto out;
        }
 
-       if (cid != CONNID_ANY && cid != CONNID_ALL && cid != 1) {
+       if (cid != SAE_CONNID_ANY && cid != SAE_CONNID_ALL && cid != 1) {
                error = EINVAL;
                goto out;
        }
@@ -2429,36 +3010,36 @@ flow_divert_control(struct socket *so, u_long cmd, caddr_t data, struct ifnet *i
        int error = 0;
 
        switch (cmd) {
-               case SIOCGCONNINFO32: {
-                       struct so_cinforeq32 cifr;
-                       bcopy(data, &cifr, sizeof (cifr));
-                       error = flow_divert_getconninfo(so, cifr.scir_cid, &cifr.scir_flags,
-                                                       &cifr.scir_ifindex, &cifr.scir_error, cifr.scir_src,
-                                                       &cifr.scir_src_len, cifr.scir_dst, &cifr.scir_dst_len,
-                                                       &cifr.scir_aux_type, cifr.scir_aux_data,
-                                                       &cifr.scir_aux_len);
-                       if (error == 0) {
-                               bcopy(&cifr, data, sizeof (cifr));
-                       }
-                       break;
+       case SIOCGCONNINFO32: {
+               struct so_cinforeq32 cifr;
+               bcopy(data, &cifr, sizeof(cifr));
+               error = flow_divert_getconninfo(so, cifr.scir_cid, &cifr.scir_flags,
+                   &cifr.scir_ifindex, &cifr.scir_error, cifr.scir_src,
+                   &cifr.scir_src_len, cifr.scir_dst, &cifr.scir_dst_len,
+                   &cifr.scir_aux_type, cifr.scir_aux_data,
+                   &cifr.scir_aux_len);
+               if (error == 0) {
+                       bcopy(&cifr, data, sizeof(cifr));
                }
+               break;
+       }
 
-               case SIOCGCONNINFO64: {
-                       struct so_cinforeq64 cifr;
-                       bcopy(data, &cifr, sizeof (cifr));
-                       error = flow_divert_getconninfo(so, cifr.scir_cid, &cifr.scir_flags,
-                                                       &cifr.scir_ifindex, &cifr.scir_error, cifr.scir_src,
-                                                       &cifr.scir_src_len, cifr.scir_dst, &cifr.scir_dst_len,
-                                                       &cifr.scir_aux_type, cifr.scir_aux_data,
-                                                       &cifr.scir_aux_len);
-                       if (error == 0) {
-                               bcopy(&cifr, data, sizeof (cifr));
-                       }
-                       break;
+       case SIOCGCONNINFO64: {
+               struct so_cinforeq64 cifr;
+               bcopy(data, &cifr, sizeof(cifr));
+               error = flow_divert_getconninfo(so, cifr.scir_cid, &cifr.scir_flags,
+                   &cifr.scir_ifindex, &cifr.scir_error, cifr.scir_src,
+                   &cifr.scir_src_len, cifr.scir_dst, &cifr.scir_dst_len,
+                   &cifr.scir_aux_type, cifr.scir_aux_data,
+                   &cifr.scir_aux_len);
+               if (error == 0) {
+                       bcopy(&cifr, data, sizeof(cifr));
                }
+               break;
+       }
 
-               default:
-                       error = EOPNOTSUPP;
+       default:
+               error = EOPNOTSUPP;
        }
 
        return error;
@@ -2489,10 +3070,10 @@ flow_divert_in6_control(struct socket *so, u_long cmd, caddr_t data, struct ifne
 }
 
 static errno_t
-flow_divert_data_out(struct socket *so, int flags, mbuf_t data, struct sockaddr *to, mbuf_t control, struct proc *p __unused)
+flow_divert_data_out(struct socket *so, int flags, mbuf_t data, struct sockaddr *to, mbuf_t control, struct proc *p)
 {
-       struct flow_divert_pcb  *fd_cb  = so->so_fd_pcb;
-       int                                             error   = 0;
+       struct flow_divert_pcb  *fd_cb  = so->so_fd_pcb;
+       int                                             error   = 0;
        struct inpcb *inp;
 
        VERIFY((so->so_flags & SOF_FLOW_DIVERT) && so->so_fd_pcb != NULL);
@@ -2512,8 +3093,9 @@ flow_divert_data_out(struct socket *so, int flags, mbuf_t data, struct sockaddr
                error = EINVAL;
                goto done; /* We don't support OOB data */
        }
-       
-       error = flow_divert_check_no_cellular(fd_cb);
+
+       error = flow_divert_check_no_cellular(fd_cb) ||
+           flow_divert_check_no_expensive(fd_cb);
        if (error) {
                goto done;
        }
@@ -2521,16 +3103,21 @@ flow_divert_data_out(struct socket *so, int flags, mbuf_t data, struct sockaddr
        /* Implicit connect */
        if (!(fd_cb->flags & FLOW_DIVERT_CONNECT_STARTED)) {
                FDLOG0(LOG_INFO, fd_cb, "implicit connect");
-               error = flow_divert_connect_out(so, to, NULL);
+               error = flow_divert_connect_out(so, to, p);
                if (error) {
                        goto done;
                }
+
+               if (so->so_flags1 & SOF1_DATA_IDEMPOTENT) {
+                       /* Open up the send window so that the data will get sent right away */
+                       fd_cb->send_window = mbuf_pkthdr_len(data);
+               }
        }
 
        FDLOG(LOG_DEBUG, fd_cb, "app wrote %lu bytes", mbuf_pkthdr_len(data));
 
        fd_cb->bytes_written_by_app += mbuf_pkthdr_len(data);
-       error = flow_divert_send_app_data(fd_cb, data);
+       error = flow_divert_send_app_data(fd_cb, data, to);
        if (error) {
                goto done;
        }
@@ -2543,7 +3130,7 @@ flow_divert_data_out(struct socket *so, int flags, mbuf_t data, struct sockaddr
 
 done:
        if (data) {
-               mbuf_free(data);
+               mbuf_freem(data);
        }
        if (control) {
                mbuf_free(control);
@@ -2551,64 +3138,26 @@ done:
        return error;
 }
 
-boolean_t
-flow_divert_is_dns_service(struct socket *so)
+static int
+flow_divert_preconnect(struct socket *so)
 {
-       uint32_t ctl_unit = 0;
-       flow_divert_check_policy(so, NULL, TRUE, &ctl_unit);
-       FDLOG(LOG_INFO, &nil_pcb, "Check for DNS resulted in %u", ctl_unit);
-       return (ctl_unit == DNS_SERVICE_GROUP_UNIT);
-}
+       struct flow_divert_pcb  *fd_cb  = so->so_fd_pcb;
+       int error = 0;
 
-errno_t
-flow_divert_check_policy(struct socket *so, proc_t p, boolean_t match_delegate, uint32_t *ctl_unit)
-{
-       int error = EPROTOTYPE;
+       if (!(fd_cb->flags & FLOW_DIVERT_CONNECT_STARTED) && fd_cb->connect_packet != NULL) {
+               FDLOG0(LOG_INFO, fd_cb, "Pre-connect read: sending saved connect packet");
+               mbuf_t connect_packet = fd_cb->connect_packet;
+               fd_cb->connect_packet = NULL;
 
-       if (ctl_unit != NULL) {
-               *ctl_unit = 0;
-       }
+               error = flow_divert_send_packet(fd_cb, connect_packet, TRUE);
+               if (error) {
+                       mbuf_freem(connect_packet);
+               }
 
-       if (SOCK_DOM(so) != PF_INET
-#if INET6
-           && SOCK_DOM(so) != PF_INET6
-#endif
-           )
-       {
-               return error;
+               fd_cb->flags |= FLOW_DIVERT_CONNECT_STARTED;
        }
 
-       if (g_signing_id_trie.root != NULL_TRIE_IDX) {
-               int release_proc = flow_divert_get_src_proc(so, &p, match_delegate);
-               if (p != PROC_NULL) {
-                       proc_lock(p);
-                       if (p->p_csflags & CS_VALID) {
-                               const char *signing_id = cs_identity_get(p);
-                               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, (const uint8_t *)signing_id);
-                                       if (result != NULL_TRIE_IDX) {
-                                               uint32_t unit = TRIE_NODE(&g_signing_id_trie, result).group_unit;
-
-                                               error = 0;
-
-                                               FDLOG(LOG_INFO, &nil_pcb, "%s matched, ctl_unit = %u", signing_id, unit);
-
-                                               if (ctl_unit != NULL) {
-                                                       *ctl_unit = unit;
-                                               }
-                                       }
-                                       lck_rw_done(&g_flow_divert_group_lck);
-                               }
-                       }
-                       proc_unlock(p);
-
-                       if (release_proc) {
-                               proc_rele(p);
-                       }
-               }
-       }
+       soclearfastopen(so);
 
        return error;
 }
@@ -2624,18 +3173,32 @@ flow_divert_set_protosw(struct socket *so)
        else {
                so->so_proto = (struct protosw *)&g_flow_divert_in6_protosw;
        }
-#endif /* INET6 */
+#endif  /* INET6 */
+}
+
+static void
+flow_divert_set_udp_protosw(struct socket *so)
+{
+       so->so_flags |= SOF_FLOW_DIVERT;
+       if (SOCK_DOM(so) == PF_INET) {
+               so->so_proto = &g_flow_divert_in_udp_protosw;
+       }
+#if INET6
+       else {
+               so->so_proto = (struct protosw *)&g_flow_divert_in6_udp_protosw;
+       }
+#endif  /* INET6 */
 }
 
 static errno_t
 flow_divert_attach(struct socket *so, uint32_t flow_id, uint32_t ctl_unit)
 {
-       int                                                                     error           = 0;
-       struct flow_divert_pcb                          *fd_cb          = NULL;
-       struct ifnet                                            *ifp            = NULL;
-       struct inpcb                                            *inp            = NULL;
-       struct socket                                           *old_so;
-       mbuf_t                                                          recv_data       = NULL;
+       int                                                                     error           = 0;
+       struct flow_divert_pcb                          *fd_cb          = NULL;
+       struct ifnet                                            *ifp            = NULL;
+       struct inpcb                                            *inp            = NULL;
+       struct socket                                           *old_so;
+       mbuf_t                                                          recv_data       = NULL;
 
        socket_unlock(so, 0);
 
@@ -2661,15 +3224,19 @@ flow_divert_attach(struct socket *so, uint32_t flow_id, uint32_t ctl_unit)
        /* Dis-associate the flow divert control block from its current socket */
        old_so = fd_cb->so;
 
-       inp = sotoinpcb(old_so); 
+       inp = sotoinpcb(old_so);
 
        VERIFY(inp != NULL);
 
        socket_lock(old_so, 0);
-       soisdisconnected(old_so);
+       flow_divert_disconnect_socket(old_so);
        old_so->so_flags &= ~SOF_FLOW_DIVERT;
        old_so->so_fd_pcb = NULL;
-       old_so->so_proto = pffindproto(SOCK_DOM(old_so), IPPROTO_TCP, SOCK_STREAM);
+       if (SOCK_TYPE(old_so) == SOCK_STREAM) {
+               old_so->so_proto = pffindproto(SOCK_DOM(old_so), IPPROTO_TCP, SOCK_STREAM);
+       } else if (SOCK_TYPE(old_so) == SOCK_DGRAM) {
+               old_so->so_proto = pffindproto(SOCK_DOM(old_so), IPPROTO_UDP, SOCK_DGRAM);
+       }
        fd_cb->so = NULL;
        /* Save the output interface */
        ifp = inp->inp_last_outifp;
@@ -2701,7 +3268,45 @@ done:
        socket_lock(so, 0);
 
        if (fd_cb != NULL) {
-               FDRELEASE(fd_cb);       /* Release the reference obtained via flow_divert_pcb_lookup */
+               FDRELEASE(fd_cb);       /* Release the reference obtained via flow_divert_pcb_lookup */
+       }
+
+       return error;
+}
+
+errno_t
+flow_divert_implicit_data_out(struct socket *so, int flags, mbuf_t data, struct sockaddr *to, mbuf_t control, struct proc *p)
+{
+       struct flow_divert_pcb  *fd_cb  = so->so_fd_pcb;
+       struct inpcb *inp;
+       int error = 0;
+
+       inp = sotoinpcb(so);
+       if (inp == NULL) {
+               return EINVAL;
+       }
+
+       if (fd_cb == NULL) {
+               uint32_t fd_ctl_unit = necp_socket_get_flow_divert_control_unit(inp);
+               if (fd_ctl_unit > 0) {
+                       error = flow_divert_pcb_init(so, fd_ctl_unit);
+                       fd_cb  = so->so_fd_pcb;
+                       if (error != 0 || fd_cb == NULL) {
+                               goto done;
+                       }
+               } else {
+                       error = ENETDOWN;
+                       goto done;
+               }
+       }
+       return flow_divert_data_out(so, flags, data, to, control, p);
+
+done:
+       if (data) {
+               mbuf_freem(data);
+       }
+       if (control) {
+               mbuf_free(control);
        }
 
        return error;
@@ -2716,7 +3321,7 @@ flow_divert_pcb_init(struct socket *so, uint32_t ctl_unit)
        if (so->so_flags & SOF_FLOW_DIVERT) {
                return EALREADY;
        }
-               
+
        fd_cb = flow_divert_pcb_create(so);
        if (fd_cb != NULL) {
                error = flow_divert_pcb_insert(fd_cb, ctl_unit);
@@ -2724,11 +3329,14 @@ flow_divert_pcb_init(struct socket *so, uint32_t ctl_unit)
                        FDLOG(LOG_ERR, fd_cb, "pcb insert failed: %d", error);
                        FDRELEASE(fd_cb);
                } else {
-                       fd_cb->log_level = LOG_NOTICE;
                        fd_cb->control_group_unit = ctl_unit;
                        so->so_fd_pcb = fd_cb;
 
-                       flow_divert_set_protosw(so);
+                       if (SOCK_TYPE(so) == SOCK_STREAM) {
+                               flow_divert_set_protosw(so);
+                       } else if (SOCK_TYPE(so) == SOCK_DGRAM) {
+                               flow_divert_set_udp_protosw(so);
+                       }
 
                        FDLOG0(LOG_INFO, fd_cb, "Created");
                }
@@ -2742,11 +3350,12 @@ flow_divert_pcb_init(struct socket *so, uint32_t ctl_unit)
 errno_t
 flow_divert_token_set(struct socket *so, struct sockopt *sopt)
 {
-       uint32_t                                        ctl_unit                = 0;
-       uint32_t                                        key_unit                = 0;
-       uint32_t                                        flow_id                 = 0;
-       int                                                     error                   = 0;
-       mbuf_t                                          token                   = NULL;
+       uint32_t                                        ctl_unit                = 0;
+       uint32_t                                        key_unit                = 0;
+       uint32_t                                        flow_id                 = 0;
+       int                                                     error                   = 0;
+       int                                                     hmac_error              = 0;
+       mbuf_t                                          token                   = NULL;
 
        if (so->so_flags & SOF_FLOW_DIVERT) {
                error = EALREADY;
@@ -2759,37 +3368,43 @@ flow_divert_token_set(struct socket *so, struct sockopt *sopt)
                goto done;
        }
 
-       if (SOCK_TYPE(so) != SOCK_STREAM ||
-           SOCK_PROTO(so) != IPPROTO_TCP ||
+       if ((SOCK_TYPE(so) != SOCK_STREAM && SOCK_TYPE(so) != SOCK_DGRAM) ||
+           (SOCK_PROTO(so) != IPPROTO_TCP && SOCK_PROTO(so) != IPPROTO_UDP) ||
            (SOCK_DOM(so) != PF_INET
 #if INET6
-            && SOCK_DOM(so) != PF_INET6
+           && SOCK_DOM(so) != PF_INET6
 #endif
-               ))
-       {
+           )) {
                error = EINVAL;
                goto done;
        } else {
-               struct tcpcb *tp = sototcpcb(so);
-               if (tp == NULL || tp->t_state != TCPS_CLOSED) {
-                       error = EINVAL;
-                       goto done;
+               if (SOCK_TYPE(so) == SOCK_STREAM && SOCK_PROTO(so) == IPPROTO_TCP) {
+                       struct tcpcb *tp = sototcpcb(so);
+                       if (tp == NULL || tp->t_state != TCPS_CLOSED) {
+                               error = EINVAL;
+                               goto done;
+                       }
                }
        }
 
        error = soopt_getm(sopt, &token);
        if (error) {
+               token = NULL;
                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);
+               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;
@@ -2812,11 +3427,12 @@ flow_divert_token_set(struct socket *so, struct sockopt *sopt)
        }
 
        socket_unlock(so, 0);
-       error = flow_divert_packet_verify_hmac(token, (key_unit != 0 ? key_unit : ctl_unit));
+       hmac_error = flow_divert_packet_verify_hmac(token, (key_unit != 0 ? key_unit : ctl_unit));
        socket_lock(so, 0);
 
-       if (error) {
-               FDLOG(LOG_ERR, &nil_pcb, "HMAC verfication failed: %d", error);
+       if (hmac_error && hmac_error != ENOENT) {
+               FDLOG(LOG_ERR, &nil_pcb, "HMAC verfication failed: %d", hmac_error);
+               error = hmac_error;
                goto done;
        }
 
@@ -2833,7 +3449,7 @@ flow_divert_token_set(struct socket *so, struct sockopt *sopt)
                        int log_level = LOG_NOTICE;
 
                        error = flow_divert_packet_get_tlv(token, 0, FLOW_DIVERT_TLV_LOG_LEVEL,
-                                                              sizeof(log_level), &log_level, NULL);
+                           sizeof(log_level), &log_level, NULL);
                        if (error == 0) {
                                fd_cb->log_level = log_level;
                        }
@@ -2846,6 +3462,13 @@ flow_divert_token_set(struct socket *so, struct sockopt *sopt)
                error = flow_divert_attach(so, flow_id, ctl_unit);
        }
 
+       if (hmac_error == 0) {
+               struct flow_divert_pcb *fd_cb = so->so_fd_pcb;
+               if (fd_cb != NULL) {
+                       fd_cb->flags |= FLOW_DIVERT_HAS_HMAC;
+               }
+       }
+
 done:
        if (token != NULL) {
                mbuf_freem(token);
@@ -2857,12 +3480,12 @@ done:
 errno_t
 flow_divert_token_get(struct socket *so, struct sockopt *sopt)
 {
-       uint32_t                                        ctl_unit;
-       int                                                     error                                           = 0;
-       uint8_t                                         hmac[SHA_DIGEST_LENGTH];
-       struct flow_divert_pcb          *fd_cb                                          = so->so_fd_pcb;
-       mbuf_t                                          token                                           = NULL;
-       struct flow_divert_group        *control_group                          = NULL;
+       uint32_t                                        ctl_unit;
+       int                                                     error                                           = 0;
+       uint8_t                                         hmac[SHA_DIGEST_LENGTH];
+       struct flow_divert_pcb          *fd_cb                                          = so->so_fd_pcb;
+       mbuf_t                                          token                                           = NULL;
+       struct flow_divert_group        *control_group                          = NULL;
 
        if (!(so->so_flags & SOF_FLOW_DIVERT)) {
                error = EINVAL;
@@ -2894,12 +3517,18 @@ flow_divert_token_get(struct socket *so, struct sockopt *sopt)
                goto done;
        }
 
+       if (fd_cb->app_data != NULL) {
+               error = flow_divert_packet_append_tlv(token, FLOW_DIVERT_TLV_APP_DATA, fd_cb->app_data_length, fd_cb->app_data);
+               if (error) {
+                       goto done;
+               }
+       }
+
        socket_unlock(so, 0);
        lck_rw_lock_shared(&g_flow_divert_group_lck);
 
        if (g_flow_divert_groups != NULL && g_active_group_count > 0 &&
-           fd_cb->control_group_unit > 0 && fd_cb->control_group_unit < GROUP_COUNT_MAX)
-       {
+           fd_cb->control_group_unit > 0 && fd_cb->control_group_unit < GROUP_COUNT_MAX) {
                control_group = g_flow_divert_groups[fd_cb->control_group_unit];
        }
 
@@ -2927,9 +3556,15 @@ flow_divert_token_get(struct socket *so, struct sockopt *sopt)
                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 */
+               token = NULL;   /* For some reason, soopt_mcopyout() frees the mbuf if it fails */
                goto done;
        }
 
@@ -2944,8 +3579,8 @@ done:
 static errno_t
 flow_divert_kctl_connect(kern_ctl_ref kctlref __unused, struct sockaddr_ctl *sac, void **unitinfo)
 {
-       struct flow_divert_group        *new_group;
-       int                             error           = 0;
+       struct flow_divert_group        *new_group      = NULL;
+       int                             error           = 0;
 
        if (sac->sc_unit >= GROUP_COUNT_MAX) {
                error = EINVAL;
@@ -2966,15 +3601,16 @@ 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);
+       new_group->signing_id_trie.root = NULL_TRIE_IDX;
 
        lck_rw_lock_exclusive(&g_flow_divert_group_lck);
 
        if (g_flow_divert_groups == NULL) {
                MALLOC(g_flow_divert_groups,
-                      struct flow_divert_group **,
-                      GROUP_COUNT_MAX * sizeof(struct flow_divert_group *),
-                      M_TEMP,
-                      M_WAITOK | M_ZERO);
+                   struct flow_divert_group **,
+                   GROUP_COUNT_MAX * sizeof(struct flow_divert_group *),
+                   M_TEMP,
+                   M_WAITOK | M_ZERO);
        }
 
        if (g_flow_divert_groups == NULL) {
@@ -3000,9 +3636,8 @@ done:
 static errno_t
 flow_divert_kctl_disconnect(kern_ctl_ref kctlref __unused, uint32_t unit, void *unitinfo)
 {
-       struct flow_divert_group        *group  = NULL;
-       errno_t                                         error   = 0;
-       uint16_t                                        node    = 0;
+       struct flow_divert_group        *group  = NULL;
+       errno_t                                         error   = 0;
 
        if (unit >= GROUP_COUNT_MAX) {
                return EINVAL;
@@ -3014,7 +3649,7 @@ flow_divert_kctl_disconnect(kern_ctl_ref kctlref __unused, uint32_t unit, void *
 
        if (g_flow_divert_groups == NULL || g_active_group_count == 0) {
                panic("flow divert group %u is disconnecting, but no groups are active (groups = %p, active count = %u", unit,
-                     g_flow_divert_groups, g_active_group_count);
+                   g_flow_divert_groups, g_active_group_count);
        }
 
        group = g_flow_divert_groups[unit];
@@ -3031,6 +3666,14 @@ flow_divert_kctl_disconnect(kern_ctl_ref kctlref __unused, uint32_t unit, void *
                        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--;
@@ -3043,13 +3686,6 @@ flow_divert_kctl_disconnect(kern_ctl_ref kctlref __unused, uint32_t unit, void *
                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;
@@ -3064,11 +3700,11 @@ flow_divert_kctl_send(kern_ctl_ref kctlref __unused, uint32_t unit __unused, voi
 static void
 flow_divert_kctl_rcvd(kern_ctl_ref kctlref __unused, uint32_t unit __unused, void *unitinfo, int flags __unused)
 {
-       struct flow_divert_group        *group  = (struct flow_divert_group *)unitinfo;
+       struct flow_divert_group        *group  = (struct flow_divert_group *)unitinfo;
 
        if (!OSTestAndClear(GROUP_BIT_CTL_ENQUEUE_BLOCKED, &group->atomic_bits)) {
-               struct flow_divert_pcb                  *fd_cb;
-               SLIST_HEAD(, flow_divert_pcb)   tmp_list;
+               struct flow_divert_pcb                  *fd_cb;
+               SLIST_HEAD(, flow_divert_pcb)   tmp_list;
 
                lck_rw_lock_shared(&g_flow_divert_group_lck);
                lck_rw_lock_exclusive(&group->lck);
@@ -3117,13 +3753,13 @@ flow_divert_kctl_rcvd(kern_ctl_ref kctlref __unused, uint32_t unit __unused, voi
 static int
 flow_divert_kctl_init(void)
 {
-       struct kern_ctl_reg     ctl_reg;
-       int                     result;
+       struct kern_ctl_reg     ctl_reg;
+       int                     result;
 
        memset(&ctl_reg, 0, sizeof(ctl_reg));
 
-       strncpy(ctl_reg.ctl_name, FLOW_DIVERT_CONTROL_NAME, sizeof(ctl_reg.ctl_name));
-       ctl_reg.ctl_name[sizeof(ctl_reg.ctl_name)-1] = '\0';
+       strlcpy(ctl_reg.ctl_name, FLOW_DIVERT_CONTROL_NAME, sizeof(ctl_reg.ctl_name));
+       ctl_reg.ctl_name[sizeof(ctl_reg.ctl_name) - 1] = '\0';
        ctl_reg.ctl_flags = CTL_FLAG_PRIVILEGED | CTL_FLAG_REG_EXTENDED;
        ctl_reg.ctl_sendsize = FD_CTL_SENDBUFF_SIZE;
        ctl_reg.ctl_recvsize = FD_CTL_RCVBUFF_SIZE;
@@ -3147,7 +3783,7 @@ void
 flow_divert_init(void)
 {
        memset(&nil_pcb, 0, sizeof(nil_pcb));
-       nil_pcb.log_level = LOG_INFO;
+       nil_pcb.log_level = LOG_NOTICE;
 
        g_tcp_protosw = pffindproto(AF_INET, IPPROTO_TCP, SOCK_STREAM);
 
@@ -3166,6 +3802,7 @@ flow_divert_init(void)
        g_flow_divert_in_usrreqs.pru_send = flow_divert_data_out;
        g_flow_divert_in_usrreqs.pru_shutdown = flow_divert_shutdown;
        g_flow_divert_in_usrreqs.pru_sockaddr = flow_divert_getsockaddr;
+       g_flow_divert_in_usrreqs.pru_preconnect = flow_divert_preconnect;
 
        g_flow_divert_in_protosw.pr_usrreqs = &g_flow_divert_in_usrreqs;
        g_flow_divert_in_protosw.pr_ctloutput = flow_divert_ctloutput;
@@ -3180,6 +3817,40 @@ flow_divert_init(void)
        g_flow_divert_in_protosw.pr_filter_head.tqh_last =
            (struct socket_filter **)(uintptr_t)0xdeadbeefdeadbeef;
 
+       /* UDP */
+       g_udp_protosw = pffindproto(AF_INET, IPPROTO_UDP, SOCK_DGRAM);
+       VERIFY(g_udp_protosw != NULL);
+
+       memcpy(&g_flow_divert_in_udp_protosw, g_udp_protosw, sizeof(g_flow_divert_in_udp_protosw));
+       memcpy(&g_flow_divert_in_udp_usrreqs, g_udp_protosw->pr_usrreqs, sizeof(g_flow_divert_in_udp_usrreqs));
+
+       g_flow_divert_in_udp_usrreqs.pru_connect = flow_divert_connect_out;
+       g_flow_divert_in_udp_usrreqs.pru_connectx = flow_divert_connectx_out;
+       g_flow_divert_in_udp_usrreqs.pru_control = flow_divert_in_control;
+       g_flow_divert_in_udp_usrreqs.pru_disconnect = flow_divert_close;
+       g_flow_divert_in_udp_usrreqs.pru_disconnectx = flow_divert_disconnectx;
+       g_flow_divert_in_udp_usrreqs.pru_peeraddr = flow_divert_getpeername;
+       g_flow_divert_in_udp_usrreqs.pru_rcvd = flow_divert_rcvd;
+       g_flow_divert_in_udp_usrreqs.pru_send = flow_divert_data_out;
+       g_flow_divert_in_udp_usrreqs.pru_shutdown = flow_divert_shutdown;
+       g_flow_divert_in_udp_usrreqs.pru_sockaddr = flow_divert_getsockaddr;
+       g_flow_divert_in_udp_usrreqs.pru_sosend_list = pru_sosend_list_notsupp;
+       g_flow_divert_in_udp_usrreqs.pru_soreceive_list = pru_soreceive_list_notsupp;
+       g_flow_divert_in_udp_usrreqs.pru_preconnect = flow_divert_preconnect;
+
+       g_flow_divert_in_udp_protosw.pr_usrreqs = &g_flow_divert_in_usrreqs;
+       g_flow_divert_in_udp_protosw.pr_ctloutput = flow_divert_ctloutput;
+
+       /*
+        * Socket filters shouldn't attach/detach to/from this protosw
+        * since pr_protosw is to be used instead, which points to the
+        * real protocol; if they do, it is a bug and we should panic.
+        */
+       g_flow_divert_in_udp_protosw.pr_filter_head.tqh_first =
+           (struct socket_filter *)(uintptr_t)0xdeadbeefdeadbeef;
+       g_flow_divert_in_udp_protosw.pr_filter_head.tqh_last =
+           (struct socket_filter **)(uintptr_t)0xdeadbeefdeadbeef;
+
 #if INET6
        g_tcp6_protosw = (struct ip6protosw *)pffindproto(AF_INET6, IPPROTO_TCP, SOCK_STREAM);
 
@@ -3198,6 +3869,7 @@ flow_divert_init(void)
        g_flow_divert_in6_usrreqs.pru_send = flow_divert_data_out;
        g_flow_divert_in6_usrreqs.pru_shutdown = flow_divert_shutdown;
        g_flow_divert_in6_usrreqs.pru_sockaddr = flow_divert_getsockaddr;
+       g_flow_divert_in6_usrreqs.pru_preconnect = flow_divert_preconnect;
 
        g_flow_divert_in6_protosw.pr_usrreqs = &g_flow_divert_in6_usrreqs;
        g_flow_divert_in6_protosw.pr_ctloutput = flow_divert_ctloutput;
@@ -3210,7 +3882,41 @@ flow_divert_init(void)
            (struct socket_filter *)(uintptr_t)0xdeadbeefdeadbeef;
        g_flow_divert_in6_protosw.pr_filter_head.tqh_last =
            (struct socket_filter **)(uintptr_t)0xdeadbeefdeadbeef;
-#endif /* INET6 */
+
+       /* UDP6 */
+       g_udp6_protosw = (struct ip6protosw *)pffindproto(AF_INET6, IPPROTO_UDP, SOCK_DGRAM);
+
+       VERIFY(g_udp6_protosw != NULL);
+
+       memcpy(&g_flow_divert_in6_udp_protosw, g_udp6_protosw, sizeof(g_flow_divert_in6_udp_protosw));
+       memcpy(&g_flow_divert_in6_udp_usrreqs, g_udp6_protosw->pr_usrreqs, sizeof(g_flow_divert_in6_udp_usrreqs));
+
+       g_flow_divert_in6_udp_usrreqs.pru_connect = flow_divert_connect_out;
+       g_flow_divert_in6_udp_usrreqs.pru_connectx = flow_divert_connectx6_out;
+       g_flow_divert_in6_udp_usrreqs.pru_control = flow_divert_in6_control;
+       g_flow_divert_in6_udp_usrreqs.pru_disconnect = flow_divert_close;
+       g_flow_divert_in6_udp_usrreqs.pru_disconnectx = flow_divert_disconnectx;
+       g_flow_divert_in6_udp_usrreqs.pru_peeraddr = flow_divert_getpeername;
+       g_flow_divert_in6_udp_usrreqs.pru_rcvd = flow_divert_rcvd;
+       g_flow_divert_in6_udp_usrreqs.pru_send = flow_divert_data_out;
+       g_flow_divert_in6_udp_usrreqs.pru_shutdown = flow_divert_shutdown;
+       g_flow_divert_in6_udp_usrreqs.pru_sockaddr = flow_divert_getsockaddr;
+       g_flow_divert_in6_udp_usrreqs.pru_sosend_list = pru_sosend_list_notsupp;
+       g_flow_divert_in6_udp_usrreqs.pru_soreceive_list = pru_soreceive_list_notsupp;
+       g_flow_divert_in6_udp_usrreqs.pru_preconnect = flow_divert_preconnect;
+
+       g_flow_divert_in6_udp_protosw.pr_usrreqs = &g_flow_divert_in6_udp_usrreqs;
+       g_flow_divert_in6_udp_protosw.pr_ctloutput = flow_divert_ctloutput;
+       /*
+        * Socket filters shouldn't attach/detach to/from this protosw
+        * since pr_protosw is to be used instead, which points to the
+        * real protocol; if they do, it is a bug and we should panic.
+        */
+       g_flow_divert_in6_udp_protosw.pr_filter_head.tqh_first =
+           (struct socket_filter *)(uintptr_t)0xdeadbeefdeadbeef;
+       g_flow_divert_in6_udp_protosw.pr_filter_head.tqh_last =
+           (struct socket_filter **)(uintptr_t)0xdeadbeefdeadbeef;
+#endif  /* INET6 */
 
        flow_divert_grp_attr = lck_grp_attr_alloc_init();
        if (flow_divert_grp_attr == NULL) {
@@ -3240,9 +3946,6 @@ flow_divert_init(void)
 
        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) {