Merge pull request #3900 from keszybz/fix-3607

Fix 3607

3 files changed, 173 insertions, 197 deletions

diff --git a/src/core/service.c b/src/core/service.c
index eb125cb999..3c9455a5f8 100644
--- a/src/core/service.c
+++ b/src/core/service.c
@@ -761,6 +761,11 @@ static void service_dump(Unit *u, FILE *f, const char *prefix) {
                         prefix, s->bus_name,
                         prefix, yes_no(s->bus_name_good));
 
+        if (UNIT_ISSET(s->accept_socket))
+                fprintf(f,
+                        "%sAccept Socket: %s\n",
+                        prefix, UNIT_DEREF(s->accept_socket)->id);
+
         kill_context_dump(&s->kill_context, f, prefix);
         exec_context_dump(&s->exec_context, f, prefix);
 
@@ -1036,6 +1041,20 @@ static int service_coldplug(Unit *u) {
         if (!IN_SET(s->deserialized_state, SERVICE_DEAD, SERVICE_FAILED, SERVICE_AUTO_RESTART))
                 (void) unit_setup_dynamic_creds(u);
 
+        if (UNIT_ISSET(s->accept_socket)) {
+                Socket* socket = SOCKET(UNIT_DEREF(s->accept_socket));
+
+                if (socket->max_connections_per_source > 0) {
+                        SocketPeer *peer;
+
+                        /* Make a best-effort attempt at bumping the connection count */
+                        if (socket_acquire_peer(socket, s->socket_fd, &peer) > 0) {
+                                socket_peer_unref(s->peer);
+                                s->peer = peer;
+                        }
+                }
+        }
+
         service_set_state(s, s->deserialized_state);
         return 0;
 }
