Many changes and additions to pub/sub APIs

fuzzing/EpollEchoServer.dict

@@ -0,0 +1,12 @@
+"get"
+"post"
+"get /"
+"http/1.1"
+"upgrade: websocket"
+"\x0D\x0A"
+"sec-websocket-key: dGhlIHNhbXBsZSBub25jZQ=="
+"sec-websocket-version: 13"
+"get / http/1.1"
+"sec-websocket-extensions: permessage-deflate"
+"sec-websocket-protocol: "
+" "

src/TopicTree.h

@@ -65,6 +65,12 @@ struct Topic {
     std::map<unsigned int, std::pair<std::string, std::string>> messages;
 
     std::set<Subscriber *> subs;
+
+    /* Locked or not, used only when iterating over a Subscriber's topics */
+    bool locked = false;
+
+    /* Full name is used when iterating topcis */
+    std::string fullName;
 };
 
 struct Hole {
@@ -149,6 +155,26 @@ private:
     int numTriggeredTopics = 0;
     Subscriber *min = (Subscriber *) UINTPTR_MAX;
 
+    /* Returns Topic, or nullptr. Topic can be root if empty string given. */
+    Topic *lookupTopic(std::string_view topic) {
+        /* Lookup exact Topic ptr from string */
+        Topic *iterator = root;
+        for (size_t start = 0, stop = 0; stop != std::string::npos; start = stop + 1) {
+            stop = topic.find('/', start);
+            std::string_view segment = topic.substr(start, stop - start);
+
+            std::map<std::string_view, Topic *>::iterator it = iterator->children.find(segment);
+            if (it == iterator->children.end()) {
+                /* This topic does not even exist */
+                return nullptr;
+            }
+
+            iterator = it->second;
+        }
+
+        return iterator;
+    }
+
     /* Cull or trim unused Topic nodes from leaf to root */
     void trimTree(Topic *topic) {
         while (!topic->subs.size() && !topic->children.size() && !topic->terminatingWildcardChild && !topic->wildcardChild) {
@@ -192,8 +218,11 @@ private:
         }
     }
 
-    /* Should be getData and commit? */
+    /* Publishes to all matching topics and wildcards. Returns whether at least one topic was a match. */
-    void publish(Topic *iterator, size_t start, size_t stop, std::string_view topic, std::pair<std::string_view, std::string_view> message) {
+    bool publish(Topic *iterator, size_t start, size_t stop, std::string_view topic, std::pair<std::string_view, std::string_view> message) {
+
+        /* Whether we matched with at least one topic */
+        bool didMatch = false;
 
         /* Iterate over all segments in given topic */
         for (; stop != std::string::npos; start = stop + 1) {
@@ -207,7 +236,8 @@ private:
              * instace the error is found late while iterating the topic segments. */
             if (segment.length() == 1) {
                 if (segment[0] == '+' || segment[0] == '#') {
-                    return;
+                    /* "Fail" here, but not necessarily for the entire publish */
+                    return didMatch;
                 }
             }
 
@@ -225,17 +255,19 @@ private:
                     triggeredTopics[numTriggeredTopics++] = iterator->terminatingWildcardChild;
                     iterator->terminatingWildcardChild->triggered = true;
                 }
+
+                didMatch = true;
             }
 
             /* Do we have a wildcard child? */
             if (iterator->wildcardChild) {
-                publish(iterator->wildcardChild, stop + 1, stop, topic, message);
+                didMatch |= publish(iterator->wildcardChild, stop + 1, stop, topic, message);
             }
 
             std::map<std::string_view, Topic *>::iterator it = iterator->children.find(segment);
             if (it == iterator->children.end()) {
                 /* Stop trying to match by exact string */
-                return;
+                return didMatch;
             }
 
             iterator = it->second;
@@ -254,6 +286,9 @@ private:
             triggeredTopics[numTriggeredTopics++] = iterator;
             iterator->triggered = true;
         }
+
+        /* We obviously matches exactly here */
+        return true;
     }
 
 public:
@@ -302,6 +337,12 @@ public:
                 newTopic->wildcardChild = nullptr;
                 memcpy(newTopic->name, segment.data(), segment.length());
 
+                /* Set fullname as parent's name plus our name */
+                newTopic->fullName.reserve(newTopic->parent->fullName.length() + 1 + segment.length());
+                newTopic->fullName.append(newTopic->parent->fullName);
+                newTopic->fullName.append("/");
+                newTopic->fullName.append(segment);
+
                 /* For simplicity we do insert wildcards with text */
                 iterator->children.insert(lb, {std::string_view(newTopic->name, segment.length()), newTopic});
 
