@oculum/scanner 1.0.14 → 1.0.15

package/src/__tests__/taint/source-classifier.test.ts

@@ -0,0 +1,367 @@
+/**
+ * Unit tests for AST-based source classifier
+ */
+
+import Parser from 'tree-sitter'
+import { createHash } from 'crypto'
+// eslint-disable-next-line @typescript-eslint/no-require-imports
+const { typescript: TypeScript, tsx: TSX } = require('tree-sitter-typescript')
+import type { ParsedAST } from '../../parse/ast'
+import { classifySources } from '../../taint/source-classifier'
+import type { TaintSource, TaintSourceType } from '../../taint/types'
+
+// Dedicated parsers for this test file — kept at module scope to prevent native GC
+const tsParser = new Parser()
+tsParser.setLanguage(TypeScript)
+const tsxParser = new Parser()
+tsxParser.setLanguage(TSX)
+
+let counter = 0
+function getSources(code: string, filePath = 'test.ts'): TaintSource[] {
+  const fp = filePath === 'test.ts' ? `source-test-${++counter}.ts` : filePath
+  const parser = fp.endsWith('.tsx') ? tsxParser : tsParser
+  const tree = parser.parse(code)
+  const hash = createHash('sha1').update(code).digest('hex')
+  const ast: ParsedAST = { tree, language: fp.endsWith('.tsx') ? 'tsx' : 'typescript', filePath: fp, hash }
+  return classifySources(ast)
+}
+
+function findSourceByType(sources: TaintSource[], type: TaintSourceType): TaintSource | undefined {
+  return sources.find(s => s.sourceType === type)
+}
+
+describe('classifySources', () => {
+  // ====================================================================
+  // Express / generic HTTP sources
+  // ====================================================================
+  describe('Express / generic HTTP', () => {
+    it('detects req.body', () => {
+      const sources = getSources(`
+        app.post('/api', (req, res) => {
+          const data = req.body
+        })
+      `)
+      const s = findSourceByType(sources, 'http_body')
+      expect(s).toBeDefined()
+      expect(s!.confidence).toBe('high')
+      expect(s!.framework).toBe('express')
+      expect(s!.taintKinds.has('sql')).toBe(true)
+      expect(s!.taintKinds.has('xss')).toBe(true)
+    })
+
+    it('detects req.query', () => {
+      const sources = getSources(`
+        const search = req.query.search
+      `)
+      const s = findSourceByType(sources, 'http_query')
+      expect(s).toBeDefined()
+      expect(s!.confidence).toBe('high')
+    })
+
+    it('detects req.params', () => {
+      const sources = getSources(`
+        const id = req.params.id
+      `)
+      const s = findSourceByType(sources, 'http_params')
+      expect(s).toBeDefined()
+    })
+
+    it('detects req.headers', () => {
+      const sources = getSources(`
+        const auth = req.headers.authorization
+      `)
+      const s = findSourceByType(sources, 'http_headers')
+      expect(s).toBeDefined()
+      // Headers carry fewer taint kinds
+      expect(s!.taintKinds.has('xss')).toBe(true)
+      expect(s!.taintKinds.has('command')).toBe(true)
+    })
+
+    it('detects req.cookies', () => {
+      const sources = getSources(`
+        const token = req.cookies.session
+      `)
+      const s = findSourceByType(sources, 'http_headers')
+      expect(s).toBeDefined()
+    })
+
+    it('detects deep property access: req.body.user.email', () => {
+      const sources = getSources(`
+        const email = req.body.user.email
+      `)
+      const s = findSourceByType(sources, 'http_body')
+      expect(s).toBeDefined()
+      expect(s!.expression).toContain('req.body')
+    })
+  })
+
+  // ====================================================================
+  // Destructuring
+  // ====================================================================
+  describe('destructuring', () => {
+    it('detects const { email } = req.body', () => {
+      const sources = getSources(`
+        const { email, name } = req.body
+      `)
+      const s = findSourceByType(sources, 'http_body')
+      expect(s).toBeDefined()
+    })
+
+    it('detects destructured with rename: const { email: e } = req.body', () => {
+      const sources = getSources(`
+        const { email: e } = req.body
+      `)
+      const s = findSourceByType(sources, 'http_body')
+      expect(s).toBeDefined()
+    })
+  })
+
+  // ====================================================================
+  // Async patterns (Next.js App Router)
+  // ====================================================================
+  describe('async patterns', () => {
+    it('detects const data = await req.json()', () => {
+      const sources = getSources(`
+        export async function POST(req: Request) {
+          const data = await req.json()
+        }
+      `)
+      const s = findSourceByType(sources, 'http_body')
+      expect(s).toBeDefined()
+      expect(s!.framework).toBe('nextjs')
+    })
+
+    it('detects await request.text()', () => {
+      const sources = getSources(`
+        async function handler(request: Request) {
+          const body = await request.text()
+        }
+      `)
+      const s = findSourceByType(sources, 'http_body')
+      expect(s).toBeDefined()
+    })
+
+    it('detects await request.formData()', () => {
+      const sources = getSources(`
+        async function handler(request: Request) {
+          const form = await request.formData()
+        }
+      `)
+      const s = findSourceByType(sources, 'form_data')
+      expect(s).toBeDefined()
+    })
+  })
+
+  // ====================================================================
+  // Type-based source detection
+  // ====================================================================
+  describe('type-based detection', () => {
+    it('detects parameter typed as NextRequest', () => {
+      const sources = getSources(`
+        export async function POST(req: NextRequest) {
+          const data = await req.json()
+        }
+      `)
+      // Should find both: the type-based source AND the member access source
+      expect(sources.length).toBeGreaterThanOrEqual(1)
+      // The type-based one marks the whole parameter as tainted
+      const typeBased = sources.find(s => s.expression.includes('NextRequest'))
+      expect(typeBased).toBeDefined()
+      expect(typeBased!.confidence).toBe('high')
+      expect(typeBased!.framework).toBe('nextjs')
+    })
+
+    it('detects parameter typed as Request', () => {
+      const sources = getSources(`
+        export function GET(request: Request) {
+          return new Response('hello')
+        }
+      `)
+      const s = sources.find(s => s.expression.includes('Request'))
+      expect(s).toBeDefined()
+    })
+
+    it('detects parameter typed as NextApiRequest', () => {
+      const sources = getSources(`
+        export default function handler(req: NextApiRequest, res: NextApiResponse) {
+          const data = req.body
+        }
+      `)
+      const typeBased = sources.find(s => s.expression.includes('NextApiRequest'))
+      expect(typeBased).toBeDefined()
+      expect(typeBased!.framework).toBe('nextjs')
+    })
+  })
+
+  // ====================================================================
+  // LLM output sources
+  // ====================================================================
+  describe('LLM output', () => {
+    it('detects OpenAI response.choices[0].message.content', () => {
+      const sources = getSources(`
+        const reply = response.choices[0].message.content
+      `)
+      const s = findSourceByType(sources, 'llm_output')
+      expect(s).toBeDefined()
+      expect(s!.taintKinds.has('command')).toBe(true)
+      expect(s!.taintKinds.has('sql')).toBe(true)
+      expect(s!.taintKinds.has('xss')).toBe(true)
+      expect(s!.taintKinds.has('prompt')).toBe(true)
+    })
+
+    it('detects Anthropic response.content[0].text', () => {
+      const sources = getSources(`
+        const text = response.content[0].text
+      `)
+      const s = findSourceByType(sources, 'llm_output')
+      expect(s).toBeDefined()
+    })
+
+    it('detects completion.choices[0].message.content', () => {
+      const sources = getSources(`
+        const content = completion.choices[0].message.content
+      `)
+      const s = findSourceByType(sources, 'llm_output')
+      expect(s).toBeDefined()
+    })
+
+    it('detects .generated_text', () => {
+      const sources = getSources(`
+        const text = result.generated_text
+      `)
+      const s = findSourceByType(sources, 'llm_output')
+      expect(s).toBeDefined()
+    })
+  })
+
+  // ====================================================================
+  // AI-specific sources
+  // ====================================================================
+  describe('AI-specific sources', () => {
+    it('detects tool output', () => {
+      const sources = getSources(`
+        const result = toolResult.output
+      `)
+      const s = findSourceByType(sources, 'tool_output')
+      expect(s).toBeDefined()
+      // Tool output carries all taint kinds
+      expect(s!.taintKinds.size).toBe(6)
+    })
+
+    it('detects tool_call.result', () => {
+      const sources = getSources(`
+        const data = tool_call.result
+      `)
+      const s = findSourceByType(sources, 'tool_output')
+      expect(s).toBeDefined()
+    })
+  })
+
+  // ====================================================================
+  // URL Search Params
+  // ====================================================================
+  describe('URL search params', () => {
+    it('detects searchParams.get()', () => {
+      const sources = getSources(`
+        const q = searchParams.get('query')
+      `)
+      const s = findSourceByType(sources, 'url_search_params')
+      expect(s).toBeDefined()
+    })
+  })
+
+  // ====================================================================
+  // Other frameworks
+  // ====================================================================
+  describe('other frameworks', () => {
+    it('detects Koa ctx.request.body', () => {
+      const sources = getSources(`
+        const data = ctx.request.body
+      `)
+      const s = findSourceByType(sources, 'http_body')
+      expect(s).toBeDefined()
+      expect(s!.framework).toBe('koa')
+    })
+
+    it('detects AWS Lambda event.body', () => {
+      const sources = getSources(`
+        exports.handler = async (event) => {
+          const body = event.body
+        }
+      `)
+      const s = findSourceByType(sources, 'http_body')
+      expect(s).toBeDefined()
+      expect(s!.framework).toBe('aws-lambda')
+    })
+
+    it('detects process.argv (CLI)', () => {
+      const sources = getSources(`
+        const args = process.argv
+      `)
+      const s = findSourceByType(sources, 'cli_args')
+      expect(s).toBeDefined()
+      expect(s!.confidence).toBe('low')
+    })
+
+    it('detects process.env', () => {
+      const sources = getSources(`
+        const dbUrl = process.env.DATABASE_URL
+      `)
+      const s = findSourceByType(sources, 'env_var')
+      expect(s).toBeDefined()
+      expect(s!.confidence).toBe('low')
+      expect(s!.taintKinds.has('ssrf')).toBe(false)
+      expect(s!.taintKinds.has('command')).toBe(true)
+      // env_var no longer carries 'path' taint — env vars are operator-controlled, not user input
+      expect(s!.taintKinds.has('path')).toBe(false)
+    })
+  })
+
+  // ====================================================================
+  // Negative cases
+  // ====================================================================
+  describe('negative cases', () => {
+    it('does NOT flag fetchConfig() as user input', () => {
+      const sources = getSources(`
+        const config = fetchConfig()
+      `)
+      expect(sources).toHaveLength(0)
+    })
+
+    it('does NOT flag Math.random() as user input', () => {
+      const sources = getSources(`
+        const val = Math.random()
+      `)
+      expect(sources).toHaveLength(0)
+    })
+
+    it('does NOT flag string literals as user input', () => {
+      const sources = getSources(`
+        const url = "https://example.com"
+      `)
+      expect(sources).toHaveLength(0)
+    })
+
+    it('does NOT flag local function calls', () => {
+      const sources = getSources(`
+        function getData() { return { name: 'test' } }
+        const data = getData()
+      `)
+      expect(sources).toHaveLength(0)
+    })
+  })
+
+  // ====================================================================
+  // Deduplication
+  // ====================================================================
+  describe('deduplication', () => {
+    it('does not double-count nested access like req.body.email', () => {
+      const sources = getSources(`
+        const email = req.body.email
+      `)
+      // Should find req.body and skip counting req.body.email as a separate source
+      const bodySources = sources.filter(s => s.sourceType === 'http_body')
+      expect(bodySources).toHaveLength(1)
+    })
+  })
+})