wikijs/server/helpers/brute-knex.js

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'
  }
}