From 573b1f6ac2cc88817ce8a0cc8d05d164b7530eb7 Mon Sep 17 00:00:00 2001 From: Hailey Date: Thu, 21 Aug 2025 10:12:15 -0700 Subject: [PATCH] start using database --- bskylink/package.json | 1 + bskylink/src/cache/cache.ts | 131 +-------------------- bskylink/src/cache/ozoneAgent.ts | 49 -------- bskylink/src/cache/rule.ts | 20 +--- bskylink/src/cache/safelinkClient.ts | 167 +++++++++++++++++++++++++++ bskylink/src/context.ts | 3 + bskylink/src/db/schema.ts | 19 +++ bskylink/src/routes/redirect.ts | 29 ++--- bskylink/yarn.lock | 5 + src/env/helper.ts | 27 +++++ 10 files changed, 235 insertions(+), 216 deletions(-) delete mode 100644 bskylink/src/cache/ozoneAgent.ts create mode 100644 bskylink/src/cache/safelinkClient.ts create mode 100644 src/env/helper.ts diff --git a/bskylink/package.json b/bskylink/package.json index 334742065b..3df3361391 100644 --- a/bskylink/package.json +++ b/bskylink/package.json @@ -17,6 +17,7 @@ "http-terminator": "^3.2.0", "i18n": "^0.15.1", "kysely": "^0.27.3", + "lru-cache": "^11.1.0", "pg": "^8.12.0", "pino": "^9.2.0", "uhtml": "^4.7.1", diff --git a/bskylink/src/cache/cache.ts b/bskylink/src/cache/cache.ts index 1f8ee7b8ea..521b4c5730 100644 --- a/bskylink/src/cache/cache.ts +++ b/bskylink/src/cache/cache.ts @@ -2,8 +2,7 @@ import {ToolsOzoneSafelinkDefs} from '@atproto/api' import {type ServiceConfig} from '../config.js' import {redirectLogger} from '../logger.js' -import {OzoneAgent} from './ozoneAgent.js' -let cacheCursor: string | undefined +import {OzoneAgent} from './safelinkClient.js' export class EventCache { private rules = new Map() @@ -18,134 +17,6 @@ export class EventCache { return this.cfg } - insert(key: string, evt: ToolsOzoneSafelinkDefs.Event) { - const existing = this.rules.get(key) - if (!existing || new Date(evt.createdAt) > new Date(existing.createdAt)) { - this.rules.set(key, evt) - } - } - - smartUpdateDomain(event: ToolsOzoneSafelinkDefs.Event) { - let domain: string - try { - domain = new URL(event.url).hostname - } catch (error) { - redirectLogger.error( - `[EventCache:smartUpdateDomain] Invalid URL: ${event.url}, error: ${error}`, - ) - throw new Error( - `[EventCache:smartUpdateDomain] Error parsing domain from URL: ${error}`, - ) - } - event.url = domain - - try { - redirectLogger.info( - `[EventCache] smartUpdateDomain called for domain: ${domain}, action: ${event.action}`, - ) - - if (event.action) { - this.insert(domain, event) - redirectLogger.info( - `[EventCache] rule updated or inserted for: ${domain}`, - ) - return - } - } catch (error) { - redirectLogger.error( - `[EventCache:smartUpdateDomain] Error updating rule for domain: ${domain}, error: ${error}`, - ) - throw new Error( - `[EventCache:smartUpdateDomain] Error processing domain event: ${error}`, - ) - } - } - - smartUpdateUrl(event: ToolsOzoneSafelinkDefs.Event) { - let url: string - try { - url = new URL(event.url).toString() - } catch (error) { - redirectLogger.error( - `[EventCache:smartUpdateUrl] Invalid URL: ${event.url}, error: ${error}`, - ) - throw new Error(`[EventCache:smartUpdateUrl] Error parsing URL: ${error}`) - } - event.url = url - - try { - redirectLogger.info( - `[EventCache] smartUpdateUrl called for url: ${url}, action: ${event.action}`, - ) - - if (event.action) { - this.insert(url, event) - redirectLogger.info( - `[EventCache] rule updated or inserted for url: ${url}`, - ) - return - } - } catch (error) { - redirectLogger.error( - `[EventCache:smartUpdateUrl] Error updating rule for url: ${url}, error: ${error}`, - ) - throw new Error( - `[EventCache:smartUpdateUrl] Error processing url event: ${error}`, - ) - } - } - - smartUpdate(event: ToolsOzoneSafelinkDefs.Event) { - if (event.pattern === ToolsOzoneSafelinkDefs.DOMAIN) { - redirectLogger.info( - `[EventCache] smartUpdate called for domain event: ${event.url}, performing ${event.action}`, - ) - return this.smartUpdateDomain(event) - } - if (event.pattern === ToolsOzoneSafelinkDefs.URL) { - redirectLogger.info( - `[EventCache] smartUpdate called for url event: ${event.url}`, - ) - return this.smartUpdateUrl(event) - } - throw new Error('[EventCache] Unknown event pattern') - } - - /** - * Attempts to retrieve an event for the given URL. - * Checks in order: domain, domain+path, then full URL. - */ - smartGet(url: string): ToolsOzoneSafelinkDefs.Event | undefined { - const parsedUrl = new URL(url) - const domain = parsedUrl.hostname - const domainAndPath = domain + parsedUrl.pathname - - const byDomain = this.rules.get(domain) - if (byDomain) { - return byDomain - } - - const byDomainAndPath = this.rules.get(domainAndPath) - if (byDomainAndPath) { - return byDomainAndPath - } - - return this.rules.get(url) - } - - delete(event: ToolsOzoneSafelinkDefs.Event) { - this.rules.delete(event.url) - } - - get(url: string): ToolsOzoneSafelinkDefs.Event | undefined { - const event = this.rules.get(url) - return event - } - - list(): ToolsOzoneSafelinkDefs.Event[] { - return Array.from(this.rules.values()) - } - // Adaptive polling: slow down if no new events, speed up if updates found async adaptiveFetchAndUpdate() { const prevCursor = cacheCursor diff --git a/bskylink/src/cache/ozoneAgent.ts b/bskylink/src/cache/ozoneAgent.ts deleted file mode 100644 index 80ae41d289..0000000000 --- a/bskylink/src/cache/ozoneAgent.ts +++ /dev/null @@ -1,49 +0,0 @@ -import {AtpAgent, CredentialSession} from '@atproto/api' - -import {type ServiceConfig} from '../config' - -export class OzoneAgent { - public session: CredentialSession - public agent: AtpAgent - private cfg: ServiceConfig - - constructor(cfg: ServiceConfig) { - this.cfg = cfg - this.session = new CredentialSession( - new URL(cfg.ozoneUrl || 'http://localhost:2583'), - ) - this.agent = new AtpAgent(this.session) - } - - public async getSession(): Promise { - if (!this.session.hasSession) { - await this.getAgent() - } - return this.session - } - - public async getAgent(): Promise { - if (!this.cfg.ozoneAgentHandle && !this.cfg.ozoneAgentPass) { - throw new Error( - 'OZONE_AGENT_HANDLE and OZONE_AGENT_PASS environment variables must be set', - ) - } - - const identifier = this.cfg.ozoneAgentHandle || 'did:plc:invalid' - const password = this.cfg.ozoneAgentPass || 'invalid' - - if (!this.session.hasSession) { - await this.session.login({identifier, password}) - } - - try { - await this.agent.com.atproto.server.getSession() - } catch (err) { - if ((err as any).status === 401) { - await this.session.login({identifier, password}) - } - } - - return this.agent - } -} diff --git a/bskylink/src/cache/rule.ts b/bskylink/src/cache/rule.ts index d5831da794..012afa6995 100644 --- a/bskylink/src/cache/rule.ts +++ b/bskylink/src/cache/rule.ts @@ -1,22 +1,6 @@ -export type RuleEventType = '#addRule' | '#updateRule' | '#removeRule' -export type RulePatternType = '#domain' | '#url' -export type RuleActionType = '#block' | '#warn' | '#whitelist' -export type RuleReasonType = '#csam' | '#spam' | '#phishing' | '#none' +import {SafelinkRule} from '../db/schema' -export interface Rule { - id: number - eventType: RuleEventType - url: string - pattern: RulePatternType - action: RuleActionType - reason: RuleReasonType - createdBy: string // DID format - createdAt: string // ISO datetime string - comment?: string -} - -// Example Rule object -export const exampleRule: Rule = { +export const exampleRule: SafelinkRule = { id: 1, eventType: '#addRule', url: 'https://malicious.example.com/phishing', diff --git a/bskylink/src/cache/safelinkClient.ts b/bskylink/src/cache/safelinkClient.ts new file mode 100644 index 0000000000..c9c87d527c --- /dev/null +++ b/bskylink/src/cache/safelinkClient.ts @@ -0,0 +1,167 @@ +import {Agent, AtpAgent, CredentialSession} from '@atproto/api' +import {LRUCache} from 'lru-cache' + +import {type ServiceConfig} from '../config' +import {SafelinkRule, RulePatternType} from '../db/schema' +import Database from '../db' +import {redirectLogger} from '../logger' + +export class SafelinkClient { + private domainCache: LRUCache + private urlCache: LRUCache + + private db: Database + + constructor({ + db, + }: { + identifier: string + password: string + pdsHost: string + db: Database + }) { + this.domainCache = new LRUCache({ + max: 10000, + }) + + this.urlCache = new LRUCache({ + max: 25000, + }) + + this.db = db + } + + public async tryFindRule(link: string): Promise { + const u = new URL(link) + u.search = '' + u.hash = '' + + const d = new URL(u.href) + d.pathname = '' + + const urlRule = this.urlCache.get(u.href) + if (urlRule) { + return urlRule + } + + const domainRule = this.domainCache.get(d.href) + if (domainRule) { + return domainRule + } + + try { + const maybeUrlRule = await this.getRule(u.href, '#url') + this.urlCache.set(u.href, maybeUrlRule) + return maybeUrlRule + } catch (e) { + this.urlCache.set(u.href, 'ok') + } + + try { + const maybeDomainRule = await this.getRule(u.href, '#domain') + this.domainCache.set(d.href, maybeDomainRule) + return maybeDomainRule + } catch (e) { + this.domainCache.set(d.href, 'ok') + } + + return 'ok' + } + + private getRule(url: string, pattern: RulePatternType) { + return this.db.db + .selectFrom('safelink_rule') + .where('url', '=', url) + .where('pattern', '=', pattern) + .executeTakeFirstOrThrow() + } + + private addRule(rule: SafelinkRule) { + this.db.db + .insertInto('safelink_rule') + .values(rule) + .execute() + .catch(err => { + redirectLogger.error( + {error: err, rule}, + 'failed to add rule to database', + ) + }) + + if (rule.pattern === '#domain') { + this.domainCache.set(rule.url, rule) + } else { + this.urlCache.set(rule.url, rule) + } + } + + private async removeRule(rule: SafelinkRule) { + await this.db.db + .deleteFrom('safelink_rule') + .where('pattern', '=', '#domain') + .where('url', '=', rule.url) + .execute() + .catch(err => { + redirectLogger.error( + {error: err, rule}, + 'failed to remove rule from database', + ) + }) + + if (rule.pattern === '#domain') { + this.domainCache.delete(rule.url) + } else { + this.urlCache.delete(rule.url) + } + } + + public run() { + // poll and add/remove rules as needed + } +} + +export class OzoneAgent { + public session: CredentialSession + public agent: AtpAgent + private cfg: ServiceConfig + + constructor(cfg: ServiceConfig) { + this.cfg = cfg + this.session = new CredentialSession( + new URL(cfg.ozoneUrl || 'http://localhost:2583'), + ) + this.agent = new AtpAgent(this.session) + } + + public async getSession(): Promise { + if (!this.session.hasSession) { + await this.getAgent() + } + return this.session + } + + public async getAgent(): Promise { + if (!this.cfg.ozoneAgentHandle && !this.cfg.ozoneAgentPass) { + throw new Error( + 'OZONE_AGENT_HANDLE and OZONE_AGENT_PASS environment variables must be set', + ) + } + + const identifier = this.cfg.ozoneAgentHandle || 'did:plc:invalid' + const password = this.cfg.ozoneAgentPass || 'invalid' + + if (!this.session.hasSession) { + await this.session.login({identifier, password}) + } + + try { + await this.agent.com.atproto.server.getSession() + } catch (err) { + if ((err as any).status === 401) { + await this.session.login({identifier, password}) + } + } + + return this.agent + } +} diff --git a/bskylink/src/context.ts b/bskylink/src/context.ts index 7e6f2f34e8..f7243f0b20 100644 --- a/bskylink/src/context.ts +++ b/bskylink/src/context.ts @@ -1,3 +1,4 @@ +import {SafelinkClient} from './cache/safelinkClient.js' import {Config} from './config.js' import Database from './db/index.js' @@ -9,11 +10,13 @@ export type AppContextOptions = { export class AppContext { cfg: Config db: Database + safelinkClient: SafelinkClient abortController = new AbortController() constructor(private opts: AppContextOptions) { this.cfg = this.opts.cfg this.db = this.opts.db + this.safelinkClient = new SafelinkClient({db: this.opts.db}) } static async fromConfig(cfg: Config, overrides?: Partial) { diff --git a/bskylink/src/db/schema.ts b/bskylink/src/db/schema.ts index 8d97f58005..d8f0e9d10f 100644 --- a/bskylink/src/db/schema.ts +++ b/bskylink/src/db/schema.ts @@ -2,6 +2,7 @@ import {Selectable} from 'kysely' export type DbSchema = { link: Link + safelink_rule: SafelinkRule } export interface Link { @@ -14,4 +15,22 @@ export enum LinkType { StarterPack = 1, } +export type RuleEventType = '#addRule' | '#updateRule' | '#removeRule' +export type RulePatternType = '#domain' | '#url' +export type RuleActionType = '#block' | '#warn' | '#whitelist' +export type RuleReasonType = '#csam' | '#spam' | '#phishing' | '#none' + +export interface SafelinkRule { + id: number + eventType: RuleEventType + url: string + pattern: RulePatternType + action: RuleActionType + reason: RuleReasonType + createdBy: string + createdAt: string + comment?: string +} + export type LinkEntry = Selectable +export type RuleEntry = Selectable diff --git a/bskylink/src/routes/redirect.ts b/bskylink/src/routes/redirect.ts index 9a311da942..6a99153f6a 100644 --- a/bskylink/src/routes/redirect.ts +++ b/bskylink/src/routes/redirect.ts @@ -1,6 +1,5 @@ import assert from 'node:assert' -import {ToolsOzoneSafelinkDefs} from '@atproto/api' import {DAY, SECOND} from '@atproto/common' import {type Express} from 'express' import {type Hole} from 'uhtml' @@ -52,20 +51,14 @@ export default function (ctx: AppContext, app: Express) { let hole: Hole | undefined if (ctx.cfg.service.safelinkEnabled) { - const rulePresent: ToolsOzoneSafelinkDefs.Event | undefined = - ctx.cfg.eventCache.smartGet(link) + const rule = await ctx.safelinkClient.tryFindRule(link) - if (rulePresent) { - switch (rulePresent.action) { - case ToolsOzoneSafelinkDefs.WHITELIST: - redirectLogger.info( - `Whitelist rule matched for ${rulePresent.url}`, - ) + if (rule !== 'ok') { + switch (rule.action) { + case '#whitelist': + redirectLogger.info(`Whitelist rule matched for ${rule.url}`) break - case ToolsOzoneSafelinkDefs.REMOVERULE: - redirectLogger.info(`Remove rule matched for ${rulePresent.url}`) - break - case ToolsOzoneSafelinkDefs.BLOCK: + case '#block': hole = linkWarningLayout( 'Blocked Link Warning', linkWarningContents(req, { @@ -74,9 +67,9 @@ export default function (ctx: AppContext, app: Express) { }), ) res.setHeader('Cache-Control', 'no-store') - redirectLogger.info(`Block rule matched for ${rulePresent.url}`) + redirectLogger.info(`Block rule matched for ${rule.url}`) break - case ToolsOzoneSafelinkDefs.WARN: + case '#warn': hole = linkWarningLayout( 'Malicious Link Warning', linkWarningContents(req, { @@ -85,15 +78,13 @@ export default function (ctx: AppContext, app: Express) { }), ) res.setHeader('Cache-Control', 'no-store') - redirectLogger.info(`Warn rule matched for ${rulePresent.url}`) + redirectLogger.info(`Warn rule matched for ${rule.url}`) break default: redirectLogger.warn( - `${rulePresent.action} rule (an unknown rule) matched for ${rulePresent.url}`, + `${rule.action} rule (an unknown rule) matched for ${rule.url}`, ) } - } else { - redirectLogger.info(`No rule present for ${rulePresent.url}`) } } diff --git a/bskylink/yarn.lock b/bskylink/yarn.lock index dc34ff8434..4d463ec091 100644 --- a/bskylink/yarn.lock +++ b/bskylink/yarn.lock @@ -620,6 +620,11 @@ kysely@^0.27.3: resolved "https://registry.yarnpkg.com/kysely/-/kysely-0.27.3.tgz#6cc6c757040500b43c4ac596cdbb12be400ee276" integrity sha512-lG03Ru+XyOJFsjH3OMY6R/9U38IjDPfnOfDgO3ynhbDr+Dz8fak+X6L62vqu3iybQnj+lG84OttBuU9KY3L9kA== +lru-cache@^11.1.0: + version "11.1.0" + resolved "https://registry.yarnpkg.com/lru-cache/-/lru-cache-11.1.0.tgz#afafb060607108132dbc1cf8ae661afb69486117" + integrity sha512-QIXZUBJUx+2zHUdQujWejBkcD9+cs94tLn0+YL8UrCh+D5sCXZ4c7LaEH48pNwRY3MLDgqUFyhlCyjJPf1WP0A== + make-plural@^7.0.0: version "7.4.0" resolved "https://registry.yarnpkg.com/make-plural/-/make-plural-7.4.0.tgz#fa6990dd550dea4de6b20163f74e5ed83d8a8d6d" diff --git a/src/env/helper.ts b/src/env/helper.ts new file mode 100644 index 0000000000..25e934151f --- /dev/null +++ b/src/env/helper.ts @@ -0,0 +1,27 @@ +/** + * Validates a proxy header is valid if supplied + */ +export const isValidProxyHeaderOrThrow = (val: string) => { + if (!val) { + return + } + + if (!val.startsWith('did:')) { + throw new Error( + 'Configured proxy header is invalid. Does not start with `did:`', + ) + } + + const pts = val.split('#') + if (pts.length !== 2) { + throw new Error( + 'Configured proxy header is invalid. Does not contain a single `#`', + ) + } + + if (pts[1].length < 1) { + throw new Error( + 'Configured proxy header is invalid. Does not contain a valid service after the `#`', + ) + } +}