@@ -2130,6 +2149,12 @@ static int service_serialize(Unit *u, FILE *f, FDSet *fds) {
         if (r < 0)
                 return r;
 
+        if (UNIT_ISSET(s->accept_socket)) {
+                r = unit_serialize_item(u, f, "accept-socket", UNIT_DEREF(s->accept_socket)->id);
+                if (r < 0)
+                        return r;
+        }
+
         r = unit_serialize_item_fd(u, f, fds, "socket-fd", s->socket_fd);
         if (r < 0)
                 return r;
@@ -2260,6 +2285,17 @@ static int service_deserialize_item(Unit *u, const char *key, const char *value,
                         s->control_command_id = id;
                         s->control_command = s->exec_command[id];
                 }
+        } else if (streq(key, "accept-socket")) {
+                Unit *socket;
+
+                r = manager_load_unit(u->manager, value, NULL, NULL, &socket);
+                if (r < 0)
+                        log_unit_debug_errno(u, r, "Failed to load accept-socket unit: %s", value);
+                else {
+                        unit_ref_set(&s->accept_socket, socket);
+                        SOCKET(socket)->n_connections++;
+                }
+
         } else if (streq(key, "socket-fd")) {
                 int fd;
 
diff --git a/src/core/socket.c b/src/core/socket.c
index ff55885fb3..7f3f154a16 100644
--- a/src/core/socket.c
+++ b/src/core/socket.c
@@ -59,6 +59,13 @@
 #include "user-util.h"
 #include "in-addr-util.h"
 
+struct SocketPeer {
+        unsigned n_ref;
+
+        Socket *socket;
+        union sockaddr_union peer;
+};
+
 static const UnitActiveState state_translation_table[_SOCKET_STATE_MAX] = {
         [SOCKET_DEAD] = UNIT_INACTIVE,
         [SOCKET_START_PRE] = UNIT_ACTIVATING,
@@ -78,9 +85,6 @@ static const UnitActiveState state_translation_table[_SOCKET_STATE_MAX] = {
 static int socket_dispatch_io(sd_event_source *source, int fd, uint32_t revents, void *userdata);
 static int socket_dispatch_timer(sd_event_source *source, usec_t usec, void *userdata);
 
-SocketPeer *socket_peer_new(void);
-int socket_find_peer(Socket *s, int fd, SocketPeer **p);
-
 static void socket_init(Unit *u) {
         Socket *s = SOCKET(u);
 
@@ -151,10 +155,10 @@ static void socket_done(Unit *u) {
 
         socket_free_ports(s);
 
-        while ((p = hashmap_steal_first(s->peers_by_address)))
+        while ((p = set_steal_first(s->peers_by_address)))
                 p->socket = NULL;
 
-        s->peers_by_address = hashmap_free(s->peers_by_address);
+        s->peers_by_address = set_free(s->peers_by_address);
 
         s->exec_runtime = exec_runtime_unref(s->exec_runtime);
         exec_command_free_array(s->exec_command, _SOCKET_EXEC_COMMAND_MAX);
@@ -482,10 +486,11 @@ static void peer_address_hash_func(const void *p, struct siphash *state) {
         const SocketPeer *s = p;
 
         assert(s);
+        assert(IN_SET(s->peer.sa.sa_family, AF_INET, AF_INET6));
 
         if (s->peer.sa.sa_family == AF_INET)
                 siphash24_compress(&s->peer.in.sin_addr, sizeof(s->peer.in.sin_addr), state);
-        else if (s->peer.sa.sa_family == AF_INET6)
+        else
                 siphash24_compress(&s->peer.in6.sin6_addr, sizeof(s->peer.in6.sin6_addr), state);
 }
 
@@ -503,8 +508,7 @@ static int peer_address_compare_func(const void *a, const void *b) {
         case AF_INET6:
                 return memcmp(&x->peer.in6.sin6_addr, &y->peer.in6.sin6_addr, sizeof(x->peer.in6.sin6_addr));
         }
-
-        return -1;
+        assert_not_reached("Black sheep in the family!");
 }
 
 const struct hash_ops peer_address_hash_ops = {
@@ -519,7 +523,7 @@ static int socket_load(Unit *u) {
         assert(u);
         assert(u->load_state == UNIT_STUB);
 
-        r = hashmap_ensure_allocated(&s->peers_by_address, &peer_address_hash_ops);
+        r = set_ensure_allocated(&s->peers_by_address, &peer_address_hash_ops);
         if (r < 0)
                 return r;
 
@@ -537,6 +541,87 @@ static int socket_load(Unit *u) {
         return socket_verify(s);
 }
 
+static SocketPeer *socket_peer_new(void) {
+        SocketPeer *p;
+
+        p = new0(SocketPeer, 1);
+        if (!p)
+                return NULL;
+
+        p->n_ref = 1;
+
+        return p;
+}
+
+SocketPeer *socket_peer_ref(SocketPeer *p) {
+        if (!p)
+                return NULL;
+
+        assert(p->n_ref > 0);
+        p->n_ref++;
+
+        return p;
+}
+
+SocketPeer *socket_peer_unref(SocketPeer *p) {
+        if (!p)
+                return NULL;
+
+        assert(p->n_ref > 0);
+
+        p->n_ref--;
+
+        if (p->n_ref > 0)
+                return NULL;
+
+        if (p->socket)
+                set_remove(p->socket->peers_by_address, p);
+
+        return mfree(p);
+}
+
+int socket_acquire_peer(Socket *s, int fd, SocketPeer **p) {
+        _cleanup_(socket_peer_unrefp) SocketPeer *remote = NULL;
+        SocketPeer sa = {}, *i;
+        socklen_t salen = sizeof(sa.peer);
+        int r;
+
+        assert(fd >= 0);
+        assert(s);
+
+        r = getpeername(fd, &sa.peer.sa, &salen);
+        if (r < 0)
+                return log_error_errno(errno, "getpeername failed: %m");
+
+        if (!IN_SET(sa.peer.sa.sa_family, AF_INET, AF_INET6)) {
+                *p = NULL;
+                return 0;
+        }
+
+        i = set_get(s->peers_by_address, &sa);
+        if (i) {
+                *p = socket_peer_ref(i);
+                return 1;
+        }
+
+        remote = socket_peer_new();
+        if (!remote)
+                return log_oom();
+
+        remote->peer = sa.peer;
+
+        r = set_put(s->peers_by_address, remote);
+        if (r < 0)
+                return r;
+
+        remote->socket = s;
+
+        *p = remote;
+        remote = NULL;
+
+        return 1;
+}
+
 _const_ static const char* listen_lookup(int family, int type) {
 
         if (family == AF_NETLINK)
@@ -2102,22 +2187,26 @@ static void socket_enter_running(Socket *s, int cfd) {
                 Service *service;
 
                 if (s->n_connections >= s->max_connections) {
-                        log_unit_warning(UNIT(s), "Too many incoming connections (%u), refusing connection attempt.", s->n_connections);
+                        log_unit_warning(UNIT(s), "Too many incoming connections (%u), dropping connection.",
+                                         s->n_connections);
                         safe_close(cfd);
                         return;
                 }
 
                 if (s->max_connections_per_source > 0) {
-                        r = socket_find_peer(s, cfd, &p);
+                        r = socket_acquire_peer(s, cfd, &p);
                         if (r < 0) {
                                 safe_close(cfd);
                                 return;
-                        }
+                        } else if (r > 0 && p->n_ref > s->max_connections_per_source) {
+                                _cleanup_free_ char *t = NULL;
+
+                                sockaddr_pretty(&p->peer.sa, FAMILY_ADDRESS_SIZE(p->peer.sa.sa_family), true, false, &t);
 
-                        if (p->n_ref > s->max_connections_per_source) {
-                                log_unit_warning(UNIT(s), "Too many incoming connections (%u) from source, refusing connection attempt.", p->n_ref);
+                                log_unit_warning(UNIT(s),
+                                                 "Too many incoming connections (%u) from source %s, dropping connection.",
+                                                 p->n_ref, strnull(t));
                                 safe_close(cfd);
-                                p = NULL;
                                 return;
                         }
                 }
@@ -2163,10 +2252,8 @@ static void socket_enter_running(Socket *s, int cfd) {
                 cfd = -1; /* We passed ownership of the fd to the service now. Forget it here. */
                 s->n_connections++;
 
-                if (s->max_connections_per_source > 0) {
-                        service->peer = socket_peer_ref(p);
-                        p = NULL;
-                }
+                service->peer = p; /* Pass ownership of the peer reference */
+                p = NULL;
 
                 r = manager_add_job(UNIT(s)->manager, JOB_START, UNIT(service), JOB_REPLACE, &error, NULL);
                 if (r < 0) {
@@ -2313,9 +2400,7 @@ static int socket_stop(Unit *u) {
 
 static int socket_serialize(Unit *u, FILE *f, FDSet *fds) {
         Socket *s = SOCKET(u);
-        SocketPeer *k;
         SocketPort *p;
-        Iterator i;
         int r;
 
         assert(u);
@@ -2366,19 +2451,14 @@ static int socket_serialize(Unit *u, FILE *f, FDSet *fds) {
                 }
         }
 
-        HASHMAP_FOREACH(k, s->peers_by_address, i) {
-                _cleanup_free_ char *t = NULL;
-
-                r = sockaddr_pretty(&k->peer.sa, FAMILY_ADDRESS_SIZE(k->peer.sa.sa_family), true, true, &t);
-                if (r < 0)
-                        return r;
-
-                unit_serialize_item_format(u, f, "peer", "%u %s", k->n_ref, t);
-        }
-
         return 0;
 }
 
+static void socket_port_take_fd(SocketPort *p, FDSet *fds, int fd) {
+        safe_close(p->fd);
+        p->fd = fdset_remove(fds, fd);
+}
+
 static int socket_deserialize_item(Unit *u, const char *key, const char *value, FDSet *fds) {
         Socket *s = SOCKET(u);
 
@@ -2433,18 +2513,13 @@ static int socket_deserialize_item(Unit *u, const char *key, const char *value,
 
                 if (sscanf(value, "%i %n", &fd, &skip) < 1 || fd < 0 || !fdset_contains(fds, fd))
                         log_unit_debug(u, "Failed to parse fifo value: %s", value);
-                else {
-
+                else
                         LIST_FOREACH(port, p, s->ports)
                                 if (p->type == SOCKET_FIFO &&
-                                    path_equal_or_files_same(p->path, value+skip))
+                                    path_equal_or_files_same(p->path, value+skip)) {
+                                        socket_port_take_fd(p, fds, fd);
                                         break;
-
-                        if (p) {
-                                safe_close(p->fd);
-                                p->fd = fdset_remove(fds, fd);
-                        }
-                }
+                                }
 
         } else if (streq(key, "special")) {
                 int fd, skip = 0;
@@ -2452,18 +2527,13 @@ static int socket_deserialize_item(Unit *u, const char *key, const char *value,
 
                 if (sscanf(value, "%i %n", &fd, &skip) < 1 || fd < 0 || !fdset_contains(fds, fd))
                         log_unit_debug(u, "Failed to parse special value: %s", value);
-                else {
-
+                else
                         LIST_FOREACH(port, p, s->ports)
                                 if (p->type == SOCKET_SPECIAL &&
-                                    path_equal_or_files_same(p->path, value+skip))
+                                    path_equal_or_files_same(p->path, value+skip)) {
+                                        socket_port_take_fd(p, fds, fd);
                                         break;
-
-                        if (p) {
-                                safe_close(p->fd);
-                                p->fd = fdset_remove(fds, fd);
-                        }
-                }
+                                }
 
         } else if (streq(key, "mqueue")) {
                 int fd, skip = 0;
@@ -2471,18 +2541,13 @@ static int socket_deserialize_item(Unit *u, const char *key, const char *value,
 
                 if (sscanf(value, "%i %n", &fd, &skip) < 1 || fd < 0 || !fdset_contains(fds, fd))
                         log_unit_debug(u, "Failed to parse mqueue value: %s", value);
-                else {
-
+                else
                         LIST_FOREACH(port, p, s->ports)
                                 if (p->type == SOCKET_MQUEUE &&
-                                    streq(p->path, value+skip))
+                                    streq(p->path, value+skip)) {
+                                        socket_port_take_fd(p, fds, fd);
                                         break;
-
-                        if (p) {
-                                safe_close(p->fd);
-                                p->fd = fdset_remove(fds, fd);
-                        }
-                }
+                                }
 
         } else if (streq(key, "socket")) {
                 int fd, type, skip = 0;
@@ -2490,17 +2555,12 @@ static int socket_deserialize_item(Unit *u, const char *key, const char *value,
 
                 if (sscanf(value, "%i %i %n", &fd, &type, &skip) < 2 || fd < 0 || type < 0 || !fdset_contains(fds, fd))
                         log_unit_debug(u, "Failed to parse socket value: %s", value);
-                else {
-
+                else
                         LIST_FOREACH(port, p, s->ports)
-                                if (socket_address_is(&p->address, value+skip, type))
+                                if (socket_address_is(&p->address, value+skip, type)) {
+                                        socket_port_take_fd(p, fds, fd);
                                         break;
-
-                        if (p) {
-                                safe_close(p->fd);
-                                p->fd = fdset_remove(fds, fd);
-                        }
-                }
+                                }
 
         } else if (streq(key, "netlink")) {
                 int fd, skip = 0;
@@ -2508,17 +2568,12 @@ static int socket_deserialize_item(Unit *u, const char *key, const char *value,
 
                 if (sscanf(value, "%i %n", &fd, &skip) < 1 || fd < 0 || !fdset_contains(fds, fd))
                         log_unit_debug(u, "Failed to parse socket value: %s", value);
-                else {
-
+                else
                         LIST_FOREACH(port, p, s->ports)
-                                if (socket_address_is_netlink(&p->address, value+skip))
+                                if (socket_address_is_netlink(&p->address, value+skip)) {
+                                        socket_port_take_fd(p, fds, fd);
                                         break;
-
-                        if (p) {
-                                safe_close(p->fd);
-                                p->fd = fdset_remove(fds, fd);
-                        }
-                }
+                                }
 
         } else if (streq(key, "ffs")) {
                 int fd, skip = 0;
@@ -2526,46 +2581,14 @@ static int socket_deserialize_item(Unit *u, const char *key, const char *value,
 
                 if (sscanf(value, "%i %n", &fd, &skip) < 1 || fd < 0 || !fdset_contains(fds, fd))
                         log_unit_debug(u, "Failed to parse ffs value: %s", value);
-                else {
-
+                else
                         LIST_FOREACH(port, p, s->ports)
                                 if (p->type == SOCKET_USB_FUNCTION &&
-                                    path_equal_or_files_same(p->path, value+skip))
+                                    path_equal_or_files_same(p->path, value+skip)) {
+                                        socket_port_take_fd(p, fds, fd);
                                         break;
+                                }
 
-                        if (p) {
-                                safe_close(p->fd);
-                                p->fd = fdset_remove(fds, fd);
-                        }
-                }
-
-        } else if (streq(key, "peer")) {
-                _cleanup_(socket_peer_unrefp) SocketPeer *p;
-                int n_ref, skip = 0;
-                SocketAddress a;
-                int r;
-
-                if (sscanf(value, "%u %n", &n_ref, &skip) < 1 || n_ref < 1)
-                        log_unit_debug(u, "Failed to parse socket peer value: %s", value);
-                else {
-                        r = socket_address_parse(&a, value+skip);
-                        if (r < 0)
-                                return r;
-
-                        p = socket_peer_new();
-                        if (!p)
-                                return log_oom();
-
-                        p->n_ref = n_ref;
-                        memcpy(&p->peer, &a.sockaddr, sizeof(a.sockaddr));
-                        p->socket = s;
-
-                        r = hashmap_put(s->peers_by_address, p, p);
-                        if (r < 0)
-                                return r;
-
-                        p = NULL;
-                }
         } else
                 log_unit_debug(UNIT(s), "Unknown serialization key: %s", key);
 
@@ -2662,83 +2685,6 @@ _pure_ static bool socket_check_gc(Unit *u) {
         return s->n_connections > 0;
 }
 
-SocketPeer *socket_peer_new(void) {
-        SocketPeer *p;
-
-        p = new0(SocketPeer, 1);
-        if (!p)
-                return NULL;
-
-        p->n_ref = 1;
-
-        return p;
-}
-
-SocketPeer *socket_peer_ref(SocketPeer *p) {
-        if (!p)
-                return NULL;
-
-        assert(p->n_ref > 0);
-        p->n_ref++;
-
-        return p;
-}
-
-SocketPeer *socket_peer_unref(SocketPeer *p) {
-        if (!p)
-                return NULL;
-
-        assert(p->n_ref > 0);
-
-        p->n_ref--;
-
-        if (p->n_ref > 0)
-                return NULL;
-
-        if (p->socket)
-                (void) hashmap_remove(p->socket->peers_by_address, p);
-
-        free(p);
-
-        return NULL;
-}
-
-int socket_find_peer(Socket *s, int fd, SocketPeer **p) {
-        _cleanup_free_ SocketPeer *remote = NULL;
-        SocketPeer sa, *i;
-        socklen_t salen = sizeof(sa.peer);
-        int r;
-
-        assert(fd >= 0);
-        assert(s);
-
-        r = getpeername(fd, &sa.peer.sa, &salen);
-        if (r < 0)
-                return log_error_errno(errno, "getpeername failed: %m");
-
-        i = hashmap_get(s->peers_by_address, &sa);
-        if (i) {
-                *p = i;
-                return 1;
-        }
-
-        remote = socket_peer_new();
-        if (!remote)
-                return log_oom();
-
-        memcpy(&remote->peer, &sa.peer, sizeof(union sockaddr_union));
-        remote->socket = s;
-
-        r = hashmap_put(s->peers_by_address, remote, remote);
-        if (r < 0)
-                return r;
-
-        *p = remote;
-        remote = NULL;
-
-        return 0;
-}
-
 static int socket_dispatch_io(sd_event_source *source, int fd, uint32_t revents, void *userdata) {
         SocketPort *p = userdata;
         int cfd = -1;
diff --git a/src/core/socket.h b/src/core/socket.h
index 2fe38ef2aa..89f4664510 100644
--- a/src/core/socket.h
+++ b/src/core/socket.h
@@ -80,7 +80,7 @@ struct Socket {
 
         LIST_HEAD(SocketPort, ports);
 
-        Hashmap *peers_by_address;
+        Set *peers_by_address;
 
         unsigned n_accepted;
         unsigned n_connections;
@@ -168,15 +168,9 @@ struct Socket {
         RateLimit trigger_limit;
 };
 
-struct SocketPeer {
-        unsigned n_ref;
-
-        Socket *socket;
-        union sockaddr_union peer;
-};
-
 SocketPeer *socket_peer_ref(SocketPeer *p);
 SocketPeer *socket_peer_unref(SocketPeer *p);
+int socket_acquire_peer(Socket *s, int fd, SocketPeer **p);
 
 DEFINE_TRIVIAL_CLEANUP_FUNC(SocketPeer*, socket_peer_unref);