3 files changed, 49 insertions, 9 deletions

diff --git a/src/basic/alloc-util.h b/src/basic/alloc-util.h
index 679ba7f398..ceeee519b7 100644
--- a/src/basic/alloc-util.h
+++ b/src/basic/alloc-util.h
@@ -51,25 +51,29 @@ static inline void freep(void *p) {
 
 #define _cleanup_free_ _cleanup_(freep)
 
-_malloc_  _alloc_(1, 2) static inline void *malloc_multiply(size_t a, size_t b) {
-        if (_unlikely_(b != 0 && a > ((size_t) -1) / b))
+static inline bool size_multiply_overflow(size_t size, size_t need) {
+        return _unlikely_(need != 0 && size > (SIZE_MAX / need));
+}
+
+_malloc_  _alloc_(1, 2) static inline void *malloc_multiply(size_t size, size_t need) {
+        if (size_multiply_overflow(size, need))
                 return NULL;
 
-        return malloc(a * b);
+        return malloc(size * need);
 }
 
-_alloc_(2, 3) static inline void *realloc_multiply(void *p, size_t a, size_t b) {
-        if (_unlikely_(b != 0 && a > ((size_t) -1) / b))
+_alloc_(2, 3) static inline void *realloc_multiply(void *p, size_t size, size_t need) {
+        if (size_multiply_overflow(size, need))
                 return NULL;
 
-        return realloc(p, a * b);
+        return realloc(p, size * need);
 }
 
-_alloc_(2, 3) static inline void *memdup_multiply(const void *p, size_t a, size_t b) {
-        if (_unlikely_(b != 0 && a > ((size_t) -1) / b))
+_alloc_(2, 3) static inline void *memdup_multiply(const void *p, size_t size, size_t need) {
+        if (size_multiply_overflow(size, need))
                 return NULL;
 
-        return memdup(p, a * b);
+        return memdup(p, size * need);
 }
 
 void* greedy_realloc(void **p, size_t *allocated, size_t need, size_t size);
diff --git a/src/basic/socket-util.c b/src/basic/socket-util.c
index 49e5f5b125..58512686e3 100644
--- a/src/basic/socket-util.c
+++ b/src/basic/socket-util.c
@@ -936,3 +936,37 @@ int receive_one_fd(int transport_fd, int flags) {
 
         return *(int*) CMSG_DATA(found);
 }
+
+ssize_t next_datagram_size_fd(int fd) {
+        ssize_t l;
+        int k;
+
+        /* This is a bit like FIONREAD/SIOCINQ, however a bit more powerful. The difference being: recv(MSG_PEEK) will
+         * actually cause the next datagram in the queue to be validated regarding checksums, which FIONREAD dosn't
+         * do. This difference is actually of major importance as we need to be sure that the size returned here
+         * actually matches what we will read with recvmsg() next, as otherwise we might end up allocating a buffer of
+         * the wrong size. */
+
+        l = recv(fd, NULL, 0, MSG_PEEK|MSG_TRUNC);
+        if (l < 0) {
+                if (errno == EOPNOTSUPP)
+                        goto fallback;
+
+                return -errno;
+        }
+        if (l == 0)
+                goto fallback;
+
+        return l;
+
+fallback:
+        k = 0;
+
+        /* Some sockets (AF_PACKET) do not support null-sized recv() with MSG_TRUNC set, let's fall back to FIONREAD
+         * for them. Checksums don't matter for raw sockets anyway, hence this should be fine. */
+
+        if (ioctl(fd, FIONREAD, &k) < 0)
+                return -errno;
+
+        return (ssize_t) k;
+}
diff --git a/src/basic/socket-util.h b/src/basic/socket-util.h
index 92edc1dc22..d17a2f35f8 100644
--- a/src/basic/socket-util.h
+++ b/src/basic/socket-util.h
@@ -133,5 +133,7 @@ int send_one_fd_sa(int transport_fd,
 #define send_one_fd(transport_fd, fd, flags) send_one_fd_sa(transport_fd, fd, NULL, 0, flags)
 int receive_one_fd(int transport_fd, int flags);
 
+ssize_t next_datagram_size_fd(int fd);
+
 #define CMSG_FOREACH(cmsg, mh)                                          \
         for ((cmsg) = CMSG_FIRSTHDR(mh); (cmsg); (cmsg) = CMSG_NXTHDR((mh), (cmsg)))