]> git.saurik.com Git - apple/xnu.git/blobdiff - bsd/netinet/flow_divert.c
xnu-4903.241.1.tar.gz
[apple/xnu.git] / bsd / netinet / flow_divert.c
index 697016c4996a806e54d433e36391091f817c69d4..c3bc7930cd23061355fd8c48ad1a0f277e38df2f 100644 (file)
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2012-2014 Apple Inc. All rights reserved.
+ * Copyright (c) 2012-2017 Apple Inc. All rights reserved.
  *
  * @APPLE_OSREFERENCE_LICENSE_HEADER_START@
  * 
 #include <netinet/flow_divert.h>
 #include <netinet/flow_divert_proto.h>
 #if INET6
+#include <netinet6/in6_pcb.h>
 #include <netinet6/ip6protosw.h>
 #endif /* INET6 */
 #include <dev/random/randomdev.h>
 #include <libkern/crypto/sha1.h>
 #include <libkern/crypto/crypto_internal.h>
+#include <os/log.h>
 
 #define FLOW_DIVERT_CONNECT_STARTED            0x00000001
 #define FLOW_DIVERT_READ_CLOSED                        0x00000002
 #define FLOW_DIVERT_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, ...) do {                                                                                    \
-       if (level <= (pcb)->log_level) {                                                                                                \
-               log((level > LOG_NOTICE ? LOG_NOTICE : level), "%s (%u): " format "\n", __FUNCTION__, (pcb)->hash, __VA_ARGS__);        \
-       }                                                                                                                                                               \
-} while (0)
+#define FDLOG(level, pcb, format, ...) \
+       os_log_with_type(OS_LOG_DEFAULT, flow_divert_syslog_type_to_oslog_type(level), "(%u): " format "\n", (pcb)->hash, __VA_ARGS__)
 
-#define FDLOG0(level, pcb, msg) do {                                                                                           \
-       if (level <= (pcb)->log_level) {                                                                                                \
-               log((level > LOG_NOTICE ? LOG_NOTICE : level), "%s (%u): %s\n", __FUNCTION__, (pcb)->hash, msg);                                \
-       }                                                                                                                                                               \
-} while (0)
+#define FDLOG0(level, pcb, msg) \
+       os_log_with_type(OS_LOG_DEFAULT, flow_divert_syslog_type_to_oslog_type(level), "(%u): " msg "\n", (pcb)->hash)
 
 #define FDRETAIN(pcb)                  if ((pcb) != NULL) OSIncrementAtomic(&(pcb)->ref_count)
 #define FDRELEASE(pcb)                                                                                                         \
@@ -93,7 +90,7 @@
 #define FDLOCK(pcb)                                            lck_mtx_lock(&(pcb)->mtx)
 #define FDUNLOCK(pcb)                                  lck_mtx_unlock(&(pcb)->mtx)
 
-#define FD_CTL_SENDBUFF_SIZE                   (2 * FLOW_DIVERT_CHUNK_SIZE)
+#define FD_CTL_SENDBUFF_SIZE                   (128 * 1024)
 #define FD_CTL_RCVBUFF_SIZE                            (128 * 1024)
 
 #define GROUP_BIT_CTL_ENQUEUE_BLOCKED  0
 #define GROUP_COUNT_MAX                                        32
 #define FLOW_DIVERT_MAX_NAME_SIZE              4096
 #define FLOW_DIVERT_MAX_KEY_SIZE               1024
-
-#define DNS_SERVICE_GROUP_UNIT                 (GROUP_COUNT_MAX + 1)
+#define FLOW_DIVERT_MAX_TRIE_MEMORY            (1024 * 1024)
 
 struct flow_divert_trie_node
 {
        uint16_t start;
        uint16_t length;
        uint16_t child_map;
-       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
@@ -138,7 +118,6 @@ static struct flow_divert_pcb               nil_pcb;
 decl_lck_rw_data(static, g_flow_divert_group_lck);
 static struct flow_divert_group                **g_flow_divert_groups                  = NULL;
 static uint32_t                                                g_active_group_count                    = 0;
-static struct flow_divert_trie         g_signing_id_trie;
 
 static lck_grp_attr_t                          *flow_divert_grp_attr                   = NULL;
 static lck_attr_t                                      *flow_divert_mtx_attr                   = NULL;
@@ -149,13 +128,51 @@ 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_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;
+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 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)
@@ -182,8 +199,6 @@ flow_divert_packet_type2str(uint8_t packet_type)
                        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:
@@ -210,12 +225,11 @@ 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;
+       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;
-       errno_t                                         result                                          = 0;
        int                                                     try_count                                       = 0;
 
        if (ctl_unit == 0 || ctl_unit >= GROUP_COUNT_MAX) {
@@ -277,7 +291,7 @@ flow_divert_pcb_insert(struct flow_divert_pcb *fd_cb, uint32_t ctl_unit)
                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,7 +300,7 @@ done:
        lck_rw_done(&g_flow_divert_group_lck);
        socket_lock(fd_cb->so, 0);
 
-       return result;
+       return error;
 }
 
 static struct flow_divert_pcb *
@@ -326,6 +340,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);
 }
 
@@ -371,10 +391,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 +404,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 +420,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 +455,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) {
@@ -581,6 +601,7 @@ flow_divert_add_data_statistics(struct flow_divert_pcb *fd_cb, int data_len, Boo
                INP_ADD_STAT(inp, cell, wifi, wired, rxpackets, 1);
                INP_ADD_STAT(inp, cell, wifi, wired, rxbytes, data_len);
        }
+       inp_set_activity_bitmap(inp);
 }
 
 static errno_t
@@ -778,8 +799,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;
@@ -796,6 +818,10 @@ flow_divert_trie_search(struct flow_divert_trie *trie, const uint8_t *string_byt
                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]);
                        }
@@ -806,17 +832,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();
        }
@@ -841,7 +956,7 @@ flow_divert_send_packet(struct flow_divert_pcb *fd_cb, mbuf_t packet, Boolean en
 
        if (fd_cb->group == NULL) {
                fd_cb->so->so_error = ECONNABORTED;
-               soisdisconnected(fd_cb->so);
+               flow_divert_disconnect_socket(fd_cb->so);
                return ECONNABORTED;
        }
 
@@ -870,9 +985,111 @@ 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, mbuf_t connect_packet)
+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)
 {
        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,
@@ -882,6 +1099,23 @@ flow_divert_send_connect(struct flow_divert_pcb *fd_cb, struct sockaddr *to, mbu
                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_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,
@@ -923,39 +1157,56 @@ flow_divert_send_connect(struct flow_divert_pcb *fd_cb, struct sockaddr *to, mbu
                fd_cb->connect_token = NULL;
        } else {
                uint32_t ctl_unit = htonl(fd_cb->control_group_unit);
-               int port;
 
                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;
                }
        }
 
-       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) {
+               *out_connect_packet = connect_packet;
+       } else if (connect_packet != NULL) {
+               mbuf_freem(connect_packet);
+       }
+
        return error;
 }
 
@@ -972,7 +1223,7 @@ 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;
        }
@@ -992,7 +1243,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;
@@ -1092,12 +1343,12 @@ 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;
@@ -1109,15 +1360,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);
@@ -1147,28 +1407,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) {
@@ -1178,19 +1486,14 @@ 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;
+       int     error           = 0;
 
        if (to_send > fd_cb->send_window) {
                to_send = fd_cb->send_window;
@@ -1200,57 +1503,94 @@ flow_divert_send_app_data(struct flow_divert_pcb *fd_cb, mbuf_t data)
                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;
 
-               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;
+                       }
+
+                       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;
+                       }
                }
        }
 
