Files
2026-07-13 12:06:16 +08:00

336 lines
8.2 KiB
JavaScript

const { Readable } = require('stream')
const QlobberSub = require('qlobber/aedes/qlobber-sub')
const { QlobberTrue } = require('qlobber')
const Packet = require('aedes-packet')
const QlobberOpts = {
wildcard_one: '+',
wildcard_some: '#',
separator: '/'
}
const CREATE_ON_EMPTY = true
function * multiIterables (iterables) {
for (const iter of iterables) {
yield * iter
}
}
function * retainedMessagesByPattern (retained, pattern) {
const qlobber = new QlobberTrue(QlobberOpts)
qlobber.add(pattern)
for (const [topic, packet] of retained) {
if (qlobber.test(topic)) {
yield packet
}
}
}
function * willsByBrokers (wills, brokers) {
for (const will of wills.values()) {
if (!brokers[will.brokerId]) {
yield will
}
}
}
function * clientListbyTopic (subscriptions, topic) {
for (const [clientId, topicMap] of subscriptions) {
if (topicMap.has(topic)) {
yield clientId
}
}
}
class MemoryPersistence {
// private class members start with #
#retained
#subscriptions
#outgoing
#incoming
#wills
#clientsCount
#trie
constructor () {
// using Maps for convenience and security (risk on prototype polution)
// Map ( topic -> packet )
this.#retained = new Map()
// Map ( clientId -> Map( topic -> { qos, rh, rap, nl } ))
this.#subscriptions = new Map()
// Map ( clientId > [ packet ] }
this.#outgoing = new Map()
// Map ( clientId -> { packetId -> Packet } )
this.#incoming = new Map()
// Map( clientId -> will )
this.#wills = new Map()
this.#clientsCount = 0
this.#trie = new QlobberSub(QlobberOpts)
}
storeRetained (pkt, cb) {
const packet = Object.assign({}, pkt)
if (packet.payload.length === 0) {
this.#retained.delete(packet.topic)
} else {
this.#retained.set(packet.topic, packet)
}
cb(null)
}
createRetainedStreamCombi (patterns) {
const iterables = patterns.map((p) => {
return retainedMessagesByPattern(this.#retained, p)
})
return Readable.from(multiIterables(iterables))
}
createRetainedStream (pattern) {
return Readable.from(retainedMessagesByPattern(this.#retained, pattern))
}
addSubscriptions (client, subs, cb) {
let stored = this.#subscriptions.get(client.id)
const trie = this.#trie
if (!stored) {
stored = new Map()
this.#subscriptions.set(client.id, stored)
this.#clientsCount++
}
for (const sub of subs) {
const storedSub = stored.get(sub.topic)
if (sub.qos > 0) {
trie.add(sub.topic, {
clientId: client.id,
topic: sub.topic,
qos: sub.qos,
rh: sub.rh,
rap: sub.rap,
nl: sub.nl
})
} else if (storedSub?.qos > 0) {
trie.remove(sub.topic, {
clientId: client.id,
topic: sub.topic
})
}
stored.set(sub.topic, { qos: sub.qos, rh: sub.rh, rap: sub.rap, nl: sub.nl })
}
cb(null, client)
}
removeSubscriptions (client, subs, cb) {
const stored = this.#subscriptions.get(client.id)
const trie = this.#trie
if (stored) {
for (const topic of subs) {
const storedSub = stored.get(topic)
if (storedSub !== undefined) {
if (storedSub.qos > 0) {
trie.remove(topic, { clientId: client.id, topic })
}
stored.delete(topic)
}
}
if (stored.size === 0) {
this.#clientsCount--
this.#subscriptions.delete(client.id)
}
}
cb(null, client)
}
subscriptionsByClient (client, cb) {
let subs = null
const stored = this.#subscriptions.get(client.id)
if (stored) {
subs = []
for (const [topic, storedSub] of stored) {
subs.push({ topic, ...storedSub })
}
}
cb(null, subs, client)
}
countOffline (cb) {
return cb(null, this.#trie.subscriptionsCount, this.#clientsCount)
}
subscriptionsByTopic (pattern, cb) {
cb(null, this.#trie.match(pattern))
}
cleanSubscriptions (client, cb) {
const trie = this.#trie
const stored = this.#subscriptions.get(client.id)
if (stored) {
for (const [topic, storedSub] of stored) {
if (storedSub.qos > 0) {
trie.remove(topic, { clientId: client.id, topic })
}
}
this.#clientsCount--
this.#subscriptions.delete(client.id)
}
cb(null, client)
}
#outgoingEnqueuePerSub (sub, packet) {
const id = sub.clientId
const queue = getMapRef(this.#outgoing, id, [], CREATE_ON_EMPTY)
queue[queue.length] = new Packet(packet)
}
outgoingEnqueue (sub, packet, cb) {
this.#outgoingEnqueuePerSub(sub, packet)
process.nextTick(cb)
}
outgoingEnqueueCombi (subs, packet, cb) {
for (let i = 0; i < subs.length; i++) {
this.#outgoingEnqueuePerSub(subs[i], packet)
}
process.nextTick(cb)
}
outgoingUpdate (client, packet, cb) {
const outgoing = getMapRef(this.#outgoing, client.id, [], CREATE_ON_EMPTY)
let temp
for (let i = 0; i < outgoing.length; i++) {
temp = outgoing[i]
if (temp.brokerId === packet.brokerId) {
if (temp.brokerCounter === packet.brokerCounter) {
temp.messageId = packet.messageId
return cb(null, client, packet)
}
/*
Maximum of messageId (packet identifier) is 65535 and will be rotated,
brokerCounter is to ensure the packet identifier be unique.
The for loop is going to search which packet messageId should be updated
in the #outgoing queue.
If there is a case that brokerCounter is different but messageId is same,
we need to let the loop keep searching
*/
} else if (temp.messageId === packet.messageId) {
outgoing[i] = packet
return cb(null, client, packet)
}
}
cb(new Error('no such packet'), client, packet)
}
outgoingClearMessageId (client, packet, cb) {
const outgoing = getMapRef(this.#outgoing, client.id, [], CREATE_ON_EMPTY)
let temp
for (let i = 0; i < outgoing.length; i++) {
temp = outgoing[i]
if (temp.messageId === packet.messageId) {
outgoing.splice(i, 1)
return cb(null, temp)
}
}
cb()
}
outgoingStream (client) {
// shallow clone the outgoing queue for this client to avoid race conditions
const outgoing = [].concat(getMapRef(this.#outgoing, client.id, []))
return Readable.from(outgoing)
}
incomingStorePacket (client, packet, cb) {
const id = client.id
const store = getMapRef(this.#incoming, id, {}, CREATE_ON_EMPTY)
store[packet.messageId] = new Packet(packet)
store[packet.messageId].messageId = packet.messageId
cb(null)
}
incomingGetPacket (client, packet, cb) {
const id = client.id
const store = getMapRef(this.#incoming, id, {})
let err = null
this.#incoming.set(id, store)
if (!store[packet.messageId]) {
err = new Error('no such packet')
}
cb(err, store[packet.messageId])
}
incomingDelPacket (client, packet, cb) {
const id = client.id
const store = getMapRef(this.#incoming, id, {})
const toDelete = store[packet.messageId]
let err = null
if (!toDelete) {
err = new Error('no such packet')
} else {
delete store[packet.messageId]
}
cb(err)
}
putWill (client, packet, cb) {
packet.brokerId = this.broker.id
packet.clientId = client.id
this.#wills.set(client.id, packet)
cb(null, client)
}
getWill (client, cb) {
cb(null, this.#wills.get(client.id), client)
}
delWill (client, cb) {
const will = this.#wills.get(client.id)
this.#wills.delete(client.id)
cb(null, will, client)
}
streamWill (brokers = {}) {
return Readable.from(willsByBrokers(this.#wills, brokers))
}
getClientList (topic) {
return Readable.from(clientListbyTopic(this.#subscriptions, topic))
}
destroy (cb) {
this.#retained = null
if (cb) {
cb(null)
}
}
}
function getMapRef (map, key, ifEmpty, createOnEmpty = false) {
const value = map.get(key)
if (value === undefined && createOnEmpty) {
map.set(key, ifEmpty)
}
return value || ifEmpty
}
module.exports = () => { return new MemoryPersistence() }
module.exports.Packet = Packet