1 files changed, 265 insertions, 0 deletions

diff --git a/src/event/ngx_event_udp.c b/src/event/ngx_event_udp.c
index 962b8e37e..8b0ae5539 100644
--- a/src/event/ngx_event_udp.c
+++ b/src/event/ngx_event_udp.c
@@ -12,13 +12,28 @@
 
 #if !(NGX_WIN32)
 
+struct ngx_udp_connection_s {
+    ngx_rbtree_node_t   node;
+    ngx_connection_t   *connection;
+    ngx_buf_t          *buffer;
+};
+
+
 static void ngx_close_accepted_udp_connection(ngx_connection_t *c);
+static ssize_t ngx_udp_shared_recv(ngx_connection_t *c, u_char *buf,
+    size_t size);
+static ngx_int_t ngx_insert_udp_connection(ngx_connection_t *c);
+static void ngx_delete_udp_connection(void *data);
+static ngx_connection_t *ngx_lookup_udp_connection(ngx_listening_t *ls,
+    struct sockaddr *sockaddr, socklen_t socklen,
+    struct sockaddr *local_sockaddr, socklen_t local_socklen);
 
 
 void
 ngx_event_recvmsg(ngx_event_t *ev)
 {
     ssize_t            n;
+    ngx_buf_t          buf;
     ngx_log_t         *log;
     ngx_err_t          err;
     socklen_t          socklen, local_socklen;
@@ -215,6 +230,43 @@ ngx_event_recvmsg(ngx_event_t *ev)
 
 #endif
 
+        c = ngx_lookup_udp_connection(ls, sockaddr, socklen, local_sockaddr,
+                                      local_socklen);
+
+        if (c) {
+
+#if (NGX_DEBUG)
+            if (c->log->log_level & NGX_LOG_DEBUG_EVENT) {
+                ngx_log_handler_pt  handler;
+
+                handler = c->log->handler;
+                c->log->handler = NULL;
+
+                ngx_log_debug2(NGX_LOG_DEBUG_EVENT, c->log, 0,
+                               "recvmsg: fd:%d n:%z", c->fd, n);
+
+                c->log->handler = handler;
+            }
+#endif
+
+            ngx_memzero(&buf, sizeof(ngx_buf_t));
+
+            buf.pos = buffer;
+            buf.last = buffer + n;
+
+            rev = c->read;
+
+            c->udp->buffer = &buf;
+            rev->ready = 1;
+
+            rev->handler(rev);
+
+            c->udp->buffer = NULL;
+            rev->ready = 0;
+
+            goto next;
+        }
+
 #if (NGX_STAT_STUB)
         (void) ngx_atomic_fetch_add(ngx_stat_accepted, 1);
 #endif
@@ -257,6 +309,7 @@ ngx_event_recvmsg(ngx_event_t *ev)
 
         *log = ls->log;
 
+        c->recv = ngx_udp_shared_recv;
         c->send = ngx_udp_send;
         c->send_chain = ngx_udp_send_chain;
 
@@ -344,11 +397,18 @@ ngx_event_recvmsg(ngx_event_t *ev)
         }
 #endif
 
+        if (ngx_insert_udp_connection(c) != NGX_OK) {
+            ngx_close_accepted_udp_connection(c);
+            return;
+        }
+
         log->data = NULL;
         log->handler = NULL;
 
         ls->handler(c);
 
+    next:
+
         if (ngx_event_flags & NGX_USE_KQUEUE_EVENT) {
             ev->available -= n;
         }
@@ -373,4 +433,209 @@ ngx_close_accepted_udp_connection(ngx_connection_t *c)
 #endif
 }
 
