1 files changed, 188 insertions, 131 deletions

diff --git a/src/socket-proxy/socket-proxyd.c b/src/socket-proxy/socket-proxyd.c
index ac47c85150..d54a05a087 100644
--- a/src/socket-proxy/socket-proxyd.c
+++ b/src/socket-proxy/socket-proxyd.c
@@ -33,6 +33,7 @@
 
 #include "sd-daemon.h"
 #include "sd-event.h"
+#include "sd-resolve.h"
 #include "log.h"
 #include "socket-util.h"
 #include "util.h"
@@ -44,10 +45,15 @@
 #define BUFFER_SIZE (256 * 1024)
 #define CONNECTIONS_MAX 256
 
+static const char *arg_remote_host = NULL;
+
 #define _cleanup_freeaddrinfo_ _cleanup_(freeaddrinfop)
 DEFINE_TRIVIAL_CLEANUP_FUNC(struct addrinfo *, freeaddrinfo);
 
 typedef struct Context {
+        sd_event *event;
+        sd_resolve *resolve;
+
         Set *listen;
         Set *connections;
 } Context;
@@ -63,9 +69,9 @@ typedef struct Connection {
         size_t server_to_client_buffer_size, client_to_server_buffer_size;
 
         sd_event_source *server_event_source, *client_event_source;
-} Connection;
 
