diff options
Diffstat (limited to 'src/libknot/probe/probe.c')
| -rw-r--r-- | src/libknot/probe/probe.c | 228 |
1 files changed, 228 insertions, 0 deletions
diff --git a/src/libknot/probe/probe.c b/src/libknot/probe/probe.c new file mode 100644 index 0000000..341c48b --- /dev/null +++ b/src/libknot/probe/probe.c @@ -0,0 +1,228 @@ +/* Copyright (C) 2023 CZ.NIC, z.s.p.o. <knot-dns@labs.nic.cz> + + This program is free software: you can redistribute it and/or modify + it under the terms of the GNU General Public License as published by + the Free Software Foundation, either version 3 of the License, or + (at your option) any later version. + + This program is distributed in the hope that it will be useful, + but WITHOUT ANY WARRANTY; without even the implied warranty of + MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the + GNU General Public License for more details. + + You should have received a copy of the GNU General Public License + along with this program. If not, see <https://www.gnu.org/licenses/>. + */ + +#include <fcntl.h> +#include <poll.h> +#include <stdbool.h> +#include <stdio.h> +#include <stdlib.h> +#include <sys/types.h> +#include <sys/socket.h> +#include <sys/stat.h> +#include <sys/un.h> +#include <unistd.h> + +#include "libknot/attribute.h" +#include "libknot/errcode.h" +#include "libknot/probe/probe.h" +#include "contrib/time.h" + +struct knot_probe { + struct sockaddr_un path; + uint32_t last_unconn_time; + bool consumer; + int fd; +}; + +_public_ +knot_probe_t *knot_probe_alloc(void) +{ + knot_probe_t *probe = calloc(1, sizeof(*probe)); + if (probe == NULL) { + return NULL; + } + + probe->fd = -1; + + return probe; +} + +_public_ +void knot_probe_free(knot_probe_t *probe) +{ + if (probe == NULL) { + return; + } + + close(probe->fd); + if (probe->consumer) { + (void)unlink(probe->path.sun_path); + } + free(probe); +} + +static int probe_connect(knot_probe_t *probe) +{ + return connect(probe->fd, (const struct sockaddr *)(&probe->path), + sizeof(probe->path)); +} + +static int probe_init(knot_probe_t *probe, const char *dir, uint16_t idx) +{ + if (probe == NULL || dir == NULL || idx == 0) { + return KNOT_EINVAL; + } + + probe->path.sun_family = AF_UNIX; + int ret = snprintf(probe->path.sun_path, sizeof(probe->path.sun_path), + "%s/probe%02u.sock", dir, idx); + if (ret < 0 || ret >= sizeof(probe->path.sun_path)) { + return KNOT_ERANGE; + } + + probe->fd = socket(AF_UNIX, SOCK_DGRAM, 0); + if (probe->fd < 0) { + return knot_map_errno(); + } + + if (fcntl(probe->fd, F_SETFL, O_NONBLOCK) == -1) { + close(probe->fd); + probe->fd = -1; + return knot_map_errno(); + } + + return KNOT_EOK; +} + +_public_ +int knot_probe_set_producer(knot_probe_t *probe, const char *dir, uint16_t idx) +{ + int ret = probe_init(probe, dir, idx); + if (ret != KNOT_EOK) { + return ret; + } + + ret = probe_connect(probe); + if (ret != 0) { + return KNOT_ECONN; + } + + return KNOT_EOK; +} + +_public_ +int knot_probe_set_consumer(knot_probe_t *probe, const char *dir, uint16_t idx) +{ + int ret = probe_init(probe, dir, idx); + if (ret != KNOT_EOK) { + return ret; + } + + probe->consumer = true; + + (void)unlink(probe->path.sun_path); + + ret = bind(probe->fd, (const struct sockaddr *)(&probe->path), + sizeof(probe->path)); + if (ret != 0) { + return knot_map_errno(); + } + + if (chmod(probe->path.sun_path, S_IWUSR | S_IWGRP | S_IWOTH) != 0) { + close(probe->fd); + return knot_map_errno(); + } + + return KNOT_EOK; +} + +_public_ +int knot_probe_fd(knot_probe_t *probe) +{ + if (probe == NULL) { + return -1; + } + + return probe->fd; +} + +_public_ +int knot_probe_produce(knot_probe_t *probe, const knot_probe_data_t *data, uint8_t count) +{ + if (probe == NULL || data == NULL || count != 1) { + return KNOT_EINVAL; + } + + size_t used_len = sizeof(*data) - KNOT_DNAME_MAXLEN + data->query.qname_len; + if (send(probe->fd, data, used_len, 0) == -1) { + struct timespec now = time_now(); + if (now.tv_sec - probe->last_unconn_time > 2) { + probe->last_unconn_time = now.tv_sec; + if ((errno == ENOTCONN || errno == ECONNREFUSED) && + probe_connect(probe) == 0 && + send(probe->fd, data, used_len, 0) > 0) { + return KNOT_EOK; + } + } + return knot_map_errno(); + } + + return KNOT_EOK; +} + +_public_ +int knot_probe_consume(knot_probe_t *probe, knot_probe_data_t *data, uint8_t count, + int timeout_ms) +{ + if (probe == NULL || data == NULL || count == 0) { + return KNOT_EINVAL; + } + +#ifdef ENABLE_RECVMMSG + struct mmsghdr msgs[count]; + struct iovec iovecs[count]; + + memset(msgs, 0, sizeof(msgs)); + for (int i = 0; i < count; i++) { + iovecs[i].iov_base = &(data[i]); + iovecs[i].iov_len = sizeof(*data); + msgs[i].msg_hdr.msg_iov = &iovecs[i]; + msgs[i].msg_hdr.msg_iovlen = 1; + } +#else + struct iovec iov = { + .iov_base = data, + .iov_len = sizeof(*data) + }; + struct msghdr msg = { + .msg_iov = &iov, + .msg_iovlen = 1 + }; +#endif + + struct pollfd pfd = { .fd = probe->fd, .events = POLLIN }; + int ret = poll(&pfd, 1, timeout_ms); + if (ret == -1) { + return knot_map_errno(); + } else if ((pfd.revents & POLLIN) == 0) { + return 0; + } + +#ifdef ENABLE_RECVMMSG + ret = recvmmsg(probe->fd, msgs, count, 0, NULL); +#else + ret = recvmsg(probe->fd, &msg, 0); +#endif + if (ret == -1) { + return knot_map_errno(); + } + +#ifdef ENABLE_RECVMMSG + return ret; +#else + return (ret > 0 ? 1 : 0); +#endif +} |