@@ -1331,6 +1671,7 @@ flow_divert_handle_connect_result(struct flow_divert_pcb *fd_cb, mbuf_t packet,
        int                                                     out_if_index            = 0;
        struct sockaddr_storage         remote_address;
        uint32_t                                        send_window;
+       uint32_t                                        app_data_length         = 0;
 
        memset(&local_address, 0, sizeof(local_address));
        memset(&remote_address, 0, sizeof(remote_address));
@@ -1351,32 +1692,37 @@ flow_divert_handle_connect_result(struct flow_divert_pcb *fd_cb, mbuf_t packet,
 
        error = flow_divert_packet_get_tlv(packet, offset, FLOW_DIVERT_TLV_CTL_UNIT, sizeof(ctl_unit), &ctl_unit, NULL);
        if (error) {
-               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");
        }
 
+       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) {
@@ -1408,14 +1754,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) {
@@ -1428,6 +1775,27 @@ flow_divert_handle_connect_result(struct flow_divert_pcb *fd_cb, mbuf_t packet,
                        goto set_socket_state;
                }
 
+               if (app_data_length > 0) {
+                       uint8_t *app_data = NULL;
+                       MALLOC(app_data, uint8_t *, app_data_length, M_TEMP, M_WAITOK);
+                       if (app_data != NULL) {
+                               error = flow_divert_packet_get_tlv(packet, offset, FLOW_DIVERT_TLV_APP_DATA, app_data_length, app_data, NULL);
+                               if (error == 0) {
+                                       FDLOG(LOG_INFO, fd_cb, "Got %u bytes of app data from the connect result", app_data_length);
+                                       if (fd_cb->app_data != NULL) {
+                                               FREE(fd_cb->app_data, M_TEMP);
+                                       }
+                                       fd_cb->app_data = app_data;
+                                       fd_cb->app_data_length = app_data_length;
+                               } else {
+                                       FDLOG(LOG_ERR, fd_cb, "Failed to copy %u bytes of application data from the connect result packet", app_data_length);
+                                       FREE(app_data, M_TEMP);
+                               }
+                       } else {
+                               FDLOG(LOG_ERR, fd_cb, "Failed to allocate a buffer of size %u to hold the application data from the connect result", app_data_length);
+                       }
+               }
+
                ifnet_head_lock_shared();
                if (out_if_index > 0 && out_if_index <= if_index) {
                        ifp = ifindex2ifnet[out_if_index];
@@ -1449,23 +1817,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) {
@@ -1482,8 +1851,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);
                }
 
@@ -1528,7 +1898,7 @@ flow_divert_handle_close(struct flow_divert_pcb *fd_cb, mbuf_t packet, int offse
                
                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) {
@@ -1540,49 +1910,117 @@ 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) {
+                       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) != 0) {
+                       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_check_no_expensive(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
@@ -1597,7 +2035,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) {
@@ -1613,8 +2051,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) {
@@ -1623,7 +2062,7 @@ flow_divert_handle_group_init(struct flow_divert_group *group, mbuf_t packet, in
        }
 
        if (key_size == 0 || key_size > FLOW_DIVERT_MAX_KEY_SIZE) {
-               FDLOG(LOG_ERR, &nil_pcb, "Invalid key size: %lu", key_size);
+               FDLOG(LOG_ERR, &nil_pcb, "Invalid key size: %u", key_size);
                return;
        }
 
@@ -1646,6 +2085,11 @@ 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);
 }
 
@@ -1656,6 +2100,7 @@ flow_divert_handle_properties_update(struct flow_divert_pcb *fd_cb, mbuf_t packe
        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");
 
@@ -1664,32 +2109,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);
                }
 
@@ -1697,18 +2146,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);
        }
@@ -1716,7 +2196,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;
@@ -1727,34 +2207,40 @@ 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;
+               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;
        }
 
@@ -1768,10 +2254,18 @@ flow_divert_handle_app_map_create(mbuf_t packet, int offset)
        child_maps_mem_size = (sizeof(*new_trie.child_maps) * CHILD_MAP_SIZE * new_trie.child_maps_count);
        bytes_mem_size = (sizeof(*new_trie.bytes) * new_trie.bytes_count);
 
-       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);
+               lck_rw_done(&group->lck);
                return;
        }
 
@@ -1795,23 +2289,13 @@ flow_divert_handle_app_map_create(mbuf_t packet, int offset)
             cursor >= 0;
             cursor = flow_divert_packet_find_tlv(packet, cursor, FLOW_DIVERT_TLV_SIGNING_ID, &error, 1))
        {
-               size_t sid_size = 0;
+               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_INFO, &nil_pcb, "Setting group unit for %s to %d", FLOW_DIVERT_DNS_SERVICE_SIGNING_ID, DNS_SERVICE_GROUP_UNIT);
-                                       TRIE_NODE(&new_trie, new_node_idx).group_unit = DNS_SERVICE_GROUP_UNIT;
-                               }
-                       } else {
+                       if (new_node_idx == NULL_TRIE_IDX) {
                                insert_error = EINVAL;
                                break;
                        }
@@ -1823,72 +2307,12 @@ 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
@@ -1904,6 +2328,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);
@@ -1919,10 +2349,7 @@ flow_divert_input(mbuf_t packet, struct flow_divert_group *group)
                                flow_divert_handle_group_init(group, packet, sizeof(hdr));
                                break;
                        case FLOW_DIVERT_PKT_APP_MAP_CREATE:
-                               flow_divert_handle_app_map_create(packet, sizeof(hdr));
-                               break;
-                       case FLOW_DIVERT_PKT_APP_MAP_UPDATE:
-                               flow_divert_handle_app_map_update(group, packet, sizeof(hdr));
+                               flow_divert_handle_app_map_create(group, packet, sizeof(hdr));
                                break;
                        default:
                                FDLOG(LOG_WARNING, &nil_pcb, "got an unknown message type: %d", hdr.packet_type);
@@ -1963,7 +2390,7 @@ flow_divert_input(mbuf_t packet, struct flow_divert_group *group)
        FDRELEASE(fd_cb);
 
 done:
-       mbuf_free(packet);
+       mbuf_freem(packet);
        return error;
 }
 
@@ -1995,6 +2422,7 @@ flow_divert_close_all(struct flow_divert_group *group)
                        flow_divert_pcb_remove(fd_cb);
                        flow_divert_update_closed_state(fd_cb, SHUT_RDWR, TRUE);
                        fd_cb->so->so_error = ECONNABORTED;
+                       flow_divert_disconnect_socket(fd_cb->so);
                        socket_unlock(fd_cb->so, 0);
                }
                FDUNLOCK(fd_cb);
@@ -2018,6 +2446,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);
        }
@@ -2040,8 +2470,10 @@ flow_divert_close(struct socket *so)
 
        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);
@@ -2054,9 +2486,10 @@ 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) {
+       if (aid != SAE_ASSOCID_ANY && aid != SAE_ASSOCID_ALL) {
                return (EINVAL);
        }
 
@@ -2108,6 +2541,106 @@ 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)
@@ -2145,6 +2678,25 @@ 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)
 {
@@ -2201,8 +2753,7 @@ flow_divert_connect_out(struct socket *so, struct sockaddr *to, proc_t p)
        struct inpcb                    *inp    = sotoinpcb(so);
        struct sockaddr_in              *sinp;
        mbuf_t                                  connect_packet = NULL;
-       char                                    *signing_id = NULL;
-       int                                             free_signing_id = 0;
+       int                                             do_send = 1;
 
        VERIFY((so->so_flags & SOF_FLOW_DIVERT) && so->so_fd_pcb != NULL);
 
@@ -2224,12 +2775,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;
@@ -2244,125 +2789,61 @@ flow_divert_connect_out(struct socket *so, struct sockaddr *to, proc_t p)
                }
        }
 
-       error = flow_divert_packet_init(fd_cb, FLOW_DIVERT_PKT_CONNECT, &connect_packet);
-       if (error) {
-               goto done;
-       }
-
-       error = EPERM;
+       FDLOG0(LOG_INFO, fd_cb, "Connecting");
 
-       if (fd_cb->connect_token != NULL) {
-               size_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;
-                       }
+       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;
                }
-       }
 
-       socket_unlock(so, 0);
-       if (g_signing_id_trie.root != NULL_TRIE_IDX) {
-               proc_t src_proc = p;
-               int release_proc = 0;
-                       
-               if (signing_id == NULL) {
-                       release_proc = flow_divert_get_src_proc(so, &src_proc, FALSE);
-                       if (src_proc != PROC_NULL) {
-                               proc_lock(src_proc);
-                               if (src_proc->p_csflags & CS_VALID) {
-                                       signing_id = (char *)cs_identity_get(src_proc);
-                               } 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;
+               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 (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);
-                       lck_rw_done(&g_flow_divert_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");
+               error = flow_divert_create_connect_packet(fd_cb, to, so, p, &connect_packet);
+               if (error) {
+                       goto done;
                }
 
-               if (src_proc != PROC_NULL) {
-                       proc_unlock(src_proc);
-                       if (release_proc) {
-                               proc_rele(src_proc);
-                       }
+               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_WARNING, fd_cb, "The signing ID trie is empty");
-       }
-       socket_lock(so, 0);
-
-       if (free_signing_id) {
-               FREE(signing_id, M_TEMP);
-       }
-
-       if (error) {
-               goto done;
+               FDLOG0(LOG_INFO, fd_cb, "Sending saved connect packet");
+               connect_packet = fd_cb->connect_packet;
+               fd_cb->connect_packet = NULL;
        }
 
-       FDLOG0(LOG_INFO, fd_cb, "Connecting");
+       if (do_send) {
+               error = flow_divert_send_packet(fd_cb, connect_packet, TRUE);
+               if (error) {
+                       goto done;
+               }
 
-       error = flow_divert_send_connect(fd_cb, to, connect_packet);
-       if (error) {
-               goto done;
+               fd_cb->flags |= FLOW_DIVERT_CONNECT_STARTED;
+       } else {
+               fd_cb->connect_packet = connect_packet;
+               connect_packet = NULL;
        }
 
-       fd_cb->flags |= FLOW_DIVERT_CONNECT_STARTED;
-
        soisconnecting(so);
 
 done:
        if (error && connect_packet != NULL) {
-               mbuf_free(connect_packet);
+               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;
 
@@ -2370,20 +2851,42 @@ flow_divert_connectx_out_common(struct socket *so, int af,
                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);
 
-       error = flow_divert_connect_out(so, dst_se->se_addr, p);
+               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);
+               }
+
+               /*
+                * 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 */
@@ -2393,29 +2896,27 @@ flow_divert_connectx_out_common(struct socket *so, int af,
 }
 
 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,
+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)
@@ -2432,7 +2933,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;
        }
@@ -2563,7 +3064,7 @@ 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;
@@ -2596,16 +3097,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;
        }
