Skip to content

Commit

Permalink
Allow Remote Config acknowledgements to be async (#4638)
Browse files Browse the repository at this point in the history
Remote Config updates coming from the Agent needs to be acknowledged. This PR
changes the existing behavior by optionally allowing the acknowledgement to be
async. This is achieved by:

- Changing the way you register products from using the `EventEmitter` API to a
  custom API that only allows a single handler per product (event emitters can
  have more than one listener, which is problematic as we can't support more
  than one of the listeners to acknowledge the config). The new API is as
  follows:
  - `rc.setProductHandler(product, handler)` (prevously `rc.on(product,
    handler)`)
  - `rc.removeProductHandler(product)` (prevously `rc.off(product, handler)`)

- The new product handler is called similar to the old event listener, with the
  following exception: A new optional 4th argument is supplied called `ack`.
  This is a callback which can be called (sync or async) either without any
  arguments (to set the state to `ACKNOWLEDGED`) or with a single error argument
  (to set the state to `ERROR`).
  - If the handler function signature takes less than 4 arguments or doesn't use
    the rest operator (`...args`), and...
    - ...the handler doesn't return a `Promise`: The state is set to
      `ACKNOWLEDGED` right away without the handler having to do anything
      (existing behavior)
    - ...the handler returns a `Promise`, we wait until the promise is resolved
      or rejected and set the state accordingly (new behavior)
  - If the handler function signature takes 4 or more arguments or use the rest
    operator (`...args`): The state is left as `UNACKNOWLEDGED` until the `ack`
    callback is called (new behavior)
  - In any case, the state is still set to `ERROR` if the handler throws
    (existing behavior)

The `RemoteConfigManager` is still an `EventEmitter` however because the
`kPreUpdate` symbol is still being emitted.
  • Loading branch information
watson authored Sep 4, 2024
1 parent 4b565ac commit 52c2993
Show file tree
Hide file tree
Showing 7 changed files with 232 additions and 147 deletions.
15 changes: 8 additions & 7 deletions packages/dd-trace/src/appsec/remote_config/index.js
Original file line number Diff line number Diff line change
Expand Up @@ -28,7 +28,7 @@ function enable (config, appsec) {
rc.updateCapabilities(RemoteConfigCapabilities.ASM_API_SECURITY_SAMPLE_RATE, true)
}

rc.on('ASM_FEATURES', (action, rcConfig) => {
rc.setProductHandler('ASM_FEATURES', (action, rcConfig) => {
if (!rcConfig) return

if (activation === Activation.ONECLICK) {
Expand Down Expand Up @@ -76,9 +76,10 @@ function enableWafUpdate (appsecConfig) {
rc.updateCapabilities(RemoteConfigCapabilities.ASM_CUSTOM_BLOCKING_RESPONSE, true)
rc.updateCapabilities(RemoteConfigCapabilities.ASM_TRUSTED_IPS, true)

rc.on('ASM_DATA', noop)
rc.on('ASM_DD', noop)
rc.on('ASM', noop)
// TODO: delete noop handlers and kPreUpdate and replace with batched handlers
rc.setProductHandler('ASM_DATA', noop)
rc.setProductHandler('ASM_DD', noop)
rc.setProductHandler('ASM', noop)

rc.on(RemoteConfigManager.kPreUpdate, RuleManager.updateWafFromRC)
}
Expand All @@ -98,9 +99,9 @@ function disableWafUpdate () {
rc.updateCapabilities(RemoteConfigCapabilities.ASM_CUSTOM_BLOCKING_RESPONSE, false)
rc.updateCapabilities(RemoteConfigCapabilities.ASM_TRUSTED_IPS, false)

rc.off('ASM_DATA', noop)
rc.off('ASM_DD', noop)
rc.off('ASM', noop)
rc.removeProductHandler('ASM_DATA')
rc.removeProductHandler('ASM_DD')
rc.removeProductHandler('ASM')

rc.off(RemoteConfigManager.kPreUpdate, RuleManager.updateWafFromRC)
}
Expand Down
140 changes: 89 additions & 51 deletions packages/dd-trace/src/appsec/remote_config/manager.js
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,7 @@ const clientId = uuid()
const DEFAULT_CAPABILITY = Buffer.alloc(1).toString('base64') // 0x00

const kPreUpdate = Symbol('kPreUpdate')
const kSupportsAckCallback = Symbol('kSupportsAckCallback')

// There MUST NOT exist separate instances of RC clients in a tracer making separate ClientGetConfigsRequest
// with their own separated Client.ClientState.
Expand All @@ -32,14 +33,26 @@ class RemoteConfigManager extends EventEmitter {
port: config.port
}))

this._handlers = new Map()
const appliedConfigs = this.appliedConfigs = new Map()

this.scheduler = new Scheduler((cb) => this.poll(cb), pollInterval)

this.state = {
client: {
state: { // updated by `parseConfig()`
state: { // updated by `parseConfig()` and `poll()`
root_version: 1,
targets_version: 0,
config_states: [],
// Use getter so `apply_*` can be updated async and still affect the content of `config_states`
get config_states () {
return Array.from(appliedConfigs.values()).map((conf) => ({
id: conf.id,
version: conf.version,
product: conf.product,
apply_state: conf.apply_state,
apply_error: conf.apply_error
}))
},
has_error: false,
error: '',
backend_client_state: ''
Expand All @@ -60,8 +73,6 @@ class RemoteConfigManager extends EventEmitter {
},
cached_target_files: [] // updated by `parseConfig()`
}

this.appliedConfigs = new Map()
}

updateCapabilities (mask, value) {
Expand All @@ -82,32 +93,24 @@ class RemoteConfigManager extends EventEmitter {
this.state.client.capabilities = Buffer.from(str, 'hex').toString('base64')
}

on (event, listener) {
super.on(event, listener)

setProductHandler (product, handler) {
this._handlers.set(product, handler)
this.updateProducts()

if (this.state.client.products.length) {
if (this.state.client.products.length === 1) {
this.scheduler.start()
}

return this
}

off (event, listener) {
super.off(event, listener)

removeProductHandler (product) {
this._handlers.delete(product)
this.updateProducts()

if (!this.state.client.products.length) {
if (this.state.client.products.length === 0) {
this.scheduler.stop()
}

return this
}

updateProducts () {
this.state.client.products = this.eventNames().filter(e => typeof e === 'string')
this.state.client.products = Array.from(this._handlers.keys())
}

getPayload () {
Expand Down Expand Up @@ -228,24 +231,11 @@ class RemoteConfigManager extends EventEmitter {
this.dispatch(toApply, 'apply')
this.dispatch(toModify, 'modify')

this.state.client.state.config_states = []
this.state.cached_target_files = []

for (const conf of this.appliedConfigs.values()) {
this.state.client.state.config_states.push({
id: conf.id,
version: conf.version,
product: conf.product,
apply_state: conf.apply_state,
apply_error: conf.apply_error
})

this.state.cached_target_files.push({
path: conf.path,
length: conf.length,
hashes: Object.entries(conf.hashes).map((entry) => ({ algorithm: entry[0], hash: entry[1] }))
})
}
this.state.cached_target_files = Array.from(this.appliedConfigs.values()).map((conf) => ({
path: conf.path,
length: conf.length,
hashes: Object.entries(conf.hashes).map((entry) => ({ algorithm: entry[0], hash: entry[1] }))
}))
}
}

Expand All @@ -254,20 +244,7 @@ class RemoteConfigManager extends EventEmitter {
// TODO: we need a way to tell if unapply configs were handled by kPreUpdate or not, because they're always
// emitted unlike the apply and modify configs

// in case the item was already handled by kPreUpdate
if (item.apply_state === UNACKNOWLEDGED || action === 'unapply') {
try {
// TODO: do we want to pass old and new config ?
const hadListeners = this.emit(item.product, action, item.file, item.id)

if (hadListeners) {
item.apply_state = ACKNOWLEDGED
}
} catch (err) {
item.apply_state = ERROR
item.apply_error = err.toString()
}
}
this._callHandlerFor(action, item)

if (action === 'unapply') {
this.appliedConfigs.delete(item.path)
Expand All @@ -276,6 +253,49 @@ class RemoteConfigManager extends EventEmitter {
}
}
}

_callHandlerFor (action, item) {
// in case the item was already handled by kPreUpdate
if (item.apply_state !== UNACKNOWLEDGED && action !== 'unapply') return

const handler = this._handlers.get(item.product)

if (!handler) return

try {
if (supportsAckCallback(handler)) {
// If the handler accepts an `ack` callback, expect that to be called and set `apply_state` accordinly
// TODO: do we want to pass old and new config ?
handler(action, item.file, item.id, (err) => {
if (err) {
item.apply_state = ERROR
item.apply_error = err.toString()
} else if (item.apply_state !== ERROR) {
item.apply_state = ACKNOWLEDGED
}
})
} else {
// If the handler doesn't accept an `ack` callback, assume `apply_state` is `ACKNOWLEDGED`,
// unless it returns a promise, in which case we wait for the promise to be resolved or rejected.
// TODO: do we want to pass old and new config ?
const result = handler(action, item.file, item.id)
if (result instanceof Promise) {
result.then(
() => { item.apply_state = ACKNOWLEDGED },
(err) => {
item.apply_state = ERROR
item.apply_error = err.toString()
}
)
} else {
item.apply_state = ACKNOWLEDGED
}
}
} catch (err) {
item.apply_state = ERROR
item.apply_error = err.toString()
}
}
}

function fromBase64JSON (str) {
Expand All @@ -299,4 +319,22 @@ function parseConfigPath (configPath) {
}
}

function supportsAckCallback (handler) {
if (kSupportsAckCallback in handler) return handler[kSupportsAckCallback]

const numOfArgs = handler.length
let result = false

if (numOfArgs >= 4) {
result = true
} else if (numOfArgs !== 0) {
const source = handler.toString()
result = source.slice(0, source.indexOf(')')).includes('...')
}

handler[kSupportsAckCallback] = result

return result
}

module.exports = RemoteConfigManager
6 changes: 3 additions & 3 deletions packages/dd-trace/src/proxy.js
Original file line number Diff line number Diff line change
Expand Up @@ -83,7 +83,7 @@ class Tracer extends NoopProxy {
if (config.remoteConfig.enabled && !config.isCiVisibility) {
const rc = remoteConfig.enable(config, this._modules.appsec)

rc.on('APM_TRACING', (action, conf) => {
rc.setProductHandler('APM_TRACING', (action, conf) => {
if (action === 'unapply') {
config.configure({}, true)
} else {
Expand All @@ -92,7 +92,7 @@ class Tracer extends NoopProxy {
this._enableOrDisableTracing(config)
})

rc.on('AGENT_CONFIG', (action, conf) => {
rc.setProductHandler('AGENT_CONFIG', (action, conf) => {
if (!conf?.name?.startsWith('flare-log-level.')) return

if (action === 'unapply') {
Expand All @@ -103,7 +103,7 @@ class Tracer extends NoopProxy {
}
})

rc.on('AGENT_TASK', (action, conf) => {
rc.setProductHandler('AGENT_TASK', (action, conf) => {
if (action === 'unapply' || !conf) return
if (conf.task_type !== 'tracer_flare' || !conf.args) return

Expand Down
3 changes: 3 additions & 0 deletions packages/dd-trace/test/appsec/index.spec.js
Original file line number Diff line number Diff line change
Expand Up @@ -1182,6 +1182,7 @@ describe('IP blocking', function () {
}).then(() => {
throw new Error('Not expected')
}).catch((err) => {
expect(err.message).to.not.equal('Not expected')
expect(err.response.status).to.be.equal(500)
expect(err.response.data).to.deep.equal(jsonDefaultContent)
})
Expand All @@ -1196,6 +1197,7 @@ describe('IP blocking', function () {
}).then(() => {
throw new Error('Not expected')
}).catch((err) => {
expect(err.message).to.not.equal('Not expected')
expect(err.response.status).to.be.equal(500)
expect(err.response.data).to.deep.equal(htmlDefaultContent)
})
Expand Down Expand Up @@ -1241,6 +1243,7 @@ describe('IP blocking', function () {
}).then(() => {
throw new Error('Not resolve expected')
}).catch((err) => {
expect(err.message).to.not.equal('Not resolve expected')
expect(err.response.status).to.be.equal(301)
expect(err.response.headers.location).to.be.equal('/error')
})
Expand Down
Loading

0 comments on commit 52c2993

Please sign in to comment.