From 6bf2e7558961ff4790098474018ef8456c20ebac Mon Sep 17 00:00:00 2001 From: Alex Hultman Date: Sat, 20 Feb 2021 15:58:20 +0100 Subject: [PATCH] Many changes and additions to pub/sub APIs --- fuzzing/EpollEchoServer.dict | 12 +++++++ src/TopicTree.h | 61 +++++++++++++++++++++++++++++++----- src/WebSocket.h | 59 ++++++++++++++++++++++++++++------ src/WebSocketContextData.h | 14 ++++++--- 4 files changed, 124 insertions(+), 22 deletions(-) create mode 100644 fuzzing/EpollEchoServer.dict diff --git a/fuzzing/EpollEchoServer.dict b/fuzzing/EpollEchoServer.dict new file mode 100644 index 0000000..48b49eb --- /dev/null +++ b/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: " +" " \ No newline at end of file diff --git a/src/TopicTree.h b/src/TopicTree.h index b1d6844..82d06e3 100644 --- a/src/TopicTree.h +++ b/src/TopicTree.h @@ -65,6 +65,12 @@ struct Topic { std::map> messages; std::set 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::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? */ - void publish(Topic *iterator, size_t start, size_t stop, std::string_view topic, std::pair message) { + /* Publishes to all matching topics and wildcards. Returns whether at least one topic was a match. */ + bool publish(Topic *iterator, size_t start, size_t stop, std::string_view topic, std::pair 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::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 message, Subscriber *sender = nullptr) { + bool publish(std::string_view topic, std::pair 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) { diff --git a/src/WebSocket.h b/src/WebSocket.h index 1415d23..4d0b8fc 100644 --- a/src/WebSocket.h +++ b/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 */ - void unsubscribeAll() { + /* Returns whether this socket is subscribed to the specified topic */ + bool isSubscribed(std::string_view topic) { WebSocketContextData *webSocketContextData = (WebSocketContextData *) 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 */ - void publish(std::string_view topic, std::string_view message, OpCode opCode = OpCode::TEXT, bool compress = false) { + /* Returns number of subscribers for this topic, or 0 for failure */ + unsigned int numSubscribers(std::string_view topic) { WebSocketContextData *webSocketContextData = (WebSocketContextData *) 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 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 *webSocketContextData = (WebSocketContextData *) 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); } }; diff --git a/src/WebSocketContextData.h b/src/WebSocketContextData.h index f1f16a3..61fbea8 100644 --- a/src/WebSocketContextData.h +++ b/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(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(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; } };