@@ -339,15 +380,16 @@ public:
         return {(unsigned int) iterator->subs.size(), false};
     }
 
-    void publish(std::string_view topic, std::pair<std::string_view, std::string_view> message, Subscriber *sender = nullptr) {
+    bool publish(std::string_view topic, std::pair<std::string_view, std::string_view> message, Subscriber *sender = nullptr) {
         /* Add a hole for the sender if one */
         if (sender) {
             senderHoles[sender].push_back(messageId);
         }
 
-        publish(root, 0, 0, topic, message);
+        auto ret = publish(root, 0, 0, topic, message);
         /* MessageIDs are reset on drain - this should be fine since messages itself are cleared on drain */
         messageId++;
+        return ret;
     }
 
     /* Returns a pair of numSubscribers after operation, and whether we were subscribed prior */
@@ -369,6 +411,11 @@ public:
                 iterator = it->second;
             }
 
+            /* Is this topic locked? If so, we cannot unsubscribe from it */
+            if (iterator->locked) {
+                return {iterator->subs.size(), false};
+            }
+
             /* Try and remove this topic from our list */
             for (auto it = subscriber->subscriptions.begin(); it != subscriber->subscriptions.end(); it++) {
                 if (*it == iterator) {

src/WebSocket.h

@@ -1,5 +1,5 @@
 /*
- * Authored by Alex Hultman, 2018-2020.
+ * Authored by Alex Hultman, 2018-2021.
  * Intellectual property of third-party.
 
  * Licensed under the Apache License, Version 2.0 (the "License");
@@ -223,31 +223,70 @@ public:
         return webSocketContextData->topicTree.unsubscribe(topic, webSocketData->subscriber, nonStrict);
     }
 
-    /* Unsubscribe from all topics you might be subscribed to */
+    /* Returns whether this socket is subscribed to the specified topic */
-    void unsubscribeAll() {
+    bool isSubscribed(std::string_view topic) {
         WebSocketContextData<SSL> *webSocketContextData = (WebSocketContextData<SSL> *) us_socket_context_ext(SSL,
             (us_socket_context_t *) us_socket_context(SSL, (us_socket_t *) this)
         );
 
-        WebSocketData *webSocketData = (WebSocketData *) us_socket_ext(SSL, (us_socket_t *) this);
+        Topic *t = webSocketContextData->lookupTopic(topic);
+        if (t) {
+            return t->subs.find(this) != t->subs.end();
+        }
 
-        webSocketContextData->topicTree.unsubscribeAll(webSocketData->subscriber);
+        return false;
     }
 
-    /* Publish a message to a topic according to MQTT rules and syntax */
+    /* Returns number of subscribers for this topic, or 0 for failure */
-    void publish(std::string_view topic, std::string_view message, OpCode opCode = OpCode::TEXT, bool compress = false) {
+    unsigned int numSubscribers(std::string_view topic) {
         WebSocketContextData<SSL> *webSocketContextData = (WebSocketContextData<SSL> *) us_socket_context_ext(SSL,
             (us_socket_context_t *) us_socket_context(SSL, (us_socket_t *) this)
         );
 
-        /* Make us a subscriber if we aren't yet (important for allocating a sender address) */
+        Topic *t = webSocketContextData->lookupTopic(topic);
+        if (t) {
+            return t->subs.size();
+        }
+
+        return 0;
+    }
+
+    /* Iterates all topics of this WebSocket. Every topic is represented by [name, numSubscribers].
+     * Can be called in close handler. It is possible to modify the subscription list while
+     * inside the callback ONLY IF not modifying the topic passed to the callback.
+     * Topic names are valid only for the duration of the callback. */
+    void iterateTopics(MoveOnlyFunction<void(std::string_view, unsigned int)> cb) {
+        WebSocketData *webSocketData = (WebSocketData *) us_socket_ext(SSL, (us_socket_t *) this);
+
+        if (webSocketData->subscriber) {
+            for (Topic *t : webSocketData->subscriber->subscriptions) {
+                /* Lock this topic so that nobody may unsubscribe from it during this callback */
+                t->locked = true;
+
+                cb(t->fullName, (unsigned int) t->subs.size());
+
+                t->locked = false;
+            }
+        }
+    }
+
+    /* Publish a message to a topic according to MQTT rules and syntax. Returns [numSubscribers, success].
+     * We, the WebSocket, must be subscribed to the topic itself and if so - no message will be sent to ourselves.
+     * Use App::publish for an unconditional publish that simply publishes to whomever might be subscribed. */
+    bool publish(std::string_view topic, std::string_view message, OpCode opCode = OpCode::TEXT, bool compress = false) {
+        WebSocketContextData<SSL> *webSocketContextData = (WebSocketContextData<SSL> *) us_socket_context_ext(SSL,
+            (us_socket_context_t *) us_socket_context(SSL, (us_socket_t *) this)
+        );
+
+        /* We cannot be a subscriber of this topic if we are not a subscriber of anything */
         WebSocketData *webSocketData = (WebSocketData *) us_socket_ext(SSL, (us_socket_t *) this);
         if (!webSocketData->subscriber) {
-            webSocketData->subscriber = new Subscriber(this);
+            /* Failure, but still do return the number of subscribers */
+            return false;
         }
 
         /* Publish as sender, does not receive its own messages even if subscribed to relevant topics */
-        webSocketContextData->publish(topic, message, opCode, compress, webSocketData->subscriber);
+        return webSocketContextData->publish(topic, message, opCode, compress, webSocketData->subscriber);
     }
 };
 

src/WebSocketContextData.h

@@ -224,7 +224,9 @@ public:
     }
 
     /* Helper for topictree publish, common path from app and ws */
-    void publish(std::string_view topic, std::string_view message, OpCode opCode, bool compress, Subscriber *sender = nullptr) {
+    bool publish(std::string_view topic, std::string_view message, OpCode opCode, bool compress, Subscriber *sender = nullptr) {
+        bool didMatch = false;
+
         /* We frame the message right here and only pass raw bytes to the pub/subber */
         char *dst = (char *) malloc(protocol::messageFrameSize(message.size()));
         size_t dst_length = protocol::formatMessage<true>(dst, message.data(), message.length(), opCode, message.length(), false);
@@ -232,7 +234,7 @@ public:
         /* If compression is disabled */
         if (compression == DISABLED) {
             /* Leave second field empty as nobody will ever read it */
-            topicTree.publish(topic, {std::string_view(dst, dst_length), {}}, sender);
+            didMatch |= topicTree.publish(topic, {std::string_view(dst, dst_length), {}}, sender);
         } else {
             /* DEDICATED_COMPRESSOR always takes the same path as must always have MessageMetadata as head */
             if (compress || compression != SHARED_COMPRESSOR) {
@@ -249,7 +251,7 @@ public:
                     size_t dst_compressed_length = protocol::formatMessage<true>(dst_compressed, compressedMessage.data(), compressedMessage.length(), opCode, compressedMessage.length(), true);
 
                     /* Always publish the shortest one in any case */
-                    topicTree.publish(topic, {std::string_view(dst, dst_length), dst_compressed_length >= dst_length ? std::string_view(dst, dst_length) : std::string_view(dst_compressed, dst_compressed_length)}, sender);
+                    didMatch |= topicTree.publish(topic, {std::string_view(dst, dst_length), dst_compressed_length >= dst_length ? std::string_view(dst, dst_length) : std::string_view(dst_compressed, dst_compressed_length)}, sender);
 
                     /* We don't care for allocation here */
                     ::free(dst_compressed);
@@ -267,7 +269,7 @@ public:
                     memcpy(dst_compressed + sizeof(MessageMetadata), message.data(), message.length());
 
                     /* Interpretation of compressed data depends on what compressor we use */
-                    topicTree.publish(topic, {
+                    didMatch |= topicTree.publish(topic, {
                         std::string_view(dst, dst_length),
                         std::string_view(dst_compressed, message.length() + sizeof(MessageMetadata))
                     }, sender);
@@ -277,11 +279,13 @@ public:
             } else {
                 /* If not compressing, put same message on both tracks (only valid for SHARED_COMPRESSOR).
                  * DEDICATED_COMPRESSOR_xKB must never end up here as we don't put a proper head here. */
-                topicTree.publish(topic, {std::string_view(dst, dst_length), std::string_view(dst, dst_length)}, sender);
+                didMatch |= topicTree.publish(topic, {std::string_view(dst, dst_length), std::string_view(dst, dst_length)}, sender);
             }
         }
 
         ::free(dst);
+
+        return didMatch;
     }
 };