wikijs-fork/server/helpers/brute-knex.js
2020-04-21 23:16:13 -04:00

149 lines
3.9 KiB
JavaScript

const AbstractClientStore = require('express-brute/lib/AbstractClientStore')
const KnexStore = module.exports = function (options) {
options = options || Object.create(null)
AbstractClientStore.apply(this, arguments)
this.options = Object.assign(Object.create(null), KnexStore.defaults, options)
if (this.options.knex) {
this.knex = this.options.knex
} else {
this.knex = require('knex')(KnexStore.defaultsKnex)
}
if (options.createTable === false) {
this.ready = Promise.resolve()
} else {
this.ready = this.knex.schema.hasTable(this.options.tablename)
.then((exists) => {
if (exists) {
return
}
return this.knex.schema.createTable(this.options.tablename, (table) => {
table.string('key')
table.bigInteger('firstRequest').nullable()
table.bigInteger('lastRequest').nullable()
table.bigInteger('lifetime').nullable()
table.integer('count')
})
})
}
}
KnexStore.prototype = Object.create(AbstractClientStore.prototype)
KnexStore.prototype.set = async function (key, value, lifetime, callback) {
try {
lifetime = lifetime || 0
await this.ready
const resp = await this.knex.transaction((trx) => {
return trx
.select('*')
.forUpdate()
.from(this.options.tablename)
.where('key', '=', key)
.then((foundKeys) => {
if (foundKeys.length === 0) {
return trx.from(this.options.tablename)
.insert({
key: key,
lifetime: new Date(Date.now() + lifetime * 1000).getTime(),
lastRequest: new Date(value.lastRequest).getTime(),
firstRequest: new Date(value.firstRequest).getTime(),
count: value.count
})
} else {
return trx(this.options.tablename)
.where('key', '=', key)
.update({
lifetime: new Date(Date.now() + lifetime * 1000).getTime(),
count: value.count,
lastRequest: new Date(value.lastRequest).getTime()
})
}
})
})
callback(null, resp)
} catch (err) {
callback(err, null)
}
}
KnexStore.prototype.get = async function (key, callback) {
try {
await this.ready
await this.clearExpired()
const resp = await this.knex.select('*')
.from(this.options.tablename)
.where('key', '=', key)
let o = null
if (resp[0]) {
o = {}
o.lastRequest = new Date(resp[0].lastRequest)
o.firstRequest = new Date(resp[0].firstRequest)
o.count = resp[0].count
}
callback(null, o)
} catch (err) {
callback(err, null)
}
}
KnexStore.prototype.reset = async function (key, callback) {
try {
await this.ready
const resp = await this.knex(this.options.tablename)
.where('key', '=', key)
.del()
callback(null, resp)
} catch (err) {
callback(err, null)
}
}
KnexStore.prototype.increment = async function (key, lifetime, callback) {
try {
const result = await this.get(key)
let resp = null
if (result) {
resp = await this.knex(this.options.tablename)
.increment('count', 1)
.where('key', '=', key)
} else {
resp = await this.knex(this.options.tablename)
.insert({
key: key,
firstRequest: new Date().getTime(),
lastRequest: new Date().getTime(),
lifetime: new Date(Date.now() + lifetime * 1000).getTime(),
count: 1
})
}
callback(null, resp)
} catch (err) {
callback(err, null)
}
}
KnexStore.prototype.clearExpired = async function (callback) {
await this.ready
return this.knex(this.options.tablename)
.del()
.where('lifetime', '<', new Date().getTime())
}
KnexStore.defaults = {
tablename: 'brute',
createTable: true
}
KnexStore.defaultsKnex = {
client: 'sqlite3',
// debug: true,
connection: {
filename: './brute-knex.sqlite'
}
}