int nn_trie_subscribe (struct nn_trie *self, const uint8_t *data, size_t size)
{
int i;
struct nn_trie_node **node;
struct nn_trie_node **n;
struct nn_trie_node *ch;
struct nn_trie_node *old_node;
int pos;
uint8_t c;
uint8_t c2;
uint8_t new_min;
uint8_t new_max;
int old_children;
int new_children;
int inserted;
int more_nodes;
/* Step 1 -- Traverse the trie. */
node = &self->root;
pos = 0;
while (1) {
/* If there are no more nodes on the path, go to step 4. */
if (!*node)
goto step4;
/* Check whether prefix matches the new subscription. */
pos = nn_node_check_prefix (*node, data, size);
data += pos;
size -= pos;
/* If only part of the prefix matches, go to step 2. */
if (pos < (*node)->prefix_len)
goto step2;
/* Even if whole prefix matches and there's no more data to match,
go directly to step 5. */
if (!size)
goto step5;
/* Move to the next node. If it is not present, go to step 3. */
n = nn_node_next (*node, *data);
if (!n || !*n)
goto step3;
node = n;
++data;
--size;
}
/* Step 2 -- Split the prefix into two parts if required. */
step2:
ch = *node;
*node = nn_alloc (sizeof (struct nn_trie_node) +
sizeof (struct nn_trie_node*), "trie node");
assert (*node);
(*node)->refcount = 0;
(*node)->prefix_len = pos;
(*node)->type = 1;
memcpy ((*node)->prefix, ch->prefix, pos);
(*node)->u.sparse.children [0] = ch->prefix [pos];
ch->prefix_len -= (pos + 1);
memmove (ch->prefix, ch->prefix + pos + 1, ch->prefix_len);
ch = nn_node_compact (ch);
*nn_node_child (*node, 0) = ch;
pos = (*node)->prefix_len;
/* Step 3 -- Adjust the child array to accommodate the new character. */
step3:
/* If there are no more data in the subscription, there's nothing to
adjust in the child array. Proceed directly to the step 5. */
if (!size)
goto step5;
/* If the new branch fits into sparse array... */
if ((*node)->type < NN_TRIE_SPARSE_MAX) {
*node = nn_realloc (*node, sizeof (struct nn_trie_node) +
((*node)->type + 1) * sizeof (struct nn_trie_node*));
assert (*node);
(*node)->u.sparse.children [(*node)->type] = *data;
++(*node)->type;
node = nn_node_child (*node, (*node)->type - 1);
*node = NULL;
++data;
--size;
goto step4;
}
/* If the node is already a dense array, resize it to fit the next
character. */
if ((*node)->type == NN_TRIE_DENSE_TYPE) {
c = *data;
if (c < (*node)->u.dense.min || c > (*node)->u.dense.max) {
new_min = (*node)->u.dense.min < c ? (*node)->u.dense.min : c;
new_max = (*node)->u.dense.max > c ? (*node)->u.dense.max : c;
*node = nn_realloc (*node, sizeof (struct nn_trie_node) +
(new_max - new_min + 1) * sizeof (struct nn_trie_node*));
assert (*node);
old_children = (*node)->u.dense.max - (*node)->u.dense.min + 1;
new_children = new_max - new_min + 1;
if ((*node)->u.dense.min != new_min) {
inserted = (*node)->u.dense.min - new_min;
memmove (nn_node_child (*node, inserted),
nn_node_child (*node, 0),
old_children * sizeof (struct nn_trie_node*));
memset (nn_node_child (*node, 0), 0,
inserted * sizeof (struct nn_trie_node*));
}
else {
memset (nn_node_child (*node, old_children), 0,
(new_children - old_children) *
sizeof (struct nn_trie_node*));
}
(*node)->u.dense.min = new_min;
(*node)->u.dense.max = new_max;
}
++(*node)->u.dense.nbr;
node = nn_node_child (*node, c - (*node)->u.dense.min);
++data;
--size;
goto step4;
}
/* This is a sparse array, but no more children can be added to it.
We have to convert it into a dense array. */
{
/* First, determine the range of children. */
new_min = 255;
new_max = 0;
for (i = 0; i != (*node)->type; ++i) {
c2 = (*node)->u.sparse.children [i];
new_min = new_min < c2 ? new_min : c2;
new_max = new_max > c2 ? new_max : c2;
}
new_min = new_min < *data ? new_min : *data;
new_max = new_max > *data ? new_max : *data;
/* Create a new mode, while keeping the old one for a while. */
old_node = *node;
*node = (struct nn_trie_node*) nn_alloc (sizeof (struct nn_trie_node) +
(new_max - new_min + 1) * sizeof (struct nn_trie_node*),
"trie node");
assert (*node);
/* Fill in the new node. */
(*node)->refcount = 0;
(*node)->prefix_len = old_node->prefix_len;
(*node)->type = NN_TRIE_DENSE_TYPE;
memcpy ((*node)->prefix, old_node->prefix, old_node->prefix_len);
(*node)->u.dense.min = new_min;
(*node)->u.dense.max = new_max;
(*node)->u.dense.nbr = old_node->type + 1;
memset (*node + 1, 0, (new_max - new_min + 1) *
sizeof (struct nn_trie_node*));
for (i = 0; i != old_node->type; ++i)
*nn_node_child (*node, old_node->u.sparse.children [i] - new_min) =
*nn_node_child (old_node, i);
node = nn_node_next (*node, *data);
++data;
--size;
/* Get rid of the obsolete old node. */
nn_free (old_node);
}
/* Step 4 -- Create new nodes for remaining part of the subscription. */
step4:
assert (!*node);
while (1) {
/* Create a new node to hold the next part of the subscription. */
more_nodes = size > NN_TRIE_PREFIX_MAX;
*node = nn_alloc (sizeof (struct nn_trie_node) +
(more_nodes ? sizeof (struct nn_trie_node*) : 0), "trie node");
assert (*node);
/* Fill in the new node. */
(*node)->refcount = 0;
(*node)->type = more_nodes ? 1 : 0;
(*node)->prefix_len = size < (size_t) NN_TRIE_PREFIX_MAX ?
size : (size_t) NN_TRIE_PREFIX_MAX;
memcpy ((*node)->prefix, data, (*node)->prefix_len);
data += (*node)->prefix_len;
size -= (*node)->prefix_len;
if (!more_nodes)
break;
(*node)->u.sparse.children [0] = *data;
node = nn_node_child (*node, 0);
++data;
--size;
}
/* Step 5 -- Create the subscription as such. */
step5:
++(*node)->refcount;
/* Return 1 in case of a fresh subscription. */
return (*node)->refcount == 1 ? 1 : 0;
}
static int nn_node_unsubscribe (struct nn_trie_node **self,
const uint8_t *data, size_t size)
{
int i;
int j;
int index;
int new_min;
struct nn_trie_node **ch;
struct nn_trie_node *new_node;
struct nn_trie_node *ch2;
if (!size)
goto found;
/* If prefix does not match the data, return. */
if (nn_node_check_prefix (*self, data, size) != (*self)->prefix_len)
return 0;
/* Skip the prefix. */
data += (*self)->prefix_len;
size -= (*self)->prefix_len;
if (!size)
goto found;
/* Move to the next node. */
ch = nn_node_next (*self, *data);
if (!ch)
return 0; /* TODO: This should be an error. */
/* Recursive traversal of the trie happens here. If the subscription
wasn't really removed, nothing have changed in the trie and
no additional pruning is needed. */
if (nn_node_unsubscribe (ch, data + 1, size - 1) == 0)
return 0;
/* Subscription removal is already done. Now we are going to compact
the trie. However, if the following node remains in place, there's
nothing to compact here. */
if (*ch)
return 1;
/* Sparse array. */
if ((*self)->type < NN_TRIE_DENSE_TYPE) {
/* Get the indices of the removed child. */
for (index = 0; index != (*self)->type; ++index)
if ((*self)->u.sparse.children [index] == *data)
break;
assert (index != (*self)->type);
/* Remove the destroyed child from both lists of children. */
memmove (
(*self)->u.sparse.children + index,
(*self)->u.sparse.children + index + 1,
(*self)->type - index - 1);
memmove (
nn_node_child (*self, index),
nn_node_child (*self, index + 1),
((*self)->type - index - 1) * sizeof (struct nn_trie_node*));
--(*self)->type;
*self = nn_realloc (*self, sizeof (struct nn_trie_node) +
((*self)->type * sizeof (struct nn_trie_node*)));
assert (*self);
/* If there are no more children and no refcount, we can delete
the node altogether. */
if (!(*self)->type && !nn_node_has_subscribers (*self)) {
nn_free (*self);
*self = NULL;
return 1;
}
/* Try to merge the node with the following node. */
*self = nn_node_compact (*self);
return 1;
}
/* Dense array. */
/* In this case the array stays dense. We have to adjust the limits of
the array, if appropriate. */
if ((*self)->u.dense.nbr > NN_TRIE_SPARSE_MAX + 1) {
/* If the removed item is the leftmost one, trim the array from
the left side. */
if (*data == (*self)->u.dense.min) {
for (i = 0; i != (*self)->u.dense.max - (*self)->u.dense.min + 1;
++i)
if (*nn_node_child (*self, i))
break;
new_min = i + (*self)->u.dense.min;
memmove (nn_node_child (*self, 0), nn_node_child (*self, i),
((*self)->u.dense.max - new_min + 1) *
sizeof (struct nn_trie_node*));
(*self)->u.dense.min = new_min;
--(*self)->u.dense.nbr;
*self = nn_realloc (*self, sizeof (struct nn_trie_node) +
((*self)->u.dense.max - new_min + 1) *
sizeof (struct nn_trie_node*));
assert (*self);
return 1;
}
/* If the removed item is the rightmost one, trim the array from
the right side. */
if (*data == (*self)->u.dense.max) {
for (i = (*self)->u.dense.max - (*self)->u.dense.min; i != 0; --i)
if (*nn_node_child (*self, i))
break;
(*self)->u.dense.max = i + (*self)->u.dense.min;
--(*self)->u.dense.nbr;
*self = nn_realloc (*self, sizeof (struct nn_trie_node) +
((*self)->u.dense.max - (*self)->u.dense.min + 1) *
sizeof (struct nn_trie_node*));
assert (*self);
return 1;
}
/* If the item is removed from the middle of the array, do nothing. */
--(*self)->u.dense.nbr;
return 1;
}
/* Convert dense array into sparse array. */
{
new_node = nn_alloc (sizeof (struct nn_trie_node) +
NN_TRIE_SPARSE_MAX * sizeof (struct nn_trie_node*), "trie node");
assert (new_node);
new_node->refcount = 0;
new_node->prefix_len = (*self)->prefix_len;
memcpy (new_node->prefix, (*self)->prefix, new_node->prefix_len);
new_node->type = NN_TRIE_SPARSE_MAX;
j = 0;
for (i = 0; i != (*self)->u.dense.max - (*self)->u.dense.min + 1;
++i) {
ch2 = *nn_node_child (*self, i);
if (ch2) {
new_node->u.sparse.children [j] = i + (*self)->u.dense.min;
*nn_node_child (new_node, j) = ch2;
++j;
}
}
assert (j == NN_TRIE_SPARSE_MAX);
nn_free (*self);
*self = new_node;
return 1;
}
found:
/* We are at the end of the subscription here. */
/* Subscription doesn't exist. */
if (nn_slow (!*self || !nn_node_has_subscribers (*self)))
return -EINVAL;
/* Subscription exists. Unsubscribe. */
--(*self)->refcount;
/* If reference count has dropped to zero we can try to compact
the node. */
if (!(*self)->refcount) {
/* If there are no children, we can delete the node altogether. */
if (!(*self)->type) {
nn_free (*self);
*self = NULL;
return 1;
}
/* Try to merge the node with the following node. */
*self = nn_node_compact (*self);
return 1;
}
return 0;
}
/* Main body of the daemon. */
static void nn_tcpmuxd_routine (void *arg)
{
int rc;
struct nn_tcpmuxd_ctx *ctx;
int conn;
int pos;
char service [256];
struct nn_tcpmuxd_conn *tc = 0;
size_t sz;
ssize_t ssz;
int i;
struct nn_list_item *it;
unsigned char buf [2];
struct timeval tv;
ctx = (struct nn_tcpmuxd_ctx*) arg;
while (1) {
/* Wait for events. */
rc = (int32_t)poll (ctx->pfd, (int32_t)ctx->pfd_size, -1);
errno_assert (rc >= 0);
nn_assert (rc != 0);
/* There's an incoming TCP connection. */
if (ctx->pfd [0].revents & POLLIN) {
/* Accept the connection. */
conn = accept (ctx->tcp_listener, NULL, NULL);
if (conn < 0 && errno == ECONNABORTED)
continue;
errno_assert (conn >= 0);
/* Set timeouts to prevent malevolent client blocking the service.
Note that these options are not supported on Solaris. */
tv.tv_sec = 0;
tv.tv_usec = 100000;
rc = setsockopt (conn, SOL_SOCKET, SO_RCVTIMEO, &tv, sizeof (tv));
errno_assert (rc == 0 || (rc < 0 && errno == ENOPROTOOPT));
rc = setsockopt (conn, SOL_SOCKET, SO_SNDTIMEO, &tv, sizeof (tv));
errno_assert (rc == 0 || (rc < 0 && errno == ENOPROTOOPT));
/* Read TCPMUX header. */
pos = 0;
while (1) {
nn_assert (pos < sizeof (service));
ssz = recv (conn, &service [pos], 1, 0);
if (ssz < 0 && errno == EAGAIN) {
close (conn);
continue;
}
errno_assert (ssz >= 0);
nn_assert (ssz == 1);
service [pos] = tolower ((uint32_t)service [pos]);
if (pos > 0 && service [pos - 1] == 0x0d &&
service [pos] == 0x0a)
break;
++pos;
}
service [pos - 1] = 0;
/* Check whether specified service is listening. */
for (it = nn_list_begin (&ctx->conns);
it != nn_list_end (&ctx->conns);
it = nn_list_next (&ctx->conns, it)) {
tc = nn_cont (it, struct nn_tcpmuxd_conn, item);
if (strcmp (service, tc->service) == 0)
break;
}
/* If no one is listening, tear down the connection. */
if (it == nn_list_end (&ctx->conns)) {
ssz = send (conn, "-\x0d\x0a", 3, 0);
if (ssz < 0 && errno == EAGAIN) {
close (conn);
continue;
}
errno_assert (ssz >= 0);
nn_assert (ssz == 3);
close (conn);
continue;
}
/* Send TCPMUX reply. */
ssz = send (conn, "+\x0d\x0a", 3, 0);
if (ssz < 0 && errno == EAGAIN) {
close (conn);
continue;
}
errno_assert (ssz >= 0);
nn_assert (ssz == 3);
nn_assert (tc != 0);
/* Pass the file descriptor to the listening process. */
rc = nn_tcpmuxd_send_fd (tc->fd, conn);
errno_assert (rc == 0);
}
/* There's an incoming IPC connection. */
if (ctx->pfd [1].revents & POLLIN) {
/* Accept the connection. */
conn = accept (ctx->ipc_listener, NULL, NULL);
if (conn < 0 && errno == ECONNABORTED)
continue;
errno_assert (conn >= 0);
/* Create new connection entry. */
tc = nn_alloc (sizeof (struct nn_tcpmuxd_conn), "tcpmuxd_conn");
nn_assert (tc);
tc->fd = conn;
nn_list_item_init (&tc->item);
/* Adjust the pollset. We will poll for errors only. */
ctx->pfd_size++;
if (ctx->pfd_size > ctx->pfd_capacity) {
ctx->pfd_capacity *= 2;
ctx->pfd = nn_realloc (ctx->pfd,
sizeof (struct pollfd) * ctx->pfd_capacity);
alloc_assert (ctx->pfd);
}
ctx->pfd [ctx->pfd_size - 1].fd = conn;
ctx->pfd [ctx->pfd_size - 1].events = 0;
ctx->pfd [ctx->pfd_size - 1].revents = 0;
/* Read the connection header. */
ssz = recv (conn, buf, 2, 0);
errno_assert (ssz >= 0);
nn_assert (ssz == 2);
sz = nn_gets (buf);
tc->service = nn_alloc (sz + 1, "tcpmuxd_conn.service");
nn_assert (tc->service);
ssz = recv (conn, tc->service, sz, 0);
errno_assert (ssz >= 0);
nn_assert (ssz == sz);
for (i = 0; i != sz; ++i)
tc->service [i] = tolower ((uint32_t)tc->service [i]);
tc->service [sz] = 0;
/* Add the entry to the IPC connections list. */
nn_list_insert (&ctx->conns, &tc->item, nn_list_end (&ctx->conns));
}
for (i = 2; i < ctx->pfd_size; ++i) {
if (ctx->pfd [i].revents & POLLERR ||
ctx->pfd [i].revents & POLLHUP) {
nn_tcpmuxd_disconnect (ctx, i);
i--;
}
}
}
