Skip to content

Commit

Permalink
rate_limit: replace callback with async/await
Browse files Browse the repository at this point in the history
  • Loading branch information
@msimerson
msimerson committed May 26, 2022
1 parent 79716ca commit 4d0b9be
Show file tree
Hide file tree
Showing 2 changed files with 34 additions and 63 deletions.
77 changes: 28 additions & 49 deletions index.js
View file
Original file line number Diff line number Diff line change
Expand Up @@ -396,33 +396,34 @@ function getLimit (value) {
return parseInt(match[1], 10);
}

exports.rate_limit = async function (connection, key, value, cb) {
exports.rate_limit = async function (connection, key, value) {

if (value === 0) { // Limit disabled for this host
connection.loginfo(this, `rate limit disabled for: ${key}`);
return cb(null, false);
return false
}

// CAUTION: !value would match that 0 value -^
if (!key || !value) return cb();
if (!this.db) return cb();
if (!key || !value) return
if (!this.db) return

const limit = getLimit(value);
const ttl = getTTL(value);

if (!limit || ! ttl) {
return cb(new Error(`syntax error: key=${key} value=${value}`));
connection.results.add(this, { err: `syntax error: key=${key} value=${value}` });
return
}

connection.logdebug(this, `key=${key} limit=${limit} ttl=${ttl}`);

try {
const newval = await this.db.incr(key)
if (newval === 1) this.db.expire(key, ttl);
cb(null, parseInt(newval, 10) > limit); // boolean
return parseInt(newval, 10) > limit // boolean
}
catch (err) {
cb(err);
connection.results.add(this, { err: `${key}:${err}` });
}
}

Expand Down Expand Up @@ -531,69 +532,47 @@ exports.rate_conn_enforce = async function (next, connection) {
}
}

exports.rate_rcpt_sender = function (next, connection, params) {
const plugin = this;

const [ key, value ] = plugin.get_mail_key('rate_rcpt_sender', connection.transaction.mail_from)

plugin.rate_limit(connection, `rate_rcpt_sender:${key}`, value, (err, over) => {
if (err) {
connection.results.add(plugin, { err: `rate_rcpt_sender:${err}` });
return next();
}
exports.rate_rcpt_sender = async function (next, connection, params) {

connection.results.add(plugin, { rate_rcpt_sender: value });
const [ key, value ] = this.get_mail_key('rate_rcpt_sender', connection.transaction.mail_from)
connection.results.add(this, { rate_rcpt_sender: value });

if (!over) return next();
const over = await this.rate_limit(connection, `rate_rcpt_sender:${key}`, value)
if (!over) return next();

connection.results.add(plugin, { fail: 'rate_rcpt_sender' });
plugin.penalize(connection, false, 'rcpt rate limit exceeded', next);
})
connection.results.add(this, { fail: 'rate_rcpt_sender' });
this.penalize(connection, false, 'rcpt rate limit exceeded', next);
}

exports.rate_rcpt_null = function (next, connection, params) {
const plugin = this;
exports.rate_rcpt_null = async function (next, connection, params) {

if (!params) return next();
if (Array.isArray(params)) params = params[0];
if (params.user) return next();

// Message from the null sender
const [ key, value ] = plugin.get_mail_key('rate_rcpt_null', params)

plugin.rate_limit(connection, `rate_rcpt_null:${key}`, value, (err, over) => {
if (err) {
connection.results.add(plugin, { err: `rate_rcpt_null:${err}` });
return next();
}
const [ key, value ] = this.get_mail_key('rate_rcpt_null', params)
connection.results.add(this, { rate_rcpt_null: value });

connection.results.add(plugin, { rate_rcpt_null: value });
const over = await this.rate_limit(connection, `rate_rcpt_null:${key}`, value)
if (!over) return next();

if (!over) return next();

connection.results.add(plugin, { fail: 'rate_rcpt_null' });
plugin.penalize(connection, false, 'null recip rate limit', next);
})
connection.results.add(this, { fail: 'rate_rcpt_null' });
this.penalize(connection, false, 'null recip rate limit', next);
}

exports.rate_rcpt = function (next, connection, params) {
const plugin = this;
if (Array.isArray(params)) params = params[0];
const [ key, value ] = plugin.get_mail_key('rate_rcpt', params)

plugin.rate_limit(connection, `rate_rcpt:${key}`, value, (err, over) => {
if (err) {
connection.results.add(plugin, { err: `rate_rcpt:${err}` });
return next();
}

connection.results.add(plugin, { rate_rcpt: value });
const [ key, value ] = plugin.get_mail_key('rate_rcpt', params)
connection.results.add(plugin, { rate_rcpt: value });

if (!over) return next();
const over = plugin.rate_limit(connection, `rate_rcpt:${key}`, value)
if (!over) return next();

connection.results.add(plugin, { fail: 'rate_rcpt' });
plugin.penalize(connection, false, 'rate limit exceeded', next);
})
connection.results.add(plugin, { fail: 'rate_rcpt' });
plugin.penalize(connection, false, 'rate limit exceeded', next);
}

/*
Expand Down
20 changes: 6 additions & 14 deletions test/rate_limit.js
View file
Original file line number Diff line number Diff line change
Expand Up @@ -70,22 +70,14 @@ describe('rate_limit', function () {
server);
})

it('no limit', function (done) {
this.plugin.rate_limit(this.connection, 'key', 0, function (err, is_limited) {
// console.log(arguments);
assert.equal(err, undefined);
assert.equal(is_limited, false);
done();
})
it('no limit', async function () {
const is_limited = await this.plugin.rate_limit(this.connection, 'key', 0)
assert.equal(is_limited, false);
})

it('below 50/5m limit', function (done) {
this.plugin.rate_limit(this.connection, 'key', '50/5m', function (err, is_limited) {
// console.log(arguments);
assert.equal(err, undefined);
assert.equal(is_limited, false);
done();
})
it('below 50/5m limit', async function () {
const is_limited = await this.plugin.rate_limit(this.connection, 'key', '50/5m')
assert.equal(is_limited, false);
})
})

Expand Down

0 comments on commit 4d0b9be

Please sign in to comment.