all files / wasm-metering/ index.js

100% Statements 131/131
100% Branches 58/58
100% Functions 7/7
100% Lines 122/122
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242     203×   203×   203× 200× 28× 16×   172× 14× 158× 38× 107× 107× 28×     120× 113×     203×       33×   64×     14×   14× 28×         14× 14× 14× 14×   14×   14× 20×     20× 64× 64×   64× 64× 20×         20×   19× 19×         20×   20× 20×     14× 14×                         81×     11× 11× 30× 30× 30× 19×   11×       11×           13× 13× 13×   13×     13× 13× 13× 13×     13× 13×   13×         13×         13×   13× 85× 85×     13×   13× 13×     10× 10×   13×       13× 13×   10× 12× 12× 10×     10×             10× 14× 14× 14× 14×   14×   10×   26× 13× 16×             13×         13×                            
const toolkit = require('wasm-json-toolkit')
const leb128 = require('leb128/unsigned')
const text2json = toolkit.text2json
const SECTION_IDS = require('wasm-json-toolkit/json2wasm').SECTION_IDS
const defaultCostTable = require('./defaultCostTable.json')
 
// gets the cost of an operation for entry in a section from the cost table
function getCost (json, costTable = {}, defaultCost = 0) {
  let cost = 0
  // finds the default cost
  defaultCost = costTable['DEFAULT'] !== undefined ? costTable['DEFAULT'] : 0
 
  if (typeof costTable === 'function') {
    cost = costTable(json)
  } else if (Array.isArray(json)) {
    json.forEach(el => {
      cost += getCost(el, costTable)
    })
  } else if (typeof costTable === 'number') {
    cost = costTable
  } else if (typeof json === 'object') {
    for (const propName in json) {
      const propCost = costTable[propName]
      if (propCost) {
        cost += getCost(json[propName], propCost, defaultCost)
      }
    }
  } else if (costTable[json] === undefined) {
    cost = defaultCost
  } else {
    cost = costTable[json]
  }
  return cost
}
 
// meters a single code entrie
function meterCodeEntry (entry, costTable, meterFuncIndex, meterType, cost = 0) {
  function meteringStatement (cost, meteringImportIndex) {
    return text2json(`${meterType}.const ${cost} call ${meteringImportIndex}`)
  }
  function remapOp (op, funcIndex) {
    if (op.name === 'call' && op.immediates >= funcIndex) {
      op.immediates = (++op.immediates).toString()
    }
  }
  function meterTheMeteringStatement () {
    const code = meteringStatement(0, 0)
    // sum the operations cost
    return code.reduce(
      (sum, op) => sum + getCost(op.name, costTable.code)
      , 0)
  }
 
  // operations that can possible cause a branch
  const branchingOps = new Set(['end', 'br', 'br_table', 'br_if', 'if', 'else', 'return', 'loop'])
  const meteringOverHead = meterTheMeteringStatement()
  let code = entry.code.slice()
  let meteredCode = []
 
  cost += getCost(entry.locals, costTable.local)
 
  while (code.length) {
    let i = 0
 
    // meters a segment of wasm code
    while (true) {
      const op = code[i++]
      remapOp(op, meterFuncIndex)
 
      cost += getCost(op.name, costTable.code)
      if (branchingOps.has(op.name)) {
        break
      }
    }
 
    // add the metering statement
    if (cost !== 0) {
      // add the cost of metering
      cost += meteringOverHead
      meteredCode = meteredCode
        .concat(meteringStatement(cost, meterFuncIndex))
    }
 
    // start a new segment
    meteredCode = meteredCode
      .concat(code.slice(0, i))
    code = code.slice(i)
    cost = 0
  }
 
  entry.code = meteredCode
  return entry
}
 
/**
 * Injects metering into a JSON output of [wasm2json](https://github.com/ewasm/wasm-json-toolkit#wasm2json)
 * @param {Object} json the json tobe metered
 * @param {Object} opts
 * @param {Object} [opts.costTable=defaultTable] the cost table to meter with. See these notes about the default.
 * @param {String} [opts.moduleStr='metering'] the import string for the metering function
 * @param {String} [opts.fieldStr='usegas'] the field string for the metering function
 * @param {String} [opts.meterType='i64'] the regerster type that is used to meter. Can be `i64`, `i32`, `f64`, `f32`
 * @return {Object} the metered json
 */