@@ -2618,7 +3124,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);
@@ -2626,6 +3132,30 @@ done:
        return error;
 }
 
+static int
+flow_divert_preconnect(struct socket *so)
+{
+       struct flow_divert_pcb  *fd_cb  = so->so_fd_pcb;
+       int error = 0;
+
+       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;
+
+               error = flow_divert_send_packet(fd_cb, connect_packet, TRUE);
+               if (error) {
+                       mbuf_freem(connect_packet);
+               }
+
+               fd_cb->flags |= FLOW_DIVERT_CONNECT_STARTED;
+       }
+
+       soclearfastopen(so);
+
+       return error;
+}
+
 static void
 flow_divert_set_protosw(struct socket *so)
 {
@@ -2640,6 +3170,20 @@ flow_divert_set_protosw(struct socket *so)
 #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)
 {
@@ -2679,10 +3223,14 @@ flow_divert_attach(struct socket *so, uint32_t flow_id, uint32_t ctl_unit)
        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;
@@ -2720,6 +3268,44 @@ done:
        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;
+}
+
 errno_t
 flow_divert_pcb_init(struct socket *so, uint32_t ctl_unit)
 {
@@ -2737,11 +3323,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");
                }
@@ -2759,6 +3348,7 @@ flow_divert_token_set(struct socket *so, struct sockopt *sopt)
        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) {
@@ -2772,8 +3362,8 @@ 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
@@ -2783,26 +3373,33 @@ flow_divert_token_set(struct socket *so, struct sockopt *sopt)
                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;
@@ -2825,11 +3422,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;
        }
 
@@ -2859,6 +3457,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);
@@ -2907,6 +3512,13 @@ 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);
 
@@ -2940,6 +3552,12 @@ 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 */
@@ -2957,7 +3575,7 @@ 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;
+       struct flow_divert_group        *new_group      = NULL;
        int                             error           = 0;
 
        if (sac->sc_unit >= GROUP_COUNT_MAX) {
@@ -2979,6 +3597,7 @@ flow_divert_kctl_connect(kern_ctl_ref kctlref __unused, struct sockaddr_ctl *sac
        RB_INIT(&new_group->pcb_tree);
        new_group->ctl_unit = sac->sc_unit;
        MBUFQ_INIT(&new_group->send_queue);
+       new_group->signing_id_trie.root = NULL_TRIE_IDX;
 
        lck_rw_lock_exclusive(&g_flow_divert_group_lck);
 
@@ -3015,7 +3634,6 @@ flow_divert_kctl_disconnect(kern_ctl_ref kctlref __unused, uint32_t unit, void *
 {
        struct flow_divert_group        *group  = NULL;
        errno_t                                         error   = 0;
-       uint16_t                                        node    = 0;
 
        if (unit >= GROUP_COUNT_MAX) {
                return EINVAL;
@@ -3044,6 +3662,14 @@ flow_divert_kctl_disconnect(kern_ctl_ref kctlref __unused, uint32_t unit, void *
                        group->token_key = NULL;
                        group->token_key_size = 0;
                }
+
+               /* 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--;
@@ -3056,13 +3682,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;
@@ -3179,6 +3798,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;
@@ -3193,6 +3813,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);
 
@@ -3211,6 +3865,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;
@@ -3223,6 +3878,40 @@ 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;
+
+       /* 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();
@@ -3253,9 +3942,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) {