-static const char *arg_remote_host = NULL;
+        sd_resolve_query *resolve_query;
+} Connection;
 
 static void connection_free(Connection *c) {
         assert(c);
@@ -82,6 +88,8 @@ static void connection_free(Connection *c) {
         safe_close_pair(c->server_to_client_buffer);
         safe_close_pair(c->client_to_server_buffer);
 
+        sd_resolve_query_unref(c->resolve_query);
+
         free(c);
 }
 
@@ -99,66 +107,9 @@ static void context_free(Context *context) {
 
         set_free(context->listen);
         set_free(context->connections);
-}
-
-static int get_remote_sockaddr(union sockaddr_union *sa, socklen_t *salen) {
-        int r;
-
-        assert(sa);
-        assert(salen);
-
-        if (path_is_absolute(arg_remote_host)) {
-                sa->un.sun_family = AF_UNIX;
-                strncpy(sa->un.sun_path, arg_remote_host, sizeof(sa->un.sun_path)-1);
-                sa->un.sun_path[sizeof(sa->un.sun_path)-1] = 0;
-
-                *salen = offsetof(union sockaddr_union, un.sun_path) + strlen(sa->un.sun_path);
-
-        } else if (arg_remote_host[0] == '@') {
-                sa->un.sun_family = AF_UNIX;
-                sa->un.sun_path[0] = 0;
-                strncpy(sa->un.sun_path+1, arg_remote_host+1, sizeof(sa->un.sun_path)-2);
-                sa->un.sun_path[sizeof(sa->un.sun_path)-1] = 0;
-
-                *salen = offsetof(union sockaddr_union, un.sun_path) + 1 + strlen(sa->un.sun_path + 1);
-
-        } else {
-                _cleanup_freeaddrinfo_ struct addrinfo *result = NULL;
-                const char *node, *service;
-
-                struct addrinfo hints = {
-                        .ai_family = AF_UNSPEC,
-                        .ai_socktype = SOCK_STREAM,
-                        .ai_flags = AI_ADDRCONFIG
-                };
-
-                service = strrchr(arg_remote_host, ':');
-                if (service) {
-                        node = strndupa(arg_remote_host, service - arg_remote_host);
-                        service ++;
-                } else {
-                        node = arg_remote_host;
-                        service = "80";
-                }
-
-                log_debug("Looking up address info for %s:%s", node, service);
-                r = getaddrinfo(node, service, &hints, &result);
-                if (r != 0) {
-                        log_error("Failed to resolve host %s:%s: %s", node, service, gai_strerror(r));
-                        return -EHOSTUNREACH;
-                }
 
-                assert(result);
-                if (result->ai_addrlen > sizeof(union sockaddr_union)) {
-                        log_error("Address too long.");
-                        return -E2BIG;
-                }
-
-                memcpy(sa, result->ai_addr, result->ai_addrlen);
-                *salen = result->ai_addrlen;
-        }
-
-        return 0;
+        sd_event_unref(context->event);
+        sd_resolve_unref(context->resolve);
 }
 
 static int connection_create_pipes(Connection *c, int buffer[2], size_t *sz) {
@@ -247,7 +198,7 @@ static int connection_shovel(
         return 0;
 }
 
-static int connection_enable_event_sources(Connection *c, sd_event *event);
+static int connection_enable_event_sources(Connection *c);
 
 static int traffic_cb(sd_event_source *s, int fd, uint32_t revents, void *userdata) {
         Connection *c = userdata;
@@ -283,7 +234,7 @@ static int traffic_cb(sd_event_source *s, int fd, uint32_t revents, void *userda
         if (c->client_fd == -1 && c->client_to_server_buffer_full <= 0)
                 goto quit;
 
-        r = connection_enable_event_sources(c, sd_event_source_get_event(s));
+        r = connection_enable_event_sources(c);
         if (r < 0)
                 goto quit;
 
@@ -294,12 +245,11 @@ quit:
         return 0; /* ignore errors, continue serving */
 }
 
-static int connection_enable_event_sources(Connection *c, sd_event *event) {
+static int connection_enable_event_sources(Connection *c) {
         uint32_t a = 0, b = 0;
         int r;
 
         assert(c);
-        assert(event);
 
         if (c->server_to_client_buffer_full > 0)
                 b |= EPOLLOUT;
@@ -314,7 +264,7 @@ static int connection_enable_event_sources(Connection *c, sd_event *event) {
         if (c->server_event_source)
                 r = sd_event_source_set_io_events(c->server_event_source, a);
         else if (c->server_fd >= 0)
-                r = sd_event_add_io(event, &c->server_event_source, c->server_fd, a, traffic_cb, c);
+                r = sd_event_add_io(c->context->event, &c->server_event_source, c->server_fd, a, traffic_cb, c);
         else
                 r = 0;
 
@@ -326,7 +276,7 @@ static int connection_enable_event_sources(Connection *c, sd_event *event) {
         if (c->client_event_source)
                 r = sd_event_source_set_io_events(c->client_event_source, b);
         else if (c->client_fd >= 0)
-                r = sd_event_add_io(event, &c->client_event_source, c->client_fd, b, traffic_cb, c);
+                r = sd_event_add_io(c->context->event, &c->client_event_source, c->client_fd, b, traffic_cb, c);
         else
                 r = 0;
 
@@ -338,6 +288,30 @@ static int connection_enable_event_sources(Connection *c, sd_event *event) {
         return 0;
 }
 
+static int connection_complete(Connection *c) {
+        int r;
+
+        assert(c);
+
+        r = connection_create_pipes(c, c->server_to_client_buffer, &c->server_to_client_buffer_size);
+        if (r < 0)
+                goto fail;
+
+        r = connection_create_pipes(c, c->client_to_server_buffer, &c->client_to_server_buffer_size);
+        if (r < 0)
+                goto fail;
+
+        r = connection_enable_event_sources(c);
+        if (r < 0)
+                goto fail;
+
+        return 0;
+
+fail:
+        connection_free(c);
+        return 0; /* ignore errors, continue serving */
+}
+
 static int connect_cb(sd_event_source *s, int fd, uint32_t revents, void *userdata) {
         Connection *c = userdata;
         socklen_t solen;
@@ -361,17 +335,49 @@ static int connect_cb(sd_event_source *s, int fd, uint32_t revents, void *userda
 
         c->client_event_source = sd_event_source_unref(c->client_event_source);
 
-        r = connection_create_pipes(c, c->server_to_client_buffer, &c->server_to_client_buffer_size);
-        if (r < 0)
-                goto fail;
+        return connection_complete(c);
 
-        r = connection_create_pipes(c, c->client_to_server_buffer, &c->client_to_server_buffer_size);
-        if (r < 0)
-                goto fail;
+fail:
+        connection_free(c);
+        return 0; /* ignore errors, continue serving */
+}
 
-        r = connection_enable_event_sources(c, sd_event_source_get_event(s));
-        if (r < 0)
+static int connection_start(Connection *c, struct sockaddr *sa, socklen_t salen) {
+        int r;
+
+        assert(c);
+        assert(sa);
+        assert(salen);
+
+        c->client_fd = socket(sa->sa_family, SOCK_STREAM|SOCK_NONBLOCK|SOCK_CLOEXEC, 0);
+        if (c->client_fd < 0) {
+                log_error("Failed to get remote socket: %m");
                 goto fail;
+        }
+
+        r = connect(c->client_fd, sa, salen);
+        if (r < 0) {
+                if (errno == EINPROGRESS) {
+                        r = sd_event_add_io(c->context->event, &c->client_event_source, c->client_fd, EPOLLOUT, connect_cb, c);
+                        if (r < 0) {
+                                log_error("Failed to add connection socket: %s", strerror(-r));
+                                goto fail;
+                        }
+
+                        r = sd_event_source_set_enabled(c->client_event_source, SD_EVENT_ONESHOT);
+                        if (r < 0) {
+                                log_error("Failed to enable oneshot event source: %s", strerror(-r));
+                                goto fail;
+                        }
+                } else {
+                        log_error("Failed to connect to remote host: %m");
+                        goto fail;
+                }
+        } else {
+                r = connection_complete(c);
+                if (r < 0)
+                        goto fail;
+        }
 
         return 0;
 
@@ -380,14 +386,88 @@ fail:
         return 0; /* ignore errors, continue serving */
 }
 
-static int add_connection_socket(Context *context, sd_event *event, int fd) {
+static int resolve_cb(sd_resolve_query *q, int ret, const struct addrinfo *ai, void *userdata) {
+        Connection *c = userdata;
+
+        assert(q);
+        assert(c);
+
+        if (ret != 0) {
+                log_error("Failed to resolve host: %s", gai_strerror(ret));
+                goto fail;
+        }
+
+        c->resolve_query = sd_resolve_query_unref(c->resolve_query);
+
+        return connection_start(c, ai->ai_addr, ai->ai_addrlen);
+
+fail:
+        connection_free(c);
+        return 0; /* ignore errors, continue serving */
+}
+
+static int resolve_remote(Connection *c) {
+
+        static const struct addrinfo hints = {
+                .ai_family = AF_UNSPEC,
+                .ai_socktype = SOCK_STREAM,
+                .ai_flags = AI_ADDRCONFIG
+        };
+
         union sockaddr_union sa = {};
+        const char *node, *service;
         socklen_t salen;
+        int r;
+
+        if (path_is_absolute(arg_remote_host)) {
+                sa.un.sun_family = AF_UNIX;
+                strncpy(sa.un.sun_path, arg_remote_host, sizeof(sa.un.sun_path)-1);
+                sa.un.sun_path[sizeof(sa.un.sun_path)-1] = 0;
+
+                salen = offsetof(union sockaddr_union, un.sun_path) + strlen(sa.un.sun_path);
+
+                return connection_start(c, &sa.sa, salen);
+        }
+
+        if (arg_remote_host[0] == '@') {
+                sa.un.sun_family = AF_UNIX;
+                sa.un.sun_path[0] = 0;
+                strncpy(sa.un.sun_path+1, arg_remote_host+1, sizeof(sa.un.sun_path)-2);
+                sa.un.sun_path[sizeof(sa.un.sun_path)-1] = 0;
+
+                salen = offsetof(union sockaddr_union, un.sun_path) + 1 + strlen(sa.un.sun_path + 1);
+
+                return connection_start(c, &sa.sa, salen);
+        }
+
+        service = strrchr(arg_remote_host, ':');
+        if (service) {
+                node = strndupa(arg_remote_host, service - arg_remote_host);
+                service ++;
+        } else {
+                node = arg_remote_host;
+                service = "80";
+        }
+
+        log_debug("Looking up address info for %s:%s", node, service);
+        r = sd_resolve_getaddrinfo(c->context->resolve, &c->resolve_query, node, service, &hints, resolve_cb, c);
+        if (r < 0) {
+                log_error("Failed to resolve remote host: %s", strerror(-r));
+                goto fail;
+        }
+
+        return 0;
+
+fail:
+        connection_free(c);
+        return 0; /* ignore errors, continue serving */
+}
+
+static int add_connection_socket(Context *context, int fd) {
         Connection *c;
         int r;
 
         assert(context);
-        assert(event);
         assert(fd >= 0);
 
         if (set_size(context->connections) > CONNECTIONS_MAX) {
@@ -397,12 +477,16 @@ static int add_connection_socket(Context *context, sd_event *event, int fd) {
         }
 
         r = set_ensure_allocated(&context->connections, trivial_hash_func, trivial_compare_func);
-        if (r < 0)
-                return log_oom();
+        if (r < 0) {
+                log_oom();
+                return 0;
+        }
 
         c = new0(Connection, 1);
-        if (!c)
-                return log_oom();
+        if (!c) {
+                log_oom();
+                return 0;
+        }
 
         c->context = context;
         c->server_fd = fd;
@@ -413,48 +497,11 @@ static int add_connection_socket(Context *context, sd_event *event, int fd) {
         r = set_put(context->connections, c);
         if (r < 0) {
                 free(c);
-                return log_oom();
-        }
-
-        r = get_remote_sockaddr(&sa, &salen);
-        if (r < 0)
-                goto fail;
-
-        c->client_fd = socket(sa.sa.sa_family, SOCK_STREAM|SOCK_NONBLOCK|SOCK_CLOEXEC, 0);
-        if (c->client_fd < 0) {
-                log_error("Failed to get remote socket: %m");
-                goto fail;
-        }
-
-        r = connect(c->client_fd, &sa.sa, salen);
-        if (r < 0) {
-                if (errno == EINPROGRESS) {
-                        r = sd_event_add_io(event, &c->client_event_source, c->client_fd, EPOLLOUT, connect_cb, c);
-                        if (r < 0) {
-                                log_error("Failed to add connection socket: %s", strerror(-r));
-                                goto fail;
-                        }
-
-                        r = sd_event_source_set_enabled(c->client_event_source, SD_EVENT_ONESHOT);
-                        if (r < 0) {
-                                log_error("Failed to enable oneshot event source: %s", strerror(-r));
-                                goto fail;
-                        }
-                } else {
-                        log_error("Failed to connect to remote host: %m");
-                        goto fail;
-                }
-        } else {
-                r = connection_enable_event_sources(c, event);
-                if (r < 0)
-                        goto fail;
+                log_oom();
+                return 0;
         }
 
-        return 0;
-
-fail:
-        connection_free(c);
-        return 0; /* ignore non-OOM errors, continue serving */
+        return resolve_remote(c);
 }
 
 static int accept_cb(sd_event_source *s, int fd, uint32_t revents, void *userdata) {
@@ -475,7 +522,7 @@ static int accept_cb(sd_event_source *s, int fd, uint32_t revents, void *userdat
                 getpeername_pretty(nfd, &peer);
                 log_debug("New connection from %s", strna(peer));
 
-                r = add_connection_socket(context, sd_event_source_get_event(s), nfd);
+                r = add_connection_socket(context, nfd);
                 if (r < 0) {
                         log_error("Failed to accept connection, ignoring: %s", strerror(-r));
                         safe_close(fd);
@@ -485,19 +532,18 @@ static int accept_cb(sd_event_source *s, int fd, uint32_t revents, void *userdat
         r = sd_event_source_set_enabled(s, SD_EVENT_ONESHOT);
         if (r < 0) {
                 log_error("Error while re-enabling listener with ONESHOT: %s", strerror(-r));
-                sd_event_exit(sd_event_source_get_event(s), r);
+                sd_event_exit(context->event, r);
                 return r;
         }
 
         return 1;
 }
 
-static int add_listen_socket(Context *context, sd_event *event, int fd) {
+static int add_listen_socket(Context *context, int fd) {
         sd_event_source *source;
         int r;
 
         assert(context);
-        assert(event);
         assert(fd >= 0);
 
         r = set_ensure_allocated(&context->listen, trivial_hash_func, trivial_compare_func);
@@ -522,7 +568,7 @@ static int add_listen_socket(Context *context, sd_event *event, int fd) {
                 return r;
         }
 
-        r = sd_event_add_io(event, &source, fd, EPOLLIN, accept_cb, context);
+        r = sd_event_add_io(context->event, &source, fd, EPOLLIN, accept_cb, context);
         if (r < 0) {
                 log_error("Failed to add event source: %s", strerror(-r));
                 return r;
@@ -612,7 +658,6 @@ static int parse_argv(int argc, char *argv[]) {
 }
 
 int main(int argc, char *argv[]) {
-        _cleanup_event_unref_ sd_event *event = NULL;
         Context context = {};
         int r, n, fd;
 
@@ -623,13 +668,25 @@ int main(int argc, char *argv[]) {
         if (r <= 0)
                 goto finish;
 
-        r = sd_event_default(&event);
+        r = sd_event_default(&context.event);
         if (r < 0) {
                 log_error("Failed to allocate event loop: %s", strerror(-r));
                 goto finish;
         }
 
-        sd_event_set_watchdog(event, true);
+        r = sd_resolve_default(&context.resolve);
+        if (r < 0) {
+                log_error("Failed to allocate resolver: %s", strerror(-r));
+                goto finish;
+        }
+
+        r = sd_resolve_attach_event(context.resolve, context.event, 0);
+        if (r < 0) {
+                log_error("Failed to attach resolver: %s", strerror(-r));
+                goto finish;
+        }
+
+        sd_event_set_watchdog(context.event, true);
 
         n = sd_listen_fds(1);
         if (n < 0) {
@@ -643,12 +700,12 @@ int main(int argc, char *argv[]) {
         }
 
         for (fd = SD_LISTEN_FDS_START; fd < SD_LISTEN_FDS_START + n; fd++) {
-                r = add_listen_socket(&context, event, fd);
+                r = add_listen_socket(&context, fd);
                 if (r < 0)
                         goto finish;
         }
 
-        r = sd_event_loop(event);
+        r = sd_event_loop(context.event);
         if (r < 0) {
                 log_error("Failed to run event loop: %s", strerror(-r));
                 goto finish;