File size: 6,022 Bytes
5f2247e | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 | """Rivet Discipline Gate — pre-flight check on every suggestion.
Before any code suggestion reaches the user, it passes through this gate.
The gate checks for dangerous patterns, flags security implications, and
assigns a confidence level. If a suggestion fails the gate, it's blocked
or annotated with warnings.
"""
import re
from dataclasses import dataclass, field
from enum import Enum
class Confidence(str, Enum):
HIGH = "HIGH" # Traced the full code path
MEDIUM = "MEDIUM" # Read the code but not run it
LOW = "LOW" # Reasoning from architecture, not source
class Severity(str, Enum):
BLOCK = "BLOCK" # Do not suggest this
WARN = "WARN" # Suggest with prominent warning
INFO = "INFO" # Note for awareness
@dataclass
class GateResult:
passed: bool = True
flags: list = field(default_factory=list)
confidence: Confidence = Confidence.MEDIUM
requires_review: list = field(default_factory=list)
def add_flag(self, severity: Severity, message: str):
self.flags.append({"severity": severity.value, "message": message})
if severity == Severity.BLOCK:
self.passed = False
def format_warnings(self) -> str:
if not self.flags:
return ""
lines = ["\n⚠️ **Discipline Gate Flags:**"]
for f in self.flags:
icon = "🛑" if f["severity"] == "BLOCK" else "⚠️" if f["severity"] == "WARN" else "ℹ️"
lines.append(f" {icon} [{f['severity']}] {f['message']}")
if self.requires_review:
lines.append(f"\n 📋 Requires review from: {', '.join(self.requires_review)}")
lines.append(f"\n 📊 Confidence: {self.confidence.value}")
return "\n".join(lines)
# Pattern matchers for dangerous operations
DESTRUCTIVE_MIGRATION_PATTERNS = [
r'\bDROP\s+(TABLE|COLUMN|INDEX|CONSTRAINT)',
r'\bALTER\s+TABLE\s+\w+\s+DROP',
r'\bALTER\s+TABLE\s+\w+\s+ALTER\s+COLUMN\s+\w+\s+TYPE',
r'\bTRUNCATE\s+TABLE',
r'\bDELETE\s+FROM\s+\w+\s*;', # Unqualified DELETE
]
AUTH_PATTERNS = [
r'\bJWT_SECRET\b',
r'\bverifyToken\b',
r'\brequireAdmin\b',
r'\brequireModerator\b',
r'\bauth\s*middleware\b',
r'\bsession\s*cookie\b',
r'\brejectIfIneligible\b',
r'\bBearer\b',
]
WEBHOOK_BYPASS_PATTERNS = [
r"process\.env\.\w+\s*\|\|\s*['\"]", # env fallback to empty string
r"if\s*\(\s*!secret\s*\)\s*return", # skip on missing secret
r"if\s*\(\s*!.*SECRET.*\)\s*return",
]
TRANSACTION_GAPS = [
r'await\s+\w+\.(debit|credit|transfer|purchase|delete)\(', # Multi-step without tx
]
def check_migration_safety(content: str) -> list:
"""Check for destructive migration patterns."""
flags = []
for pattern in DESTRUCTIVE_MIGRATION_PATTERNS:
matches = re.findall(pattern, content, re.IGNORECASE)
if matches:
flags.append((Severity.BLOCK,
f"Destructive migration detected: {pattern}. "
f"Staging and prod share the database — additive only."))
return flags
def check_auth_impact(content: str) -> list:
"""Flag anything touching authentication."""
flags = []
for pattern in AUTH_PATTERNS:
if re.search(pattern, content, re.IGNORECASE):
flags.append((Severity.WARN,
"This touches authentication. Review with security before merging."))
break
return flags
def check_webhook_safety(content: str) -> list:
"""Check for webhook signature bypass patterns (Audit C2)."""
flags = []
for pattern in WEBHOOK_BYPASS_PATTERNS:
if re.search(pattern, content):
flags.append((Severity.WARN,
"Webhook signature bypass pattern detected (Audit C2). "
"Never skip verification when env var is missing — reject instead."))
break
return flags
def check_transaction_safety(content: str) -> list:
"""Flag multi-step DB operations that may need transaction wrapping (Audit H3)."""
flags = []
matches = re.findall(r'await\s+\w+\.\w+\(', content)
if len(matches) >= 2:
for pattern in TRANSACTION_GAPS:
if re.search(pattern, content):
flags.append((Severity.WARN,
"Multi-step DB operation detected. Ensure these are wrapped "
"in a single transaction (Audit H3: gem debit without tx protection)."))
break
return flags
def run_gate(suggestion: str, context: str = "") -> GateResult:
"""Run the full discipline gate on a suggestion.
Args:
suggestion: The code/text being suggested to the user
context: Optional surrounding context (the question, file being edited)
Returns:
GateResult with pass/fail, flags, and confidence
"""
result = GateResult()
full_content = f"{context}\n{suggestion}"
# Run all checks
for severity, message in check_migration_safety(full_content):
result.add_flag(severity, message)
for severity, message in check_auth_impact(full_content):
result.add_flag(severity, message)
result.requires_review.append("security")
for severity, message in check_webhook_safety(full_content):
result.add_flag(severity, message)
result.requires_review.append("security")
for severity, message in check_transaction_safety(full_content):
result.add_flag(severity, message)
return result
if __name__ == "__main__":
# Quick test
test_migration = "ALTER TABLE students DROP COLUMN legacy_score;"
result = run_gate(test_migration)
print(f"Migration test: passed={result.passed}")
print(result.format_warnings())
test_webhook = """
const secret = process.env.STRIPE_SECRET || ''
if (!secret) return next()
"""
result = run_gate(test_webhook)
print(f"\nWebhook test: passed={result.passed}")
print(result.format_warnings())
|