start using database

This commit is contained in:
Hailey
2025-08-21 10:12:15 -07:00
parent b06abba6b2
commit 573b1f6ac2
10 changed files with 235 additions and 216 deletions
+1
View File
@@ -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",
+1 -130
View File
@@ -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
-49
View File
@@ -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
}
}
+2 -18
View File
@@ -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',
+167
View File
@@ -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
}
}
+3
View File
@@ -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>) {
+19
View File
@@ -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>
+10 -19
View File
@@ -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}`)
}
}
+5
View File
@@ -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"
+27
View File
@@ -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 `#`',
)
}
}