exports.meterJSON = (json, opts) => {
  function findSection (module, sectionName) {
    return module.find(section => section.name === sectionName)
  }
 
  function createSection (module, sectionName) {
    const newSectionId = SECTION_IDS[sectionName]
    for (let index in module) {
      const section = module[index]
      const sectionId = SECTION_IDS[section.name]
      if (sectionId) {
        if (newSectionId < sectionId) {
          // inject a new section
          module.splice(index, 0, {
            name: sectionName,
            entries: []
          })
          return
        }
      }
    }
  }
 
  let funcIndex = 0
  let initialCost = 0
  let functionModule, typeModule
 
  let {costTable, moduleStr, fieldStr, meterType} = opts
 
  // set defaults
  if (!costTable) costTable = defaultCostTable
  if (!moduleStr) moduleStr = 'metering'
  if (!fieldStr) fieldStr = 'usegas'
  if (!meterType) meterType = 'i64'
 
  // add nessicarry sections iff they don't exist
  if (!findSection(json, 'type')) createSection(json, 'type')
  if (!findSection(json, 'import')) createSection(json, 'import')
 
  const importJson = {
    'moduleStr': moduleStr,
    'fieldStr': fieldStr,
    'kind': 'function'
  }
  const importType = {
    'form': 'func',
    'params': [meterType]
  }
 
  json = json.slice(0)
 
  for (let section of json) {
    section = Object.assign(section)
    switch (section.name) {
      case 'type':
        // mark the import index
        importJson.type = section.entries.push(importType) - 1
        // save for use for the code section
        typeModule = section
        break
      case 'function':
        // save for use for the code section
        functionModule = section
        break
      case 'import':
        for (const entry of section.entries) {
          initialCost += getCost(entry, costTable.import)
          if (entry.kind === 'function') {
            funcIndex++
          }
        }
        // append the metering import
        section.entries.push(importJson)
        break
      case 'export':
        for (const entry of section.entries) {
          initialCost += getCost(entry, costTable.global)
          if (entry.kind === 'function' && entry.index >= funcIndex) {
            entry.index++
          }
        }
        break
      case 'element':
        for (const entry of section.entries) {
          initialCost += getCost(entry, costTable.element)
          // remap elements indices
          entry.elements = entry.elements.map(el => el >= funcIndex ? ++el : el)
        }
        break
      case 'start':
        initialCost += getCost(section, costTable.start)
        // remap start index
        if (section.index >= funcIndex) section.index++
        break
      case 'code':
        for (const i in section.entries) {
          const entry = section.entries[i]
          const typeIndex = functionModule.entries[i]
          const type = typeModule.entries[typeIndex]
          const cost = getCost(type, costTable.type)
 
          meterCodeEntry(entry, costTable.code, funcIndex, meterType, cost)
        }
        break
      default:
        if (section.entries) {
          for (const entry of section.entries) {
            initialCost += getCost(entry, costTable[section.name])
          }
        }
    }
  }
 
  // add custom section with the initial cost
  json.splice(1, 0, {
    'name': 'custom',
    'sectionName': 'initCost',
    'payload': leb128.encode(initialCost)
  })
  return json
}
 
/**
 * Injects metering into a webassembly binary
 * @param {Object} json the json tobe metered
 * @param {Object} opts
 * @param {Object} [opts.costTable=defaultTable] the cost table to meter with. See these notes about the default.
 * @param {String} [opts.moduleStr='metering'] the import string for the metering function
 * @param {String} [opts.fieldStr='usegas'] the field string for the metering function
 * @param {String} [opts.meterType='i64'] the regerster type that is used to meter. Can be `i64`, `i32`, `f64`, `f32`
 * @return {Buffer}
 */
exports.meterWASM = (wasm, opts = {}) => {
  let json = toolkit.wasm2json(wasm)
  json = exports.meterJSON(json, opts)
  return toolkit.json2wasm(json)
}