+
+static ssize_t
+ngx_udp_shared_recv(ngx_connection_t *c, u_char *buf, size_t size)
+{
+    ssize_t     n;
+    ngx_buf_t  *b;
+
+    if (c->udp == NULL || c->udp->buffer == NULL) {
+        return NGX_AGAIN;
+    }
+
+    b = c->udp->buffer;
+
+    n = ngx_min(b->last - b->pos, (ssize_t) size);
+
+    ngx_memcpy(buf, b->pos, n);
+
+    c->udp->buffer = NULL;
+    c->read->ready = 0;
+
+    return n;
+}
+
+
+void
+ngx_udp_rbtree_insert_value(ngx_rbtree_node_t *temp,
+    ngx_rbtree_node_t *node, ngx_rbtree_node_t *sentinel)
+{
+    ngx_int_t               rc;
+    ngx_connection_t       *c, *ct;
+    ngx_rbtree_node_t     **p;
+    ngx_udp_connection_t   *udp, *udpt;
+
+    for ( ;; ) {
+
+        if (node->key < temp->key) {
+
+            p = &temp->left;
+
+        } else if (node->key > temp->key) {
+
+            p = &temp->right;
+
+        } else { /* node->key == temp->key */
+
+            udp = (ngx_udp_connection_t *) node;
+            c = udp->connection;
+
+            udpt = (ngx_udp_connection_t *) temp;
+            ct = udpt->connection;
+
+            rc = ngx_cmp_sockaddr(c->sockaddr, c->socklen,
+                                  ct->sockaddr, ct->socklen, 1);
+
+            if (rc == 0 && c->listening->wildcard) {
+                rc = ngx_cmp_sockaddr(c->local_sockaddr, c->local_socklen,
+                                      ct->local_sockaddr, ct->local_socklen, 1);
+            }
+
+            p = (rc < 0) ? &temp->left : &temp->right;
+        }
+
+        if (*p == sentinel) {
+            break;
+        }
+
+        temp = *p;
+    }
+
+    *p = node;
+    node->parent = temp;
+    node->left = sentinel;
+    node->right = sentinel;
+    ngx_rbt_red(node);
+}
+
+
+static ngx_int_t
+ngx_insert_udp_connection(ngx_connection_t *c)
+{
+    uint32_t               hash;
+    ngx_pool_cleanup_t    *cln;
+    ngx_udp_connection_t  *udp;
+
+    if (c->udp) {
+        return NGX_OK;
+    }
+
+    udp = ngx_pcalloc(c->pool, sizeof(ngx_udp_connection_t));
+    if (udp == NULL) {
+        return NGX_ERROR;
+    }
+
+    udp->connection = c;
+
+    ngx_crc32_init(hash);
+    ngx_crc32_update(&hash, (u_char *) c->sockaddr, c->socklen);
+
+    if (c->listening->wildcard) {
+        ngx_crc32_update(&hash, (u_char *) c->local_sockaddr, c->local_socklen);
+    }
+
+    ngx_crc32_final(hash);
+
+    udp->node.key = hash;
+
+    cln = ngx_pool_cleanup_add(c->pool, 0);
+    if (cln == NULL) {
+        return NGX_ERROR;
+    }
+
+    cln->data = c;
+    cln->handler = ngx_delete_udp_connection;
+
+    ngx_rbtree_insert(&c->listening->rbtree, &udp->node);
+
+    c->udp = udp;
+
+    return NGX_OK;
+}
+
+
+static void
+ngx_delete_udp_connection(void *data)
+{
+    ngx_connection_t  *c = data;
+
+    ngx_rbtree_delete(&c->listening->rbtree, &c->udp->node);
+}
+
+
+static ngx_connection_t *
+ngx_lookup_udp_connection(ngx_listening_t *ls, struct sockaddr *sockaddr,
+    socklen_t socklen, struct sockaddr *local_sockaddr, socklen_t local_socklen)
+{
+    uint32_t               hash;
+    ngx_int_t              rc;
+    ngx_connection_t      *c;
+    ngx_rbtree_node_t     *node, *sentinel;
+    ngx_udp_connection_t  *udp;
+
+#if (NGX_HAVE_UNIX_DOMAIN)
+
+    if (sockaddr->sa_family == AF_UNIX) {
+        struct sockaddr_un *saun = (struct sockaddr_un *) sockaddr;
+
+        if (socklen <= (socklen_t) offsetof(struct sockaddr_un, sun_path)
+            || saun->sun_path[0] == '\0')
+        {
+            ngx_log_debug0(NGX_LOG_DEBUG_EVENT, ngx_cycle->log, 0,
+                           "unbound unix socket");
+            return NULL;
+        }
+    }
+
+#endif
+
+    node = ls->rbtree.root;
+    sentinel = ls->rbtree.sentinel;
+
+    ngx_crc32_init(hash);
+    ngx_crc32_update(&hash, (u_char *) sockaddr, socklen);
+
+    if (ls->wildcard) {
+        ngx_crc32_update(&hash, (u_char *) local_sockaddr, local_socklen);
+    }
+
+    ngx_crc32_final(hash);
+
+    while (node != sentinel) {
+
+        if (hash < node->key) {
+            node = node->left;
+            continue;
+        }
+
+        if (hash > node->key) {
+            node = node->right;
+            continue;
+        }
+
+        /* hash == node->key */
+
+        udp = (ngx_udp_connection_t *) node;
+
+        c = udp->connection;
+
+        rc = ngx_cmp_sockaddr(sockaddr, socklen,
+                              c->sockaddr, c->socklen, 1);
+
+        if (rc == 0 && ls->wildcard) {
+            rc = ngx_cmp_sockaddr(local_sockaddr, local_socklen,
+                                  c->local_sockaddr, c->local_socklen, 1);
+        }
+
+        if (rc == 0) {
+            return c;
+        }
+
+        node = (rc < 0) ? node->left : node->right;
+    }
+
+    return NULL;
+}
+
 #endif