diff --git a/cmd/pgproxy/main.go b/cmd/pgproxy/main.go index 320de88..8c652cd 100644 --- a/cmd/pgproxy/main.go +++ b/cmd/pgproxy/main.go @@ -6,11 +6,15 @@ import ( "crypto/tls" "crypto/x509" "database/sql" + _ "github.com/microsoft/go-mssqldb" "log/slog" "net/http" "net/url" "os" "os/signal" + "pgproxy/internal/pipeline" + "pgproxy/internal/sqlnorm" + "pgproxy/internal/wire/tds" "strings" "syscall" "time" @@ -139,22 +143,22 @@ func main() { shuntSSLMode := parseSSLMode(cfg.ShuntNodeURL) proxyCfg := proxy.Config{ - ListenAddr: cfg.ListenAddr, - BackendAddr: parseBackendAddr(cfg.BackendURL), - BackendSSLMode: backendSSLMode, - BackendUser: backendUser, - BackendPassword: backendPassword, - BackendDatabase: backendDatabase, - ShuntAddr: parseBackendAddr(cfg.ShuntNodeURL), - ShuntSSLMode: shuntSSLMode, - ShuntUser: shuntUser, - ShuntPassword: shuntPassword, - ShuntDatabase: shuntDatabase, - Mode: cfg.Mode, - Analyzer: analyzer, - Cache: approvalCache, - Reporter: rep, - MetricsReporter: metricsRep, + ListenAddr: cfg.ListenAddr, + BackendAddr: parseBackendAddr(cfg.BackendURL), + BackendSSLMode: backendSSLMode, + BackendUser: backendUser, + BackendPassword: backendPassword, + BackendDatabase: backendDatabase, + ShuntAddr: parseBackendAddr(cfg.ShuntNodeURL), + ShuntSSLMode: shuntSSLMode, + ShuntUser: shuntUser, + ShuntPassword: shuntPassword, + ShuntDatabase: shuntDatabase, + Mode: cfg.Mode, + Analyzer: analyzer, + Cache: approvalCache, + Reporter: rep, + MetricsReporter: metricsRep, TLSConfig: tlsConfig, BackendTLSConfig: backendTLSConfig, RequireClientTLS: cfg.RequireClientTLS, @@ -234,6 +238,77 @@ func main() { errChan <- p.ListenAndServe() }() + // Experimental: TDS passthrough frontend (multi-protocol milestone 2, + // see docs/multiprotocol-design.md). Enabled only when both env vars + // are set, e.g. TDS_LISTEN_ADDR=:1434 TDS_BACKEND_ADDR=localhost:1433. + if tdsListen := os.Getenv("TDS_LISTEN_ADDR"); tdsListen != "" { + tdsBackend := os.Getenv("TDS_BACKEND_ADDR") + if tdsBackend == "" { + slog.Error("TDS_LISTEN_ADDR set but TDS_BACKEND_ADDR missing; TDS frontend disabled") + } else { + tdsSrv := &tds.Server{ListenAddr: tdsListen, BackendAddr: tdsBackend} + if certPath := os.Getenv("TDS_TLS_CERT_PATH"); certPath != "" { + cert, cerr := tls.LoadX509KeyPair(certPath, os.Getenv("TDS_TLS_KEY_PATH")) + if cerr != nil { + slog.Error("TDS TLS disabled: certificate load failed", "error", cerr) + } else { + tdsSrv.TLSConfig = tds.NewTLSConfig(cert) + slog.Info("TDS strict TLS enabled (TDS 8.0)", "cert", certPath) + } + } else if os.Getenv("TDS_TLS_SELF_SIGNED") == "true" { + tlsCfg, terr := tds.SelfSignedTLSConfig("localhost", "127.0.0.1") + if terr != nil { + slog.Error("TDS TLS disabled: self-signed generation failed", "error", terr) + } else { + tdsSrv.TLSConfig = tlsCfg + slog.Warn("TDS strict TLS enabled with SELF-SIGNED certificate (dev only)") + } + } + if cfg.AnalysisEnabled { + tdsMode := cfg.Mode + if tdsMode == config.ModeShunt { + slog.Warn("shunt mode not yet supported on the TDS frontend; using blocking") + tdsMode = config.ModeBlocking + } + tdsLLM, err := agent.NewBedrockClient(ctx) + if err != nil { + slog.Error("TDS analysis disabled: Bedrock client init failed", "error", err) + } else { + var tsqlTools *agent.TSQLTools + if dsn := os.Getenv("TDS_BACKEND_DSN"); dsn != "" { + tdsDB, derr := sql.Open("sqlserver", dsn) + if derr != nil { + slog.Error("TDS tools disabled: backend DSN open failed", "error", derr) + } else { + tsqlTools = agent.NewTSQLTools(tdsDB) + slog.Info("TDS agent tools enabled (DMV + SHOWPLAN)") + } + } else { + slog.Warn("TDS_BACKEND_DSN not set; tsql agent runs without engine tools") + } + tsqlPipeline := pipeline.New(pipeline.Config{ + Analyzer: agent.NewTSQLAgent(tdsLLM, tsqlTools), + Cache: approvalCache, + Mode: tdsMode, + SkipAnalysis: tds.SkipAnalysis, + Fingerprint: sqlnorm.Fingerprint, + AnalysisTimeout: 60 * time.Second, + CacheTTL: 24 * time.Hour, + }) + tdsSrv.Decide = func(ctx context.Context, sql string) pipeline.Verdict { + return tsqlPipeline.Decide(pipeline.Statement{Database: "tsql", SQL: sql}) + } + slog.Info("TDS analysis enabled", "mode", tdsMode, "persona", "tsql") + } + } + go func() { + if err := tdsSrv.ListenAndServe(ctx); err != nil { + errChan <- err + } + }() + } + } + // Wait for signal or error select { case sig := <-sigChan: @@ -380,9 +455,9 @@ func initMetricsReporter(ctx context.Context, cfg *config.Config) reporter.Metri func initCache(ctx context.Context, cfg *config.Config) (cache.ApprovalCache, masking.MetadataCache, masking.InferredMaskingCache, redis.UniversalClient) { if cfg.ValkeyURL != "" { vc, err := cache.NewValkeyApprovalCache(ctx, cache.ValkeyConfig{ - URL: cfg.ValkeyURL, - KeyPrefix: "pgproxy:", - ClusterMode: cfg.ValkeyClusterMode, + URL: cfg.ValkeyURL, + KeyPrefix: "pgproxy:", + ClusterMode: cfg.ValkeyClusterMode, IAMAuth: cfg.ValkeyIAMAuth, IAMUsername: cfg.ValkeyIAMUsername, CacheName: cfg.ValkeyCacheName, @@ -510,10 +585,10 @@ func initMasking(ctx context.Context, cfg *config.Config, db *sql.DB, metadataCa svc, err := masking.NewService(ctx, db, masking.ServiceConfig{ Enabled: true, ConfigPath: cfg.MaskingConfigPath, - Cache: nil, // Query rewrite cache: nil triggers default in-memory cache (caches by query fingerprint only, not database-aware) - MetadataCache: metadataCache, // Cache table column info to avoid slow information_schema queries - InferenceCache: inferenceCache, // Cache LLM-inferred masking decisions - DatabaseName: dbName, // Masking operates on single database extracted from BackendURL + Cache: nil, // Query rewrite cache: nil triggers default in-memory cache (caches by query fingerprint only, not database-aware) + MetadataCache: metadataCache, // Cache table column info to avoid slow information_schema queries + InferenceCache: inferenceCache, // Cache LLM-inferred masking decisions + DatabaseName: dbName, // Masking operates on single database extracted from BackendURL MetadataCacheTTL: cfg.MetadataCacheTTL, }, cfg) if err != nil { diff --git a/docs/multiprotocol-design.md b/docs/multiprotocol-design.md new file mode 100644 index 0000000..f766480 --- /dev/null +++ b/docs/multiprotocol-design.md @@ -0,0 +1,253 @@ +# Multi-Protocol Design: One Guardian, Two Wires (PostgreSQL + SQL Server) + +Status: PROPOSAL. Written against capitec/pg-proxy @ 25b5e64. + +## 1. The goal, stated precisely + +Per the zero-trust framing (see the PGProxy launch post): the proxy is an +**exchange that mediates between identities and data**, and it is a control +only when it is the **only permitted route** to the operational database. +Most estates run PostgreSQL/Aurora/Redshift *and* SQL Server side by side. +Therefore the goal is: + +> The same guardian binary, fronting either PostgreSQL or SQL Server, +> with identical identity, governance, masking, and audit semantics, +> selected by configuration. + +The protocol is a door. The exchange behind the doors is written once. + +## 2. Non-goals + +- **No dialect translation.** A SQL Server client talks T-SQL to a real SQL + Server. We never convert T-SQL to Postgres SQL (that is Babelfish's job, and + the reason Babelfish is not our starting point: it is a server-side Postgres + extension for *leaving* SQL Server, not a proxy for guarding it). +- No NTLM/Kerberos termination in v1 (see auth, section 6). +- No result-set rewriting. Masking remains query rewriting, on both wires. + +## 3. What must change in the current code + +Today `internal/proxy` owns both the Postgres wire protocol AND the decision +pipeline (fingerprint, cache lookup, analyze, mode decision, mask, forward). +`DATABASE_TYPE: postgresql | redshift` conflates two axes that must split: + +- **wire**: how bytes frame messages (pgwire vs TDS) +- **dialect**: how SQL is parsed, judged, and error-formatted + (postgres | redshift | tsql) + +Aurora/Redshift proved the dialect axis works. TDS forces the wire axis. + +## 4. Target architecture + +``` +cmd/proxy/ one binary; config picks frontend(s) +internal/wire/ + wire.go the Frontend contract (below) + postgres/ moved from internal/proxy, behavior-identical + tds/ new +internal/pipeline/ extracted decision core, protocol-neutral +internal/dialect/ registry: parser, persona, tools, error mapper +internal/agent/ + tsql persona and tool set +internal/masking/ parser becomes pluggable (libpg_query | ScriptDom) +internal/{cache,reporter,config,auth,metadata} unchanged +``` + +### The Frontend contract + +Small on purpose. Everything above it is written once. + +```go +// A Frontend owns one client connection end-to-end. +type Frontend interface { + // Handshake runs the protocol's startup/auth phase. + // It returns the authenticated Identity (from JWT/FedAuth/etc) + // and a connected Backend for the session. + Handshake(ctx context.Context) (Identity, Backend, error) + + // Next blocks until the client sends an interceptable statement. + // Non-statement traffic is forwarded transparently inside Next. + Next(ctx context.Context) (Statement, error) + + // Deliver applies a verdict in protocol-native form: + // forward (possibly rewritten), or reject with reason/detail/fix. + Deliver(ctx context.Context, s Statement, v Verdict) error +} + +type Statement struct { + SQL string + Kind StatementKind // batch | prepared | exec-prepared + Params []Param // TDS RPC / pgwire Bind: already separated + Fingerprint string // computed by pipeline, cached here +} +``` + +### The pipeline (extracted, unchanged in behavior) + +fingerprint → skip-list → cache → (agent | fail-open) → mode decision +(block / shunt / observe) → mask rewrite → forward → report. + +Invariant during extraction: **the existing BDD suite stays green with zero +scenario edits.** The Postgres frontend after refactor must be +bit-for-bit compatible on the wire. + +### The dialect registry + +```go +type Dialect struct { + Name string // postgres | redshift | tsql + Parser SQLParser // parse/deparse for masking + write detect + Persona AgentPersona // system prompt + Tools []AgentTool // explain/stats per engine + ErrorMapper func(Verdict) WireError // hint placement differs per wire +} +``` + +## 5. TDS frontend: the actual work + +TDS (MS-TDS, public spec) maps cleanly onto the existing two-phase design: + +1. **Framing**: 8-byte packet headers (type, status incl. EOM bit, length, + SPID). A token-stream parser is needed only for the message types we touch. +2. **PRELOGIN**: the one genuinely weird part. TLS is negotiated *inside* + TDS packets: the TLS handshake records are encapsulated in PRELOGIN-phase + packets, after which the stream switches to normal TLS framing. The proxy + must terminate client TLS and originate backend TLS (it already rewrites + auth, so it is a TLS MITM by design, same as the pgwire path with + REQUIRE_BACKEND_TLS). This is the SCRAM-war-story equivalent; budget for it. +3. **LOGIN7 + auth** (section 6). +4. **Query phase**: intercept exactly three client packet types: + - `SQLBatch` (type 0x01): ad-hoc T-SQL. The pgwire Query analog. + - `RPC` (0x03): parameterized calls (`sp_executesql`, prepared handles + `sp_prepare`/`sp_execute`). The Parse/Bind analog. **Parameters arrive + separated from SQL text** — see 8a for why this is a gift. + - `TransactionManager` (0x0E): BEGIN/COMMIT/ROLLBACK arrive as their own + packet type — cleaner than pgwire, where we string-match. Shunt mode's + "writes and transactions blocked" gets *more* rigorous on TDS. + Everything else (attention, bulk load in v1) passes through or is refused + by policy (bulk load = write; blocked in shunt, analyzable later). +5. **Server→client**: pass through unchanged except when injecting a rejection. + +### Backend leg +`microsoft/go-mssqldb` provides client-side TDS. Where its abstraction is too +high for raw relay, vendor the packet layer (MIT) rather than reimplement. + +## 6. Identity on TDS: better than the password hack + +The pgwire door smuggles the JWT through the password field. TDS is nicer: +**FEDAUTH** (MS-TDS feature ext) carries an OAuth2/Entra access token natively +— every modern SQL Server client (`sqlcmd -G`, JDBC/ODBC `Authentication= +ActiveDirectory*`, SSMS) already knows how to send a token. So on the SQL +Server door, the zero-trust identity story is *native protocol*, not a trick: + +- v1 auth modes: (a) SQL auth passthrough; (b) JWT-in-password (parity with + pgwire, works everywhere); (c) **FEDAUTH token validation** against the same + OIDC config (Entra) — proxy validates, then authenticates to the backend + with the service credential, exactly the pgwire pattern. +- Attribution: pgwire injects `application_name = jwt:`. TDS equivalents: + LOGIN7 `AppName`/`HostName` rewrite (visible in `sys.dm_exec_sessions + .program_name`) — same DBA-visible audit trail. +- NTLM/Kerberos (integrated auth) termination is explicitly out of scope for + v1; OIDC/Entra is the strategic direction, and FEDAUTH covers it natively. + +## 7. Agent dialect: tsql + +- Persona: a career SQL Server engine specialist (pick the character later; + the Tom Lane trick, applied to the SQL Server world). +- Tools (DMV-backed, read-only, same fail-open rules): + - `get_schema`: sys.tables/columns/types + - `get_indexes`: sys.indexes + sys.dm_db_index_usage_stats + - `get_table_stats`: sys.dm_db_partition_stats (rowcounts, sizes) + - `get_query_stats`: sys.dm_exec_query_stats top offenders + - `explain_query`: `SET SHOWPLAN_XML ON` on a **dedicated session** + (SHOWPLAN is session-sticky and must wrap the statement; keep a pooled + analysis connection, never the client's) +- Reject criteria tuned per engine (missing join predicates, scans on large + rowcounts, implicit conversions killing sargability, SELECT * on wide + tables, missing TOP/WHERE on huge estimates). +- Bedrock/Converse, temperature 0, MaxTokens, caching: unchanged. + +## 8. Masking on T-SQL + +- **Parser**: Microsoft **ScriptDom** — their production T-SQL parser, + open source, parse + script-generate (the libpg_query/Deparse analog). + It is .NET: run as a **sidecar** (small HTTP/JSON service, containerized, + stateless) behind the existing `QueryMasker` interface. Contract: + `POST /rewrite {sql, maskingMetadata} -> {sql', maskedColumns, unknownTables}`. + The Go proxy keeps the cache, config, and behavior policy; the sidecar only + parses and rewrites. Fail-open on sidecar unavailability, same as LLM path. +- **Mask expressions** in T-SQL: CONCAT/LEFT/RIGHT/STUFF for partial; + email/phone via STRING_SPLIT-free scalar expressions; hash mask via + `CONVERT(varchar(16), HASHBYTES('SHA2_256', CONCAT(@key, col)), 2)` — + which also delivers the blog's roadmap item (keyed HMAC-style hashing, + retiring MD5-truncation) — implement on the Postgres side simultaneously + (pgcrypto hmac) so both dialects land the upgrade together. +- Sensitivity lineage (CTEs, SELECT *, aggregates-not-sensitive, UNION + any-branch) re-implemented over ScriptDom's AST; port the rule table and + the test corpus, not the code. + +### 8a. The blog's #1 roadmap item, largely free on TDS +"Normalize literals before sending query text to Bedrock." On the RPC path, +parameters are *already separated* — the SQL text contains `@p1`, values ride +in typed fields the proxy never needs to forward to the agent. For SQLBatch +(and pgwire simple queries), add a literal-stripping pass (both parsers can do +this precisely; today's regex fingerprint normalizer is the fallback) so the +agent sees shapes, not values. This lands as part of milestone 1 since it +lives in the shared pipeline and improves the Postgres product immediately. + +## 9. Error delivery per wire + +- pgwire: ErrorResponse, fix in **Hint** (unchanged). +- TDS: ERROR token — number (pick a stable custom number, e.g. 50999), + severity 16, state, message = reason + newline + `Fix: `; sqlcmd, + SSMS, and drivers all render message text. Same content, native envelope. + +## 10. Configuration + +``` +LISTENERS: pg://:5432/postgres, tds://:1433/tsql # protocol + port + dialect +# DATABASE_TYPE deprecated -> split: +WIRE_PROTOCOL: postgres | tds (per listener) +SQL_DIALECT: postgres | redshift | tsql +REAL_DB_URL / SHUNT_DB_URL (sqlserver:// URLs for tds) +``` +One process may run both listeners (dev convenience); production runs one +guardian per backend, per the only-permitted-route rule. + +## 11. Testing strategy + +- **Invariant**: existing Postgres BDD suite green, unmodified, at every + milestone. This is the contract that the refactor changed nothing. +- TDS integration: SQL Server 2022 container (`mcr.microsoft.com/mssql/server`; + on Apple Silicon runs under Rosetta emulation in Rancher — slow but fine for + CI-of-one; azure-sql-edge is the arm64 fallback with reduced surface). +- BDD scenarios mirrored per frontend from one scenario source where possible + (godog step defs parameterized by client: pgx vs go-mssqldb). +- Golden sets for the tsql agent (must-reject / must-accept + hint-repairs-it, + per the eval philosophy). +- Mutation testing continues to gate the masking rule table. + +## 12. Milestones (each independently shippable) + +1. **Pipeline extraction + literal normalization** — pure refactor plus the + blog's top roadmap item. Risk lives here; BDD is the net. No new features. +2. **TDS passthrough** — sqlcmd → proxy → SQL Server container, raw relay, + TLS-in-PRELOGIN handled. (The "pgadmin working" commit of this project.) +3. **TDS interception + blocking** — SQLBatch/RPC into the shared pipeline, + tsql persona + DMV tools, ERROR-token delivery. Shunt + observe come almost + free (routing logic is pipeline-side; TM packets make write-blocking exact). +4. **T-SQL masking** — ScriptDom sidecar + HMAC mask type on both dialects. +5. **FEDAUTH** — native Entra token auth on the TDS door. + +## 13. Risks, honestly + +- **PRELOGIN TLS encapsulation** is fiddly; mitigate by writing it against + captured packet traces from real clients (sqlcmd, JDBC, SSMS) first. +- **go-mssqldb abstraction fit** unknown until milestone 2; fallback is + vendoring its packet layer. +- **ScriptDom sidecar** adds an operational component; contained by strict + fail-open and statelessness. +- **SHOWPLAN needs its own session** — a pooled analysis connection is new + machinery; without it the agent's best tool is missing. +- **Naming**: pg-proxy fronting SQL Server reads oddly. Defer; a module + rename is cheap once the frontend proves out. +``` diff --git a/go.mod b/go.mod index e4b314a..568e32b 100644 --- a/go.mod +++ b/go.mod @@ -18,13 +18,16 @@ require ( github.com/jackc/pgx/v5 v5.9.2 github.com/joho/godotenv v1.5.1 github.com/lib/pq v1.12.3 + github.com/microsoft/go-mssqldb v1.10.0 github.com/pganalyze/pg_query_go/v6 v6.2.2 github.com/redis/go-redis/v9 v9.18.0 golang.org/x/crypto v0.50.0 + golang.org/x/time v0.15.0 gopkg.in/yaml.v3 v3.0.1 ) require ( + github.com/aws/aws-sdk-go-v2/aws/protocol/eventstream v1.7.9 // indirect github.com/aws/aws-sdk-go-v2/feature/ec2/imds v1.18.22 // indirect github.com/aws/aws-sdk-go-v2/internal/configsources v1.4.22 // indirect github.com/aws/aws-sdk-go-v2/internal/endpoints/v2 v2.7.22 // indirect @@ -40,6 +43,8 @@ require ( github.com/cucumber/messages/go/v21 v21.0.1 // indirect github.com/dgryski/go-rendezvous v0.0.0-20200823014737-9f7001d12a5f // indirect github.com/gofrs/uuid v4.3.1+incompatible // indirect + github.com/golang-sql/civil v0.0.0-20220223132316-b832511892a9 // indirect + github.com/golang-sql/sqlexp v0.1.0 // indirect github.com/hashicorp/go-immutable-radix v1.3.1 // indirect github.com/hashicorp/go-memdb v1.3.4 // indirect github.com/hashicorp/golang-lru v0.5.4 // indirect @@ -48,11 +53,11 @@ require ( github.com/jackc/puddle/v2 v2.2.2 // indirect github.com/kr/text v0.2.0 // indirect github.com/rogpeppe/go-internal v1.14.1 // indirect + github.com/shopspring/decimal v1.4.0 // indirect github.com/spf13/pflag v1.0.7 // indirect github.com/yuin/gopher-lua v1.1.1 // indirect go.uber.org/atomic v1.11.0 // indirect golang.org/x/sync v0.20.0 // indirect golang.org/x/text v0.36.0 // indirect - golang.org/x/time v0.15.0 // indirect google.golang.org/protobuf v1.36.7 // indirect ) diff --git a/go.sum b/go.sum index 016822e..ec0d608 100644 --- a/go.sum +++ b/go.sum @@ -1,11 +1,25 @@ +github.com/Azure/azure-sdk-for-go/sdk/azcore v1.21.1 h1:jHb/wfvRikGdxMXYV3QG/SzUOPYN9KEUUuC0Yd0/vC0= +github.com/Azure/azure-sdk-for-go/sdk/azcore v1.21.1/go.mod h1:pzBXCYn05zvYIrwLgtK8Ap8QcjRg+0i76tMQdWN6wOk= +github.com/Azure/azure-sdk-for-go/sdk/azidentity v1.13.1 h1:Hk5QBxZQC1jb2Fwj6mpzme37xbCDdNTxU7O9eb5+LB4= +github.com/Azure/azure-sdk-for-go/sdk/azidentity v1.13.1/go.mod h1:IYus9qsFobWIc2YVwe/WPjcnyCkPKtnHAqUYeebc8z0= +github.com/Azure/azure-sdk-for-go/sdk/internal v1.12.0 h1:fhqpLE3UEXi9lPaBRpQ6XuRW0nU7hgg4zlmZZa+a9q4= +github.com/Azure/azure-sdk-for-go/sdk/internal v1.12.0/go.mod h1:7dCRMLwisfRH3dBupKeNCioWYUZ4SS09Z14H+7i8ZoY= +github.com/Azure/azure-sdk-for-go/sdk/security/keyvault/azkeys v1.4.0 h1:E4MgwLBGeVB5f2MdcIVD3ELVAWpr+WD6MUe1i+tM/PA= +github.com/Azure/azure-sdk-for-go/sdk/security/keyvault/azkeys v1.4.0/go.mod h1:Y2b/1clN4zsAoUd/pgNAQHjLDnTis/6ROkUfyob6psM= +github.com/Azure/azure-sdk-for-go/sdk/security/keyvault/internal v1.2.0 h1:nCYfgcSyHZXJI8J0IWE5MsCGlb2xp9fJiXyxWgmOFg4= +github.com/Azure/azure-sdk-for-go/sdk/security/keyvault/internal v1.2.0/go.mod h1:ucUjca2JtSZboY8IoUqyQyuuXvwbMBVwFOm0vdQPNhA= github.com/Azure/go-ansiterm v0.0.0-20230124172434-306776ec8161 h1:L/gRVlceqvL25UVaW/CKtUDjefjrs0SPonmDGUVOYP0= github.com/Azure/go-ansiterm v0.0.0-20230124172434-306776ec8161/go.mod h1:xomTg63KZ2rFqZQzSB4Vz2SUXa1BpHTVz9L5PTmPC4E= +github.com/AzureAD/microsoft-authentication-library-for-go v1.6.0 h1:XRzhVemXdgvJqCH0sFfrBUTnUJSBrBf7++ypk+twtRs= +github.com/AzureAD/microsoft-authentication-library-for-go v1.6.0/go.mod h1:HKpQxkWaGLJ+D/5H8QRpyQXA1eKjxkFlOMwck5+33Jk= github.com/Microsoft/go-winio v0.6.2 h1:F2VQgta7ecxGYO8k3ZZz3RS8fVIXVxONVUPlNERoyfY= github.com/Microsoft/go-winio v0.6.2/go.mod h1:yd8OoFMLzJbo9gZq8j5qaps8bJ9aShtEA8Ipt1oGCvU= github.com/alicebob/miniredis/v2 v2.37.0 h1:RheObYW32G1aiJIj81XVt78ZHJpHonHLHW7OLIshq68= github.com/alicebob/miniredis/v2 v2.37.0/go.mod h1:TcL7YfarKPGDAthEtl5NBeHZfeUQj6OXMm/+iu5cLMM= github.com/aws/aws-sdk-go-v2 v1.41.6 h1:1AX0AthnBQzMx1vbmir3Y4WsnJgiydmnJjiLu+LvXOg= github.com/aws/aws-sdk-go-v2 v1.41.6/go.mod h1:dy0UzBIfwSeot4grGvY1AqFWN5zgziMmWGzysDnHFcQ= +github.com/aws/aws-sdk-go-v2/aws/protocol/eventstream v1.7.9 h1:adBsCIIpLbLmYnkQU+nAChU5yhVTvu5PerROm+/Kq2A= +github.com/aws/aws-sdk-go-v2/aws/protocol/eventstream v1.7.9/go.mod h1:uOYhgfgThm/ZyAuJGNQ5YgNyOlYfqnGpTHXvk3cpykg= github.com/aws/aws-sdk-go-v2/config v1.32.16 h1:Q0iQ7quUgJP0F/SCRTieScnaMdXr9h/2+wze1u3cNeM= github.com/aws/aws-sdk-go-v2/config v1.32.16/go.mod h1:duCCnJEFqpt2RC6no1iK6q+8HpwOAkiUua0pY507dQc= github.com/aws/aws-sdk-go-v2/credentials v1.19.15 h1:fyvgWTszojq8hEnMi8PPBTvZdTtEVmAVyo+NFLHBhH4= @@ -88,6 +102,10 @@ github.com/golang-jwt/jwt/v5 v5.3.1 h1:kYf81DTWFe7t+1VvL7eS+jKFVWaUnK9cB1qbwn63Y github.com/golang-jwt/jwt/v5 v5.3.1/go.mod h1:fxCRLWMO43lRc8nhHWY6LGqRcf+1gQWArsqaEUEa5bE= github.com/golang-migrate/migrate/v4 v4.19.1 h1:OCyb44lFuQfYXYLx1SCxPZQGU7mcaZ7gH9yH4jSFbBA= github.com/golang-migrate/migrate/v4 v4.19.1/go.mod h1:CTcgfjxhaUtsLipnLoQRWCrjYXycRz/g5+RWDuYgPrE= +github.com/golang-sql/civil v0.0.0-20220223132316-b832511892a9 h1:au07oEsX2xN0ktxqI+Sida1w446QrXBRJ0nee3SNZlA= +github.com/golang-sql/civil v0.0.0-20220223132316-b832511892a9/go.mod h1:8vg3r2VgvsThLBIFL93Qb5yWzgyZWhEmBwUJWevAkK0= +github.com/golang-sql/sqlexp v0.1.0 h1:ZCD6MBpcuOVfGVqsEmY5/4FtYiKz6tSyUv9LPEDei6A= +github.com/golang-sql/sqlexp v0.1.0/go.mod h1:J4ad9Vo8ZCWQ2GMrC4UCQy1JpCbwU9m3EOqtpKwwwHI= github.com/google/go-cmp v0.5.5 h1:Khx7svrCpmxxtHBq5j2mp/xVjsi8hQMfNLvJFAlrGgU= github.com/google/go-cmp v0.5.5/go.mod h1:v8dTdLbMG2kIc/vJvl+f65V22dbkXbowE6jgT/gNBxE= github.com/google/uuid v1.6.0 h1:NIvaJDMOsjHA8n1jAhLSgzrAzy1Hgr+hNrb57e+94F0= @@ -98,8 +116,9 @@ github.com/hashicorp/go-immutable-radix v1.3.1/go.mod h1:0y9vanUI8NX6FsYoO3zeMjh github.com/hashicorp/go-memdb v1.3.4 h1:XSL3NR682X/cVk2IeV0d70N4DZ9ljI885xAEU8IoK3c= github.com/hashicorp/go-memdb v1.3.4/go.mod h1:uBTr1oQbtuMgd1SSGoR8YV27eT3sBHbYiNm53bMpgSg= github.com/hashicorp/go-uuid v1.0.0/go.mod h1:6SBZvOh/SIDV7/2o3Jml5SYk/TvGqwFJ/bN7x4byOro= -github.com/hashicorp/go-uuid v1.0.2 h1:cfejS+Tpcp13yd5nYHWDI6qVCny6wyX2Mt5SGur2IGE= github.com/hashicorp/go-uuid v1.0.2/go.mod h1:6SBZvOh/SIDV7/2o3Jml5SYk/TvGqwFJ/bN7x4byOro= +github.com/hashicorp/go-uuid v1.0.3 h1:2gKiV6YVmrJ1i2CKKa9obLvRieoRGviZFL26PcT/Co8= +github.com/hashicorp/go-uuid v1.0.3/go.mod h1:6SBZvOh/SIDV7/2o3Jml5SYk/TvGqwFJ/bN7x4byOro= github.com/hashicorp/golang-lru v0.5.0/go.mod h1:/m3WP610KZHVQ1SGc6re/UDhFvYD7pJ4Ao+sR/qLZy8= github.com/hashicorp/golang-lru v0.5.4 h1:YDjusn29QI/Das2iO9M0BHnIbxPeyuCHsjMW+lJfyTc= github.com/hashicorp/golang-lru v0.5.4/go.mod h1:iADmTwqILo4mZ8BN3D2Q6+9jd8WM5uGBxy+E8yxSoD4= @@ -123,8 +142,12 @@ github.com/kr/pty v1.1.1/go.mod h1:pFQYn66WHrOpPYNljwOMqo10TkYh1fy3cYio2l3bCsQ= github.com/kr/text v0.1.0/go.mod h1:4Jbv+DJW3UT/LiOwJeYQe1efqtUx/iVham/4vfdArNI= github.com/kr/text v0.2.0 h1:5Nx0Ya0ZqY2ygV366QzturHI13Jq95ApcVaJBhpS+AY= github.com/kr/text v0.2.0/go.mod h1:eLer722TekiGuMkidMxC/pM04lWEeraHUUmBw8l2grE= +github.com/kylelemons/godebug v1.1.0 h1:RPNrshWIDI6G2gRW9EHilWtl7Z6Sb1BR0xunSBf0SNc= +github.com/kylelemons/godebug v1.1.0/go.mod h1:9/0rRGxNHcop5bhtWyNeEfOS8JIWk580+fNqagV/RAw= github.com/lib/pq v1.12.3 h1:tTWxr2YLKwIvK90ZXEw8GP7UFHtcbTtty8zsI+YjrfQ= github.com/lib/pq v1.12.3/go.mod h1:/p+8NSbOcwzAEI7wiMXFlgydTwcgTr3OSKMsD2BitpA= +github.com/microsoft/go-mssqldb v1.10.0 h1:pHEt+Qz6YFPWqREq10mqSE524QQo+/QremwTCQht7TY= +github.com/microsoft/go-mssqldb v1.10.0/go.mod h1:mnG7lGa9iYJbzJqGCXyuQCegStKMr3kogDLD6+bmggg= github.com/moby/docker-image-spec v1.3.1 h1:jMKff3w6PgbfSa69GfNg+zN/XLhfXJGnEx3Nl2EsFP0= github.com/moby/docker-image-spec v1.3.1/go.mod h1:eKmb5VW8vQEh/BAr2yvVNvuiJuY6UIocYsFu/DxxRpo= github.com/moby/term v0.5.0 h1:xt8Q1nalod/v7BqbG21f8mQPqH+xAaC9C3N3wfWbVP0= @@ -137,6 +160,8 @@ github.com/opencontainers/image-spec v1.1.0 h1:8SG7/vwALn54lVB/0yZ/MMwhFrPYtpEHQ github.com/opencontainers/image-spec v1.1.0/go.mod h1:W4s4sFTMaBeK1BQLXbG4AdM2szdn85PY75RI83NrTrM= github.com/pganalyze/pg_query_go/v6 v6.2.2 h1:O0L6zMC226R82RF3X5n0Ki6HjytDsoAzuzp4ATVAHNo= github.com/pganalyze/pg_query_go/v6 v6.2.2/go.mod h1:Cn6+j4870kJz3iYNsb0VsNG04vpSWgEvBwc590J4qD0= +github.com/pkg/browser v0.0.0-20240102092130-5ac0b6a4141c h1:+mdjkGKdHQG3305AYmdv1U2eRNDiU2ErMBj1gwrq8eQ= +github.com/pkg/browser v0.0.0-20240102092130-5ac0b6a4141c/go.mod h1:7rwL4CYBLnjLxUqIJNnCWiEdr3bn6IUYi15bNlnbCCU= github.com/pkg/errors v0.9.1 h1:FEBLx1zS214owpjy7qsBeixbURkuhQAwrK5UwLGTwt4= github.com/pkg/errors v0.9.1/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0= github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4= @@ -147,6 +172,8 @@ github.com/redis/go-redis/v9 v9.18.0/go.mod h1:k3ufPphLU5YXwNTUcCRXGxUoF1fqxnhFQ github.com/rogpeppe/go-internal v1.14.1 h1:UQB4HGPB6osV0SQTLymcB4TgvyWu6ZyliaW0tI/otEQ= github.com/rogpeppe/go-internal v1.14.1/go.mod h1:MaRKkUm5W0goXpeCfT7UZI6fk/L7L7so1lCWt35ZSgc= github.com/russross/blackfriday/v2 v2.1.0/go.mod h1:+Rmxgy9KzJVeS9/2gXHxylqXiyQDYRxCVz55jmeOWTM= +github.com/shopspring/decimal v1.4.0 h1:bxl37RwXBklmTi0C79JfXCEBD1cqqHt0bbgBAGFp81k= +github.com/shopspring/decimal v1.4.0/go.mod h1:gawqmDU56v4yIKSwfBSFip1HdCCXN8/+DMd9qYNcwME= github.com/spf13/cobra v1.7.0/go.mod h1:uLxZILRyS/50WlhOIKD7W6V5bgeIt+4sICxh6uRMrb0= github.com/spf13/pflag v1.0.5/go.mod h1:McXfInJRrz4CZXVZOBLb0bTZqETkiAhM9Iw0y3An2Bg= github.com/spf13/pflag v1.0.7 h1:vN6T9TfwStFPFM5XzjsvmzZkLuaLX+HS+0SeFLRgU6M= @@ -180,6 +207,8 @@ go.uber.org/atomic v1.11.0 h1:ZvwS0R+56ePWxUNi+Atn9dWONBPp/AUETXlHW0DxSjE= go.uber.org/atomic v1.11.0/go.mod h1:LUxbIzbOniOlMKjJjyPfpl4v+PKK2cNJn91OQbhoJI0= golang.org/x/crypto v0.50.0 h1:zO47/JPrL6vsNkINmLoo/PH1gcxpls50DNogFvB5ZGI= golang.org/x/crypto v0.50.0/go.mod h1:3muZ7vA7PBCE6xgPX7nkzzjiUq87kRItoJQM1Yo8S+Q= +golang.org/x/net v0.53.0 h1:d+qAbo5L0orcWAr0a9JweQpjXF19LMXJE8Ey7hwOdUA= +golang.org/x/net v0.53.0/go.mod h1:JvMuJH7rrdiCfbeHoo3fCQU24Lf5JJwT9W3sJFulfgs= golang.org/x/sync v0.20.0 h1:e0PTpb7pjO8GAtTs2dQ6jYa5BWYlMuX047Dco/pItO4= golang.org/x/sync v0.20.0/go.mod h1:9xrNwdLfx4jkKbNva9FpL6vEN7evnE43NNNJQ2LF3+0= golang.org/x/sys v0.43.0 h1:Rlag2XtaFTxp19wS8MXlJwTvoh8ArU6ezoyFsMyCTNI= diff --git a/internal/agent/agent.go b/internal/agent/agent.go index 8e08c58..8065649 100644 --- a/internal/agent/agent.go +++ b/internal/agent/agent.go @@ -6,6 +6,7 @@ import ( "encoding/json" "fmt" "log/slog" + "pgproxy/internal/sqlnorm" "pgproxy/internal/config" "pgproxy/internal/metadata" @@ -185,11 +186,17 @@ func (s *PostgreSQLAgent) AnalyzeQuery(ctx context.Context, database string, que if err != nil { return nil, fmt.Errorf("failed to create PostgreSQL tools: %w", err) } + // Pin the trusted original so explain_query runs the real text while the + // model only ever sees a literal-redacted shape (zero-trust: no query + // literals leave the proxy toward the LLM). + tools.PinQuery(query) toolDefs := tools.GetToolDefinitions() + redacted := sqlnorm.NormalizeLiterals(query) + // Build initial message messages := []ConversationMessage{ - GetUserTextMessage(fmt.Sprintf("Analyze this SQL query for efficiency:\n\n```sql\n%s\n```\n\nUse the available tools to examine the schema, indexes, and query plan. Then provide your decision as JSON.", query)), + GetUserTextMessage(fmt.Sprintf("Analyze this SQL query for efficiency:\n\n```sql\n%s\n```\n\nNote: literal values have been redacted as ? or '?' for privacy; the structure is unchanged and explain_query runs against the real query. Use the available tools to examine the schema, indexes, and query plan. Then provide your decision as JSON.", redacted)), } // Agentic loop - keep calling until we get a final answer diff --git a/internal/agent/agent_redshift.go b/internal/agent/agent_redshift.go index 617bf4c..b07e3f1 100644 --- a/internal/agent/agent_redshift.go +++ b/internal/agent/agent_redshift.go @@ -6,6 +6,7 @@ import ( "encoding/json" "fmt" "log/slog" + "pgproxy/internal/sqlnorm" "pgproxy/internal/config" "pgproxy/internal/metadata" @@ -121,15 +122,20 @@ func (s *RedshiftAgent) AnalyzeQuery(ctx context.Context, database string, query // Use pre-initialized tools (avoids recreation on every query) toolDefs := s.tools.GetToolDefinitions() + // Pin the trusted original so explain_query runs the real text; the model + // only sees a literal-redacted shape (no query literals reach the LLM). + s.tools.PinQuery(query) + redacted := sqlnorm.NormalizeLiterals(query) + // Pre-fetch schema context for tables referenced in the query so the LLM has // the information it needs without requiring a tool round-trip. schemaCtx := s.tools.BuildSchemaContext(query) var userMsg string if schemaCtx != "" { - userMsg = fmt.Sprintf("Analyze this Redshift SQL query for efficiency:\n\n```sql\n%s\n```\n\n%s\nUse the available tools to run EXPLAIN or fetch additional metadata if needed. Then provide your decision as JSON.", query, schemaCtx) + userMsg = fmt.Sprintf("Analyze this Redshift SQL query for efficiency:\n\n```sql\n%s\n```\n\n%s\nNote: literal values are redacted as ? for privacy; explain_query runs against the real query. Use the available tools to run EXPLAIN or fetch additional metadata if needed. Then provide your decision as JSON.", redacted, schemaCtx) } else { - userMsg = fmt.Sprintf("Analyze this Redshift SQL query for efficiency:\n\n```sql\n%s\n```\n\nUse the available tools to examine the schema, distribution keys, sort keys, and query plan. Then provide your decision as JSON.", query) + userMsg = fmt.Sprintf("Analyze this Redshift SQL query for efficiency:\n\n```sql\n%s\n```\n\nNote: literal values are redacted as ? for privacy; explain_query runs against the real query. Use the available tools to examine the schema, distribution keys, sort keys, and query plan. Then provide your decision as JSON.", redacted) } // Build initial message diff --git a/internal/agent/agent_test.go b/internal/agent/agent_test.go index 77a4c08..d1abc5e 100644 --- a/internal/agent/agent_test.go +++ b/internal/agent/agent_test.go @@ -175,7 +175,6 @@ func TestAnalyzeQuery_DirectApproval(t *testing.T) { }, } - agent := NewPostgreSQLAgentWithDeps(llm, nil, nil) decision, err := agent.AnalyzeQuery(context.Background(), "testdb", "SELECT 1") @@ -212,7 +211,6 @@ func TestAnalyzeQuery_WithToolCalls(t *testing.T) { }, } - agent := NewPostgreSQLAgentWithDeps(llm, nil, nil) decision, err := agent.AnalyzeQuery(context.Background(), "testdb", "SELECT * FROM users WHERE email = 'test'") @@ -254,7 +252,6 @@ func TestAnalyzeQuery_MultipleToolCalls(t *testing.T) { }, } - agent := NewPostgreSQLAgentWithDeps(llm, nil, nil) decision, err := agent.AnalyzeQuery(context.Background(), "testdb", "SELECT * FROM users") @@ -281,7 +278,6 @@ func TestAnalyzeQuery_ToolError(t *testing.T) { }, } - agent := NewPostgreSQLAgentWithDeps(llm, nil, nil) decision, err := agent.AnalyzeQuery(context.Background(), "testdb", "SELECT * FROM nonexistent") @@ -301,7 +297,6 @@ func TestAnalyzeQuery_LLMError(t *testing.T) { errors: []error{errors.New("LLM service unavailable")}, } - agent := NewPostgreSQLAgentWithDeps(llm, nil, nil) _, err := agent.AnalyzeQuery(context.Background(), "testdb", "SELECT 1") @@ -329,7 +324,6 @@ func TestAnalyzeQuery_MaxIterations(t *testing.T) { llm := &mockConverser{responses: responses} - agent := NewPostgreSQLAgentWithDeps(llm, nil, nil) _, err := agent.AnalyzeQuery(context.Background(), "testdb", "SELECT 1") @@ -362,7 +356,6 @@ func TestAnalyzeQuery_ParallelToolCalls(t *testing.T) { }, } - agent := NewPostgreSQLAgentWithDeps(llm, nil, nil) decision, err := agent.AnalyzeQuery(context.Background(), "testdb", "SELECT * FROM users") diff --git a/internal/agent/agent_tsql.go b/internal/agent/agent_tsql.go new file mode 100644 index 0000000..5e57116 --- /dev/null +++ b/internal/agent/agent_tsql.go @@ -0,0 +1,101 @@ +package agent + +import ( + "context" + "fmt" + "log/slog" + + "pgproxy/internal/sqlnorm" +) + +// tsqlSystemPrompt is the T-SQL persona: a career SQL Server engine +// specialist reviewing queries as a gatekeeper. Milestone 3 ships the +// persona without engine tools; DMV/SHOWPLAN tools follow (see +// docs/multiprotocol-design.md §7), so the prompt instructs the agent +// to judge on query shape alone and to stay conservative. +const tsqlSystemPrompt = `You are a veteran SQL Server database engineer with decades of production experience: query store archaeology, execution plan forensics, and too many 3 AM callouts caused by preventable T-SQL. + +Your task: review T-SQL batches as a gatekeeper, deciding whether they are efficient and safe enough to run in production. + +You have tools to examine schemas, indexes, table statistics, and the estimated execution plan (SHOWPLAN_XML). Use them methodically before judging: +1. Identify the tables involved and their sizes (get_table_stats); tiny tables get a pass. +2. Check the indexes the predicates could use (get_indexes). +3. Get the estimated plan (explain_query) and read it like you wrote the optimizer: scans on large rowcounts, nested loops with huge outer inputs, implicit conversions (CONVERT_IMPLICIT), missing index warnings, spools. + +If the tools fail or return nothing, judge on the query text alone and be conservative: only reject when the text alone proves a serious problem. When in doubt, approve. + +REJECT when the evidence shows: +- Cartesian products: comma-joins or JOINs with no ON condition between large-sounding tables +- SELECT * with no WHERE and no TOP against non-trivial tables +- UPDATE or DELETE with no WHERE clause +- Non-sargable predicates that defeat any index: functions wrapped around columns in WHERE (e.g. WHERE UPPER(name) = ..., WHERE CONVERT(varchar, date_col) = ...) +- Leading-wildcard LIKE ('%term') on what is clearly a large-table scan +- Obvious implicit-conversion traps (N'...' compared to varchar keys and similar) + +Response format, JSON only: +{"approved": true/false, "reason": "one line, direct", "details": "specifics from the query text", "suggested_fix": "corrected T-SQL (required when rejecting)"} + +When rejecting you MUST provide suggested_fix containing corrected T-SQL the caller can run immediately. Note: literal values in the query you receive have been redacted as ? or '?' for privacy; keep the placeholders in your suggested fix. You are not here to be pedantic; you are catching the queries that page someone at 3 AM.` + +// TSQLAgent judges T-SQL batches via a Converser (Bedrock direct or +// Portkey). It implements pipeline.Analyzer. +type TSQLAgent struct { + llm Converser + tools *TSQLTools // optional; nil = judge on query shape alone +} + +// NewTSQLAgent creates a T-SQL analysis agent. tools may be nil. +func NewTSQLAgent(llm Converser, tools *TSQLTools) *TSQLAgent { + return &TSQLAgent{llm: llm, tools: tools} +} + +// AnalyzeQuery judges one T-SQL batch. The database parameter is +// carried for interface compatibility; without engine tools it is +// informational only. Literals are redacted before transmission. +func (s *TSQLAgent) AnalyzeQuery(ctx context.Context, database string, query string) (*QueryDecision, error) { + redacted := sqlnorm.NormalizeLiterals(query) + slog.Debug("analyzing query", "agent", "tsql", "database", database) + + var toolDefs []ToolDefinition + if s.tools != nil { + s.tools.PinQuery(query) + toolDefs = s.tools.GetToolDefinitions() + } + + messages := []ConversationMessage{ + GetUserTextMessage(fmt.Sprintf("Analyze this T-SQL batch for efficiency and safety:\n\n```sql\n%s\n```\n\nLiteral values are redacted as ? for privacy; the structure is unchanged and explain_query runs against the real query. Provide your decision as JSON.", redacted)), + } + + for i := 0; i < MaxIterations; i++ { + resp, err := s.llm.Converse(ctx, tsqlSystemPrompt, messages, toolDefs) + if err != nil { + return nil, fmt.Errorf("converse failed: %w", err) + } + if len(resp.ToolCalls) == 0 { + return parseDecision(resp.TextContent) + } + results := make(map[string]string, len(resp.ToolCalls)) + for _, call := range resp.ToolCalls { + if s.tools == nil { + results[call.ID] = "Error: no tools are available in this context" + continue + } + slog.Debug("calling tool", "agent", "tsql", "tool", call.Name) + out, terr := s.tools.CallTool(ctx, call.Name, call.Input) + if terr != nil { + out = fmt.Sprintf("Error: %v", terr) + } + results[call.ID] = out + } + messages = append(messages, GetAssistantMessageWithToolUse(resp.ToolCalls)) + messages = append(messages, GetUserMessageWithToolResults(results)) + } + return nil, fmt.Errorf("tsql agent: no decision after %d iterations", MaxIterations) +} + +// WarmupDatabase satisfies agent.Agent; the T-SQL agent has no +// metadata cache yet. +func (s *TSQLAgent) WarmupDatabase(ctx context.Context, database string) error { return nil } + +// Close satisfies agent.Agent. +func (s *TSQLAgent) Close() error { return nil } diff --git a/internal/agent/bedrock.go b/internal/agent/bedrock.go new file mode 100644 index 0000000..108485f --- /dev/null +++ b/internal/agent/bedrock.go @@ -0,0 +1,135 @@ +package agent + +import ( + "context" + "fmt" + "os" + + "github.com/aws/aws-sdk-go-v2/aws" + awsconfig "github.com/aws/aws-sdk-go-v2/config" + "github.com/aws/aws-sdk-go-v2/service/bedrockruntime" + "github.com/aws/aws-sdk-go-v2/service/bedrockruntime/document" + "github.com/aws/aws-sdk-go-v2/service/bedrockruntime/types" +) + +// BedrockClient is a direct AWS Bedrock Converser (no gateway). It is +// the original pg-proxy LLM path, retained for deployments that call +// Bedrock directly rather than through Portkey. +type BedrockClient struct { + client *bedrockruntime.Client + modelID string +} + +// NewBedrockClient creates a Converser backed by the Bedrock Converse +// API. Region and model come from AWS_REGION and BEDROCK_MODEL_ID with +// the historical defaults. +func NewBedrockClient(ctx context.Context) (*BedrockClient, error) { + region := os.Getenv("AWS_REGION") + if region == "" { + region = "eu-west-1" + } + modelID := os.Getenv("BEDROCK_MODEL_ID") + if modelID == "" { + modelID = "eu.anthropic.claude-sonnet-4-20250514-v1:0" + } + + cfg, err := awsconfig.LoadDefaultConfig(ctx, awsconfig.WithRegion(region)) + if err != nil { + return nil, fmt.Errorf("failed to load AWS config: %w", err) + } + return &BedrockClient{ + client: bedrockruntime.NewFromConfig(cfg), + modelID: modelID, + }, nil +} + +// Converse implements the Converser interface against Bedrock. +func (bc *BedrockClient) Converse( + ctx context.Context, + systemPrompt string, + messages []ConversationMessage, + tools []ToolDefinition, +) (*ConverseResponse, error) { + system := []types.SystemContentBlock{ + &types.SystemContentBlockMemberText{Value: systemPrompt}, + } + + var convMessages []types.Message + for _, msg := range messages { + convMessages = append(convMessages, types.Message{ + Role: msg.Role, + Content: msg.Content, + }) + } + + var toolConfig *types.ToolConfiguration + if len(tools) > 0 { + var toolSpecs []types.Tool + for _, tool := range tools { + toolSpecs = append(toolSpecs, &types.ToolMemberToolSpec{ + Value: types.ToolSpecification{ + Name: aws.String(tool.Name), + Description: aws.String(tool.Description), + InputSchema: &types.ToolInputSchemaMemberJson{ + Value: buildInputSchemaDoc(tool), + }, + }, + }) + } + toolConfig = &types.ToolConfiguration{Tools: toolSpecs} + } + + output, err := bc.client.Converse(ctx, &bedrockruntime.ConverseInput{ + ModelId: aws.String(bc.modelID), + System: system, + Messages: convMessages, + ToolConfig: toolConfig, + InferenceConfig: &types.InferenceConfiguration{ + MaxTokens: aws.Int32(4096), + Temperature: aws.Float32(0.0), + }, + }) + if err != nil { + return nil, fmt.Errorf("converse failed: %w", err) + } + + response := &ConverseResponse{StopReason: string(output.StopReason)} + msgOutput, ok := output.Output.(*types.ConverseOutputMemberMessage) + if !ok { + return response, nil + } + for _, block := range msgOutput.Value.Content { + switch b := block.(type) { + case *types.ContentBlockMemberText: + response.TextContent = b.Value + case *types.ContentBlockMemberToolUse: + response.ToolCalls = append(response.ToolCalls, ToolCall{ + ID: aws.ToString(b.Value.ToolUseId), + Name: aws.ToString(b.Value.Name), + Input: documentToMap(b.Value.Input), + }) + } + } + return response, nil +} + +func buildInputSchemaDoc(tool ToolDefinition) document.Interface { + props := make(map[string]any) + for name, param := range tool.Parameters { + propType := "string" + if param.Type == "boolean" { + propType = "boolean" + } else if param.Type == "integer" { + propType = "integer" + } + props[name] = map[string]any{ + "type": propType, + "description": param.Description, + } + } + schema := map[string]any{"type": "object", "properties": props} + if len(tool.Required) > 0 { + schema["required"] = tool.Required + } + return document.NewLazyDocument(schema) +} diff --git a/internal/agent/redaction_test.go b/internal/agent/redaction_test.go new file mode 100644 index 0000000..a1fcd14 --- /dev/null +++ b/internal/agent/redaction_test.go @@ -0,0 +1,70 @@ +package agent + +import ( + types "github.com/aws/aws-sdk-go-v2/service/bedrockruntime/types" + + "context" + "strings" + "testing" +) + +// TestAnalyzeQuery_RedactsLiteralsFromPrompt verifies the zero-trust +// guarantee that no query literal values are transmitted to the LLM: +// the prompt must contain the query's shape, never its values. +func TestAnalyzeQuery_RedactsLiteralsFromPrompt(t *testing.T) { + llm := &mockConverser{ + responses: []*ConverseResponse{ + {TextContent: `{"approved": true, "reason": "fine"}`}, + }, + } + + agent := NewPostgreSQLAgentWithDeps(llm, nil, nil) + query := "SELECT balance FROM accounts WHERE id_number = '8001015009087' AND acc_no = 62001234567" + if _, err := agent.AnalyzeQuery(context.Background(), "testdb", query); err != nil { + t.Fatalf("AnalyzeQuery() error = %v", err) + } + + if len(llm.calls) != 1 { + t.Fatalf("expected 1 LLM call, got %d", len(llm.calls)) + } + var prompt strings.Builder + for _, msg := range llm.calls[0].Messages { + for _, c := range msg.Content { + if tb, ok := c.(*types.ContentBlockMemberText); ok { + prompt.WriteString(tb.Value) + } + } + } + sent := prompt.String() + + for _, secret := range []string{"8001015009087", "62001234567"} { + if strings.Contains(sent, secret) { + t.Errorf("literal %q was sent to the LLM; prompt: %s", secret, sent) + } + } + // Structure must survive so the agent can still reason about the query. + for _, structural := range []string{"accounts", "id_number", "acc_no", "SELECT"} { + if !strings.Contains(sent, structural) { + t.Errorf("expected structural element %q in prompt; got: %s", structural, sent) + } + } +} + +// TestPinQuery_OverridesModelSuppliedExplainTarget verifies that once a +// query is pinned, the explain tool ignores model-supplied SQL. +func TestPinQuery_OverridesModelSuppliedExplainTarget(t *testing.T) { + tools, err := NewPostgreSQLTools(nil, nil) + if err != nil { + t.Fatal(err) + } + tools.PinQuery("SELECT 1") + if tools.pinnedQuery != "SELECT 1" { + t.Errorf("pinnedQuery = %q, want %q", tools.pinnedQuery, "SELECT 1") + } + // With no DB the tool errors before executing, but the pin path is + // exercised in ExecuteTool; integration coverage exercises the rest. + _, err = tools.CallTool(context.Background(), "explain_query", map[string]any{"query": "DROP TABLE users"}) + if err == nil { + t.Error("expected error with nil db") + } +} diff --git a/internal/agent/tools.go b/internal/agent/tools.go index 49eb82b..3c079c2 100644 --- a/internal/agent/tools.go +++ b/internal/agent/tools.go @@ -14,6 +14,12 @@ import ( type PostgreSQLTools struct { db *sql.DB cache *metadata.Cache + // pinnedQuery is the trusted original SQL under analysis. When set, + // explain_query always runs this text, regardless of what the model + // supplies: the model only ever sees a literal-redacted version of the + // query (see internal/sqlnorm), so it cannot reproduce the original, + // and pinning also prevents the model from EXPLAINing arbitrary SQL. + pinnedQuery string } // NewPostgreSQLTools creates a new PostgreSQLTools instance. @@ -26,6 +32,11 @@ func NewPostgreSQLTools(db *sql.DB, cache *metadata.Cache) (*PostgreSQLTools, er }, nil } +// PinQuery sets the trusted original query for this analysis session. +func (t *PostgreSQLTools) PinQuery(query string) { + t.pinnedQuery = query +} + // ToolDefinition represents a tool that can be called by the agent type ToolDefinition struct { Name string @@ -99,7 +110,7 @@ func (t *PostgreSQLTools) GetToolDefinitions() []ToolDefinition { }, { Name: "explain_query", - Description: "Run EXPLAIN (FORMAT JSON) on a query to get the execution plan. Never executes the query — runs inside an aborted read-only transaction.", + Description: "Run EXPLAIN (FORMAT JSON) on the query under analysis to get the execution plan. The proxy supplies the real query text server-side. Never executes the query.", Parameters: map[string]ParameterDef{ "query": { Type: "string", @@ -137,6 +148,11 @@ func (t *PostgreSQLTools) CallTool(ctx context.Context, name string, args map[st return "", fmt.Errorf("no database connection available") } query, _ := args["query"].(string) + if t.pinnedQuery != "" { + // Literals are redacted before the model sees the query, so the + // model cannot supply the real text; always explain the original. + query = t.pinnedQuery + } return t.explainQuery(ctx, query) default: return "", fmt.Errorf("unknown tool: %s", name) diff --git a/internal/agent/tools_redshift.go b/internal/agent/tools_redshift.go index 4165cbf..4243c8c 100644 --- a/internal/agent/tools_redshift.go +++ b/internal/agent/tools_redshift.go @@ -17,8 +17,9 @@ var tableRefRe = regexp.MustCompile(`(?i)(?:FROM|JOIN)\s+([a-zA-Z_][a-zA-Z0-9_]* // RedshiftTools provides Redshift-specific analysis tools for the agent type RedshiftTools struct { - db *sql.DB - cache *metadata.RedshiftCache + pinnedQuery string + db *sql.DB + cache *metadata.RedshiftCache } // NewRedshiftTools creates a new RedshiftTools instance @@ -29,6 +30,11 @@ func NewRedshiftTools(db *sql.DB, cache *metadata.RedshiftCache) *RedshiftTools } } +// PinQuery sets the trusted original query for this analysis session. +func (t *RedshiftTools) PinQuery(query string) { + t.pinnedQuery = query +} + // GetToolDefinitions returns all available Redshift-specific tool definitions func (t *RedshiftTools) GetToolDefinitions() []ToolDefinition { return []ToolDefinition{ @@ -121,6 +127,9 @@ func (t *RedshiftTools) CallTool(ctx context.Context, name string, args map[stri return t.getTableStats(table) case "explain_query": query, _ := args["query"].(string) + if t.pinnedQuery != "" { + query = t.pinnedQuery + } return t.explainQuery(ctx, query) default: return "", fmt.Errorf("unknown tool: %s", name) diff --git a/internal/agent/tools_tsql.go b/internal/agent/tools_tsql.go new file mode 100644 index 0000000..5cbf103 --- /dev/null +++ b/internal/agent/tools_tsql.go @@ -0,0 +1,231 @@ +package agent + +import ( + "context" + "database/sql" + "encoding/json" + "fmt" + "log/slog" + "strings" +) + +// TSQLTools provides SQL Server analysis tools backed by DMVs and +// SHOWPLAN (docs/multiprotocol-design.md §7). All tools are read-only. +type TSQLTools struct { + db *sql.DB + // pinnedQuery is the trusted original SQL under analysis; the model + // only sees a literal-redacted shape, so explain always runs this. + pinnedQuery string +} + +// NewTSQLTools creates SQL Server tools over a go-mssqldb pool. +func NewTSQLTools(db *sql.DB) *TSQLTools { + return &TSQLTools{db: db} +} + +// PinQuery sets the trusted original query for this analysis session. +func (t *TSQLTools) PinQuery(query string) { t.pinnedQuery = query } + +// GetToolDefinitions lists the tools offered to the model. +func (t *TSQLTools) GetToolDefinitions() []ToolDefinition { + return []ToolDefinition{ + { + Name: "get_schema", + Description: "Get columns and types for a table (sys.columns). Pass empty table to list all user tables.", + Parameters: map[string]ParameterDef{ + "table": {Type: "string", Description: "Table name, optionally schema-qualified. Empty lists all tables."}, + }, + Required: []string{}, + }, + { + Name: "get_indexes", + Description: "Get indexes for a table with usage statistics (sys.indexes + sys.dm_db_index_usage_stats).", + Parameters: map[string]ParameterDef{ + "table": {Type: "string", Description: "Table name, optionally schema-qualified."}, + }, + Required: []string{"table"}, + }, + { + Name: "get_table_stats", + Description: "Get approximate row count and size for a table (sys.dm_db_partition_stats).", + Parameters: map[string]ParameterDef{ + "table": {Type: "string", Description: "Table name, optionally schema-qualified."}, + }, + Required: []string{"table"}, + }, + { + Name: "explain_query", + Description: "Get the estimated execution plan (SHOWPLAN_XML) for the query under analysis. The proxy supplies the real query text server-side. Never executes the query.", + Parameters: map[string]ParameterDef{}, + Required: []string{}, + }, + } +} + +// CallTool dispatches a tool call. +func (t *TSQLTools) CallTool(ctx context.Context, name string, args map[string]any) (string, error) { + if t.db == nil { + return "", fmt.Errorf("no database connection available") + } + table, _ := args["table"].(string) + switch name { + case "get_schema": + return t.getSchema(ctx, table) + case "get_indexes": + return t.getIndexes(ctx, table) + case "get_table_stats": + return t.getTableStats(ctx, table) + case "explain_query": + return t.explainQuery(ctx) + default: + return "", fmt.Errorf("unknown tool: %s", name) + } +} + +// splitTable returns (schema, name) with a dbo default. +func splitTable(table string) (string, string) { + if i := strings.IndexByte(table, '.'); i >= 0 { + return table[:i], table[i+1:] + } + return "dbo", table +} + +func rowsToJSON(rows *sql.Rows) (string, error) { + cols, err := rows.Columns() + if err != nil { + return "", err + } + var out []map[string]any + for rows.Next() { + vals := make([]any, len(cols)) + ptrs := make([]any, len(cols)) + for i := range vals { + ptrs[i] = &vals[i] + } + if err := rows.Scan(ptrs...); err != nil { + return "", err + } + m := make(map[string]any, len(cols)) + for i, c := range cols { + if b, ok := vals[i].([]byte); ok { + m[c] = string(b) + } else { + m[c] = vals[i] + } + } + out = append(out, m) + } + if err := rows.Err(); err != nil { + return "", err + } + b, err := json.Marshal(out) + return string(b), err +} + +func (t *TSQLTools) getSchema(ctx context.Context, table string) (string, error) { + if table == "" { + rows, err := t.db.QueryContext(ctx, ` + SELECT s.name AS [schema], t.name AS [table] + FROM sys.tables t JOIN sys.schemas s ON t.schema_id = s.schema_id + ORDER BY s.name, t.name`) + if err != nil { + return "", err + } + defer rows.Close() + return rowsToJSON(rows) + } + schema, name := splitTable(table) + rows, err := t.db.QueryContext(ctx, ` + SELECT c.name AS [column], ty.name AS [type], c.max_length, c.is_nullable + FROM sys.columns c + JOIN sys.types ty ON c.user_type_id = ty.user_type_id + JOIN sys.tables t ON c.object_id = t.object_id + JOIN sys.schemas s ON t.schema_id = s.schema_id + WHERE s.name = @p1 AND t.name = @p2 + ORDER BY c.column_id`, schema, name) + if err != nil { + return "", err + } + defer rows.Close() + return rowsToJSON(rows) +} + +func (t *TSQLTools) getIndexes(ctx context.Context, table string) (string, error) { + schema, name := splitTable(table) + rows, err := t.db.QueryContext(ctx, ` + SELECT i.name AS [index], i.type_desc, i.is_unique, + us.user_seeks, us.user_scans, us.user_lookups + FROM sys.indexes i + JOIN sys.tables t ON i.object_id = t.object_id + JOIN sys.schemas s ON t.schema_id = s.schema_id + LEFT JOIN sys.dm_db_index_usage_stats us + ON us.object_id = i.object_id AND us.index_id = i.index_id + AND us.database_id = DB_ID() + WHERE s.name = @p1 AND t.name = @p2`, schema, name) + if err != nil { + return "", err + } + defer rows.Close() + return rowsToJSON(rows) +} + +func (t *TSQLTools) getTableStats(ctx context.Context, table string) (string, error) { + schema, name := splitTable(table) + rows, err := t.db.QueryContext(ctx, ` + SELECT SUM(ps.row_count) AS row_count, + SUM(ps.used_page_count) * 8 AS used_kb + FROM sys.dm_db_partition_stats ps + JOIN sys.tables t ON ps.object_id = t.object_id + JOIN sys.schemas s ON t.schema_id = s.schema_id + WHERE s.name = @p1 AND t.name = @p2 AND ps.index_id IN (0, 1)`, schema, name) + if err != nil { + return "", err + } + defer rows.Close() + return rowsToJSON(rows) +} + +// explainQuery returns the estimated plan for the pinned query using a +// dedicated session: SHOWPLAN_XML is session-sticky, so it must never +// run on a pooled connection shared with real traffic. +func (t *TSQLTools) explainQuery(ctx context.Context) (string, error) { + if t.pinnedQuery == "" { + return "", fmt.Errorf("no query pinned for analysis") + } + conn, err := t.db.Conn(ctx) + if err != nil { + return "", err + } + defer conn.Close() + + if _, err := conn.ExecContext(ctx, "SET SHOWPLAN_XML ON"); err != nil { + return "", fmt.Errorf("SHOWPLAN_XML ON: %w", err) + } + // Under SHOWPLAN the statement is compiled, not executed; the plan + // XML comes back as a single-column result set. + rows, err := conn.QueryContext(ctx, t.pinnedQuery) + if err != nil { + conn.ExecContext(ctx, "SET SHOWPLAN_XML OFF") + return "", fmt.Errorf("showplan compile: %w", err) + } + var plan strings.Builder + for rows.Next() { + var chunk string + if err := rows.Scan(&chunk); err != nil { + rows.Close() + conn.ExecContext(ctx, "SET SHOWPLAN_XML OFF") + return "", err + } + plan.WriteString(chunk) + } + rows.Close() + if _, err := conn.ExecContext(ctx, "SET SHOWPLAN_XML OFF"); err != nil { + slog.Warn("SHOWPLAN_XML OFF failed; discarding session", "error", err) + } + const maxPlan = 30000 + out := plan.String() + if len(out) > maxPlan { + out = out[:maxPlan] + "\n" + } + return out, nil +} diff --git a/internal/agent/tools_tsql_integration_test.go b/internal/agent/tools_tsql_integration_test.go new file mode 100644 index 0000000..1aff5ee --- /dev/null +++ b/internal/agent/tools_tsql_integration_test.go @@ -0,0 +1,63 @@ +package agent + +import ( + "context" + "database/sql" + "fmt" + "os" + "strings" + "testing" + + _ "github.com/microsoft/go-mssqldb" +) + +// Requires a running SQL Server / Azure SQL Edge: +// +// TSQL_IT_DSN='sqlserver://sa:pass@localhost:11433?encrypt=disable' go test ./internal/agent/ -run TSQLTools -v +func TestTSQLToolsAgainstRealBackend(t *testing.T) { + dsn := os.Getenv("TSQL_IT_DSN") + if dsn == "" { + t.Skip("TSQL_IT_DSN not set; skipping live T-SQL tools test") + } + db, err := sql.Open("sqlserver", dsn) + if err != nil { + t.Fatal(err) + } + defer db.Close() + ctx := context.Background() + + db.ExecContext(ctx, "DROP TABLE IF EXISTS dbo.guardian_it") + if _, err := db.ExecContext(ctx, "CREATE TABLE dbo.guardian_it (id INT PRIMARY KEY, email NVARCHAR(200))"); err != nil { + t.Fatal(err) + } + defer db.ExecContext(ctx, "DROP TABLE IF EXISTS dbo.guardian_it") + for i := 0; i < 5; i++ { + db.ExecContext(ctx, fmt.Sprintf("INSERT INTO dbo.guardian_it VALUES (%d, 'u%d@example.com')", i, i)) + } + + tools := NewTSQLTools(db) + + schema, err := tools.CallTool(ctx, "get_schema", map[string]any{"table": "guardian_it"}) + if err != nil || !strings.Contains(schema, "email") { + t.Fatalf("get_schema: %v %s", err, schema) + } + stats, err := tools.CallTool(ctx, "get_table_stats", map[string]any{"table": "guardian_it"}) + if err != nil || !strings.Contains(stats, "row_count") { + t.Fatalf("get_table_stats: %v %s", err, stats) + } + idx, err := tools.CallTool(ctx, "get_indexes", map[string]any{"table": "guardian_it"}) + if err != nil { + t.Fatalf("get_indexes: %v", err) + } + t.Logf("indexes: %.120s", idx) + + tools.PinQuery("SELECT * FROM dbo.guardian_it WHERE email = 'u1@example.com'") + plan, err := tools.CallTool(ctx, "explain_query", nil) + if err != nil { + t.Fatalf("explain_query: %v", err) + } + if !strings.Contains(plan, "ShowPlanXML") { + t.Fatalf("expected ShowPlanXML, got: %.200s", plan) + } + t.Logf("plan head: %.150s", plan) +} diff --git a/internal/pipeline/pipeline.go b/internal/pipeline/pipeline.go new file mode 100644 index 0000000..7d49214 --- /dev/null +++ b/internal/pipeline/pipeline.go @@ -0,0 +1,328 @@ +// Package pipeline contains the protocol-neutral query decision core: +// skip-list -> filter -> cache -> agent analysis -> cache write -> report +// -> mode application. +// +// It is extracted from the PostgreSQL wire handler so additional wire +// frontends (e.g. TDS for SQL Server) can share identical governance +// semantics. See docs/multiprotocol-design.md. +// +// Behavioral contract: this package reproduces the proxy's original +// decision logic verbatim, including reason-string suffixes such as +// "(cached)" and "(logged, non-blocking)", fail-open on agent errors, +// and best-effort cache/report handling. +package pipeline + +import ( + "context" + "log/slog" + "strings" + "time" + + "pgproxy/internal/agent" + "pgproxy/internal/cache" + "pgproxy/internal/config" + "pgproxy/internal/reporter" +) + +// Statement is one client statement awaiting a verdict. +type Statement struct { + Database string + SQL string + ClientIP string + Username string +} + +// Verdict is the pipeline's decision about a Statement. +type Verdict struct { + Approved bool + Shunted bool // shunt mode only: statement routes to the shunt node + Reason string + Details string + SuggestedFix string + LLMLatencyMs *int +} + +// ReportEntry captures an audit record for the reporter. +type ReportEntry struct { + Statement Statement + Fingerprint string + Reason string + Details string + SuggestedFix string + Approved bool + WasBlocked bool + CacheHit bool + RoutedTo reporter.RoutedTo +} + +// Filter optionally bypasses analysis for matching statements. +type Filter interface { + ShouldFilter(query string) bool +} + +// Analyzer is the narrow slice of agent.Agent the pipeline needs. +type Analyzer interface { + AnalyzeQuery(ctx context.Context, database string, query string) (*agent.QueryDecision, error) +} + +// ApprovalCache is the narrow slice of cache.ApprovalCache the pipeline needs. +type ApprovalCache interface { + Get(ctx context.Context, fingerprint string) (*cache.ApprovedQuery, error) + Set(ctx context.Context, fingerprint string, approval *cache.ApprovedQuery, ttl time.Duration) error +} + +// Config wires the pipeline's collaborators. Analyzer, Cache, Filter and +// Report are optional; the pipeline degrades exactly as the proxy did +// (passthrough without an analyzer, no caching without a cache, etc). +type Config struct { + Analyzer Analyzer + Cache ApprovalCache + Mode config.ProxyMode + Filter Filter + + // SkipAnalysis reports whether a statement is exempt from analysis + // (system/catalog traffic). Wire-frontend specific. + SkipAnalysis func(query string) bool + // Fingerprint produces the cache key for a statement's shape. + Fingerprint func(query string) string + // Report records an audit entry (best effort; may be nil). + Report func(ctx context.Context, e ReportEntry) + + AnalysisTimeout time.Duration + CacheTTL time.Duration +} + +// Pipeline is safe for concurrent use. +type Pipeline struct { + cfg Config +} + +func New(cfg Config) *Pipeline { + return &Pipeline{cfg: cfg} +} + +func (p *Pipeline) report(ctx context.Context, e ReportEntry) { + if p.cfg.Report != nil { + p.cfg.Report(ctx, e) + } +} + +func truncate(s string, maxLen int) string { + if len(s) > maxLen { + return s[:maxLen] + "..." + } + return s +} + +// Decide evaluates a statement under blocking / non-blocking semantics. +// This is the original Proxy.analyzeQuery, verbatim. +func (p *Pipeline) Decide(stmt Statement) Verdict { + if p.cfg.Analyzer == nil { + return Verdict{Approved: true, Reason: "passthrough mode"} + } + + if strings.TrimSpace(stmt.SQL) == "" { + slog.Debug("query skipped (empty)", "database", stmt.Database) + return Verdict{Approved: true, Reason: "empty query"} + } + + if p.cfg.SkipAnalysis != nil && p.cfg.SkipAnalysis(stmt.SQL) { + slog.Debug("query skipped (system/catalog query)", + "database", stmt.Database, + "query_preview", truncate(stmt.SQL, 80), + ) + return Verdict{Approved: true, Reason: "system query"} + } + + if p.cfg.Filter != nil && p.cfg.Filter.ShouldFilter(stmt.SQL) { + slog.Info("query matched filter, bypassing analysis", + "query_prefix", truncate(stmt.SQL, 80)) + return Verdict{Approved: true, Reason: "matched filter pattern"} + } + + ctx, cancel := context.WithTimeout(context.Background(), p.cfg.AnalysisTimeout) + defer cancel() + + fingerprint := p.cfg.Fingerprint(stmt.SQL) + + if p.cfg.Cache != nil { + cached, err := p.cfg.Cache.Get(ctx, fingerprint) + if err != nil { + slog.Error("cache error", "error", err) + } else if cached != nil { + slog.Debug("cache hit", "fingerprint", fingerprint[:8], "status", cached.Status, "reason", cached.Reason) + approved := cached.Status == cache.StatusApproved + + wasBlocked := !approved && p.cfg.Mode == config.ModeBlocking + p.report(ctx, ReportEntry{ + Statement: stmt, Fingerprint: fingerprint, + Reason: cached.Reason, Details: cached.Details, + Approved: approved, WasBlocked: wasBlocked, CacheHit: true, + RoutedTo: reporter.RoutedToNone, + }) + + if !approved && p.cfg.Mode == config.ModeNonBlocking { + return Verdict{Approved: true, Reason: cached.Reason + " (logged, non-blocking)", Details: cached.Details} + } + return Verdict{Approved: approved, Reason: cached.Reason + " (cached)", Details: cached.Details} + } + slog.Debug("cache miss", "fingerprint", fingerprint[:8]) + } + + slog.Info("analyzing query for database", "database", stmt.Database, "query_preview", truncate(stmt.SQL, 80)) + + llmStart := time.Now() + decision, err := p.cfg.Analyzer.AnalyzeQuery(ctx, stmt.Database, stmt.SQL) + llmLatencyMs := int(time.Since(llmStart).Milliseconds()) + + if err != nil { + slog.Error("agent error, allowing query", "database", stmt.Database, "error", err) + return Verdict{Approved: true, Reason: "analysis error - fail open", LLMLatencyMs: &llmLatencyMs} + } + + if p.cfg.Cache != nil { + status := cache.StatusApproved + if !decision.Approved { + status = cache.StatusRejected + } + approval := &cache.ApprovedQuery{ + Query: stmt.SQL, + Fingerprint: fingerprint, + Status: status, + Reason: decision.Reason, + Details: decision.Details, + DecidedAt: time.Now(), + } + if err := p.cfg.Cache.Set(ctx, fingerprint, approval, p.cfg.CacheTTL); err != nil { + slog.Error("cache set error", "error", err) + } else { + slog.Debug("cache set", "fingerprint", fingerprint[:8], "status", status) + } + } + + if decision.Approved { + slog.Info("query approved", "reason", decision.Reason) + if decision.Details != "" { + slog.Debug("agent details", "details", decision.Details) + } + } else { + slog.Warn("query rejected", "reason", decision.Reason) + if decision.Details != "" { + slog.Debug("agent details", "details", decision.Details) + } + if decision.SuggestedFix != "" { + slog.Info("suggested fix", "fix", decision.SuggestedFix) + } + } + + wasBlocked := !decision.Approved && p.cfg.Mode == config.ModeBlocking + p.report(ctx, ReportEntry{ + Statement: stmt, Fingerprint: fingerprint, + Reason: decision.Reason, Details: decision.Details, SuggestedFix: decision.SuggestedFix, + Approved: decision.Approved, WasBlocked: wasBlocked, CacheHit: false, + RoutedTo: reporter.RoutedToNone, + }) + + if !decision.Approved && p.cfg.Mode == config.ModeNonBlocking { + return Verdict{Approved: true, Reason: decision.Reason + " (logged, non-blocking)", Details: decision.Details, SuggestedFix: decision.SuggestedFix, LLMLatencyMs: &llmLatencyMs} + } + + return Verdict{Approved: decision.Approved, Reason: decision.Reason, Details: decision.Details, SuggestedFix: decision.SuggestedFix, LLMLatencyMs: &llmLatencyMs} +} + +// DecideShunt evaluates a statement under shunt-mode semantics: approved +// statements route to the fast node, rejected statements to the shunt +// node. This is the original Proxy.analyzeQueryShunt, verbatim. +func (p *Pipeline) DecideShunt(stmt Statement) Verdict { + if p.cfg.Analyzer == nil { + return Verdict{Approved: true, Reason: "passthrough mode"} + } + + if p.cfg.SkipAnalysis != nil && p.cfg.SkipAnalysis(stmt.SQL) { + return Verdict{Approved: true, Reason: "system query"} + } + + ctx, cancel := context.WithTimeout(context.Background(), p.cfg.AnalysisTimeout) + defer cancel() + + fingerprint := p.cfg.Fingerprint(stmt.SQL) + + if p.cfg.Cache != nil { + cached, err := p.cfg.Cache.Get(ctx, fingerprint) + if err != nil { + slog.Error("cache error", "error", err) + } else if cached != nil { + slog.Debug("cache hit (shunt)", "fingerprint", fingerprint[:8], "status", cached.Status, "shunted", cached.Shunted) + + routedTo := reporter.RoutedToFast + if cached.Shunted { + routedTo = reporter.RoutedToShunt + } + p.report(ctx, ReportEntry{ + Statement: stmt, Fingerprint: fingerprint, + Reason: cached.Reason, Details: cached.Details, + Approved: cached.Status == cache.StatusApproved, WasBlocked: false, CacheHit: true, + RoutedTo: routedTo, + }) + + if cached.Shunted { + return Verdict{Approved: false, Shunted: true, Reason: cached.Reason + " (cached, shunted)"} + } + return Verdict{Approved: cached.Status == cache.StatusApproved, Reason: cached.Reason + " (cached)"} + } + slog.Debug("cache miss (shunt)", "fingerprint", fingerprint[:8]) + } + + slog.Info("analyzing query for database (shunt mode)", "database", stmt.Database, "query_preview", truncate(stmt.SQL, 80)) + + llmStart := time.Now() + decision, err := p.cfg.Analyzer.AnalyzeQuery(ctx, stmt.Database, stmt.SQL) + _ = int(time.Since(llmStart).Milliseconds()) // TODO: metrics tracking for shunt mode + + if err != nil { + slog.Error("agent error, routing to fast node", "database", stmt.Database, "error", err) + return Verdict{Approved: true, Reason: "analysis error - fail open"} + } + + if p.cfg.Cache != nil { + status := cache.StatusApproved + shuntedFlag := false + if !decision.Approved { + status = cache.StatusRejected + shuntedFlag = true // in shunt mode, rejected = shunted + } + approval := &cache.ApprovedQuery{ + Query: stmt.SQL, + Fingerprint: fingerprint, + Status: status, + Reason: decision.Reason, + Details: decision.Details, + DecidedAt: time.Now(), + Shunted: shuntedFlag, + } + if err := p.cfg.Cache.Set(ctx, fingerprint, approval, p.cfg.CacheTTL); err != nil { + slog.Error("cache set error", "error", err) + } else { + slog.Debug("cache set (shunt)", "fingerprint", fingerprint[:8], "status", status, "shunted", shuntedFlag) + } + } + + routedTo := reporter.RoutedToFast + if !decision.Approved { + routedTo = reporter.RoutedToShunt + } + p.report(ctx, ReportEntry{ + Statement: stmt, Fingerprint: fingerprint, + Reason: decision.Reason, Details: decision.Details, SuggestedFix: decision.SuggestedFix, + Approved: decision.Approved, WasBlocked: false, CacheHit: false, + RoutedTo: routedTo, + }) + + if decision.Approved { + slog.Info("query approved (shunt mode) -> fast node", "reason", decision.Reason) + return Verdict{Approved: true, Reason: decision.Reason} + } + + slog.Warn("query rejected (shunt mode) -> shunt node", "reason", decision.Reason) + return Verdict{Approved: false, Shunted: true, Reason: decision.Reason} +} diff --git a/internal/pipeline/pipeline_test.go b/internal/pipeline/pipeline_test.go new file mode 100644 index 0000000..6f96039 --- /dev/null +++ b/internal/pipeline/pipeline_test.go @@ -0,0 +1,165 @@ +package pipeline + +import ( + "context" + "errors" + "testing" + "time" + + "pgproxy/internal/agent" + "pgproxy/internal/cache" + "pgproxy/internal/config" +) + +type fakeAnalyzer struct { + decision *agent.QueryDecision + err error + calls int +} + +func (f *fakeAnalyzer) AnalyzeQuery(ctx context.Context, database, query string) (*agent.QueryDecision, error) { + f.calls++ + return f.decision, f.err +} + +type fakeCache struct { + entry *cache.ApprovedQuery + sets []*cache.ApprovedQuery +} + +func (f *fakeCache) Get(ctx context.Context, fp string) (*cache.ApprovedQuery, error) { + return f.entry, nil +} +func (f *fakeCache) Set(ctx context.Context, fp string, q *cache.ApprovedQuery, ttl time.Duration) error { + f.sets = append(f.sets, q) + return nil +} + +func newPipeline(a Analyzer, c ApprovalCache, mode config.ProxyMode) *Pipeline { + cfg := Config{ + Mode: mode, + SkipAnalysis: func(q string) bool { return q == "SELECT 1" }, + Fingerprint: func(q string) string { return "aabbccddeeff0011" }, + AnalysisTimeout: 5 * time.Second, + CacheTTL: time.Hour, + } + if a != nil { + cfg.Analyzer = a + } + if c != nil { + cfg.Cache = c + } + return New(cfg) +} + +func stmt(sql string) Statement { + return Statement{Database: "testdb", SQL: sql, ClientIP: "127.0.0.1", Username: "leon"} +} + +func TestDecide_PassthroughWithoutAnalyzer(t *testing.T) { + v := newPipeline(nil, nil, config.ModeBlocking).Decide(stmt("SELECT * FROM t")) + if !v.Approved || v.Reason != "passthrough mode" { + t.Errorf("got %+v", v) + } +} + +func TestDecide_EmptyAndSystemQueriesSkip(t *testing.T) { + a := &fakeAnalyzer{} + p := newPipeline(a, nil, config.ModeBlocking) + if v := p.Decide(stmt(" ")); !v.Approved || v.Reason != "empty query" { + t.Errorf("empty: %+v", v) + } + if v := p.Decide(stmt("SELECT 1")); !v.Approved || v.Reason != "system query" { + t.Errorf("skip: %+v", v) + } + if a.calls != 0 { + t.Errorf("analyzer called %d times for exempt queries", a.calls) + } +} + +func TestDecide_BlockingRejects(t *testing.T) { + a := &fakeAnalyzer{decision: &agent.QueryDecision{Approved: false, Reason: "cartesian", SuggestedFix: "add ON"}} + v := newPipeline(a, nil, config.ModeBlocking).Decide(stmt("SELECT * FROM a, b")) + if v.Approved || v.Reason != "cartesian" || v.SuggestedFix != "add ON" { + t.Errorf("got %+v", v) + } + if v.LLMLatencyMs == nil { + t.Error("expected latency recorded") + } +} + +func TestDecide_NonBlockingLogsButApproves(t *testing.T) { + a := &fakeAnalyzer{decision: &agent.QueryDecision{Approved: false, Reason: "cartesian"}} + v := newPipeline(a, nil, config.ModeNonBlocking).Decide(stmt("SELECT * FROM a, b")) + if !v.Approved || v.Reason != "cartesian (logged, non-blocking)" { + t.Errorf("got %+v", v) + } +} + +func TestDecide_FailOpenOnAgentError(t *testing.T) { + a := &fakeAnalyzer{err: errors.New("bedrock down")} + v := newPipeline(a, nil, config.ModeBlocking).Decide(stmt("SELECT * FROM a, b")) + if !v.Approved || v.Reason != "analysis error - fail open" { + t.Errorf("fail-open violated: %+v", v) + } +} + +func TestDecide_CacheHitSuffixAndNoAgentCall(t *testing.T) { + a := &fakeAnalyzer{decision: &agent.QueryDecision{Approved: true, Reason: "should not be used"}} + c := &fakeCache{entry: &cache.ApprovedQuery{Status: cache.StatusRejected, Reason: "cartesian"}} + v := newPipeline(a, c, config.ModeBlocking).Decide(stmt("SELECT * FROM a, b")) + if v.Approved || v.Reason != "cartesian (cached)" { + t.Errorf("got %+v", v) + } + if a.calls != 0 { + t.Error("agent consulted despite cache hit") + } +} + +func TestDecide_CachesDecision(t *testing.T) { + a := &fakeAnalyzer{decision: &agent.QueryDecision{Approved: false, Reason: "bad"}} + c := &fakeCache{} + newPipeline(a, c, config.ModeBlocking).Decide(stmt("SELECT * FROM a, b")) + if len(c.sets) != 1 || c.sets[0].Status != cache.StatusRejected { + t.Errorf("cache writes: %+v", c.sets) + } +} + +func TestDecideShunt_RoutesRejectedToShunt(t *testing.T) { + a := &fakeAnalyzer{decision: &agent.QueryDecision{Approved: false, Reason: "bad"}} + c := &fakeCache{} + v := newPipeline(a, c, config.ModeShunt).DecideShunt(stmt("SELECT * FROM a, b")) + if v.Approved || !v.Shunted || v.Reason != "bad" { + t.Errorf("got %+v", v) + } + if len(c.sets) != 1 || !c.sets[0].Shunted { + t.Errorf("shunted flag not cached: %+v", c.sets) + } +} + +func TestDecideShunt_CachedShuntedSuffix(t *testing.T) { + c := &fakeCache{entry: &cache.ApprovedQuery{Status: cache.StatusRejected, Reason: "bad", Shunted: true}} + a := &fakeAnalyzer{} + v := newPipeline(a, c, config.ModeShunt).DecideShunt(stmt("SELECT * FROM a, b")) + if !v.Shunted || v.Reason != "bad (cached, shunted)" { + t.Errorf("got %+v", v) + } +} + +func TestReportEntriesEmitted(t *testing.T) { + var entries []ReportEntry + a := &fakeAnalyzer{decision: &agent.QueryDecision{Approved: false, Reason: "bad"}} + cfg := Config{ + Analyzer: a, + Mode: config.ModeBlocking, + SkipAnalysis: func(string) bool { return false }, + Fingerprint: func(string) string { return "aabbccddeeff0011" }, + AnalysisTimeout: 5 * time.Second, + CacheTTL: time.Hour, + Report: func(ctx context.Context, e ReportEntry) { entries = append(entries, e) }, + } + New(cfg).Decide(stmt("SELECT * FROM a, b")) + if len(entries) != 1 || !entries[0].WasBlocked || entries[0].CacheHit { + t.Errorf("entries: %+v", entries) + } +} diff --git a/internal/proxy/protocol.go b/internal/proxy/protocol.go index 9f17205..c08bf53 100644 --- a/internal/proxy/protocol.go +++ b/internal/proxy/protocol.go @@ -6,6 +6,7 @@ import ( "encoding/hex" "fmt" "log/slog" + "pgproxy/internal/sqlnorm" "regexp" "strings" @@ -131,15 +132,8 @@ func normalizeInClauses(query string) string { // fallbackFingerprintQuery provides regex-based normalization as a fallback // when pg_query_go parsing fails (e.g., for PostgreSQL-specific extensions or invalid SQL). func fallbackFingerprintQuery(query string) string { - // Normalize whitespace - normalized := strings.TrimSpace(query) - normalized = whitespaceRe.ReplaceAllString(normalized, " ") - - // Replace string literals with placeholder - normalized = stringLitRe.ReplaceAllString(normalized, "'?'") - - // Replace numeric literals with placeholder - normalized = numericLitRe.ReplaceAllString(normalized, "?") + // Normalize via the shared literal normalizer (see internal/sqlnorm) + normalized := strings.TrimSpace(sqlnorm.NormalizeLiterals(query)) return normalized } diff --git a/internal/proxy/proxy.go b/internal/proxy/proxy.go index 184614d..1bf84cc 100644 --- a/internal/proxy/proxy.go +++ b/internal/proxy/proxy.go @@ -15,6 +15,7 @@ import ( "io" "log/slog" "net" + "pgproxy/internal/pipeline" "regexp" "strings" "sync" @@ -145,9 +146,10 @@ type Proxy struct { reporter reporter.Reporter metricsReporter reporter.MetricsReporter // Optional metrics reporter for performance tracking masker QueryMasker - authenticator Authenticator // Optional JWT authenticator for client connections - allowedDatabases []string // Optional list of databases JWT users can access (empty = allow all) - queryFilter *QueryFilter // Optional filter to bypass analysis for certain patterns + authenticator Authenticator // Optional JWT authenticator for client connections + allowedDatabases []string // Optional list of databases JWT users can access (empty = allow all) + queryFilter *QueryFilter // Optional filter to bypass analysis for certain patterns + pipeline *pipeline.Pipeline // protocol-neutral decision core listener net.Listener tlsConfig *tls.Config // TLS configuration for client connections requireClientTLS bool // Require all clients to use TLS @@ -184,34 +186,34 @@ type ipLimiterEntry struct { // Config holds the configuration for creating a new Proxy. type Config struct { - ListenAddr string - BackendAddr string - BackendSSLMode string // SSL mode for backend connections - BackendUser string // Username for backend authentication (used with JWT auth) - BackendPassword string // Password for backend authentication (used with JWT auth) - BackendDatabase string // Database for proxy's metadata connection (NOT enforced on JWT clients) - ShuntAddr string // Address of shunt node for rejected queries (shunt mode only) - ShuntSSLMode string // SSL mode for shunt node connections - ShuntUser string // Username for shunt node authentication - ShuntPassword string // Password for shunt node authentication - ShuntDatabase string // Database name for shunt node - Mode config.ProxyMode - Analyzer agent.Agent - Cache cache.ApprovalCache - Reporter reporter.Reporter - MetricsReporter reporter.MetricsReporter // Optional metrics reporter for performance tracking - Masker QueryMasker - Authenticator Authenticator // Optional JWT authenticator for client connections - AllowedDatabases []string // Optional list of databases JWT users can access (empty = allow all) - QueryFilter *QueryFilter // Optional filter to bypass analysis for certain patterns - TLSConfig *tls.Config // TLS configuration for client connections (optional) - RequireClientTLS bool // Require all clients to use TLS - BackendTLSConfig *tls.Config // TLS configuration for backend connections (optional) - ShuntTLSConfig *tls.Config // TLS configuration for shunt connections (optional) - MaxConnsPerIPPerMin int // Max new connections per source IP per minute (0 = disabled) - MaxMessageSize uint32 // Max PostgreSQL message size in bytes (0 = use defaultMaxMessageSize) - StoreRawQuery bool // Store raw SQL in reporter/cache (default: fingerprint only) - AllowMD5Auth bool // Allow MD5 password authentication (deprecated, default: false) + ListenAddr string + BackendAddr string + BackendSSLMode string // SSL mode for backend connections + BackendUser string // Username for backend authentication (used with JWT auth) + BackendPassword string // Password for backend authentication (used with JWT auth) + BackendDatabase string // Database for proxy's metadata connection (NOT enforced on JWT clients) + ShuntAddr string // Address of shunt node for rejected queries (shunt mode only) + ShuntSSLMode string // SSL mode for shunt node connections + ShuntUser string // Username for shunt node authentication + ShuntPassword string // Password for shunt node authentication + ShuntDatabase string // Database name for shunt node + Mode config.ProxyMode + Analyzer agent.Agent + Cache cache.ApprovalCache + Reporter reporter.Reporter + MetricsReporter reporter.MetricsReporter // Optional metrics reporter for performance tracking + Masker QueryMasker + Authenticator Authenticator // Optional JWT authenticator for client connections + AllowedDatabases []string // Optional list of databases JWT users can access (empty = allow all) + QueryFilter *QueryFilter // Optional filter to bypass analysis for certain patterns + TLSConfig *tls.Config // TLS configuration for client connections (optional) + RequireClientTLS bool // Require all clients to use TLS + BackendTLSConfig *tls.Config // TLS configuration for backend connections (optional) + ShuntTLSConfig *tls.Config // TLS configuration for shunt connections (optional) + MaxConnsPerIPPerMin int // Max new connections per source IP per minute (0 = disabled) + MaxMessageSize uint32 // Max PostgreSQL message size in bytes (0 = use defaultMaxMessageSize) + StoreRawQuery bool // Store raw SQL in reporter/cache (default: fingerprint only) + AllowMD5Auth bool // Allow MD5 password authentication (deprecated, default: false) } // New creates a new Proxy with the given configuration. @@ -237,16 +239,16 @@ func New(cfg Config) *Proxy { authenticator: cfg.Authenticator, allowedDatabases: cfg.AllowedDatabases, queryFilter: cfg.QueryFilter, - tlsConfig: cfg.TLSConfig, - requireClientTLS: cfg.RequireClientTLS, - backendTLSConfig: cfg.BackendTLSConfig, - shuntTLSConfig: cfg.ShuntTLSConfig, - maxConns: maxConcurrentConnections, - ipLimiters: make(map[string]*ipLimiterEntry), - maxConnsPerIPpm: cfg.MaxConnsPerIPPerMin, - maxMessageSize: cfg.MaxMessageSize, - storeRawQuery: cfg.StoreRawQuery, - allowMD5Auth: cfg.AllowMD5Auth, + tlsConfig: cfg.TLSConfig, + requireClientTLS: cfg.RequireClientTLS, + backendTLSConfig: cfg.BackendTLSConfig, + shuntTLSConfig: cfg.ShuntTLSConfig, + maxConns: maxConcurrentConnections, + ipLimiters: make(map[string]*ipLimiterEntry), + maxConnsPerIPpm: cfg.MaxConnsPerIPPerMin, + maxMessageSize: cfg.MaxMessageSize, + storeRawQuery: cfg.StoreRawQuery, + allowMD5Auth: cfg.AllowMD5Auth, } } @@ -2308,102 +2310,10 @@ func (p *Proxy) sendMD5PasswordMessage(conn net.Conn, user, password string, sal // In shunt mode: approved queries go to fast node, rejected queries are shunted. // database is the client's connected database for accurate metadata lookup. func (p *Proxy) analyzeQueryShunt(database, query, clientIP, username string) (approved bool, shunted bool, reason string) { - if p.analyzer == nil { - // No agent configured - all queries go to fast node - return true, false, "passthrough mode" - } - - if shouldSkipAnalysis(query) { - // System queries go to fast node - return true, false, "system query" - } - - ctx, cancel := context.WithTimeout(context.Background(), analysisTimeout) - defer cancel() - - fingerprint := fingerprintQuery(query) - - // Check cache first - if p.cache != nil { - cached, err := p.cache.Get(ctx, fingerprint) - if err != nil { - slog.Error("cache error", "error", err) - } else if cached != nil { - slog.Debug("cache hit (shunt)", "fingerprint", fingerprint[:8], "status", cached.Status, "shunted", cached.Shunted) - - // Report for audit trail with routing info - routedTo := reporter.RoutedToFast - if cached.Shunted { - routedTo = reporter.RoutedToShunt - } - p.reportIssue(ctx, query, fingerprint, cached.Reason, cached.Details, "", clientIP, database, username, - cached.Status == cache.StatusApproved, false, true, routedTo) - - if cached.Shunted { - return false, true, cached.Reason + " (cached, shunted)" - } - return cached.Status == cache.StatusApproved, false, cached.Reason + " (cached)" - } - slog.Debug("cache miss (shunt)", "fingerprint", fingerprint[:8]) - } - - // Analyze with agent (pass database for per-database metadata) - queryPreview := query - if len(query) > 80 { - queryPreview = query[:80] + "..." - } - slog.Info("analyzing query for database (shunt mode)", "database", database, "query_preview", queryPreview) - - // Track LLM analysis latency (for logging only; shunt mode doesn't use separate metrics table yet) - llmStart := time.Now() - decision, err := p.analyzer.AnalyzeQuery(ctx, database, query) - _ = int(time.Since(llmStart).Milliseconds()) // TODO: Add metrics tracking for shunt mode - - if err != nil { - // On error, allow query to fast node (fail-open) - slog.Error("agent error, routing to fast node", "database", database, "error", err) - return true, false, "analysis error - fail open" - } - - // Cache the decision - if p.cache != nil { - status := cache.StatusApproved - shuntedFlag := false - if !decision.Approved { - status = cache.StatusRejected - shuntedFlag = true // In shunt mode, rejected = shunted - } - approval := &cache.ApprovedQuery{ - Query: query, - Fingerprint: fingerprint, - Status: status, - Reason: decision.Reason, - Details: decision.Details, - DecidedAt: time.Now(), - Shunted: shuntedFlag, - } - if err := p.cache.Set(ctx, fingerprint, approval, approvalCacheTTL); err != nil { - slog.Error("cache set error", "error", err) - } else { - slog.Debug("cache set (shunt)", "fingerprint", fingerprint[:8], "status", status, "shunted", shuntedFlag) - } - } - - // Report with routing info - routedTo := reporter.RoutedToFast - if !decision.Approved { - routedTo = reporter.RoutedToShunt - } - p.reportIssue(ctx, query, fingerprint, decision.Reason, decision.Details, decision.SuggestedFix, clientIP, database, username, - decision.Approved, false, false, routedTo) - - if decision.Approved { - slog.Info("query approved (shunt mode) -> fast node", "reason", decision.Reason) - return true, false, decision.Reason - } - - slog.Warn("query rejected (shunt mode) -> shunt node", "reason", decision.Reason) - return false, true, decision.Reason + v := p.decisionPipeline().DecideShunt(pipeline.Statement{ + Database: database, SQL: query, ClientIP: clientIP, Username: username, + }) + return v.Approved, v.Shunted, v.Reason } // analyzeQuery synchronously analyzes a query and returns whether it should be allowed. @@ -2411,124 +2321,39 @@ func (p *Proxy) analyzeQueryShunt(database, query, clientIP, username string) (a // In non-blocking mode, issues are logged but queries are always approved. // database is the client's connected database for accurate metadata lookup. func (p *Proxy) analyzeQuery(database, query, clientIP, username string) (bool, string, string, string, *int) { - if p.analyzer == nil { - // No agent configured - allow query through - return true, "passthrough mode", "", "", nil - } - - if strings.TrimSpace(query) == "" { - slog.Debug("query skipped (empty)", "database", database) - return true, "empty query", "", "", nil - } - - if shouldSkipAnalysis(query) { - slog.Debug("query skipped (system/catalog query)", - "database", database, - "query_preview", truncateForLog(query, 80), - ) - return true, "system query", "", "", nil - } - - // Check query filter - queries matching patterns bypass analysis - if p.queryFilter != nil && p.queryFilter.ShouldFilter(query) { - slog.Info("query matched filter, bypassing analysis", - "query_prefix", truncateForLog(query, 80)) - return true, "matched filter pattern", "", "", nil - } - - // Create context with timeout - ctx, cancel := context.WithTimeout(context.Background(), analysisTimeout) - defer cancel() - - // Generate fingerprint for cache lookup - fingerprint := fingerprintQuery(query) - - // Check cache first - if p.cache != nil { - cached, err := p.cache.Get(ctx, fingerprint) - if err != nil { - slog.Error("cache error", "error", err) - } else if cached != nil { - slog.Debug("cache hit", "fingerprint", fingerprint[:8], "status", cached.Status, "reason", cached.Reason) - approved := cached.Status == cache.StatusApproved - - // Report cached result for audit trail (not in shunt mode, so no routing) - wasBlocked := !approved && p.mode == config.ModeBlocking - p.reportIssue(ctx, query, fingerprint, cached.Reason, cached.Details, "", clientIP, database, username, approved, wasBlocked, true, reporter.RoutedToNone) - - // In non-blocking mode, always approve but report the issue - if !approved && p.mode == config.ModeNonBlocking { - return true, cached.Reason + " (logged, non-blocking)", cached.Details, "", nil - } - return approved, cached.Reason + " (cached)", cached.Details, "", nil - } - slog.Debug("cache miss", "fingerprint", fingerprint[:8]) - } - - // Analyze with agent (pass database for per-database metadata) - queryPreview := query - if len(query) > 80 { - queryPreview = query[:80] + "..." - } - slog.Info("analyzing query for database", "database", database, "query_preview", queryPreview) - - // Track LLM analysis latency - llmStart := time.Now() - decision, err := p.analyzer.AnalyzeQuery(ctx, database, query) - llmLatencyMs := int(time.Since(llmStart).Milliseconds()) - - if err != nil { - // On error or timeout, allow query through (fail-open) - slog.Error("agent error, allowing query", "database", database, "error", err) - return true, "analysis error - fail open", "", "", &llmLatencyMs - } - - // Cache the decision - if p.cache != nil { - status := cache.StatusApproved - if !decision.Approved { - status = cache.StatusRejected - } - approval := &cache.ApprovedQuery{ - Query: query, - Fingerprint: fingerprint, - Status: status, - Reason: decision.Reason, - Details: decision.Details, - DecidedAt: time.Now(), - } - if err := p.cache.Set(ctx, fingerprint, approval, approvalCacheTTL); err != nil { - slog.Error("cache set error", "error", err) - } else { - slog.Debug("cache set", "fingerprint", fingerprint[:8], "status", status) - } - } - - if decision.Approved { - slog.Info("query approved", "reason", decision.Reason) - if decision.Details != "" { - slog.Debug("agent details", "details", decision.Details) - } - } else { - slog.Warn("query rejected", "reason", decision.Reason) - if decision.Details != "" { - slog.Debug("agent details", "details", decision.Details) - } - if decision.SuggestedFix != "" { - slog.Info("suggested fix", "fix", decision.SuggestedFix) - } - } - - // Report all analysis results (best effort audit trail, not in shunt mode so no routing) - wasBlocked := !decision.Approved && p.mode == config.ModeBlocking - p.reportIssue(ctx, query, fingerprint, decision.Reason, decision.Details, decision.SuggestedFix, clientIP, database, username, decision.Approved, wasBlocked, false, reporter.RoutedToNone) - - // In non-blocking mode, allow rejected queries through - if !decision.Approved && p.mode == config.ModeNonBlocking { - return true, decision.Reason + " (logged, non-blocking)", decision.Details, decision.SuggestedFix, &llmLatencyMs - } + v := p.decisionPipeline().Decide(pipeline.Statement{ + Database: database, SQL: query, ClientIP: clientIP, Username: username, + }) + return v.Approved, v.Reason, v.Details, v.SuggestedFix, v.LLMLatencyMs +} - return decision.Approved, decision.Reason, decision.Details, decision.SuggestedFix, &llmLatencyMs +// decisionPipeline lazily constructs the protocol-neutral decision core. +// Callbacks bind wire-specific helpers (skip list, fingerprinting) and the +// proxy's reporter; see internal/pipeline. +func (p *Proxy) decisionPipeline() *pipeline.Pipeline { + if p.pipeline != nil { + return p.pipeline + } + cfg := pipeline.Config{ + Analyzer: p.analyzer, + Cache: p.cache, + Mode: p.mode, + SkipAnalysis: shouldSkipAnalysis, + Fingerprint: fingerprintQuery, + AnalysisTimeout: analysisTimeout, + CacheTTL: approvalCacheTTL, + Report: func(ctx context.Context, e pipeline.ReportEntry) { + p.reportIssue(ctx, e.Statement.SQL, e.Fingerprint, e.Reason, e.Details, e.SuggestedFix, + e.Statement.ClientIP, e.Statement.Database, e.Statement.Username, + e.Approved, e.WasBlocked, e.CacheHit, e.RoutedTo) + }, + } + // Assign the filter only when non-nil to avoid the nil-interface gotcha. + if p.queryFilter != nil { + cfg.Filter = p.queryFilter + } + p.pipeline = pipeline.New(cfg) + return p.pipeline } // prewarmMetadataCache eagerly loads and caches metadata for a database immediately after connection diff --git a/internal/sqlnorm/sqlnorm.go b/internal/sqlnorm/sqlnorm.go new file mode 100644 index 0000000..1d45299 --- /dev/null +++ b/internal/sqlnorm/sqlnorm.go @@ -0,0 +1,47 @@ +// Package sqlnorm provides SQL literal normalization shared by query +// fingerprinting and agent prompt redaction. +// +// Motivation (see docs/multiprotocol-design.md §8a): query literals can +// contain sensitive predicates (account numbers, emails, ID numbers). +// The analysis agent only needs the *shape* of a query to judge it, so +// literals are redacted before any query text is sent to the LLM. The +// EXPLAIN tool runs server-side against the trusted original, so plan +// quality is unaffected. +package sqlnorm + +import ( + "crypto/sha256" + "encoding/hex" + "regexp" + "strings" +) + +var ( + whitespaceRe = regexp.MustCompile(`\s+`) + // String literals, including escaped quotes via adjacent literals ('O''Brien'). + stringLitRe = regexp.MustCompile(`'[^']*'`) + // Dollar-quoted literals ($$...$$ and $tag$...$tag$). + dollarLitRe = regexp.MustCompile(`\$([A-Za-z_]*)\$[\s\S]*?\$([A-Za-z_]*)\$`) + // Numeric literals (word-bounded so identifiers like table2 are untouched). + numericLitRe = regexp.MustCompile(`\b\d+\.?\d*\b`) +) + +// NormalizeLiterals replaces string, dollar-quoted, and numeric literals +// with placeholders and collapses whitespace. The result preserves the +// query's structure (tables, columns, operators, clauses) while removing +// every literal value. +func NormalizeLiterals(sql string) string { + s := dollarLitRe.ReplaceAllString(sql, "'?'") + s = stringLitRe.ReplaceAllString(s, "'?'") + s = numericLitRe.ReplaceAllString(s, "?") + return whitespaceRe.ReplaceAllString(s, " ") +} + +// Fingerprint returns a 32-hex-char cache key for a query's shape: +// SHA-256 over the trimmed, literal-normalized text. Used by wire +// frontends whose dialect has no native parser-based normalizer. +func Fingerprint(sql string) string { + normalized := strings.TrimSpace(NormalizeLiterals(sql)) + sum := sha256.Sum256([]byte(normalized)) + return hex.EncodeToString(sum[:16]) +} diff --git a/internal/sqlnorm/sqlnorm_test.go b/internal/sqlnorm/sqlnorm_test.go new file mode 100644 index 0000000..d63a1f0 --- /dev/null +++ b/internal/sqlnorm/sqlnorm_test.go @@ -0,0 +1,62 @@ +package sqlnorm + +import ( + "strings" + "testing" +) + +func TestNormalizeLiterals(t *testing.T) { + tests := []struct { + name string + in string + want string + }{ + { + name: "string literal", + in: "SELECT * FROM users WHERE email = 'leon@example.com'", + want: "SELECT * FROM users WHERE email = '?'", + }, + { + name: "numeric literal", + in: "SELECT * FROM users WHERE id = 8817", + want: "SELECT * FROM users WHERE id = ?", + }, + { + name: "identifier with digits untouched", + in: "SELECT * FROM table2 WHERE id = 5", + want: "SELECT * FROM table2 WHERE id = ?", + }, + { + name: "adjacent quoted literals", + in: "SELECT * FROM t WHERE name = 'O''Brien'", + want: "SELECT * FROM t WHERE name = '?''?'", + }, + { + name: "dollar quoted", + in: "SELECT * FROM t WHERE doc = $tag$secret value 123$tag$", + want: "SELECT * FROM t WHERE doc = '?'", + }, + { + name: "whitespace collapsed", + in: "SELECT *\n FROM t\n WHERE a = 1", + want: "SELECT * FROM t WHERE a = ?", + }, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + if got := NormalizeLiterals(tt.in); got != tt.want { + t.Errorf("NormalizeLiterals(%q) = %q, want %q", tt.in, got, tt.want) + } + }) + } +} + +func TestNoSensitiveValueSurvives(t *testing.T) { + in := "SELECT balance FROM accounts WHERE id_number = '8001015009087' AND acc = 62001234567" + got := NormalizeLiterals(in) + for _, secret := range []string{"8001015009087", "62001234567"} { + if strings.Contains(got, secret) { + t.Errorf("sensitive value %q survived normalization: %q", secret, got) + } + } +} diff --git a/internal/wire/tds/integration_test.go b/internal/wire/tds/integration_test.go new file mode 100644 index 0000000..f83b55c --- /dev/null +++ b/internal/wire/tds/integration_test.go @@ -0,0 +1,162 @@ +package tds + +import ( + "context" + "database/sql" + "encoding/pem" + "fmt" + "net" + "os" + "path/filepath" + "strings" + "testing" + "time" + + _ "github.com/microsoft/go-mssqldb" + + "pgproxy/internal/pipeline" +) + +// TestInterceptionAgainstRealBackend drives a real SQL Server client +// through the TDS frontend to a real backend, with a fake decider: +// approved queries return rows; rejected queries surface as native SQL +// errors carrying the guardian's reason and fix. +// +// Requires a running SQL Server/Azure SQL Edge; enable with: +// +// TDS_IT_BACKEND=localhost:11433 TDS_IT_PASSWORD='...' go test ./internal/wire/tds/ -run Interception -v +func TestInterceptionAgainstRealBackend(t *testing.T) { + backend := os.Getenv("TDS_IT_BACKEND") + if backend == "" { + t.Skip("TDS_IT_BACKEND not set; skipping live TDS integration test") + } + password := os.Getenv("TDS_IT_PASSWORD") + + decide := func(ctx context.Context, q string) pipeline.Verdict { + if strings.Contains(strings.ToLower(q), "forbidden_table") { + return pipeline.Verdict{ + Approved: false, + Reason: "cartesian product across 26M rows", + SuggestedFix: "SELECT c.* FROM client c JOIN transaction t ON c.client_id = t.client_id", + } + } + return pipeline.Verdict{Approved: true, Reason: "fine"} + } + + // Free port for the frontend. + l, err := net.Listen("tcp", "127.0.0.1:0") + if err != nil { + t.Fatal(err) + } + addr := l.Addr().String() + l.Close() + + srv := &Server{ListenAddr: addr, BackendAddr: backend, Decide: decide} + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + go srv.ListenAndServe(ctx) + time.Sleep(200 * time.Millisecond) + + host, port, _ := net.SplitHostPort(addr) + dsn := fmt.Sprintf("server=%s;port=%s;user id=sa;password=%s;encrypt=disable", host, port, password) + db, err := sql.Open("sqlserver", dsn) + if err != nil { + t.Fatal(err) + } + defer db.Close() + + // Approved query flows to the backend and returns data. + var two int + if err := db.QueryRow("SELECT 1+1").Scan(&two); err != nil { + t.Fatalf("approved query failed: %v", err) + } + if two != 2 { + t.Fatalf("approved query returned %d", two) + } + + // Rejected query is blocked with the guardian's error. + _, err = db.Query("SELECT * FROM forbidden_table, other_table") + if err == nil { + t.Fatal("expected rejected query to error") + } + msg := err.Error() + for _, want := range []string{"Query blocked by SQL Engineer", "cartesian product", "Fix:", "JOIN transaction t"} { + if !strings.Contains(msg, want) { + t.Errorf("error message missing %q: %s", want, msg) + } + } + + // The connection survives a rejection: next query still works. + if err := db.QueryRow("SELECT 40+2").Scan(&two); err != nil { + t.Fatalf("post-rejection query failed: %v", err) + } + if two != 42 { + t.Fatalf("post-rejection query returned %d", two) + } +} + +// TestStrictTLSInterception repeats the interception flow with TDS 8.0 +// strict encryption: the client opens with a TLS ClientHello, the +// frontend terminates it (ALPN tds/8.0) and speaks plain TDS to the +// backend. Requires TDS_IT_BACKEND as above. +func TestStrictTLSInterception(t *testing.T) { + backend := os.Getenv("TDS_IT_BACKEND") + if backend == "" { + t.Skip("TDS_IT_BACKEND not set; skipping live TDS integration test") + } + password := os.Getenv("TDS_IT_PASSWORD") + + tlsCfg, err := SelfSignedTLSConfig("localhost", "127.0.0.1") + if err != nil { + t.Fatal(err) + } + // Strict mode mandates certificate validation (trustservercertificate + // is ignored, per TDS 8.0); pin the server certificate like a real + // client would. + certPEM := pem.EncodeToMemory(&pem.Block{Type: "CERTIFICATE", Bytes: tlsCfg.Certificates[0].Certificate[0]}) + certFile := filepath.Join(t.TempDir(), "guardian.pem") + if err := os.WriteFile(certFile, certPEM, 0o600); err != nil { + t.Fatal(err) + } + + decide := func(ctx context.Context, q string) pipeline.Verdict { + if strings.Contains(strings.ToLower(q), "forbidden_table") { + return pipeline.Verdict{Approved: false, Reason: "cartesian product", SuggestedFix: "add a JOIN condition"} + } + return pipeline.Verdict{Approved: true, Reason: "fine"} + } + + l, err := net.Listen("tcp", "127.0.0.1:0") + if err != nil { + t.Fatal(err) + } + addr := l.Addr().String() + l.Close() + + srv := &Server{ListenAddr: addr, BackendAddr: backend, Decide: decide, TLSConfig: tlsCfg} + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + go srv.ListenAndServe(ctx) + time.Sleep(200 * time.Millisecond) + + host, port, _ := net.SplitHostPort(addr) + dsn := fmt.Sprintf("server=%s;port=%s;user id=sa;password=%s;encrypt=strict;certificate=%s;hostnameincertificate=localhost", host, port, password, certFile) + db, err := sql.Open("sqlserver", dsn) + if err != nil { + t.Fatal(err) + } + defer db.Close() + + var two int + if err := db.QueryRow("SELECT 1+1").Scan(&two); err != nil { + t.Fatalf("strict-TLS approved query failed: %v", err) + } + if two != 2 { + t.Fatalf("got %d", two) + } + + _, err = db.Query("SELECT * FROM forbidden_table, other_table") + if err == nil || !strings.Contains(err.Error(), "Query blocked by SQL Engineer") { + t.Fatalf("expected guardian block over strict TLS, got: %v", err) + } +} diff --git a/internal/wire/tds/message.go b/internal/wire/tds/message.go new file mode 100644 index 0000000..823bc5d --- /dev/null +++ b/internal/wire/tds/message.go @@ -0,0 +1,192 @@ +package tds + +import ( + "encoding/binary" + "fmt" + "unicode/utf16" +) + +// Message is one complete client request: all packets up to and +// including the EOM packet, sharing one packet type. +type Message struct { + Type byte + Packets []*Packet +} + +// Payload concatenates the packet payloads. +func (m *Message) Payload() []byte { + if len(m.Packets) == 1 { + return m.Packets[0].Payload + } + var total int + for _, p := range m.Packets { + total += len(p.Payload) + } + out := make([]byte, 0, total) + for _, p := range m.Packets { + out = append(out, p.Payload...) + } + return out +} + +// SQLFromBatch extracts the SQL text from a SQLBatch message payload. +// +// Per MS-TDS, a SQLBatch payload in TDS 7.2+ begins with ALL_HEADERS: a +// little-endian uint32 TotalLength (inclusive of itself) followed by +// header structures (transaction descriptor etc.), then the SQL text as +// UTF-16LE. Some legacy clients omit ALL_HEADERS; detect by sanity of +// the length prefix. +func SQLFromBatch(payload []byte) (string, error) { + if len(payload) == 0 { + return "", fmt.Errorf("tds: empty SQLBatch payload") + } + body := payload + if len(payload) >= 4 { + total := binary.LittleEndian.Uint32(payload[:4]) + // A valid ALL_HEADERS block is at least 4 bytes, no larger than + // the payload, and leaves an even number of bytes of UTF-16 text. + if total >= 4 && int(total) <= len(payload) && (len(payload)-int(total))%2 == 0 { + body = payload[total:] + } + } + if len(body)%2 != 0 { + return "", fmt.Errorf("tds: SQLBatch text has odd byte length %d", len(body)) + } + u := make([]uint16, len(body)/2) + for i := range u { + u[i] = binary.LittleEndian.Uint16(body[2*i:]) + } + return string(utf16.Decode(u)), nil +} + +// EncodeUCS2 encodes s as UTF-16LE bytes. +func EncodeUCS2(s string) []byte { + u := utf16.Encode([]rune(s)) + out := make([]byte, len(u)*2) + for i, c := range u { + binary.LittleEndian.PutUint16(out[2*i:], c) + } + return out +} + +// Token identifiers used in server responses. +const ( + tokenError byte = 0xAA + tokenDone byte = 0xFD +) + +// DONE status flags. +const doneError uint16 = 0x0002 + +// GuardianErrorNumber is the message number used for proxy rejections. +// User-defined SQL Server error numbers start at 50000. +const GuardianErrorNumber int32 = 50999 + +// ErrorResponsePacket builds a server->client TabularResult packet +// carrying an ERROR token (severity 16) followed by DONE, which is how +// SQL Server reports a failed batch. Clients (sqlcmd, SSMS, drivers) +// render msgText; the suggested fix travels inside it. +func ErrorResponsePacket(spid uint16, msgText string) *Packet { + text := EncodeUCS2(msgText) + server := EncodeUCS2("pgproxy-guardian") + + // ERROR token body (TDS 7.2+): Number int32, State byte, Class byte, + // MsgText us_varchar, ServerName b_varchar, ProcName b_varchar, + // LineNumber int32. + body := make([]byte, 0, 16+len(text)+len(server)) + var n4 [4]byte + binary.LittleEndian.PutUint32(n4[:], uint32(GuardianErrorNumber)) + body = append(body, n4[:]...) + body = append(body, 1) // state + body = append(body, 16) // class/severity: user error + var n2 [2]byte + binary.LittleEndian.PutUint16(n2[:], uint16(len(text)/2)) // char count + body = append(body, n2[:]...) + body = append(body, text...) + body = append(body, byte(len(server)/2)) // b_varchar char count + body = append(body, server...) + body = append(body, 0) // proc name: empty + binary.LittleEndian.PutUint32(n4[:], 0) + body = append(body, n4[:]...) // line number + + payload := make([]byte, 0, 3+len(body)+13) + payload = append(payload, tokenError) + binary.LittleEndian.PutUint16(n2[:], uint16(len(body))) + payload = append(payload, n2[:]...) + payload = append(payload, body...) + + // DONE token: Status uint16, CurCmd uint16, RowCount uint64 (7.2+). + payload = append(payload, tokenDone) + binary.LittleEndian.PutUint16(n2[:], doneError) + payload = append(payload, n2[:]...) + payload = append(payload, 0, 0) // CurCmd + payload = append(payload, 0, 0, 0, 0, 0, 0, 0, 0) // RowCount + + return &Packet{ + Type: PacketTabular, + Status: StatusEOM, + SPID: spid, + Payload: payload, + } +} + +// PRELOGIN option tokens and ENCRYPTION values (MS-TDS §2.2.6.5). +const ( + preloginOptionEncryption byte = 0x01 + preloginTerminator byte = 0xFF + + // EncryptNotSup declares no encryption support. + EncryptNotSup byte = 0x02 + // EncryptStrict is the TDS 8.0 strict-encryption marker. + EncryptStrict byte = 0x04 +) + +// SetPreloginEncryption rewrites the ENCRYPTION option inside a +// PRELOGIN payload in place. Returns false if the option is absent or +// the payload is malformed. Used by the strict-TLS path to translate +// between a TDS 8.0 client (inside the terminated tunnel) and a +// backend negotiating classic 7.x PRELOGIN encryption. +func SetPreloginEncryption(payload []byte, value byte) bool { + i := 0 + for i+5 <= len(payload) { + token := payload[i] + if token == preloginTerminator { + return false + } + offset := int(payload[i+1])<<8 | int(payload[i+2]) + length := int(payload[i+3])<<8 | int(payload[i+4]) + if token == preloginOptionEncryption { + if length < 1 || offset+1 > len(payload) { + return false + } + payload[offset] = value + return true + } + i += 5 + } + return false +} + +// TDS version constants as they appear in the LOGIN7 TDSVersion field. +const ( + VerTDS74 uint32 = 0x74000004 + VerTDS80 uint32 = 0x08000000 +) + +// DowngradeLogin7Version rewrites the TDSVersion field of a LOGIN7 +// payload from 8.0 to 7.4 in place. A strict-mode (TDS 8.0) client +// inside the terminated tunnel declares 8.0; a 7.x backend refuses it, +// so the frontend presents the login as 7.4. Returns true if a +// downgrade was applied. +func DowngradeLogin7Version(payload []byte) bool { + // LOGIN7 fixed header: Length uint32 LE, then TDSVersion uint32 LE. + if len(payload) < 8 { + return false + } + ver := binary.LittleEndian.Uint32(payload[4:8]) + if ver != VerTDS80 { + return false + } + binary.LittleEndian.PutUint32(payload[4:8], VerTDS74) + return true +} diff --git a/internal/wire/tds/packet.go b/internal/wire/tds/packet.go new file mode 100644 index 0000000..ded6c46 --- /dev/null +++ b/internal/wire/tds/packet.go @@ -0,0 +1,137 @@ +// Package tds implements the server side of the MS-TDS wire protocol +// (the protocol spoken by every SQL Server client: sqlcmd, SSMS, JDBC, +// ODBC, go-mssqldb) as a proxy frontend. +// +// Milestone 2 (docs/multiprotocol-design.md §5): packet-aware +// passthrough. The proxy reads whole TDS packets from the client so the +// message boundary layer exists, relays them to the backend unchanged, +// and streams backend bytes straight back. Interception of SQLBatch/RPC +// builds on this in milestone 3. +// +// References: MS-TDS open specification; go-mssqldb and pytds as +// implementation references; Wireshark's TDS dissector for traces. +package tds + +import ( + "encoding/binary" + "fmt" + "io" +) + +// Packet types (MS-TDS §2.2.3.1.1). +const ( + PacketSQLBatch byte = 0x01 + PacketOldLogin byte = 0x02 // pre-7.0 login, refused + PacketRPC byte = 0x03 + PacketTabular byte = 0x04 // server -> client results + PacketAttention byte = 0x06 + PacketBulkLoad byte = 0x07 + PacketFedAuth byte = 0x08 // federated auth token + PacketTransaction byte = 0x0E // transaction manager request + PacketLogin7 byte = 0x10 + PacketSSPI byte = 0x11 + PacketPrelogin byte = 0x12 +) + +// PacketName returns a human-readable name for logging. +func PacketName(t byte) string { + switch t { + case PacketSQLBatch: + return "SQLBatch" + case PacketOldLogin: + return "OldLogin" + case PacketRPC: + return "RPC" + case PacketTabular: + return "TabularResult" + case PacketAttention: + return "Attention" + case PacketBulkLoad: + return "BulkLoad" + case PacketFedAuth: + return "FedAuthToken" + case PacketTransaction: + return "TransactionManager" + case PacketLogin7: + return "Login7" + case PacketSSPI: + return "SSPI" + case PacketPrelogin: + return "Prelogin" + default: + return fmt.Sprintf("Unknown(0x%02x)", t) + } +} + +// Header status bits (MS-TDS §2.2.3.1.2). +const ( + StatusEOM byte = 0x01 // end of message + StatusIgnore byte = 0x02 + StatusResetConnection byte = 0x08 + StatusResetSkipTran byte = 0x10 +) + +// headerLen is the fixed TDS packet header size. +const headerLen = 8 + +// maxPacketLen bounds a single TDS packet. The negotiated packet size +// is at most 32767 in practice; 1MiB is a defensive ceiling. +const maxPacketLen = 1 << 20 + +// Packet is one TDS packet: an 8-byte header plus payload. +type Packet struct { + Type byte + Status byte + SPID uint16 + Seq byte + Payload []byte +} + +// EOM reports whether this packet ends a client message. +func (p *Packet) EOM() bool { return p.Status&StatusEOM != 0 } + +// ReadPacket reads exactly one TDS packet from r. +func ReadPacket(r io.Reader) (*Packet, error) { + var hdr [headerLen]byte + if _, err := io.ReadFull(r, hdr[:]); err != nil { + return nil, err + } + length := binary.BigEndian.Uint16(hdr[2:4]) + if int(length) < headerLen { + return nil, fmt.Errorf("tds: packet length %d shorter than header", length) + } + if int(length) > maxPacketLen { + return nil, fmt.Errorf("tds: packet length %d exceeds limit", length) + } + payload := make([]byte, int(length)-headerLen) + if _, err := io.ReadFull(r, payload); err != nil { + return nil, err + } + return &Packet{ + Type: hdr[0], + Status: hdr[1], + SPID: binary.BigEndian.Uint16(hdr[4:6]), + Seq: hdr[6], + Payload: payload, + }, nil +} + +// WritePacket writes one TDS packet to w. +func WritePacket(w io.Writer, p *Packet) error { + length := headerLen + len(p.Payload) + if length > maxPacketLen { + return fmt.Errorf("tds: packet length %d exceeds limit", length) + } + var hdr [headerLen]byte + hdr[0] = p.Type + hdr[1] = p.Status + binary.BigEndian.PutUint16(hdr[2:4], uint16(length)) + binary.BigEndian.PutUint16(hdr[4:6], p.SPID) + hdr[6] = p.Seq + hdr[7] = 0 // window, unused + if _, err := w.Write(hdr[:]); err != nil { + return err + } + _, err := w.Write(p.Payload) + return err +} diff --git a/internal/wire/tds/packet_test.go b/internal/wire/tds/packet_test.go new file mode 100644 index 0000000..d163900 --- /dev/null +++ b/internal/wire/tds/packet_test.go @@ -0,0 +1,53 @@ +package tds + +import ( + "bytes" + "io" + "testing" +) + +func TestPacketRoundTrip(t *testing.T) { + in := &Packet{ + Type: PacketSQLBatch, + Status: StatusEOM, + SPID: 52, + Seq: 1, + Payload: []byte("SELECT 1"), + } + var buf bytes.Buffer + if err := WritePacket(&buf, in); err != nil { + t.Fatal(err) + } + out, err := ReadPacket(&buf) + if err != nil { + t.Fatal(err) + } + if out.Type != in.Type || out.Status != in.Status || out.SPID != in.SPID || out.Seq != in.Seq { + t.Errorf("header mismatch: %+v vs %+v", out, in) + } + if !bytes.Equal(out.Payload, in.Payload) { + t.Errorf("payload mismatch: %q vs %q", out.Payload, in.Payload) + } + if !out.EOM() { + t.Error("EOM bit lost") + } +} + +func TestReadPacketRejectsBadLengths(t *testing.T) { + // Length shorter than header + short := []byte{0x01, 0x01, 0x00, 0x04, 0, 0, 0, 0} + if _, err := ReadPacket(bytes.NewReader(short)); err == nil { + t.Error("expected error for short length") + } + // Truncated payload + trunc := []byte{0x01, 0x01, 0x00, 0x10, 0, 0, 0, 0, 'x'} + if _, err := ReadPacket(bytes.NewReader(trunc)); err != io.ErrUnexpectedEOF { + t.Errorf("expected ErrUnexpectedEOF, got %v", err) + } +} + +func TestPacketNames(t *testing.T) { + if PacketName(PacketPrelogin) != "Prelogin" || PacketName(0xFF) != "Unknown(0xff)" { + t.Error("naming broken") + } +} diff --git a/internal/wire/tds/server.go b/internal/wire/tds/server.go new file mode 100644 index 0000000..7b53025 --- /dev/null +++ b/internal/wire/tds/server.go @@ -0,0 +1,274 @@ +package tds + +import ( + "context" + "crypto/tls" + "fmt" + "io" + "log/slog" + "net" + "sync" + "time" + + "pgproxy/internal/pipeline" +) + +// Decider evaluates one SQL batch and returns a verdict. Wired to +// pipeline.Decide in production; fakes in tests. +type Decider func(ctx context.Context, sql string) pipeline.Verdict + +// Server is the TDS frontend. Milestone 3: SQLBatch messages are +// assembled, their SQL extracted, and submitted to the decision +// pipeline; rejected batches never reach the backend and the client +// receives a native ERROR token with the reason and suggested fix. +// All other packet types relay untouched (login sequence, RPC and +// transaction manager interception land in a later milestone). +// +// Encryption: PRELOGIN is relayed untouched; this milestone targets +// unencrypted dev loops (sqlcmd -N disable). TLS termination is +// scheduled with TDS 8.0 strict-encryption support. +type Server struct { + ListenAddr string + BackendAddr string + + // Decide is optional; when nil every batch passes (pure relay). + Decide Decider + + // TLSConfig enables TDS 8.0 strict termination: clients that open + // with a TLS ClientHello are terminated here (ALPN tds/8.0) and + // speak plain TDS inside the tunnel. Legacy 7.x clients bypass it. + TLSConfig *tls.Config + + // DialTimeout bounds the backend dial. Zero means 10s. + DialTimeout time.Duration + + listener net.Listener + wg sync.WaitGroup +} + +// ListenAndServe accepts connections until ctx is cancelled. +func (s *Server) ListenAndServe(ctx context.Context) error { + l, err := net.Listen("tcp", s.ListenAddr) + if err != nil { + return fmt.Errorf("tds: listen %s: %w", s.ListenAddr, err) + } + s.listener = l + slog.Info("TDS frontend listening", "addr", s.ListenAddr, "backend", s.BackendAddr, + "interception", s.Decide != nil) + + go func() { + <-ctx.Done() + l.Close() + }() + + for { + conn, err := l.Accept() + if err != nil { + if ctx.Err() != nil { + s.wg.Wait() + return nil + } + return err + } + s.wg.Add(1) + go func() { + defer s.wg.Done() + s.handleConn(ctx, conn) + }() + } +} + +type session struct { + clientConn net.Conn + backendConn net.Conn + clientMu sync.Mutex // sole guard for client-bound writes + client string + strict bool // client is inside a terminated TDS 8.0 TLS tunnel +} + +func (sess *session) writeClient(b []byte) error { + sess.clientMu.Lock() + defer sess.clientMu.Unlock() + _, err := sess.clientConn.Write(b) + return err +} + +func (sess *session) writeClientPacket(p *Packet) error { + sess.clientMu.Lock() + defer sess.clientMu.Unlock() + return WritePacket(sess.clientConn, p) +} + +func (s *Server) handleConn(ctx context.Context, rawConn net.Conn) { + defer rawConn.Close() + + clientConn, strict, err := maybeTerminateTLS(rawConn, s.TLSConfig) + if err != nil { + slog.Error("TDS client TLS setup failed", "client", rawConn.RemoteAddr().String(), "strict", strict, "error", err) + return + } + defer clientConn.Close() + sess := &session{clientConn: clientConn, client: rawConn.RemoteAddr().String(), strict: strict} + slog.Info("TDS client connected", "client", sess.client, "strict_tls", strict) + + timeout := s.DialTimeout + if timeout == 0 { + timeout = 10 * time.Second + } + backendConn, err := net.DialTimeout("tcp", s.BackendAddr, timeout) + if err != nil { + slog.Error("TDS backend dial failed", "backend", s.BackendAddr, "error", err) + return + } + sess.backendConn = backendConn + defer backendConn.Close() + + done := make(chan struct{}) + var once sync.Once + closeAll := func() { + once.Do(func() { + close(done) + clientConn.Close() + backendConn.Close() + }) + } + defer closeAll() + + // In strict mode the PRELOGIN exchange is translated synchronously + // before the relay loops start: the tunnel already provides + // encryption, so the backend is told NOT_SUP and the client is told + // STRICT. + if sess.strict { + if err := s.preloginExchange(sess); err != nil { + slog.Error("TDS strict prelogin exchange failed", "client", sess.client, "error", err) + return + } + } + + // client -> backend: packet-aware, with SQLBatch interception. + go func() { + defer closeAll() + s.clientLoop(ctx, sess) + }() + + // backend -> client: raw byte stream (responses are not inspected). + buf := make([]byte, 32*1024) + for { + n, err := backendConn.Read(buf) + if n > 0 { + if werr := sess.writeClient(buf[:n]); werr != nil { + return + } + } + if err != nil { + if err != io.EOF { + slog.Debug("TDS backend read ended", "error", err) + } + return + } + } +} + +// preloginExchange relays one PRELOGIN round-trip, rewriting the +// ENCRYPTION option in both directions for a strict-TLS session. +func (s *Server) preloginExchange(sess *session) error { + pkt, err := ReadPacket(sess.clientConn) + if err != nil { + return fmt.Errorf("read client prelogin: %w", err) + } + if pkt.Type == PacketPrelogin { + if !SetPreloginEncryption(pkt.Payload, EncryptNotSup) { + slog.Warn("TDS strict: client PRELOGIN has no encryption option") + } + } + if err := WritePacket(sess.backendConn, pkt); err != nil { + return fmt.Errorf("forward prelogin: %w", err) + } + resp, err := ReadPacket(sess.backendConn) + if err != nil { + return fmt.Errorf("read backend prelogin response: %w", err) + } + if resp.Type == PacketTabular || resp.Type == PacketPrelogin { + SetPreloginEncryption(resp.Payload, EncryptStrict) + } + return sess.writeClientPacket(resp) +} + +func (s *Server) clientLoop(ctx context.Context, sess *session) { + var batch []*Packet // accumulating SQLBatch packets until EOM + for { + pkt, err := ReadPacket(sess.clientConn) + if err != nil { + if err != io.EOF { + slog.Debug("TDS client read ended", "client", sess.client, "error", err) + } + return + } + slog.Debug("TDS client packet", + "client", sess.client, + "type", PacketName(pkt.Type), + "len", len(pkt.Payload), + "eom", pkt.EOM(), + ) + + if sess.strict && pkt.Type == PacketLogin7 { + if DowngradeLogin7Version(pkt.Payload) { + slog.Debug("TDS strict: LOGIN7 version presented to backend as 7.4", "client", sess.client) + } + } + + intercept := s.Decide != nil && pkt.Type == PacketSQLBatch + if !intercept { + if err := WritePacket(sess.backendConn, pkt); err != nil { + slog.Debug("TDS backend write failed", "error", err) + return + } + continue + } + + batch = append(batch, pkt) + if !pkt.EOM() { + continue + } + msg := &Message{Type: PacketSQLBatch, Packets: batch} + batch = nil + if err := s.handleBatch(ctx, sess, msg); err != nil { + return + } + } +} + +func (s *Server) handleBatch(ctx context.Context, sess *session, msg *Message) error { + sql, err := SQLFromBatch(msg.Payload()) + if err != nil { + // Malformed by our reading; fail open and let the backend judge. + slog.Warn("TDS batch parse failed, relaying unjudged", "error", err) + return relayAll(sess.backendConn, msg.Packets) + } + + v := s.Decide(ctx, sql) + if v.Approved { + slog.Debug("TDS batch approved", "reason", v.Reason) + return relayAll(sess.backendConn, msg.Packets) + } + + slog.Warn("TDS batch blocked", "client", sess.client, "reason", v.Reason) + msgText := "Query blocked by SQL Engineer: " + v.Reason + if v.Details != "" { + msgText += "\n" + v.Details + } + if v.SuggestedFix != "" { + msgText += "\nFix: " + v.SuggestedFix + } + spid := msg.Packets[0].SPID + return sess.writeClientPacket(ErrorResponsePacket(spid, msgText)) +} + +func relayAll(w io.Writer, pkts []*Packet) error { + for _, p := range pkts { + if err := WritePacket(w, p); err != nil { + return err + } + } + return nil +} diff --git a/internal/wire/tds/skip.go b/internal/wire/tds/skip.go new file mode 100644 index 0000000..212fdc3 --- /dev/null +++ b/internal/wire/tds/skip.go @@ -0,0 +1,30 @@ +package tds + +import "strings" + +// SkipAnalysis reports whether a T-SQL batch is session/tooling traffic +// that should bypass analysis (the TDS analog of the pgwire system-query +// skip list). sqlcmd, SSMS and drivers emit these constantly on connect. +func SkipAnalysis(sql string) bool { + q := strings.ToLower(strings.TrimSpace(sql)) + if q == "" { + return true + } + prefixes := []string{ + "set ", // session options (sqlcmd sends several on connect) + "use ", // database switch + "declare @", // local tooling scratch + "select @@", // @@version, @@spid, @@trancount probes + "exec sp_", // sp_reset_connection and friends + "execute sp_", + "dbcc ", + "print ", + "if @@trancount", + } + for _, p := range prefixes { + if strings.HasPrefix(q, p) { + return true + } + } + return false +} diff --git a/internal/wire/tds/tls.go b/internal/wire/tds/tls.go new file mode 100644 index 0000000..084c5bc --- /dev/null +++ b/internal/wire/tds/tls.go @@ -0,0 +1,106 @@ +package tds + +import ( + "bufio" + "crypto/ecdsa" + "crypto/elliptic" + "crypto/rand" + "crypto/tls" + "crypto/x509" + "crypto/x509/pkix" + "math/big" + "net" + "time" +) + +// alpnTDS8 is the ALPN protocol id for TDS 8.0 strict encryption. +// Strict-mode clients (SSMS 20+, SqlClient 5+, ODBC 18+, go-mssqldb +// encrypt=strict) require it during the TLS handshake. +const alpnTDS8 = "tds/8.0" + +// NewTLSConfig prepares a server TLS config for TDS 8.0 strict +// termination from a certificate. +func NewTLSConfig(cert tls.Certificate) *tls.Config { + return &tls.Config{ + Certificates: []tls.Certificate{cert}, + NextProtos: []string{alpnTDS8}, + MinVersion: tls.VersionTLS12, + } +} + +// SelfSignedTLSConfig generates an in-memory self-signed certificate +// for development loops (clients connect with trustservercertificate). +// Production supplies a real certificate via TDS_TLS_CERT_PATH. +func SelfSignedTLSConfig(hosts ...string) (*tls.Config, error) { + key, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader) + if err != nil { + return nil, err + } + serial, err := rand.Int(rand.Reader, new(big.Int).Lsh(big.NewInt(1), 128)) + if err != nil { + return nil, err + } + tmpl := x509.Certificate{ + SerialNumber: serial, + Subject: pkix.Name{CommonName: "pgproxy-tds-guardian"}, + NotBefore: time.Now().Add(-time.Hour), + NotAfter: time.Now().Add(365 * 24 * time.Hour), + KeyUsage: x509.KeyUsageKeyEncipherment | x509.KeyUsageDigitalSignature, + ExtKeyUsage: []x509.ExtKeyUsage{x509.ExtKeyUsageServerAuth}, + } + for _, h := range hosts { + if ip := net.ParseIP(h); ip != nil { + tmpl.IPAddresses = append(tmpl.IPAddresses, ip) + } else { + tmpl.DNSNames = append(tmpl.DNSNames, h) + } + } + der, err := x509.CreateCertificate(rand.Reader, &tmpl, &tmpl, &key.PublicKey, key) + if err != nil { + return nil, err + } + cert := tls.Certificate{Certificate: [][]byte{der}, PrivateKey: key} + return NewTLSConfig(cert), nil +} + +// peekedConn replays bytes already buffered by a bufio.Reader before +// reading from the underlying connection. +type peekedConn struct { + net.Conn + r *bufio.Reader +} + +func (c *peekedConn) Read(p []byte) (int, error) { return c.r.Read(p) } + +// maybeTerminateTLS inspects the first byte of a fresh client +// connection. A TLS ClientHello (0x16) means a TDS 8.0 strict client: +// terminate TLS here and speak plain TDS inside. A TDS packet type +// (e.g. 0x12 Prelogin) means a legacy 7.x client: return the +// connection unwrapped for the passthrough path. +func maybeTerminateTLS(conn net.Conn, cfg *tls.Config) (net.Conn, bool, error) { + br := bufio.NewReader(conn) + first, err := br.Peek(1) + if err != nil { + return nil, false, err + } + pc := &peekedConn{Conn: conn, r: br} + if first[0] != 0x16 { + return pc, false, nil + } + if cfg == nil { + return pc, true, errNoTLSConfig + } + tconn := tls.Server(pc, cfg) + if err := tconn.Handshake(); err != nil { + return nil, true, err + } + return tconn, true, nil +} + +var errNoTLSConfig = &tlsConfigError{} + +type tlsConfigError struct{} + +func (*tlsConfigError) Error() string { + return "tds: client attempted TLS (strict mode) but no server certificate is configured; set TDS_TLS_CERT_PATH/TDS_TLS_KEY_PATH or TDS_TLS_SELF_SIGNED=true" +} diff --git a/internal/wire/wire.go b/internal/wire/wire.go new file mode 100644 index 0000000..012f234 --- /dev/null +++ b/internal/wire/wire.go @@ -0,0 +1,59 @@ +// Package wire defines the contract between protocol frontends and the +// protocol-neutral decision pipeline. See docs/multiprotocol-design.md. +// +// A frontend owns one client connection end-to-end: it runs the +// protocol's startup/auth phase, surfaces interceptable statements to +// the pipeline, and delivers verdicts in protocol-native form (a +// PostgreSQL ErrorResponse with the fix in Hint; a TDS ERROR token with +// the fix in the message text). +// +// The PostgreSQL frontend (internal/proxy) predates this contract and +// will be adapted to it as part of the TDS milestone; new frontends +// implement it directly. +package wire + +import ( + "context" + + "pgproxy/internal/pipeline" +) + +// StatementKind distinguishes ad-hoc from prepared traffic. +type StatementKind int + +const ( + KindBatch StatementKind = iota // pgwire Query / TDS SQLBatch + KindPrepare // pgwire Parse / TDS sp_prepare + KindExecPrepared // pgwire Bind+Execute / TDS sp_execute +) + +// Statement is a wire-level statement awaiting a verdict. +type Statement struct { + pipeline.Statement + Kind StatementKind +} + +// Identity is the authenticated principal for a session, as established +// during the protocol handshake (JWT, FedAuth token, or passthrough). +type Identity struct { + Username string // verified subject (email / UPN), never client-claimed + Method string // e.g. "jwt", "fedauth", "passthrough" +} + +// Frontend owns one client connection end-to-end. +type Frontend interface { + // Handshake runs the protocol's startup and auth phase, returning + // the authenticated identity once the session is established. + Handshake(ctx context.Context) (Identity, error) + + // Next blocks until the client sends an interceptable statement. + // Non-statement traffic is relayed transparently inside Next. + Next(ctx context.Context) (Statement, error) + + // Deliver applies a verdict in protocol-native form: forward the + // (possibly rewritten) statement, or reject with reason/detail/fix. + Deliver(ctx context.Context, s Statement, v pipeline.Verdict) error + + // Close releases the connection pair. + Close() error +}