start using database

2025-08-21 10:12:15 -07:00
parent b06abba6b2
commit 573b1f6ac2
10 changed files with 235 additions and 216 deletions
@@ -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",
@@ -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<string, ToolsOzoneSafelinkDefs.Event>()
@@ -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
@@ -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<CredentialSession> {
-    if (!this.session.hasSession) {
-      await this.getAgent()
-    }
-    return this.session
-  }
-
-  public async getAgent(): Promise<AtpAgent> {
-    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
-  }
-}
@@ -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',
@@ -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<string, SafelinkRule | 'ok'>
+  private urlCache: LRUCache<string, SafelinkRule | 'ok'>
+
+  private db: Database
+
+  constructor({
+    db,
+  }: {
+    identifier: string
+    password: string
+    pdsHost: string
+    db: Database
+  }) {
+    this.domainCache = new LRUCache<string, SafelinkRule | 'ok'>({
+      max: 10000,
+    })
+
+    this.urlCache = new LRUCache<string, SafelinkRule | 'ok'>({
+      max: 25000,
+    })
+
+    this.db = db
+  }
+
+  public async tryFindRule(link: string): Promise<SafelinkRule | 'ok'> {
+    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<CredentialSession> {
+    if (!this.session.hasSession) {
+      await this.getAgent()
+    }
+    return this.session
+  }
+
+  public async getAgent(): Promise<AtpAgent> {
+    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
+  }
+}
@@ -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<AppContextOptions>) {
@@ -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<Link>
+export type RuleEntry = Selectable<SafelinkRule>
@@ -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}`)
        }
      }

@@ -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"