diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index ddf4022c9..a3967a570 100644 --- a/.github/workflows/test.yml +++ b/.github/workflows/test.yml @@ -56,6 +56,12 @@ jobs: run: | mise run --output prefix test + - name: Show proxy logs on failure + if: failure() + run: | + docker logs --timestamps proxy 2>&1 || true + docker logs --timestamps proxy-tls 2>&1 || true + - uses: ./.github/actions/send-slack-notification with: channel: engineering diff --git a/Cargo.lock b/Cargo.lock index d247b5539..3d1e7f58d 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -450,6 +450,12 @@ version = "0.22.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "72b3254f16251a8381aa12e40e3c4d2f0199f8c6508fbecb9d91f575e0fbb8c6" +[[package]] +name = "base64" +version = "0.23.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ac07cdecf99051d9a5238b80f35af32cdeba5b336e55d957b318b50137e18da5" + [[package]] name = "base64ct" version = "1.7.3" @@ -622,9 +628,9 @@ checksum = "1fd0f2584146f6f2ef48085050886acf353beff7305ebd1ae69500e27c67f64b" [[package]] name = "bytes" -version = "1.11.1" +version = "1.12.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1e748733b7cbc798e1434b6ac524f0c1ff2ab456fe201501e6497c8417a4fc33" +checksum = "fc652a48c352aef3ea3aed32080501cf3ef6ed5da78602a020c991775b0aff04" dependencies = [ "serde", ] @@ -738,7 +744,7 @@ dependencies = [ "async-mutex", "async-trait", "base16ct", - "base64", + "base64 0.22.1", "base64ct", "base85", "blake3", @@ -751,7 +757,7 @@ dependencies = [ "dirs", "futures", "hex", - "hmac", + "hmac 0.12.1", "itertools 0.12.1", "lazy_static", "log", @@ -771,7 +777,7 @@ dependencies = [ "serde_cbor", "serde_json", "serdect", - "sha2", + "sha2 0.10.8", "stack-auth", "stack-profile", "static_assertions", @@ -807,12 +813,12 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1946988e0b7f9de259d85b10c9c1fd7e1327751103b808c2c6e930b9a87c25c9" dependencies = [ "getrandom 0.2.15", - "hmac", + "hmac 0.12.1", "lazy_static", "num-bigint", "rand 0.8.6", "regex", - "sha2", + "sha2 0.10.8", "thiserror 1.0.69", ] @@ -822,7 +828,6 @@ version = "3.0.1" dependencies = [ "arc-swap", "async-trait", - "aws-lc-rs", "bigdecimal", "blake3", "bytes", @@ -836,12 +841,11 @@ dependencies = [ "eql-mapper", "exitcode", "hex", - "md-5", + "md-5 0.10.6", "metrics", "metrics-exporter-prometheus", "moka", - "oid-registry", - "pg_escape", + "pg-proto", "postgres-protocol", "postgres-types", "rand 0.9.2", @@ -849,24 +853,23 @@ dependencies = [ "regex", "rust_decimal", "rustls", + "rustls-native-certs", "rustls-pki-types", "rustls-platform-verifier 0.5.1", "serde", "serde_json", - "socket2 0.5.8", + "socket2 0.6.5", "sqltk", "temp-env", "thiserror 2.0.18", "tokio", "tokio-postgres", "tokio-postgres-rustls", - "tokio-rustls", "tokio-util", "tracing", "tracing-subscriber", "uuid", "vitaminc-protected 0.1.0-pre4.2", - "x509-parser", ] [[package]] @@ -972,6 +975,12 @@ dependencies = [ "cc", ] +[[package]] +name = "cmov" +version = "0.5.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0c9ea0ac24bc397ab3c98583a3c9ba74fa56b09a4449bbe172b9b1ddb016027a" + [[package]] name = "colorchoice" version = "1.0.3" @@ -1031,6 +1040,12 @@ version = "0.9.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c2459377285ad874054d797f3ccebf984978aa39129f6eafde5cdc8315b612f8" +[[package]] +name = "const-oid" +version = "0.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a6ef517f0926dd24a1582492c791b6a4818a4d94e789a334894aa15b0d12f55c" + [[package]] name = "constant_time_eq" version = "0.3.1" @@ -1192,6 +1207,15 @@ dependencies = [ "vitaminc", ] +[[package]] +name = "ctutils" +version = "0.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7d5515a3834141de9eafb9717ad39eea8247b5674e6066c404e8c4b365d2a29e" +dependencies = [ + "cmov", +] + [[package]] name = "darling" version = "0.20.10" @@ -1248,7 +1272,7 @@ version = "0.7.9" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f55bf8e7b65898637379c1b74eb1551107c8294ed26d855ceb9fd1a09cfc9bc0" dependencies = [ - "const-oid", + "const-oid 0.9.6", "der_derive", "flagset", "zeroize", @@ -1393,7 +1417,9 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f1dd6dbb5841937940781866fa1281a1ff7bd3bf827091440879f9994983d5c2" dependencies = [ "block-buffer 0.12.1", + "const-oid 0.10.2", "crypto-common 0.2.2", + "ctutils", ] [[package]] @@ -1569,7 +1595,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "33d852cb9b869c2a9b3df2f71a3074817f01e1844f839a144f5fcef059a4eb5d" dependencies = [ "libc", - "windows-sys 0.52.0", + "windows-sys 0.59.0", ] [[package]] @@ -2006,6 +2032,15 @@ dependencies = [ "digest 0.10.7", ] +[[package]] +name = "hmac" +version = "0.13.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6303bc9732ae41b04cb554b844a762b4115a61bfaa81e3e83050991eeb56863f" +dependencies = [ + "digest 0.11.3", +] + [[package]] name = "http" version = "1.3.1" @@ -2108,7 +2143,7 @@ version = "0.1.19" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "727805d60e7938b76b826a6ef209eb70eaa1812794f9424d4a4e2d740662df5f" dependencies = [ - "base64", + "base64 0.22.1", "bytes", "futures-channel", "futures-core", @@ -2120,7 +2155,7 @@ dependencies = [ "libc", "percent-encoding", "pin-project-lite", - "socket2 0.6.1", + "socket2 0.6.5", "system-configuration", "tokio", "tower-service", @@ -2477,7 +2512,7 @@ version = "9.3.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "5a87cc7a48537badeae96744432de36f4be2b4a34a05a5ef32e9dd8a1c169dde" dependencies = [ - "base64", + "base64 0.22.1", "js-sys", "pem", "ring", @@ -2500,9 +2535,9 @@ checksum = "09edd9e8b54e49e587e4f6295a7d29c3ea94d469cb40ab8ca70b288248a81db2" [[package]] name = "libc" -version = "0.2.177" +version = "0.2.189" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2874a2af47a2325c2001a6e6fad9b16a53b802102b528163885171cf92b15976" +checksum = "3eaf3ede3fee6db1a4c2ee091bf8a8b4dccdc6d17f656fb07896ee72867612f2" [[package]] name = "libm" @@ -2592,6 +2627,16 @@ dependencies = [ "digest 0.10.7", ] +[[package]] +name = "md-5" +version = "0.11.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "69b6441f590336821bb897fb28fc622898ccceb1d6cea3fde5ea86b090c4de98" +dependencies = [ + "cfg-if", + "digest 0.11.3", +] + [[package]] name = "memchr" version = "2.7.4" @@ -2614,7 +2659,7 @@ version = "0.17.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "df88858cd28baaaf2cfc894e37789ed4184be0e1351157aec7bf3c2266c793fd" dependencies = [ - "base64", + "base64 0.22.1", "http-body-util", "hyper", "hyper-rustls", @@ -2654,7 +2699,7 @@ dependencies = [ "cfg-if", "miette-derive", "thiserror 1.0.69", - "unicode-width", + "unicode-width 0.1.14", ] [[package]] @@ -2691,13 +2736,13 @@ dependencies = [ [[package]] name = "mio" -version = "1.0.3" +version = "1.2.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2886843bf800fba2e3377cff24abf6379b4c4d5c6681eaf9ea5b0d15090450bd" +checksum = "30d65c71f1ce40ab09135ce117d742b9f8a19ff91a41a8b57ed50bc2de59c427" dependencies = [ "libc", "wasi 0.11.0+wasi-snapshot-preview1", - "windows-sys 0.52.0", + "windows-sys 0.61.2", ] [[package]] @@ -2972,7 +3017,7 @@ version = "3.0.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1d30c53c26bc5b31a98cd02d20f25a7c8567146caf63ed593a9d87b2775291be" dependencies = [ - "base64", + "base64 0.22.1", "serde_core", ] @@ -2983,45 +3028,47 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e3148f5046208a5d56bcfc03053e3ca6334e51da8dfb19b6cdc8b306fae3283e" [[package]] -name = "pg_escape" -version = "0.1.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "44c7bc82ccbe2c7ef7ceed38dcac90d7ff46681e061e9d7310cbcd409113e303" -dependencies = [ - "phf", -] - -[[package]] -name = "phf" -version = "0.11.3" +name = "pg-proto" +version = "0.11.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1fd6780a80ae0c52cc120a26a1a42c1ae51b247a253e4e06113d23d2c2edd078" +checksum = "72f5b88d40736045b71bb72d114a6f212c98095e1f43cfc998c948e429b55012" dependencies = [ - "phf_macros", - "phf_shared", + "base64 0.23.1", + "bytes", + "hmac 0.13.0", + "pg-proto-fsm", + "postgres-protocol", + "rand 0.10.2", + "rustls", + "sha2 0.11.0", + "socket2 0.6.5", + "stringprep", + "subtle", + "tokio", + "tokio-rustls", + "tokio-util", + "x509-parser", ] [[package]] -name = "phf_generator" -version = "0.11.3" +name = "pg-proto-fsm" +version = "0.11.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3c80231409c20246a13fddb31776fb942c38553c51e871f8cbd687a4cfb5843d" +checksum = "dc37a28331c337145f6500402d7857565ff807891712d79833e0fbc00567c48d" dependencies = [ - "phf_shared", - "rand 0.8.6", + "proc-macro2", + "quote", + "railroad", + "syn 3.0.3", ] [[package]] -name = "phf_macros" +name = "phf" version = "0.11.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f84ac04429c13a7ff43785d75ad27569f2951ce0ffd30a3321230db2fc727216" +checksum = "1fd6780a80ae0c52cc120a26a1a42c1ae51b247a253e4e06113d23d2c2edd078" dependencies = [ - "phf_generator", "phf_shared", - "proc-macro2", - "quote", - "syn 2.0.117", ] [[package]] @@ -3083,19 +3130,19 @@ dependencies = [ [[package]] name = "postgres-protocol" -version = "0.6.8" +version = "0.6.12" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "76ff0abab4a9b844b93ef7b81f1efc0a366062aaef2cd702c76256b5dc075c54" +checksum = "08808e3c483c46e999108051c78334f473d5adb59d78bb80a1268c7e6aa6c514" dependencies = [ - "base64", + "base64 0.22.1", "byteorder", "bytes", "fallible-iterator", - "hmac", - "md-5", + "hmac 0.13.0", + "md-5 0.11.0", "memchr", - "rand 0.9.2", - "sha2", + "rand 0.10.2", + "sha2 0.11.0", "stringprep", ] @@ -3205,9 +3252,9 @@ dependencies = [ [[package]] name = "proc-macro2" -version = "1.0.95" +version = "1.0.107" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "02b3e5e68a3a1a02aad3ec490a98007cbc13c37cbe84a3cd7b8e406d76e7f778" +checksum = "985e7ec9bb745e6ce6535b544d84d6cd6f7ad8bd711c398938ae983b91a766d9" dependencies = [ "unicode-ident", ] @@ -3325,14 +3372,14 @@ dependencies = [ "once_cell", "socket2 0.5.8", "tracing", - "windows-sys 0.52.0", + "windows-sys 0.59.0", ] [[package]] name = "quote" -version = "1.0.40" +version = "1.0.47" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1885c039570dc00dcb4ff087a89e185fd56bae234ddc7f056a945bf36467248d" +checksum = "1fbf4db142a473a8d80c26bbf18454ed458bf8d26c8219c331daecfdbd079001" dependencies = [ "proc-macro2", ] @@ -3355,6 +3402,15 @@ version = "0.7.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "dc33ff2d4973d518d823d61aa239014831e521c75da58e3df4840d3f47749d09" +[[package]] +name = "railroad" +version = "0.3.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4bf842ad92d09c4dd1e68be1507189b01898c46fe01916c0467fb405b1f9ee0a" +dependencies = [ + "unicode-width 0.2.2", +] + [[package]] name = "rand" version = "0.8.6" @@ -3474,7 +3530,7 @@ dependencies = [ "rand_chacha 0.3.1", "serde", "serde_cbor", - "sha2", + "sha2 0.10.8", "thiserror 1.0.69", "zeroize", ] @@ -3495,7 +3551,7 @@ dependencies = [ "rand_chacha 0.3.1", "serde", "serde_cbor", - "sha2", + "sha2 0.10.8", "thiserror 1.0.69", "zeroize", ] @@ -3557,7 +3613,7 @@ checksum = "2c9283685feec7d69af75fb0e858d5e7378f33fe4fc699383b2916ab9273e03c" dependencies = [ "proc-macro2", "quote", - "syn 3.0.2", + "syn 3.0.3", ] [[package]] @@ -3604,7 +3660,7 @@ version = "0.13.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "04e9018c9d814e5f30cc16a0f03271aeab3571e609612d9fe78c1aa8d11c2f62" dependencies = [ - "base64", + "base64 0.22.1", "bytes", "encoding_rs", "futures-core", @@ -3775,14 +3831,14 @@ dependencies = [ "errno", "libc", "linux-raw-sys", - "windows-sys 0.52.0", + "windows-sys 0.59.0", ] [[package]] name = "rustls" -version = "0.23.28" +version = "0.23.43" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7160e3e10bf4535308537f3c4e1641468cd0e485175d6163087c0393c7d46643" +checksum = "0283386ce02abc0151e1761d08802dfe86c173b0b494af5cbc086574e453da06" dependencies = [ "aws-lc-rs", "log", @@ -3807,11 +3863,12 @@ dependencies = [ [[package]] name = "rustls-pki-types" -version = "1.11.0" +version = "1.15.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "917ce264624a4b4db1c364dcc35bfca9ded014d0a958cd47ad3e960e988ea51c" +checksum = "2f4925028c7eb5d1fcdaf196971378ed9d2c1c4efc7dc5d011256f76c99c0a96" dependencies = [ "web-time", + "zeroize", ] [[package]] @@ -3832,7 +3889,7 @@ dependencies = [ "security-framework", "security-framework-sys", "webpki-root-certs 0.26.8", - "windows-sys 0.52.0", + "windows-sys 0.59.0", ] [[package]] @@ -3864,9 +3921,9 @@ checksum = "f87165f0995f63a9fbeea62b64d10b4d9d8e78ec6d7d51fb2125fda7bb36788f" [[package]] name = "rustls-webpki" -version = "0.103.3" +version = "0.103.13" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e4a72fe2bcf7a6ac6fd7d0b9e5cb68aeb7d4c0a0271730218b3e92d43b4eb435" +checksum = "61c429a8649f110dddef65e2a5ad240f747e85f7758a6bccc7e5777bd33f756e" dependencies = [ "aws-lc-rs", "ring", @@ -4161,6 +4218,17 @@ dependencies = [ "digest 0.10.7", ] +[[package]] +name = "sha2" +version = "0.11.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "446ba717509524cb3f22f17ecc096f10f4822d76ab5c0b9822c5f9c284e825f4" +dependencies = [ + "cfg-if", + "cpufeatures 0.3.0", + "digest 0.11.3", +] + [[package]] name = "sharded-slab" version = "0.1.7" @@ -4258,12 +4326,12 @@ dependencies = [ [[package]] name = "socket2" -version = "0.6.1" +version = "0.6.5" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "17129e116933cf371d018bb80ae557e889637989d8638274fb25622827b03881" +checksum = "c3d1e2c7f27f8d4cb10542a02c49005dbd6e93095799d6f3be745fae9f8fedd4" dependencies = [ "libc", - "windows-sys 0.60.2", + "windows-sys 0.61.2", ] [[package]] @@ -4310,7 +4378,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6d1c9c640571eba8fa5a705ccebe5c5e50aa27d4d51f4ec1614ffd1841b8a13a" dependencies = [ "aquamarine", - "base64", + "base64 0.22.1", "cts-common", "jsonwebtoken", "miette", @@ -4355,7 +4423,7 @@ dependencies = [ "cfg-if", "libc", "psm", - "windows-sys 0.52.0", + "windows-sys 0.59.0", ] [[package]] @@ -4417,9 +4485,9 @@ dependencies = [ [[package]] name = "syn" -version = "3.0.2" +version = "3.0.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a207d6d6a2b7fc470b80443726053f18a2481b7e1eee970597051596567987a3" +checksum = "53e9bae58849f64dfa4f5d5ae372c8341f7305f82a3868709269343628b659a3" dependencies = [ "proc-macro2", "quote", @@ -4626,9 +4694,9 @@ dependencies = [ [[package]] name = "tokio" -version = "1.48.0" +version = "1.53.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ff360e02eab121e0bc37a2d3b4d4dc622e6eda3a8e5253d5435ecf5bd4c68408" +checksum = "202caea871b69668250d242070849eb495be178ed697a3e98aebce5bc81a0bed" dependencies = [ "bytes", "libc", @@ -4636,20 +4704,20 @@ dependencies = [ "parking_lot", "pin-project-lite", "signal-hook-registry", - "socket2 0.6.1", + "socket2 0.6.5", "tokio-macros", "windows-sys 0.61.2", ] [[package]] name = "tokio-macros" -version = "2.6.0" +version = "2.7.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "af407857209536a95c8e56f8231ef2c2e2aff839b22e07a1ffcbc617e9db9fa5" +checksum = "78773a2a397f451582ce068015985c33193cf6dea8b74d2a639fe457b2f07b0e" dependencies = [ "proc-macro2", "quote", - "syn 2.0.117", + "syn 3.0.3", ] [[package]] @@ -4684,7 +4752,7 @@ version = "0.13.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "27d684bad428a0f2481f42241f821db42c54e2dc81d8c00db8536c506b0a0144" dependencies = [ - "const-oid", + "const-oid 0.9.6", "ring", "rustls", "tokio", @@ -4695,9 +4763,9 @@ dependencies = [ [[package]] name = "tokio-rustls" -version = "0.26.2" +version = "0.26.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8e727b36a1a0e8b74c376ac2211e40c2c8af09fb4013c60d910495810f008e9b" +checksum = "1729aa945f29d91ba541258c8df89027d5792d85a8841fb65e8bf0f4ede4ef61" dependencies = [ "rustls", "tokio", @@ -4705,15 +4773,15 @@ dependencies = [ [[package]] name = "tokio-util" -version = "0.7.14" +version = "0.7.19" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6b9590b93e6fcc1739458317cccd391ad3955e2bde8913edf6f95f9e65a8f034" +checksum = "494815d09bf52b5548659851081238f0ca39ff638363907596da739561c62c52" dependencies = [ "bytes", "futures-core", "futures-sink", "futures-util", - "hashbrown 0.14.5", + "libc", "pin-project-lite", "tokio", ] @@ -4964,6 +5032,12 @@ version = "0.1.14" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7dd6e30e90baa6f72411720665d41d89b9a3d039dc45b8faea1ddd07f617f6af" +[[package]] +name = "unicode-width" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b4ac048d71ede7ee76d585517add45da530660ef4390e49b098733c6e897f254" + [[package]] name = "unicode-xid" version = "0.2.6" @@ -5056,7 +5130,7 @@ dependencies = [ "atomic", "getrandom 0.3.2", "js-sys", - "md-5", + "md-5 0.10.6", "serde", "sha1_smol", "wasm-bindgen", @@ -5496,7 +5570,7 @@ version = "0.1.9" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "cf221c93e13a30d793f7645a0e7762c55d169dbb0a49671918a2319d289b10bb" dependencies = [ - "windows-sys 0.48.0", + "windows-sys 0.59.0", ] [[package]] @@ -5678,15 +5752,6 @@ dependencies = [ "windows-targets 0.52.6", ] -[[package]] -name = "windows-sys" -version = "0.60.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f2f500e4d28234f72040990ec9d39e3a6b950f9f22d3dba18416c35882612bcb" -dependencies = [ - "windows-targets 0.53.5", -] - [[package]] name = "windows-sys" version = "0.61.2" @@ -6100,7 +6165,7 @@ version = "0.2.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1301e935010a701ae5f8655edc0ad17c44bad3ac5ce8c39185f75453b720ae94" dependencies = [ - "const-oid", + "const-oid 0.9.6", "der", "spki", "tls_codec", @@ -6108,9 +6173,9 @@ dependencies = [ [[package]] name = "x509-parser" -version = "0.17.0" +version = "0.18.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4569f339c0c402346d4a75a9e39cf8dad310e287eef1ff56d4c68e5067f53460" +checksum = "d43b0f71ce057da06bc0851b23ee24f3f86190b07203dd8f567d0b706a185202" dependencies = [ "asn1-rs", "data-encoding", @@ -6240,7 +6305,7 @@ version = "0.12.26" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "16f731f2de99e66396928faef44b02b39288dc9a93f77a4a3e1dcdc33c1adad0" dependencies = [ - "base64", + "base64 0.22.1", "cipherstash-config", "const-hex", "cts-common", diff --git a/PG_PROTO_FOLLOWUPS.md b/PG_PROTO_FOLLOWUPS.md new file mode 100644 index 000000000..b8633cb02 --- /dev/null +++ b/PG_PROTO_FOLLOWUPS.md @@ -0,0 +1,43 @@ +# pg-proto follow-ups + +The proxy migration delegates framing, typed messages, startup/authentication +state, typed startup middleware, async runtime middleware, demultiplexing, and +compile-time checked bounded pipeline dispatch to `pg-proto` 0.5.0. The +remaining transport adapters below would be better eliminated in `pg-proto` +itself. + +The 0.5.0 version above is the historical baseline when these follow-ups were +written; the completed migration now uses pg-proto 0.10.6. + +## Preserve buffered transport state across a split + +`Buffered::into_inner()` cannot return retained inbound bytes, pending outbound +bytes, or demultiplexer state. The proxy must finish startup through +`ReadyForQuery` before splitting the bidirectional stream and then construct new +`Buffered` values for concurrent frontend/backend processing. + +An `into_parts`/`from_parts` API, or a buffer-preserving split API, would let a +proxy change transport ownership without risking loss of bytes already read past +a message boundary. It should preserve both codec buffers and the demultiplexer. + +## Transfer a demultiplexer between transport owners + +The proxy routes backend messages through pg-proto's demultiplexer, but startup +and concurrent runtime currently use separate `Buffered` owners. A supported way +to extract and restore `Demux` state would retain startup parameter status, +cancellation-key, and readiness state without application bookkeeping. + +This may naturally be solved by the buffer-preserving transport-parts API above. + +## Publish the example proxy driver as a library API + +`pg-proto` demonstrates a clean `Buffered` + `Middleware` forwarding loop in +`examples/proxy_support`, but does not expose a configurable proxy driver from +the crate. CipherStash therefore still owns connection orchestration, concurrent +forwarding, and the small amount of glue that invokes middleware. + +A library-level proxy builder should accept downstream/upstream transports, +startup and authentication policy, typed frontend/backend middleware, timeout +policy, and an output strategy. It should own framing, phase transitions, +bounded pipeline dispatch, demultiplexing, and shutdown. Applications would then +only supply policy and message transformations. diff --git a/PG_PROTO_MIGRATION_PLAN.md b/PG_PROTO_MIGRATION_PLAN.md new file mode 100644 index 000000000..5660a6125 --- /dev/null +++ b/PG_PROTO_MIGRATION_PLAN.md @@ -0,0 +1,41 @@ +# pg-proto Migration + +> Historical plan: this records the migration's initial 0.2.1 baseline. The +> completed implementation now uses pg-proto 0.10.6. + +## Summary + +- Create `/Users/jamessadler/cipherstash/proxy-pg-proto` from current `main` (`15b7f996`) on branch `refactor/pg-proto`. +- Save this migration plan as `PG_PROTO_MIGRATION_PLAN.md` in that worktree’s repository root. +- Limit the initial deliverable to the worktree, branch, and plan document; implementation follows separately. +- Target a full migration to published [`pg-proto` 0.2.1](https://crates.io/crates/pg-proto/0.2.1), covering codecs, startup/authentication, runtime protocol-state validation, and bounded pipeline orchestration. + +## Implementation Changes + +- Replace handwritten framing, startup packet parsing, message codes, and message serialization with direction-specific `pg-proto` frontend/backend codecs. +- Convert CipherStash-specific behavior into adapters over `pg-proto` messages: + - Preserve Parse/Query SQL rewriting and parameter OID mapping. + - Preserve Bind format-code semantics, nulls, parameter reshaping, and encryption. + - Preserve ParameterDescription, RowDescription, and DataRow rewriting and batched decryption. + - Retain diagnostic-response factories while emitting `pg-proto` response types. +- Use `pg-proto` pre-startup and authentication APIs for SSL negotiation, startup, cancellation, client-facing MD5 authentication, and upstream cleartext/MD5/SCRAM authentication. Continue using existing TLS configuration and certificate policy. +- Pair downstream server-role and upstream client-role runtime FSMs through `Intermediary`. Advance both sides for forwarded messages and only the affected side for locally intercepted or synthesised messages. +- Preserve concurrent client-to-server and server-to-client processing, connection timeouts, response ordering, metrics, logging, schema reloads, and row buffering. +- Track protocol state even when encryption mapping is disabled. Preserve one-to-one cancellation forwarding; do not introduce pooling or cancellation-key translation. +- Reject unknown message tags as protocol errors, matching `pg-proto`’s fail-closed behavior. +- Remove obsolete handwritten protocol modules and direct low-level dependencies once unused. Preserve the public configuration and CLI surfaces; retain existing `ProtocolError` variants for source compatibility even where `pg-proto` supersedes them. + +## Test Plan + +- Port existing message round-trip and rewrite tests to `pg-proto` message fixtures. +- Add coverage for partial/oversized frames, malformed messages, unknown-tag rejection, SSL/TLS startup, cancellation, and all supported authentication modes. +- Exercise simple queries and extended Parse/Bind/Describe/Execute/Close/Sync pipelines, including pipelining and error draining through Sync. +- Verify text/binary formats, nulls, reshaped parameters, prepared statements, portals, COPY messages, asynchronous backend messages, and buffered DataRow decryption. +- Run formatting, clippy, proxy unit tests, and TCP/TLS integration suites. Unset `CS_PROMETHEUS__ENABLED` for the baseline unit suite; its current environment value causes the otherwise unrelated Prometheus test to fail. + +## Assumptions + +- The existing untracked `.claude/worktrees/` directory remains untouched. +- The worktree path and branch are currently available. +- No compatibility feature flag or dual protocol implementation is required. +- The plan document is left as an uncommitted worktree change unless a commit is requested separately. diff --git a/docs/errors.md b/docs/errors.md index 8f9ae535c..53b844447 100644 --- a/docs/errors.md +++ b/docs/errors.md @@ -42,6 +42,20 @@ +## Internal PostgreSQL protocol error + +CipherStash Proxy encountered an unexpected internal protocol state while processing a PostgreSQL response. + +### Error message + +``` +CipherStash Proxy encountered an internal PostgreSQL protocol error. +``` + +### How to fix + +Retry the operation. If the error persists, collect the proxy logs and report the issue to CipherStash support. The logs retain the detailed internal error needed for diagnosis. + # Authentication errors diff --git a/mise.toml b/mise.toml index 95770d8d4..cca6550f2 100644 --- a/mise.toml +++ b/mise.toml @@ -67,7 +67,8 @@ run = """ # Exit immediately if any command returns non-zero. Stops proxy from starting when the build fails set -e -{% set target = arch() ~ "-unknown-linux-gnu" | replace(from="arm64", to="aarch64") | replace(from="x64", to="x86_64") %} +{% set target_arch = arch() | replace(from="arm64", to="aarch64") | replace(from="x64", to="x86_64") %} +{% set target = target_arch ~ "-unknown-linux-gnu" %} {% set docker_platform = "linux/" ~ arch() | replace(from="x64", to="amd64") %} # build a binary @@ -232,7 +233,7 @@ run = "mise run test:nextest {{arg(name='test',default='')}}" dir = "{{config_root}}" run = """ {% set default_host = get_env(name="CS_DATABASE__HOST",default="localhost") %} -{% set default_port = get_env(name="CS_DATABASE__PORT",default=5432) %} +{% set default_port = get_env(name="CS_DATABASE__PORT",default="5432") %} host={{option(name="host",default=default_host)}} port={{option(name="port",default=default_port)}} tls={{flag(name="tls")}} @@ -244,12 +245,11 @@ echo "Testing presence of Postgres at ${host}:${port} with a maximum of ${max_re postgres_ready () { # tls test if [ "$1" = "true" ]; then - response="$(echo | openssl s_client -starttls postgres -connect $host:$port -state -no_ign_eof 2>&1)" - success=$? - if [ -n "${DEBUG}" ]; then - echo "response: ${response}" - fi - return $success + # Use psql so TLS negotiation and a real PostgreSQL connection both + # complete before the service is considered ready on Linux and macOS. + PGPASSWORD="$CS_DATABASE__PASSWORD" PGCONNECT_TIMEOUT=1 PGSSLMODE=require \ + psql -h "$host" -p "$port" -d "$CS_DATABASE__NAME" \ + -U "$CS_DATABASE__USERNAME" -c 'SELECT 1' >/dev/null 2>&1 else # plaintext test # We blackhole stderr (2>/dev/null) below because nc is echoing to stderr, @@ -653,7 +653,7 @@ fi export CIPHERSTASH_CLIENT_SECONDARY_USER_AGENT="cipherstash-proxy/${BUILD_VERSION}" # cross-compile -rustup update +rustup toolchain install stable --profile minimal --no-self-update rustup target add --toolchain stable {{ target }} cargo build --locked --target {{ target }} --release --package cipherstash-proxy diff --git a/packages/cipherstash-proxy-integration/src/common.rs b/packages/cipherstash-proxy-integration/src/common.rs index d9422c526..9b13eabad 100644 --- a/packages/cipherstash-proxy-integration/src/common.rs +++ b/packages/cipherstash-proxy-integration/src/common.rs @@ -208,6 +208,11 @@ pub fn connection_config(port: u16) -> tokio_postgres::Config { } pub async fn connect_with_tls(port: u16) -> Client { + let (client, _connection_task) = connect_with_tls_and_task(port).await; + client +} + +pub async fn connect_with_tls_and_task(port: u16) -> (Client, tokio::task::JoinHandle<()>) { let tls_config = configure_test_client(); let tls = tokio_postgres_rustls::MakeRustlsConnect::new(tls_config); @@ -217,12 +222,12 @@ pub async fn connect_with_tls(port: u16) -> Client { .await .expect("connection to database to succeed"); - tokio::spawn(async move { + let connection_task = tokio::spawn(async move { if let Err(e) = connection.await { eprintln!("connection error: {e}"); } }); - client + (client, connection_task) } pub async fn connect(port: u16) -> Client { diff --git a/packages/cipherstash-proxy-integration/src/diagnostics.rs b/packages/cipherstash-proxy-integration/src/diagnostics.rs index 25edc4926..8415d2f2f 100644 --- a/packages/cipherstash-proxy-integration/src/diagnostics.rs +++ b/packages/cipherstash-proxy-integration/src/diagnostics.rs @@ -1,11 +1,10 @@ #[cfg(test)] mod tests { - use crate::common::{clear, connect_with_tls, PROXY, PROXY_METRICS_PORT}; + use crate::common::{clear, connect_with_tls_and_task, PROXY, PROXY_METRICS_PORT}; /// Maximum number of retry attempts for fetching metrics. - /// 5 retries with 200ms delay gives ~1 second total wait time, - /// sufficient for Prometheus scrape interval in CI environments. - const METRICS_FETCH_MAX_RETRIES: u32 = 5; + /// 25 retries with 200ms delay gives ~5 seconds total wait time. + const METRICS_FETCH_MAX_RETRIES: u32 = 25; /// Delay between retry attempts in milliseconds. /// 200ms provides a reasonable balance between responsiveness and allowing @@ -13,9 +12,14 @@ mod tests { const METRICS_FETCH_RETRY_DELAY_MS: u64 = 200; /// Fetch metrics with retry logic to handle CI timing variability. - async fn fetch_metrics_with_retry(max_retries: u32, delay_ms: u64) -> String { - let url = format!("http://localhost:{}/metrics", *PROXY_METRICS_PORT); + async fn fetch_metrics_with_retry( + max_retries: u32, + delay_ms: u64, + expected: &[&str], + ) -> String { + let url = format!("http://127.0.0.1:{}/metrics", *PROXY_METRICS_PORT); let mut last_error = None; + let mut last_body = None; for attempt in 0..max_retries { if attempt > 0 { @@ -24,23 +28,26 @@ mod tests { match reqwest::get(&url).await { Ok(response) => match response.text().await { - Ok(body) => return body, + Ok(body) if expected.iter().all(|value| body.contains(value)) => return body, + Ok(body) => last_body = Some(body), Err(e) => last_error = Some(format!("Failed to read response: {}", e)), }, Err(e) => last_error = Some(format!("Failed to fetch metrics: {}", e)), } } - panic!( - "Failed to fetch metrics after {} retries: {}", - max_retries, - last_error.unwrap_or_else(|| "unknown error".to_string()) - ); + last_body.unwrap_or_else(|| { + panic!( + "Failed to fetch metrics after {} retries: {}", + max_retries, + last_error.unwrap_or_else(|| "unknown error".to_string()) + ) + }) } #[tokio::test] async fn metrics_include_statement_labels() { - let client = connect_with_tls(*PROXY).await; + let (client, connection_task) = connect_with_tls_and_task(*PROXY).await; clear().await; @@ -59,14 +66,30 @@ mod tests { .await .unwrap(); + // Closing the client gives the background connection driver a + // deterministic completion boundary before the metrics scrape. + drop(client); + connection_task + .await + .expect("PostgreSQL connection task should finish"); + // Fetch metrics with retry logic for CI robustness - let body = - fetch_metrics_with_retry(METRICS_FETCH_MAX_RETRIES, METRICS_FETCH_RETRY_DELAY_MS).await; + let body = fetch_metrics_with_retry( + METRICS_FETCH_MAX_RETRIES, + METRICS_FETCH_RETRY_DELAY_MS, + &[ + "statement_type=\"insert\"", + "statement_type=\"select\"", + "multi_statement=\"false\"", + ], + ) + .await; // Assert that the metrics include the expected labels assert!( body.contains("statement_type=\"insert\""), - "Metrics should include insert statement_type label" + "Metrics should include insert statement_type label. Found: {}", + body ); assert!( body.contains("statement_type=\"select\""), @@ -80,7 +103,7 @@ mod tests { #[tokio::test] async fn slow_statement_metrics_and_logs() { - let client = connect_with_tls(*PROXY).await; + let (client, connection_task) = connect_with_tls_and_task(*PROXY).await; clear().await; @@ -88,9 +111,23 @@ mod tests { // We use pg_sleep(2.1) to ensure it's considered slow client.query("SELECT pg_sleep(2.1)", &[]).await.unwrap(); + // Ensure the completed statement has been observed by the proxy's + // connection driver before polling its metrics endpoint. + drop(client); + connection_task + .await + .expect("PostgreSQL connection task should finish"); + // Fetch metrics with retry logic - let body = - fetch_metrics_with_retry(METRICS_FETCH_MAX_RETRIES, METRICS_FETCH_RETRY_DELAY_MS).await; + let body = fetch_metrics_with_retry( + METRICS_FETCH_MAX_RETRIES, + METRICS_FETCH_RETRY_DELAY_MS, + &[ + "cipherstash_proxy_slow_statements_total", + "cipherstash_proxy_statements_session_duration_seconds", + ], + ) + .await; // Assert that the slow statements counter is present and non-zero assert!( diff --git a/packages/cipherstash-proxy-integration/src/eql_regression.rs b/packages/cipherstash-proxy-integration/src/eql_regression.rs index e8e3580de..3099578ab 100644 --- a/packages/cipherstash-proxy-integration/src/eql_regression.rs +++ b/packages/cipherstash-proxy-integration/src/eql_regression.rs @@ -7,13 +7,13 @@ //! ## How to use these tests: //! //! 1. **Generate fixtures from main branch:** -//! ``` +//! ```text //! git checkout main //! CS_GENERATE_EQL_FIXTURES=1 cargo nextest run -p cipherstash-proxy-integration eql_regression::generate //! ``` //! //! 2. **Run regression tests on new branch:** -//! ``` +//! ```text //! git checkout //! cargo nextest run -p cipherstash-proxy-integration eql_regression::regression //! ``` diff --git a/packages/cipherstash-proxy-integration/src/multitenant/set_keyset_id.rs b/packages/cipherstash-proxy-integration/src/multitenant/set_keyset_id.rs index f76f221ce..57e4f8de6 100644 --- a/packages/cipherstash-proxy-integration/src/multitenant/set_keyset_id.rs +++ b/packages/cipherstash-proxy-integration/src/multitenant/set_keyset_id.rs @@ -69,8 +69,15 @@ mod tests { let result = client.query(&sql, &[]).await; assert!(result.is_ok()); + // A decryption failure is connection-fatal, so verify tenant isolation on + // a disposable connection while preserving this connection for the + // TENANT_1 -> TENANT_2 -> TENANT_1 switch under test. + let tenant_2_client = connect_with_tls(*PROXY).await; + let sql = format!("SET CIPHERSTASH.KEYSET_ID = '{tenant_keyset_id_2}'"); + tenant_2_client.query(&sql, &[]).await.unwrap(); + // SELECT data created by TENANT_1 AS TENANT_2 - let result = client.query(select_sql, &[&tenant_1_id]).await; + let result = tenant_2_client.query(select_sql, &[&tenant_1_id]).await; assert!(result.is_err()); // -------- @@ -150,9 +157,16 @@ mod tests { let result = client.simple_query(&sql).await; assert!(result.is_ok()); + // A decryption failure is connection-fatal, so verify tenant isolation on + // a disposable connection while preserving this connection for the + // TENANT_1 -> TENANT_2 -> TENANT_1 switch under test. + let tenant_2_client = connect_with_tls(*PROXY).await; + let sql = format!("SET CIPHERSTASH.KEYSET_ID = '{tenant_keyset_id_2}'"); + tenant_2_client.simple_query(&sql).await.unwrap(); + // SELECT data created by TENANT_1 AS TENANT_2 let sql = format!("SELECT id, encrypted_text FROM encrypted WHERE id = {tenant_1_id}"); - let result = client.simple_query(&sql).await; + let result = tenant_2_client.simple_query(&sql).await; assert!(result.is_err()); // -------- diff --git a/packages/cipherstash-proxy-integration/src/multitenant/set_keyset_name.rs b/packages/cipherstash-proxy-integration/src/multitenant/set_keyset_name.rs index 0d9af1757..0448adeca 100644 --- a/packages/cipherstash-proxy-integration/src/multitenant/set_keyset_name.rs +++ b/packages/cipherstash-proxy-integration/src/multitenant/set_keyset_name.rs @@ -64,8 +64,15 @@ mod tests { let result = client.query(&sql, &[]).await; assert!(result.is_ok()); + // A decryption failure is connection-fatal, so verify tenant isolation on + // a disposable connection while preserving this connection for the + // TENANT_1 -> TENANT_2 -> TENANT_1 switch under test. + let tenant_2_client = connect_with_tls(*PROXY).await; + let sql = format!("SET CIPHERSTASH.KEYSET_NAME = '{tenant_keyset_name_2}'"); + tenant_2_client.query(&sql, &[]).await.unwrap(); + // SELECT data created by TENANT_1 AS TENANT_2 - let result = client.query(select_sql, &[&tenant_1_id]).await; + let result = tenant_2_client.query(select_sql, &[&tenant_1_id]).await; assert!(result.is_err()); // -------- @@ -146,9 +153,16 @@ mod tests { let result = client.simple_query(&sql).await; assert!(result.is_ok()); + // A decryption failure is connection-fatal, so verify tenant isolation on + // a disposable connection while preserving this connection for the + // TENANT_1 -> TENANT_2 -> TENANT_1 switch under test. + let tenant_2_client = connect_with_tls(*PROXY).await; + let sql = format!("SET CIPHERSTASH.KEYSET_NAME = '{tenant_keyset_name_2}'"); + tenant_2_client.simple_query(&sql).await.unwrap(); + // SELECT data created by TENANT_1 AS TENANT_2 let sql = format!("SELECT id, encrypted_text FROM encrypted WHERE id = {tenant_1_id}"); - let result = client.simple_query(&sql).await; + let result = tenant_2_client.simple_query(&sql).await; assert!(result.is_err()); // -------- diff --git a/packages/cipherstash-proxy/Cargo.toml b/packages/cipherstash-proxy/Cargo.toml index 73dce1b9b..f694b0f08 100644 --- a/packages/cipherstash-proxy/Cargo.toml +++ b/packages/cipherstash-proxy/Cargo.toml @@ -5,7 +5,6 @@ edition = "2021" [dependencies] async-trait = "0.1" -aws-lc-rs = "1.13.3" bigdecimal = { version = "0.4.6", features = ["serde-json"] } blake3 = "1" arc-swap = "1.7.1" @@ -29,8 +28,7 @@ md-5 = "0.10.6" metrics = "0.24.3" metrics-exporter-prometheus = "0.17" moka = { version = "0.12", features = ["future"] } -oid-registry = "0.8" -pg_escape = "0.1.1" +pg-proto = "0.11.1" postgres-protocol = "0.6.7" postgres-types = { version = "0.2.8", features = ["with-serde_json-1"] } rand = "0.9" @@ -40,10 +38,11 @@ rust_decimal = { version = "1.36.0", default-features = false, features = [ ] } rustls = { version = "0.23.20", default-features = false, features = ["std"] } rustls-platform-verifier = "0.5.0" +rustls-native-certs = "0.8" rustls-pki-types = "1.10.0" serde = "1.0" serde_json = "1.0" -socket2 = "0.5.7" +socket2 = "0.6" sqltk = { workspace = true } thiserror = { workspace = true } tokio = { workspace = true } @@ -52,13 +51,11 @@ tokio-postgres = { version = "0.7", features = [ "with-serde_json-1", ] } tokio-postgres-rustls = "0.13.0" -tokio-rustls = "0.26.0" tokio-util = { version = "0.7.13", features = ["rt"] } tracing = { workspace = true } tracing-subscriber = { workspace = true } uuid = { version = "1.11.0", features = ["serde", "v4"] } vitaminc-protected = "0.1.0-pre4.2" -x509-parser = "0.17.0" [dev-dependencies] diff --git a/packages/cipherstash-proxy/src/connect/async_stream.rs b/packages/cipherstash-proxy/src/connect/async_stream.rs deleted file mode 100644 index 83436a05b..000000000 --- a/packages/cipherstash-proxy/src/connect/async_stream.rs +++ /dev/null @@ -1,148 +0,0 @@ -use super::{configure, connect_with_retry}; -use crate::{error::Error, log::AUTHENTICATION}; -use aws_lc_rs::digest; -use core::str; -use oid_registry::{ - Oid, OID_HASH_SHA1, OID_NIST_HASH_SHA256, OID_NIST_HASH_SHA384, OID_NIST_HASH_SHA512, - OID_PKCS1_SHA1WITHRSA, OID_PKCS1_SHA256WITHRSA, OID_PKCS1_SHA384WITHRSA, - OID_PKCS1_SHA512WITHRSA, OID_SIG_ECDSA_WITH_SHA256, OID_SIG_ECDSA_WITH_SHA384, OID_SIG_ED25519, -}; -use postgres_protocol::authentication::sasl::ChannelBinding; - -use std::{ - pin::Pin, - task::{Context, Poll}, -}; -use tokio::{ - io::{split, AsyncRead, AsyncWrite, ReadBuf}, - net::{TcpListener, TcpStream}, -}; -use tokio_rustls::TlsStream; -use tracing::debug; -use x509_parser::prelude::{FromDer, X509Certificate}; - -#[derive(Debug)] -pub enum AsyncStream { - Tcp(TcpStream), - Tls(Box>), -} - -impl AsyncStream { - pub async fn accept(listener: &TcpListener) -> Result { - let (stream, _) = listener.accept().await?; - configure(&stream); - Ok(AsyncStream::Tcp(stream)) - } - - pub async fn connect(addr: &str) -> Result { - let stream = connect_with_retry(addr).await?; - configure(&stream); - Ok(AsyncStream::Tcp(stream)) - } - - pub fn split( - self, - ) -> ( - tokio::io::ReadHalf, - tokio::io::WriteHalf, - ) { - split(self) - } - - pub fn is_tls(&self) -> bool { - matches!(self, AsyncStream::Tls(_)) - } - - pub fn is_tcp(&self) -> bool { - !self.is_tls() - } - - pub fn channel_binding(&self) -> ChannelBinding { - match self { - AsyncStream::Tcp(_) => ChannelBinding::unsupported(), - AsyncStream::Tls(stream) => { - let (_, session) = stream.get_ref(); - let certs = session.peer_certificates(); - match certs { - Some(certs) if !certs.is_empty() => { - let cert_der = &certs[0]; - X509Certificate::from_der(cert_der) - .ok() - .map(|(_, cert)| get_digest(&cert.signature_algorithm.algorithm)) - .map_or_else(ChannelBinding::unsupported, |algorithm| { - let hash = digest::digest(algorithm, certs[0].as_ref()); - ChannelBinding::tls_server_end_point(hash.as_ref().into()) - }) - } - _ => { - debug!( - target: AUTHENTICATION, - msg = "Missing certificates, ChannelBinding is unsupported" - ); - ChannelBinding::unsupported() - } - } - } - } - } -} - -/// -/// Note: SHA1 is upgraded to SHA256 as per https://datatracker.ietf.org/doc/html/rfc5929#section-4.1 -/// -fn get_digest(oid: &Oid) -> &'static digest::Algorithm { - match oid { - oid if oid == &OID_HASH_SHA1 => &digest::SHA256, - oid if oid == &OID_NIST_HASH_SHA256 => &digest::SHA256, - oid if oid == &OID_PKCS1_SHA1WITHRSA => &digest::SHA256, - oid if oid == &OID_PKCS1_SHA256WITHRSA => &digest::SHA256, - oid if oid == &OID_SIG_ECDSA_WITH_SHA256 => &digest::SHA256, - oid if oid == &OID_NIST_HASH_SHA384 => &digest::SHA384, - oid if oid == &OID_PKCS1_SHA384WITHRSA => &digest::SHA384, - oid if oid == &OID_SIG_ECDSA_WITH_SHA384 => &digest::SHA384, - oid if oid == &OID_NIST_HASH_SHA512 => &digest::SHA512, - oid if oid == &OID_PKCS1_SHA512WITHRSA => &digest::SHA512, - oid if oid == &OID_SIG_ED25519 => &digest::SHA512, - _ => panic!("Unsupported OID"), - } -} - -impl AsyncRead for AsyncStream { - fn poll_read( - mut self: Pin<&mut Self>, - cx: &mut Context<'_>, - buf: &mut ReadBuf<'_>, - ) -> Poll> { - match *self { - AsyncStream::Tcp(ref mut stream) => Pin::new(stream).poll_read(cx, buf), - AsyncStream::Tls(ref mut stream) => Pin::new(stream).poll_read(cx, buf), - } - } -} - -impl AsyncWrite for AsyncStream { - fn poll_write( - mut self: Pin<&mut Self>, - cx: &mut Context<'_>, - buf: &[u8], - ) -> Poll> { - match *self { - AsyncStream::Tcp(ref mut stream) => Pin::new(stream).poll_write(cx, buf), - AsyncStream::Tls(ref mut stream) => Pin::new(stream).poll_write(cx, buf), - } - } - - fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { - match *self { - AsyncStream::Tcp(ref mut stream) => Pin::new(stream).poll_flush(cx), - AsyncStream::Tls(ref mut stream) => Pin::new(stream).poll_flush(cx), - } - } - - fn poll_shutdown(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { - match *self { - AsyncStream::Tcp(ref mut stream) => Pin::new(stream).poll_shutdown(cx), - AsyncStream::Tls(ref mut stream) => Pin::new(stream).poll_shutdown(cx), - } - } -} diff --git a/packages/cipherstash-proxy/src/connect/channel_writer.rs b/packages/cipherstash-proxy/src/connect/channel_writer.rs deleted file mode 100644 index a6dc68ffb..000000000 --- a/packages/cipherstash-proxy/src/connect/channel_writer.rs +++ /dev/null @@ -1,109 +0,0 @@ -use bytes::BytesMut; -use tokio::{ - io::{AsyncWrite, AsyncWriteExt}, - sync::mpsc::{self, UnboundedReceiver, UnboundedSender}, -}; -use tracing::{debug, error}; - -use crate::log::PROTOCOL; - -pub type Receiver = UnboundedReceiver; -pub type Sender = UnboundedSender; - -#[derive(Debug)] -pub struct ChannelWriter -where - W: AsyncWrite + Unpin, -{ - writer: W, - receiver: Receiver, - sender: Sender, - client_id: i32, -} - -impl ChannelWriter -where - W: AsyncWrite + Unpin, -{ - pub fn new(writer: W, client_id: i32) -> Self { - let (sender, receiver): (UnboundedSender, UnboundedReceiver) = - mpsc::unbounded_channel(); - - ChannelWriter { - writer, - receiver, - sender, - client_id, - } - } - - pub async fn receive(mut self) { - debug!(target: PROTOCOL, - client_id = self.client_id, - msg = "ChannelWriter task started", - ); - - // Drop our own sender so the channel can close when frontend/backend senders are dropped - // Without this, we have a circular dependency: receiver waits for all senders to drop, - // but we're holding one of them ourselves! - drop(self.sender); - - while let Some(bytes) = self.receiver.recv().await { - debug!(target: PROTOCOL, - client_id = self.client_id, - msg = "Writing", - ?bytes - ); - - match self.writer.write_all(&bytes).await { - Ok(_) => { - debug!(target: PROTOCOL, - client_id = self.client_id, - msg = "Write complete", - ); - } - Err(err) => { - error!(target: PROTOCOL, - client_id = self.client_id, - msg = "Write error", - error = ?err - ); - break; - } - } - } - - // Channel closed - shutdown the writer to properly close the connection - debug!(target: PROTOCOL, - client_id = self.client_id, - msg = "Recv loop exited - channel closed, beginning shutdown", - ); - - // Flush any pending writes before shutdown - if let Err(err) = self.writer.flush().await { - error!(target: PROTOCOL, - client_id = self.client_id, - msg = "Error flushing writer during shutdown", - error = ?err - ); - } - - // Shutdown the write half to send FIN and properly close the connection - if let Err(err) = self.writer.shutdown().await { - error!(target: PROTOCOL, - client_id = self.client_id, - msg = "Error shutting down writer", - error = ?err - ); - } - - debug!(target: PROTOCOL, - client_id = self.client_id, - msg = "Writer shutdown complete", - ); - } - - pub fn sender(&self) -> Sender { - self.sender.clone() - } -} diff --git a/packages/cipherstash-proxy/src/connect/mod.rs b/packages/cipherstash-proxy/src/connect/mod.rs index d17d97969..722145e8d 100644 --- a/packages/cipherstash-proxy/src/connect/mod.rs +++ b/packages/cipherstash-proxy/src/connect/mod.rs @@ -1,11 +1,5 @@ -mod async_stream; -mod channel_writer; - -pub use async_stream::AsyncStream; -pub use channel_writer::{ChannelWriter, Sender}; - -use crate::{config::ServerConfig, error::Error, log::DEVELOPMENT, tls, DatabaseConfig}; -use socket2::TcpKeepalive; +use crate::{config::ServerConfig, error::Error, tls, DatabaseConfig}; +use socket2::{SockRef, TcpKeepalive}; use std::time::Duration; use tokio::{ net::{TcpListener, TcpStream}, @@ -14,13 +8,66 @@ use tokio::{ use tokio_postgres::Client; use tracing::{debug, error, info, warn}; +const MAX_RETRY_DELAY: Duration = Duration::from_secs(2); +const MAX_RETRY_COUNT: u32 = 3; const TCP_USER_TIMEOUT: Duration = Duration::from_secs(10); const TCP_KEEPALIVE_INTERVAL: Duration = Duration::from_secs(5); const TCP_KEEPALIVE_TIME: Duration = Duration::from_secs(5); const TCP_KEEPALIVE_RETRIES: u32 = 5; -const MAX_RETRY_DELAY: Duration = Duration::from_secs(2); -const MAX_RETRY_COUNT: u32 = 3; +fn configure_tcp(stream: &TcpStream) { + if let Err(error) = stream.set_nodelay(true) { + warn!(msg = "Error configuring TCP_NODELAY", error = %error); + } + + let socket = SockRef::from(stream); + #[cfg(target_os = "linux")] + if let Err(error) = socket.set_tcp_user_timeout(Some(TCP_USER_TIMEOUT)) { + warn!(msg = "Error configuring TCP_USER_TIMEOUT", error = %error); + } + + if let Err(error) = socket.set_keepalive(true) { + warn!(msg = "Error enabling TCP keepalive", error = %error); + return; + } + + let keepalive = TcpKeepalive::new() + .with_interval(TCP_KEEPALIVE_INTERVAL) + .with_time(TCP_KEEPALIVE_TIME) + .with_retries(TCP_KEEPALIVE_RETRIES); + if let Err(error) = socket.set_tcp_keepalive(&keepalive) { + warn!(msg = "Error configuring TCP keepalive", error = %error); + } +} + +pub async fn accept(listener: &TcpListener) -> Result { + let (stream, _) = listener.accept().await?; + configure_tcp(&stream); + Ok(stream) +} + +pub async fn connect(address: &str) -> Result { + debug!(msg = "Connecting to database"); + let mut delay = Duration::from_millis(100); + for attempt in 0..=MAX_RETRY_COUNT { + match TcpStream::connect(address).await { + Ok(stream) => { + configure_tcp(&stream); + return Ok(stream); + } + Err(error) if attempt < MAX_RETRY_COUNT => { + warn!(msg = "Database connection failed; retrying", %error, attempt); + time::sleep(delay).await; + delay = (delay * 2).min(MAX_RETRY_DELAY); + } + Err(error) => { + error!(msg = "Could not connect to database", error = %error); + return Err(Error::DatabaseConnection); + } + } + } + unreachable!() +} pub async fn database(config: &DatabaseConfig) -> Result { let connection_config = config.to_connection_config(); @@ -79,81 +126,3 @@ pub async fn bind_with_retry(server: &ServerConfig) -> TcpListener { retry_count += 1; } } - -pub async fn connect_with_retry(addr: &str) -> Result { - let mut retry_count = 0; - - loop { - debug!(target: DEVELOPMENT, msg = "Connecting to database"); - match TcpStream::connect(&addr).await { - Ok(stream) => { - return Ok(stream); - } - Err(err) => { - if retry_count > MAX_RETRY_COUNT { - error!(msg = "Could not connect to database", retries = ?retry_count, error = err.to_string()); - return Err(Error::DatabaseConnection); - } - } - }; - let sleep_duration_ms = - (100 * 2_u64.pow(retry_count)).min(MAX_RETRY_DELAY.as_millis() as _); - time::sleep(Duration::from_millis(sleep_duration_ms)).await; - - retry_count += 1; - } -} - -/// -/// Configure the tcp socket -/// set_nodelay -/// set_keepalive -/// -/// Keepalive is not as important without connection pooling timeouts to deal with -/// -pub fn configure(stream: &TcpStream) { - let sock_ref = socket2::SockRef::from(&stream); - - stream.set_nodelay(true).unwrap_or_else(|err| { - warn!( - msg = "Error configuring nodelay for connection", - error = err.to_string() - ); - }); - - #[cfg(target_os = "linux")] - match sock_ref.set_tcp_user_timeout(Some(TCP_USER_TIMEOUT)) { - Ok(_) => (), - Err(err) => { - warn!( - msg = "Error configuring tcp_user_timeout for connection", - error = err.to_string() - ); - } - } - - match sock_ref.set_keepalive(true) { - Ok(_) => { - let params = &TcpKeepalive::new() - .with_interval(TCP_KEEPALIVE_INTERVAL) - .with_retries(TCP_KEEPALIVE_RETRIES) - .with_time(TCP_KEEPALIVE_TIME); - - match sock_ref.set_tcp_keepalive(params) { - Ok(_) => (), - Err(err) => { - warn!( - msg = "Error configuring keepalive for connection", - error = err.to_string() - ); - } - } - } - Err(err) => { - warn!( - msg = "Error configuring connection", - error = err.to_string() - ); - } - } -} diff --git a/packages/cipherstash-proxy/src/error.rs b/packages/cipherstash-proxy/src/error.rs index dffc14935..2660f2c40 100644 --- a/packages/cipherstash-proxy/src/error.rs +++ b/packages/cipherstash-proxy/src/error.rs @@ -1,12 +1,12 @@ use crate::{postgresql::Column, Identifier}; -use bytes::BytesMut; use cipherstash_client::{encryption, schema::ColumnType}; use eql_mapper::{EqlMapperError, EqlTermVariant}; use metrics_exporter_prometheus::BuildError; use std::{io, time::Duration}; use thiserror::Error; -const ERROR_DOC_BASE_URL: &str = "https://github.com/cipherstash/proxy/blob/main/docs/errors.md"; +pub(crate) const ERROR_DOC_BASE_URL: &str = + "https://github.com/cipherstash/proxy/blob/main/docs/errors.md"; const ERROR_DOC_CONFIG_URL: &str = "https://github.com/cipherstash/proxy/blob/main/docs/how-to/index.md#configuring-proxy"; @@ -55,7 +55,7 @@ pub enum Error { Unknown, #[error(transparent)] - SendError(#[from] tokio::sync::mpsc::error::SendError), + Send(#[from] tokio::sync::mpsc::error::SendError), } impl Error { @@ -474,6 +474,21 @@ pub enum ProtocolError { #[error("Client authentication failed. Check username and password. For help visit {}#authentication-failed-client", ERROR_DOC_BASE_URL)] ClientAuthenticationFailed, + #[error("A buffered PostgreSQL DataRow was not associated with an operation")] + HeldDataRowMissingOperation, + + #[error("Buffered PostgreSQL DataRows crossed Execute operation boundaries")] + HeldDataRowOperationMismatch, + + #[error("A buffered PostgreSQL response was not a DataRow")] + HeldBackendMessageNotDataRow, + + #[error("Buffered PostgreSQL DataRows were not associated with an encrypted portal")] + HeldDataRowsNotEncrypted, + + #[error("Expected {expected} DataRow columns, received {received}")] + DataRowColumnCountMismatch { expected: usize, received: usize }, + #[error("Expected {expected} parameter format codes, received {received}")] ParameterFormatCodesMismatch { expected: usize, received: usize }, diff --git a/packages/cipherstash-proxy/src/main.rs b/packages/cipherstash-proxy/src/main.rs index a11da143d..5cf22a728 100644 --- a/packages/cipherstash-proxy/src/main.rs +++ b/packages/cipherstash-proxy/src/main.rs @@ -1,5 +1,5 @@ use cipherstash_proxy::config::TandemConfig; -use cipherstash_proxy::connect::{self, AsyncStream}; +use cipherstash_proxy::connect; use cipherstash_proxy::error::{ConfigError, Error}; use cipherstash_proxy::prometheus::CLIENTS_ACTIVE_CONNECTIONS; use cipherstash_proxy::proxy::Proxy; @@ -87,7 +87,7 @@ fn main() -> Result<(), Box> { info!(msg = "Received SIGTERM"); break; }, - Ok(client_stream) = AsyncStream::accept(&listener) => { + Ok(client_stream) = connect::accept(&listener) => { client_id += 1; @@ -97,12 +97,12 @@ fn main() -> Result<(), Box> { gauge!(CLIENTS_ACTIVE_CONNECTIONS).increment(1); - match pg::handler(client_stream,context).await { + let result = pg::handler(client_stream,context).await; + gauge!(CLIENTS_ACTIVE_CONNECTIONS).decrement(1); + + match result { Ok(_) => (), Err(err) => { - - gauge!(CLIENTS_ACTIVE_CONNECTIONS).decrement(1); - match err { Error::ConnectionClosed => { info!(msg = "Database connection closed by client"); diff --git a/packages/cipherstash-proxy/src/postgresql/context/mod.rs b/packages/cipherstash-proxy/src/postgresql/context/mod.rs index d42e015cf..a4dd1ab10 100644 --- a/packages/cipherstash-proxy/src/postgresql/context/mod.rs +++ b/packages/cipherstash-proxy/src/postgresql/context/mod.rs @@ -4,11 +4,7 @@ pub mod portal; pub mod statement; pub mod statement_metadata; pub use self::{phase_timing::PhaseTiming, portal::Portal, statement::Statement}; -use super::{ - column_mapper::ColumnMapper, - messages::{describe::Describe, Name, Target}, - Column, -}; +use super::{column_mapper::ColumnMapper, rewrite::Name, Column}; use crate::{ config::TandemConfig, error::{EncryptError, Error}, @@ -22,11 +18,12 @@ use crate::{ use cipherstash_client::IdentifiedBy; use eql_mapper::{Schema, TableResolver}; use metrics::{counter, histogram}; +use pg_proto::{Describe, DescribeTarget, OperationId, TransactionStatus}; use serde_json::json; use sqltk::parser::ast::{Expr, Ident, ObjectName, ObjectNamePart, Set, Value, ValueWithSpan}; pub use statement_metadata::StatementMetadata; use std::{ - collections::{HashMap, VecDeque}, + collections::HashMap, sync::{ atomic::{AtomicU64, Ordering}, Arc, LazyLock, RwLock, @@ -37,12 +34,7 @@ use tokio::sync::oneshot; use tracing::{debug, error, warn}; use uuid::Uuid; -type DescribeQueue = Queue; -type ExecuteQueue = Queue; -type SessionMetricsQueue = Queue; -type PortalQueue = Queue>; - -#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)] +#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, Ord, PartialOrd)] pub struct SessionId(u64); #[derive(Clone, Debug, PartialEq)] @@ -67,26 +59,25 @@ where column_mapper: ColumnMapper, statements: Arc>>>, statement_sessions: Arc>>, - portals: Arc>>, - describe: Arc>, - execute: Arc>, + portals: Arc>>>, + operations: Arc>>, schema_changed: Arc>, - session_metrics: Arc>, + session_metrics: Arc>>, table_resolver: Arc, unsafe_disable_mapping: bool, keyset_id: Arc>>, session_id_counter: Arc, + transaction_status: Arc>, } /// Context for tracking an in-flight Execute operation. /// -/// Timing data is accumulated here during backend message processing because -/// the backend operates on the execute queue rather than having direct access -/// to the session metrics queue. On completion via `complete_execution()`, -/// timing is transferred to the associated SessionMetricsContext. +/// This stores only CipherStash metadata associated with pg-proto's operation; +/// pg-proto owns protocol ordering and backpressure. #[derive(Clone, Debug)] pub struct ExecuteContext { name: Name, + portal: Option>, start: Instant, session_id: Option, /// Server wait duration (time to first response byte). @@ -98,9 +89,14 @@ pub struct ExecuteContext { } impl ExecuteContext { - fn new(name: Name, session_id: Option) -> ExecuteContext { + fn new( + name: Name, + portal: Option>, + session_id: Option, + ) -> ExecuteContext { ExecuteContext { name, + portal, start: Instant::now(), session_id, server_wait_duration: None, @@ -133,6 +129,12 @@ impl ExecuteContext { } } +#[derive(Clone, Debug, Default)] +struct OperationContext { + describe_statement: Option>, + execute: Option, +} + #[derive(Clone, Debug)] pub struct SessionMetricsContext { id: SessionId, @@ -160,11 +162,6 @@ impl SessionMetricsContext { } } -#[derive(Clone, Debug)] -pub struct Queue { - pub queue: VecDeque, -} - impl Context where T: EncryptionService, @@ -183,10 +180,9 @@ where statements: Arc::new(RwLock::new(HashMap::new())), statement_sessions: Arc::new(RwLock::new(HashMap::new())), portals: Arc::new(RwLock::new(HashMap::new())), - describe: Arc::new(RwLock::from(Queue::new())), - execute: Arc::new(RwLock::from(Queue::new())), + operations: Arc::new(RwLock::new(HashMap::new())), schema_changed: Arc::new(RwLock::from(false)), - session_metrics: Arc::new(RwLock::from(Queue::new())), + session_metrics: Arc::new(RwLock::new(HashMap::new())), table_resolver: Arc::new(TableResolver::new_editable(schema)), client_id, config, @@ -197,33 +193,57 @@ where unsafe_disable_mapping: false, keyset_id: Arc::new(RwLock::new(None)), session_id_counter: Arc::new(AtomicU64::new(1)), + transaction_status: Arc::new(RwLock::new(TransactionStatus::Idle)), } } - pub fn set_describe(&mut self, describe: Describe) { + pub fn set_describe(&mut self, operation: OperationId, describe: Describe) { debug!(target: CONTEXT, client_id = self.client_id, describe = ?describe); - let _ = self.describe.write().map(|mut queue| queue.add(describe)); + let statement = match &describe { + Describe { + name, + target: DescribeTarget::Portal, + } => self.get_portal_statement(name), + Describe { + name, + target: DescribeTarget::Statement, + } => self.get_statement(name), + }; + let _ = self.operations.write().map(|mut operations| { + operations.entry(operation).or_default().describe_statement = statement; + }); } /// /// Marks the current Describe as complete /// Removes the Describe from the Queue /// - pub fn complete_describe(&mut self) { + pub fn complete_describe(&mut self, operation: OperationId) { debug!(target: CONTEXT, client_id = self.client_id, msg = "Describe complete"); - let _ = self.describe.write().map(|mut queue| queue.complete()); + let _ = self.operations.write().map(|mut operations| { + if let Some(context) = operations.get_mut(&operation) { + context.describe_statement = None; + if context.execute.is_none() { + operations.remove(&operation); + } + } + }); } pub fn start_session(&mut self) -> SessionId { let id = SessionId(self.session_id_counter.fetch_add(1, Ordering::Relaxed)); let ctx = SessionMetricsContext::new(id); - let _ = self.session_metrics.write().map(|mut queue| queue.add(ctx)); + let _ = self + .session_metrics + .write() + .map(|mut sessions| sessions.insert(id, ctx)); id } - pub fn finish_session(&mut self) { + pub fn finish_session(&mut self, session_id: Option) { debug!(target: CONTEXT, client_id = self.client_id, msg = "Session Metrics finished"); - if let Some(session) = self.get_session_metrics() { + let session = session_id.and_then(|id| self.session_metrics.write().ok()?.remove(&id)); + if let Some(session) = session { let duration = session.duration(); let metadata = &session.metadata; @@ -287,43 +307,27 @@ where ); } } - - let _ = self - .session_metrics - .write() - .map(|mut queue| queue.complete()); } - pub fn set_execute(&mut self, name: Name, session_id: Option) { + pub fn set_execute( + &mut self, + operation: OperationId, + name: Name, + session_id: Option, + ) { debug!(target: CONTEXT, client_id = self.client_id, execute = ?name); - let ctx = ExecuteContext::new(name, session_id); - - let _ = self.execute.write().map(|mut queue| queue.add(ctx)); + let portal = self.get_portal(&name); + let ctx = ExecuteContext::new(name, portal, session_id); + let _ = self.operations.write().map(|mut operations| { + operations.entry(operation).or_default().execute = Some(ctx); + }); } /// Set execute state for portal, looking up session ID internally. - pub fn set_execute_for_portal(&mut self, name: Name) { + pub fn set_execute_for_portal(&mut self, operation: OperationId, name: Name) { let session_id = self.get_portal_session_id(&name); - self.set_execute(name, session_id); - } - - /// Number of entries currently queued in the `execute` queue. - /// - /// Test-only accessor used by the BUG-300 regression tests to assert the - /// per-connection queues are drained (and do not grow unbounded). - #[cfg(test)] - pub(crate) fn execute_queue_len(&self) -> usize { - self.execute.read().unwrap().queue.len() - } - - /// Number of entries currently queued in the `session_metrics` queue. - /// - /// Test-only accessor used by the BUG-300 regression tests to assert the - /// per-connection queues are drained (and do not grow unbounded). - #[cfg(test)] - pub(crate) fn session_metrics_queue_len(&self) -> usize { - self.session_metrics.read().unwrap().queue.len() + self.set_execute(operation, name, session_id); } /// Marks the current Execution as Complete. @@ -332,9 +336,8 @@ where /// - `server_wait_duration` (time to first response byte) is recorded to the session /// - `server_response_duration` (time receiving response data) is added to the session /// - /// This two-phase timing pattern exists because the backend operates on the execute queue - /// rather than having direct access to the session. Timing is accumulated in ExecuteContext - /// during message processing, then transferred to the correct SessionMetricsContext here. + /// Timing is transferred to the session identified by this operation rather + /// than inferred from response order. /// /// If the associated portal is Unnamed, it is closed. /// @@ -343,10 +346,20 @@ where /// An unnamed portal is destroyed at the end of the transaction, or as soon as the next Bind statement specifying the unnamed portal as destination is issued /// /// https://www.postgresql.org/docs/current/protocol-flow.html#PROTOCOL-FLOW-EXT-QUERY - pub fn complete_execution(&mut self) { + pub fn complete_execution(&mut self, operation: OperationId) -> Option { debug!(target: CONTEXT, client_id = self.client_id, msg = "Execute complete"); - if let Some(execute) = self.get_execute() { + let execute = self.operations.write().ok().and_then(|mut operations| { + let execute = operations.get_mut(&operation)?.execute.take(); + if operations + .get(&operation) + .is_some_and(|context| context.describe_statement.is_none()) + { + operations.remove(&operation); + } + execute + }); + if let Some(execute) = execute { if let Some(session_id) = execute.session_id() { if let Some(wait) = execute.server_wait_duration() { self.record_server_wait_duration(session_id, wait); @@ -358,25 +371,27 @@ where } // Get labels from current session metadata - let (statement_type, protocol, mapped, multi_statement) = - if let Some(session) = self.get_session_metrics() { - let metadata = &session.metadata; - ( - metadata - .statement_type - .map(|t| t.as_label()) - .unwrap_or("unknown"), - metadata.protocol.map(|p| p.as_label()).unwrap_or("unknown"), - if metadata.encrypted { "true" } else { "false" }, - if metadata.multi_statement { - "true" - } else { - "false" - }, - ) - } else { - ("unknown", "unknown", "false", "false") - }; + let (statement_type, protocol, mapped, multi_statement) = if let Some(session) = execute + .session_id() + .and_then(|id| self.get_session_metrics(id)) + { + let metadata = &session.metadata; + ( + metadata + .statement_type + .map(|t| t.as_label()) + .unwrap_or("unknown"), + metadata.protocol.map(|p| p.as_label()).unwrap_or("unknown"), + if metadata.encrypted { "true" } else { "false" }, + if metadata.multi_statement { + "true" + } else { + "false" + }, + ) + } else { + ("unknown", "unknown", "false", "false") + }; histogram!( STATEMENTS_EXECUTION_DURATION_SECONDS, @@ -387,12 +402,12 @@ where ) .record(execute.duration()); - if execute.name.is_unnamed() { - self.close_portal(&execute.name); + if execute.name.is_empty() { + self.close_portal_if_current(&execute.name, execute.portal.as_ref()); } + return execute.session_id(); } - - let _ = self.execute.write().map(|mut queue| queue.complete()); + None } pub fn add_statement(&mut self, name: Name, statement: Statement) { @@ -406,15 +421,39 @@ where pub fn close_statement(&mut self, name: &Name) { debug!(target: CONTEXT, client_id = self.client_id, statement = ?name); - let _ = self + let statement = self .statements .write() - .map(|mut guarded| guarded.remove(name)); + .ok() + .and_then(|mut guarded| guarded.remove(name)); - let _ = self + if let Some(statement) = statement { + let _ = self.portals.write().map(|mut guarded| { + guarded.retain(|_, portal| { + !matches!(portal.as_ref(), Portal::Encrypted { statement: portal_statement, .. } if Arc::ptr_eq(portal_statement, &statement)) + }); + }); + } + + let session_id = self .statement_sessions .write() - .map(|mut guarded| guarded.remove(name)); + .ok() + .and_then(|mut guarded| guarded.remove(name)); + self.finish_session(session_id); + } + + pub fn transaction_status(&self) -> TransactionStatus { + self.transaction_status + .read() + .map(|status| *status) + .unwrap_or(TransactionStatus::Idle) + } + + pub fn set_transaction_status(&mut self, status: TransactionStatus) { + if let Ok(mut current) = self.transaction_status.write() { + *current = status; + } } /// Close both statement and its associated portal. @@ -425,12 +464,10 @@ where pub fn add_portal(&mut self, name: Name, portal: Portal) { debug!(target: CONTEXT, client_id = self.client_id, name = ?name, portal = ?portal); - let _ = self.portals.write().map(|mut portals| { - portals - .entry(name) - .or_insert_with(Queue::new) - .add(Arc::new(portal)); - }); + let _ = self + .portals + .write() + .map(|mut portals| portals.insert(name, Arc::new(portal))); } pub fn get_statement(&self, name: &Name) -> Option> { @@ -462,7 +499,7 @@ where warn!( target: CONTEXT, client_id = self.client_id, - prepared_statement = %name.as_str(), + prepared_statement = %String::from_utf8_lossy(name), msg = "Session lookup failed for prepared statement, using latest session" ); } @@ -475,10 +512,18 @@ where /// pub fn close_portal(&mut self, name: &Name) { debug!(target: CONTEXT, client_id = self.client_id, msg = "Close Portal", name = ?name); + let _ = self.portals.write().map(|mut portals| portals.remove(name)); + } + + fn close_portal_if_current(&mut self, name: &Name, expected: Option<&Arc>) { let _ = self.portals.write().map(|mut portals| { - portals - .entry(name.clone()) - .and_modify(|queue| queue.complete()); + if expected.is_some_and(|expected| { + portals + .get(name) + .is_some_and(|current| Arc::ptr_eq(current, expected)) + }) { + portals.remove(name); + } }); } @@ -486,14 +531,12 @@ where debug!(target: CONTEXT, client_id = self.client_id, src = "Get Portal", portal = ?name); let portals = self.portals.read().ok()?; - let queue = portals.get(name)?; - queue.next().cloned() + portals.get(name).cloned() } pub fn get_portal_statement(&self, name: &Name) -> Option> { let portals = self.portals.read().ok()?; - let queue = portals.get(name)?; - let portal = queue.next()?; + let portal = portals.get(name)?; debug!(target: CONTEXT, client_id = self.client_id, portal = ?portal); @@ -505,59 +548,52 @@ where pub fn get_portal_session_id(&self, name: &Name) -> Option { let portals = self.portals.read().ok()?; - let queue = portals.get(name)?; - let portal = queue.next()?; + let portal = portals.get(name)?; portal.session_id() } - pub fn get_statement_for_row_decription(&self) -> Option> { - if let Some(statement) = self.get_statement_from_describe() { + pub fn get_statement_for_operation(&self, operation: OperationId) -> Option> { + let operations = self.operations.read().ok()?; + let context = operations.get(&operation)?; + if let Some(statement) = &context.describe_statement { return Some(statement.clone()); } - - if let Some(Portal::Encrypted { statement, .. }) = self.get_portal_from_execute().as_deref() - { - return Some(statement.clone()); - }; - - None + match context.execute.as_ref()?.portal.as_deref()? { + Portal::Encrypted { statement, .. } => Some(statement.clone()), + Portal::Passthrough { .. } => None, + } } - pub fn get_statement_from_describe(&self) -> Option> { - let queue = self.describe.read().ok()?; - let describe = queue.next()?; - - debug!(target: CONTEXT, client_id = self.client_id, msg = "Get Statement", describe = ?describe); - - match describe { - Describe { - ref name, - target: Target::Portal, - } => self.get_portal_statement(name), - Describe { - ref name, - target: Target::Statement, - } => self.get_statement(name), - } + pub fn get_statement_from_describe(&self, operation: OperationId) -> Option> { + self.operations + .read() + .ok()? + .get(&operation)? + .describe_statement + .clone() } - pub fn get_portal_from_execute(&self) -> Option> { - let queue = self.execute.read().ok()?; - let execute_context = queue.next()?; - let name = &execute_context.name; - self.get_portal(name) + pub fn get_portal_from_execute(&self, operation: OperationId) -> Option> { + self.operations + .read() + .ok()? + .get(&operation)? + .execute + .as_ref()? + .portal + .clone() } - pub fn get_execute(&self) -> Option { - let queue = self.execute.read().ok()?; - let execute_context = queue.next()?; + pub fn get_execute(&self, operation: OperationId) -> Option { + let operations = self.operations.read().ok()?; + let execute_context = operations.get(&operation)?.execute.as_ref()?; debug!(target: CONTEXT, client_id = self.client_id, msg = "Get Execute", execute = ?execute_context); Some(execute_context.to_owned()) } - pub fn get_session_metrics(&self) -> Option { - let queue = self.session_metrics.read().ok()?; - let session_context = queue.next()?; + pub fn get_session_metrics(&self, session_id: SessionId) -> Option { + let sessions = self.session_metrics.read().ok()?; + let session_context = sessions.get(&session_id)?; debug!(target: CONTEXT, client_id = self.client_id, msg = "Get Session Metrics", session_metrics = ?session_context); Some(session_context.to_owned()) } @@ -897,20 +933,16 @@ where where F: FnOnce(&mut SessionMetricsContext), { - if let Ok(mut queue) = self.session_metrics.write() { - if let Some(session) = queue - .queue - .iter_mut() - .find(|session| session.id() == session_id) - { + if let Ok(mut sessions) = self.session_metrics.write() { + if let Some(session) = sessions.get_mut(&session_id) { f(session); } } } pub fn latest_session_id(&self) -> Option { - let queue = self.session_metrics.read().ok()?; - queue.queue.back().map(|session| session.id()) + let sessions = self.session_metrics.read().ok()?; + sessions.keys().max().copied() } /// Record parse phase duration for the session (first write wins) @@ -1004,77 +1036,62 @@ where } /// Record server wait for first response; otherwise accumulate response time for the current execute - pub fn record_execute_server_timing(&mut self, duration: Duration) { - if let Ok(mut queue) = self.execute.write() { - if let Some(execute) = queue.current_mut() { + pub fn record_execute_server_timing(&mut self, operation: OperationId, duration: Duration) { + if let Ok(mut operations) = self.operations.write() { + if let Some(execute) = operations + .get_mut(&operation) + .and_then(|op| op.execute.as_mut()) + { execute.record_server_wait_or_add_response(duration); } } } /// Add decrypt phase duration for the current execute session (if any) - pub fn add_decrypt_duration_for_execute(&mut self, duration: Duration) { - let session_id = self.get_execute().and_then(|execute| execute.session_id()); + pub fn add_decrypt_duration_for_execute(&mut self, operation: OperationId, duration: Duration) { + let session_id = self + .get_execute(operation) + .and_then(|execute| execute.session_id()); if let Some(session_id) = session_id { self.add_decrypt_duration(session_id, duration); } } /// Add client write duration for the current execute session (if any) - pub fn add_client_write_duration_for_execute(&mut self, duration: Duration) { - let session_id = self.get_execute().and_then(|execute| execute.session_id()); + pub fn add_client_write_duration_for_execute( + &mut self, + operation: OperationId, + duration: Duration, + ) { + let session_id = self + .get_execute(operation) + .and_then(|execute| execute.session_id()); if let Some(session_id) = session_id { self.add_client_write_duration(session_id, duration); } } } -impl Queue { - pub fn new() -> Self { - Queue { - queue: VecDeque::new(), - } - } - - pub fn complete(&mut self) { - let _ = self.queue.pop_front(); - } - - pub fn next(&self) -> Option<&T> { - self.queue.front() - } - - pub fn add(&mut self, item: T) { - self.queue.push_back(item); - } - - /// Get mutable reference to the current (first) item in the queue - pub fn current_mut(&mut self) -> Option<&mut T> { - self.queue.front_mut() - } -} - #[cfg(test)] mod tests { - use super::{Context, Describe, KeysetIdentifier, Portal, Statement}; + use super::{Context, KeysetIdentifier, Portal, Statement}; use crate::{ config::LogConfig, error::Error, log, - postgresql::{ - messages::{Name, Target}, - Column, - }, + postgresql::{rewrite::Name, Column}, proxy::{EncryptConfig, EncryptionService}, TandemConfig, }; use cipherstash_client::IdentifiedBy; use eql_mapper::Schema; + use pg_proto::TransactionStatus; use sqltk::parser::{dialect::PostgreSqlDialect, parser::Parser}; use std::sync::Arc; use tokio::sync::mpsc; use uuid::Uuid; + #[derive(Clone)] struct TestService {} #[async_trait::async_trait] @@ -1131,108 +1148,63 @@ mod tests { Portal::encrypted_with_format_codes(statement.clone(), vec![], None) } - fn get_statement(portal: Arc) -> Arc { - match portal.as_ref() { - Portal::Encrypted { statement, .. } => statement.clone(), - _ => { - panic!("Expected Encrypted Portal"); - } - } - } - #[test] - pub fn get_statement_from_describe() { - log::init(LogConfig::default()); - + fn closing_a_statement_finishes_its_metrics_session() { let mut context = create_context(); + let name = Name::default(); + let session_id = context.start_session(); + context.set_statement_session(name.clone(), session_id); - let name = Name::from("name"); + context.close_statement(&name); - context.add_statement(name.clone(), statement()); - - let statement = context.get_statement(&name).unwrap(); - - let describe = Describe { - name, - target: Target::Statement, - }; - context.set_describe(describe); - - let s = context.get_statement_from_describe().unwrap(); - - assert_eq!(s, statement) + assert!(context.get_session_metrics(session_id).is_none()); + assert!(context.get_statement_session(&name).is_none()); } #[test] - pub fn execution_flow() { - log::init(LogConfig::default()); - + fn closing_a_statement_invalidates_only_its_portals() { let mut context = create_context(); + let closed_statement_name = Name::from("closed_statement"); + let retained_statement_name = Name::from("retained_statement"); + context.add_statement(closed_statement_name.clone(), statement()); + context.add_statement(retained_statement_name.clone(), statement()); - let statement_name = Name::from("statement"); - let portal_name = Name::from("portal"); - - // Add statement to context - context.add_statement(statement_name.clone(), statement()); + let closed_statement = context.get_statement(&closed_statement_name).unwrap(); + let retained_statement = context.get_statement(&retained_statement_name).unwrap(); + let closed_portal_name = Name::from("differently_named_portal"); + let retained_portal_name = Name::from("retained_portal"); + let passthrough_portal_name = Name::from("passthrough_portal"); + context.add_portal(closed_portal_name.clone(), portal(&closed_statement)); + context.add_portal(retained_portal_name.clone(), portal(&retained_statement)); + context.add_portal(passthrough_portal_name.clone(), Portal::passthrough(None)); - // Get statement from context - let statement = context.get_statement(&statement_name).unwrap(); + context.close_statement(&closed_statement_name); - // Add portal pointing to statement to context - context.add_portal(portal_name.clone(), portal(&statement)); - - // Add statement name to execute context - context.set_execute(portal_name.clone(), None); - - // Portal statement should be the right statement - let portal = context.get_portal_from_execute().unwrap(); - - let statement = get_statement(portal); - assert_eq!(statement, statement); - - // Complete the execution - context.complete_execution(); - - // Should be no portal for execute context - let portal = context.get_portal_from_execute(); - assert!(portal.is_none()); - - // Unamed portal is closed on complete - let portal = context.get_portal(&portal_name); - assert!(portal.is_some()); + assert!(context.get_portal(&closed_portal_name).is_none()); + assert!(context.get_portal(&retained_portal_name).is_some()); + assert!(context.get_portal(&passthrough_portal_name).is_some()); } - /// Unit test for the queue-draining primitives. - /// - /// `complete_execution()` / `finish_session()` are the only drains for the - /// per-connection `execute` / `session_metrics` queues. This asserts that - /// calling them after enqueuing a session + execute leaves both queues - /// empty, across many iterations. - /// - /// Note: this exercises the primitives directly — it does *not* drive the - /// backend passthrough path that actually caused BUG-300 (that early - /// returned without calling these). The backend-level regression test for - /// BUG-300 lives in `backend.rs` - /// (`passthrough_drains_queues_on_execute_terminating_message`). #[test] - pub fn complete_execution_and_finish_session_drain_queues() { - log::init(LogConfig::default()); - + fn transaction_status_tracks_backend_ready_state() { let mut context = create_context(); + assert_eq!(context.transaction_status(), TransactionStatus::Idle); - for _ in 0..1000 { - // Frontend: a session + execute are enqueued for every statement. - let session_id = context.start_session(); - context.set_execute(Name::unnamed(), Some(session_id)); + context.set_transaction_status(TransactionStatus::InTransaction); - // Drain primitives, normally called by the backend on an - // execute-terminating message (CommandComplete / ErrorResponse / …). - context.complete_execution(); - context.finish_session(); - } + assert_eq!( + context.transaction_status(), + TransactionStatus::InTransaction + ); + } - assert_eq!(context.execute_queue_len(), 0); - assert_eq!(context.session_metrics_queue_len(), 0); + fn get_statement(portal: Arc) -> Arc { + match portal.as_ref() { + Portal::Encrypted { statement, .. } => statement.clone(), + _ => { + panic!("Expected Encrypted Portal"); + } + } } #[test] @@ -1249,9 +1221,6 @@ mod tests { context.add_statement(statement_name_1.clone(), statement()); context.add_statement(statement_name_2.clone(), statement()); - // Replicate pipelined execution - // Add multiple portals with the same name - // Pointing to different statements let portal_name = Name::from("portal"); let statement_1 = context.get_statement(&statement_name_1).unwrap(); @@ -1260,88 +1229,10 @@ mod tests { let statement_2 = context.get_statement(&statement_name_2).unwrap(); context.add_portal(portal_name.clone(), portal(&statement_2)); - // Execute both portals - context.set_execute(portal_name.clone(), None); - context.set_execute(portal_name.clone(), None); - - // Portal should point to first statement - let portal = context.get_portal_from_execute().unwrap(); - let statement = get_statement(portal); - assert_eq!(statement, statement); - - let portal = context.get_portal_from_execute().unwrap(); - let statement = get_statement(portal); - assert_eq!(statement_1, statement); - - // Complete execution - context.complete_execution(); - - // Portal should point to second statement - let portal = context.get_portal_from_execute().unwrap(); - - let statement = get_statement(portal); - assert_eq!(statement_1, statement); - } - - #[test] - pub fn pipeline_execution() { - log::init(LogConfig::default()); - - let mut context = create_context(); - - let statement_name_1 = Name::from("statement_1"); - let portal_name_1 = Name::unnamed(); - - let statement_name_2 = Name::from("statement_2"); - let portal_name_2 = Name::unnamed(); - - let statement_name_3 = Name::from("statement_3"); - let portal_name_3 = Name::from("portal_3"); - - // Add statement to context - context.add_statement(statement_name_1.clone(), statement()); - context.add_statement(statement_name_2.clone(), statement()); - context.add_statement(statement_name_3.clone(), statement()); - - // Create portals for each statement - let statement_1 = context.get_statement(&statement_name_1).unwrap(); - context.add_portal(portal_name_1.clone(), portal(&statement_1)); - - let statement_2 = context.get_statement(&statement_name_2).unwrap(); - context.add_portal(portal_name_2.clone(), portal(&statement_2)); - - let statement_3 = context.get_statement(&statement_name_3).unwrap(); - context.add_portal(portal_name_3.clone(), portal(&statement_3)); - - // Add portals to execute context - context.set_execute(portal_name_1.clone(), None); - context.set_execute(portal_name_2.clone(), None); - context.set_execute(portal_name_3.clone(), None); - - // Multiple calls return the portal for the first Execution context - let portal = context.get_portal_from_execute().unwrap(); - let statement = get_statement(portal); - assert_eq!(statement_1, statement); - - let portal = context.get_portal_from_execute().unwrap(); - let statement = get_statement(portal); - assert_eq!(statement_1, statement); - - // Complete the execution of the first portal - context.complete_execution(); - - // Returns the next portal - let portal = context.get_portal_from_execute().unwrap(); - let statement = get_statement(portal); - assert_eq!(statement_2, statement); - - // Complete the execution - context.complete_execution(); - - // Returns the next portal - let portal = context.get_portal_from_execute().unwrap(); - let statement = get_statement(portal); - assert_eq!(statement_3, statement); + let portal = context.get_portal(&portal_name).unwrap(); + assert_eq!(statement_2, get_statement(portal)); + context.close_portal(&portal_name); + assert!(context.get_portal(&portal_name).is_none()); } fn parse_statement(sql: &str) -> sqltk::parser::ast::Statement { diff --git a/packages/cipherstash-proxy/src/postgresql/data/from_sql.rs b/packages/cipherstash-proxy/src/postgresql/data/from_sql.rs index aafd74679..9ee51d606 100644 --- a/packages/cipherstash-proxy/src/postgresql/data/from_sql.rs +++ b/packages/cipherstash-proxy/src/postgresql/data/from_sql.rs @@ -1,7 +1,7 @@ use crate::{ error::{Error, MappingError}, log::ENCODING, - postgresql::{format_code::FormatCode, messages::bind::BindParam}, + postgresql::{format_code::FormatCode, rewrite::bind::BindParam}, }; use bigdecimal::BigDecimal; use bytes::BytesMut; @@ -603,7 +603,7 @@ fn decimal_from_sql( #[cfg(test)] mod binary_json_value_tests { use super::*; - use crate::postgresql::{format_code::FormatCode, messages::bind::BindParam}; + use crate::postgresql::{format_code::FormatCode, rewrite::bind::BindParam}; use bytes::BytesMut; fn binary_param(bytes: &[u8]) -> BindParam { @@ -733,7 +733,7 @@ mod tests { config::LogConfig, log, postgresql::{ - data::bind_param_from_sql, format_code::FormatCode, messages::bind::BindParam, Column, + data::bind_param_from_sql, format_code::FormatCode, rewrite::bind::BindParam, Column, }, Identifier, }; diff --git a/packages/cipherstash-proxy/src/postgresql/diagnostics.rs b/packages/cipherstash-proxy/src/postgresql/diagnostics.rs new file mode 100644 index 000000000..3bd0a8ddf --- /dev/null +++ b/packages/cipherstash-proxy/src/postgresql/diagnostics.rs @@ -0,0 +1,127 @@ +//! CipherStash diagnostic response policy over pg-proto's wire model. +use bytes::Bytes; +use pg_proto::{DiagnosticField, DiagnosticResponse}; +use regex::Regex; +use std::sync::LazyLock; + +pub const CODE_UNDEFINED_COLUMN: &str = "42703"; +pub const CODE_INVALID_PASSWORD: &str = "28P01"; +pub const CODE_RAISE_EXCEPTION: &str = "P0001"; +pub const CODE_SYNTAX_ERROR: &str = "42601"; +pub const CODE_INVALID_TEXT_REPRESENTATION: &str = "22P02"; +pub const CODE_IDLE_SESSION_TIMEOUT: &str = "57P05"; +pub const CODE_SYSTEM_ERROR: &str = "58000"; + +fn response(fields: impl IntoIterator) -> DiagnosticResponse { + DiagnosticResponse { + fields: fields + .into_iter() + .map(|(code, value)| DiagnosticField { + code, + value: Bytes::from(value), + }) + .collect(), + } +} + +fn standard(severity: &str, code: &str, message: String) -> DiagnosticResponse { + response([ + (b'S', severity.to_owned()), + (b'V', severity.to_owned()), + (b'C', code.to_owned()), + (b'M', message), + ]) +} + +pub fn connection_timeout(message: String) -> DiagnosticResponse { + standard("FATAL", CODE_IDLE_SESSION_TIMEOUT, message) +} + +pub fn invalid_password(message: String) -> DiagnosticResponse { + standard("FATAL", CODE_INVALID_PASSWORD, message) +} + +pub fn invalid_sql_statement(message: String) -> DiagnosticResponse { + let line = extract_line_from_parse_error(&message); + let position = extract_position_from_parse_error(&message); + let mut fields = vec![ + (b'S', "ERROR".to_owned()), + (b'V', "ERROR".to_owned()), + (b'C', CODE_SYNTAX_ERROR.to_owned()), + (b'M', message), + ]; + if let Some(line) = line { + fields.push((b'L', line.to_string())); + } + if let Some(position) = position { + fields.push((b'P', position.to_string())); + } + response(fields) +} + +pub fn invalid_parameter(message: String, table: &str, column: &str) -> DiagnosticResponse { + response([ + (b'S', "ERROR".to_owned()), + (b'V', "ERROR".to_owned()), + (b'C', CODE_INVALID_TEXT_REPRESENTATION.to_owned()), + (b'M', message), + (b't', table.to_owned()), + (b'c', column.to_owned()), + ]) +} + +pub fn unknown_column(message: String, table: &str, column: &str) -> DiagnosticResponse { + response([ + (b'S', "ERROR".to_owned()), + (b'V', "ERROR".to_owned()), + (b'C', CODE_UNDEFINED_COLUMN.to_owned()), + (b'M', message), + (b't', table.to_owned()), + (b'c', column.to_owned()), + (b'R', "cipherstash-proxy".to_owned()), + ]) +} + +pub fn system_error(message: String) -> DiagnosticResponse { + standard("FATAL", CODE_SYSTEM_ERROR, message) +} + +pub fn is_fatal(response: &DiagnosticResponse) -> bool { + response + .fields + .iter() + .any(|field| field.code == b'S' && field.value.as_ref() == b"FATAL") +} + +fn extract_line_from_parse_error(message: &str) -> Option { + static RE: LazyLock = LazyLock::new(|| Regex::new(r"\s*Line:\s*(\d+)").unwrap()); + RE.captures(message) + .and_then(|capture| capture.get(1)?.as_str().parse().ok()) +} + +fn extract_position_from_parse_error(message: &str) -> Option { + static RE: LazyLock = LazyLock::new(|| Regex::new(r"\s*Column:\s*(\d+)").unwrap()); + RE.captures(message) + .and_then(|capture| capture.get(1)?.as_str().parse().ok()) +} + +#[cfg(test)] +mod tests { + use super::*; + + fn field(response: &DiagnosticResponse, code: u8) -> Option<&[u8]> { + response + .fields + .iter() + .find(|field| field.code == code) + .map(|field| field.value.as_ref()) + } + + #[test] + fn sql_parse_error_includes_line_and_position() { + let response = + invalid_sql_statement("sql syntax error in blah Line: 1, Column: 2".to_owned()); + assert_eq!(field(&response, b'L'), Some(b"1".as_slice())); + assert_eq!(field(&response, b'P'), Some(b"2".as_slice())); + } +} diff --git a/packages/cipherstash-proxy/src/postgresql/driver.rs b/packages/cipherstash-proxy/src/postgresql/driver.rs new file mode 100644 index 000000000..b0d2424d7 --- /dev/null +++ b/packages/cipherstash-proxy/src/postgresql/driver.rs @@ -0,0 +1,221 @@ +use super::{middleware::CipherStashMiddlewareFactory, Context}; +use crate::{connect, error::Error, proxy::EncryptionService, tls}; +use pg_proto::{ + BackendHoldLimits, BoundedPipeline, CancellationPolicy, Client, ClientTlsConfig, + ClientTlsPolicy, ClientTlsProvider, ConnectTarget, ForwardedMessage, FrontendMessage, + InMemoryCancellationRegistry, InitialServerContext, Intermediary, Server, ServerIdentity, + ServerIdentityProvider, ServerTlsPolicy, SslMode, StartupParameters, StartupRouteResolver, + StaticClientCredentials, StaticMd5ServerCredentials, +}; +use std::{convert::Infallible, sync::Arc}; +use tokio::net::TcpStream; +use tracing::{info, warn}; + +#[derive(Clone)] +struct Route(String); + +impl StartupRouteResolver for Route { + type Error = Infallible; + async fn resolve( + &self, + _: StartupParameters, + _: InitialServerContext<'_, Peer>, + ) -> Result { + Ok(ConnectTarget::new(self.0.clone())) + } +} + +#[derive(Clone)] +struct DownstreamIdentity(ServerIdentity); +impl ServerIdentityProvider for DownstreamIdentity { + type Error = Infallible; + fn resolve(&self) -> Result { + Ok(self.0.clone()) + } +} + +#[derive(Clone)] +struct UpstreamTls { + server_name: rustls_pki_types::ServerName<'static>, + verify: bool, +} +impl ClientTlsProvider for UpstreamTls { + type Error = Error; + async fn resolve(&self, _: &ConnectTarget) -> Result { + let name = self.server_name.clone(); + let result = rustls_native_certs::load_native_certs(); + for error in result.errors { + warn!(msg = "Could not load a native TLS certificate", %error); + } + let mut roots = rustls::RootCertStore::empty(); + for certificate in result.certs { + roots + .add(certificate) + .map_err(|error| std::io::Error::new(std::io::ErrorKind::InvalidData, error))?; + } + if self.verify && roots.is_empty() { + return Err(crate::error::ConfigError::from( + crate::error::TlsConfigError::InvalidCertificate, + ) + .into()); + } + Ok(ClientTlsConfig::new(name, roots)) + } +} + +pub async fn handler(client_stream: TcpStream, context: Context) -> Result<(), Error> +where + S: EncryptionService + Clone, +{ + let address = context.database_socket_address(); + let downstream_auth = StaticMd5ServerCredentials::new( + context.database_username().to_owned(), + context.database_password(), + ); + let upstream_auth = StaticClientCredentials::new( + context.database_username().to_owned(), + context.database_password(), + ); + + macro_rules! run { + ($server:expr, $client:expr) => {{ + let connection_timeout = context.connection_timeout(); + let intermediary = Intermediary::builder() + .server($server) + .client($client) + .startup_resolver(Route(address.clone())) + .cancellation(CancellationPolicy::Forward) + .cancellation_registry(InMemoryCancellationRegistry::default()) + .pipeline(BoundedPipeline::new(256).expect("non-zero pipeline bound")) + .backend_batching( + BackendHoldLimits::new(4096, 64 * 1024 * 1024) + .expect("non-zero backend hold limits"), + ) + .middleware(CipherStashMiddlewareFactory(context.clone())) + .build() + .map_err(invalid_data)?; + let accept = intermediary.accept(client_stream, (), ()); + let accepted = match connection_timeout { + Some(duration) => tokio::time::timeout(duration, accept) + .await + .map_err(|_| Error::ConnectionTimeout { duration })?, + None => accept.await, + } + .map_err(invalid_data)?; + let mut session = match accepted { + pg_proto::IntermediaryAccept::Session(session) => session, + pg_proto::IntermediaryAccept::CancellationForwarded => return Ok(()), + }; + loop { + let forward = session.forward_next(); + let forwarded = match connection_timeout { + Some(duration) => tokio::time::timeout(duration, forward) + .await + .map_err(|_| Error::ConnectionTimeout { duration })?, + None => forward.await, + }; + let forwarded = match forwarded { + Ok(forwarded) => forwarded, + Err(pg_proto::ForwardError::Frontend( + pg_proto::FrontendProjectionError::Capacity(_), + )) => continue, + Err(pg_proto::ForwardError::Middleware(error)) => return Err(error), + Err(error) => return Err(invalid_data(error)), + }; + if matches!( + forwarded, + ForwardedMessage::Frontend(FrontendMessage::Terminate) + ) { + break; + } + } + Ok(()) + }}; + } + + macro_rules! run_client { + ($server:expr) => {{ + if context.database_tls_disabled() { + let client = Client::builder() + .connector(|target: &ConnectTarget| { + let address = target.name().to_owned(); + async move { connect::connect(&address).await } + }) + .tls(ClientTlsPolicy::Disabled) + .authentication(upstream_auth.clone()) + .build() + .map_err(invalid_data)?; + run!($server, client) + } else { + let provider = UpstreamTls { + server_name: context.config().database.server_name()?.to_owned(), + verify: context.config().database.with_tls_verification, + }; + let mode = upstream_ssl_mode(context.config()); + let client = Client::builder() + .connector(|target: &ConnectTarget| { + let address = target.name().to_owned(); + async move { connect::connect(&address).await } + }) + .tls(ClientTlsPolicy::libpq(mode, provider)) + .authentication(upstream_auth.clone()) + .build() + .map_err(invalid_data)?; + run!($server, client) + } + }}; + } + + info!( + msg = "Client connected", + database = address, + client_id = context.client_id + ); + if let Some(tls_config) = context.tls_config() { + let (config, leaf) = tls::configure_server_with_leaf(tls_config)?; + let identity = DownstreamIdentity(ServerIdentity::new(Arc::new(config), leaf)); + if context.require_tls() { + return run_client!(Server::builder() + .tls(ServerTlsPolicy::Required(identity)) + .authentication(downstream_auth) + .build() + .map_err(invalid_data)?); + } + return run_client!(Server::builder() + .tls(ServerTlsPolicy::Optional(identity)) + .authentication(downstream_auth) + .build() + .map_err(invalid_data)?); + } + + run_client!(Server::builder() + .tls(ServerTlsPolicy::Disabled) + .authentication(downstream_auth) + .build() + .map_err(invalid_data)?) +} + +fn invalid_data(error: impl std::fmt::Display) -> Error { + std::io::Error::new(std::io::ErrorKind::InvalidData, error.to_string()).into() +} + +fn upstream_ssl_mode(config: &crate::TandemConfig) -> SslMode { + if config.database.with_tls_verification { + SslMode::VerifyFull + } else { + // Preserve the proxy's historical opportunistic-TLS policy: attempt + // SSL, but continue in plaintext when the database rejects it. + SslMode::Prefer + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn upstream_tls_without_verification_is_opportunistic() { + let config = crate::TandemConfig::for_testing(); + assert_eq!(upstream_ssl_mode(&config), SslMode::Prefer); + } +} diff --git a/packages/cipherstash-proxy/src/postgresql/error_handler.rs b/packages/cipherstash-proxy/src/postgresql/error_handler.rs index 4de24e03f..2478f7b17 100644 --- a/packages/cipherstash-proxy/src/postgresql/error_handler.rs +++ b/packages/cipherstash-proxy/src/postgresql/error_handler.rs @@ -4,10 +4,10 @@ /// components, ensuring that all errors are properly converted to PostgreSQL /// ErrorResponse messages and sent to clients in a protocol-compliant manner. use crate::{ - connect::Sender, - error::{EncryptError, Error, MappingError}, - postgresql::messages::error_response::ErrorResponse, + error::{EncryptError, Error, MappingError, ProtocolError, ERROR_DOC_BASE_URL}, + postgresql::diagnostics, }; +use pg_proto::DiagnosticResponse; /// Trait for components that can send PostgreSQL error responses to clients. /// @@ -15,99 +15,79 @@ use crate::{ /// frontend and backend components, providing consistent error conversion /// and client communication. pub trait PostgreSqlErrorHandler { - /// Get the client sender for this component - fn client_sender(&mut self) -> &mut Sender; - /// Get the client ID for logging purposes fn client_id(&self) -> i32; - /// Convert various error types into appropriate PostgreSQL ErrorResponse messages. + /// Convert various error types into PostgreSQL `DiagnosticResponse` messages. /// /// # Error Type Mapping /// - /// - `MappingError` -> InvalidSqlStatement error + /// - `MappingError::InvalidParameter` -> Invalid parameter error + /// - Other `MappingError` values -> Invalid SQL statement error /// - `EncryptError::UnknownColumn` -> Unknown column error - /// - `EncryptError::CouldNotRetrieveKey` -> Key retrieval error + /// - `EncryptError::CouldNotDecryptDataForKeyset` -> System error + /// - `EncryptError::UnknownKeysetIdentifier` -> System error + /// - `Error::ConnectionTimeout` -> Idle session timeout error /// - All others -> System error /// /// # Arguments /// - /// * `err` - The error to be converted to a PostgreSQL ErrorResponse - fn error_to_response(&self, err: Error) -> ErrorResponse { + /// * `err` - The error to be converted to a `DiagnosticResponse` + fn error_to_response(&self, err: Error) -> DiagnosticResponse { match err { Error::Mapping(MappingError::InvalidParameter(ref column)) => { - ErrorResponse::invalid_parameter( + diagnostics::invalid_parameter( err.to_string(), &column.table_name(), &column.column_name(), ) } - Error::Mapping(err) => ErrorResponse::invalid_sql_statement(err.to_string()), + Error::Mapping(err) => diagnostics::invalid_sql_statement(err.to_string()), Error::Encrypt(EncryptError::UnknownColumn { ref table, ref column, - }) => ErrorResponse::unknown_column(err.to_string(), table, column), + }) => diagnostics::unknown_column(err.to_string(), table, column), Error::Encrypt(EncryptError::CouldNotDecryptDataForKeyset { .. }) => { - ErrorResponse::system_error(err.to_string()) + diagnostics::system_error(err.to_string()) } Error::Encrypt(EncryptError::UnknownKeysetIdentifier { .. }) => { - ErrorResponse::system_error(err.to_string()) + diagnostics::system_error(err.to_string()) } - Error::ConnectionTimeout { .. } => ErrorResponse::connection_timeout(err.to_string()), - _ => ErrorResponse::system_error(err.to_string()), + Error::ConnectionTimeout { .. } => diagnostics::connection_timeout(err.to_string()), + Error::Protocol( + ProtocolError::HeldDataRowMissingOperation + | ProtocolError::HeldDataRowOperationMismatch + | ProtocolError::HeldBackendMessageNotDataRow + | ProtocolError::HeldDataRowsNotEncrypted, + ) => diagnostics::system_error(format!( + "CipherStash Proxy encountered an internal PostgreSQL protocol error. For help visit {ERROR_DOC_BASE_URL}#protocol-internal-error" + )), + _ => diagnostics::system_error(err.to_string()), } } - - /// Send an ErrorResponse message to the client. - /// - /// Converts the error to a PostgreSQL ErrorResponse and sends it - /// to the client via the component's sender channel. - /// - /// # Arguments - /// - /// * `error_response` - The ErrorResponse to send to the client - fn send_error_response(&mut self, err: Error) -> Result<(), Error>; } #[cfg(test)] mod tests { use super::*; - use crate::postgresql::messages::error_response::{ - ErrorResponseCode, CODE_IDLE_SESSION_TIMEOUT, CODE_SYSTEM_ERROR, - }; + use crate::postgresql::diagnostics::{CODE_IDLE_SESSION_TIMEOUT, CODE_SYSTEM_ERROR}; use std::time::Duration; /// Minimal implementation of PostgreSqlErrorHandler for testing the default method. struct TestHandler; impl PostgreSqlErrorHandler for TestHandler { - fn client_sender(&mut self) -> &mut Sender { - unimplemented!("not needed for error_to_response tests") - } - fn client_id(&self) -> i32 { 0 } - - fn send_error_response(&mut self, _err: Error) -> Result<(), Error> { - unimplemented!("not needed for error_to_response tests") - } } - fn error_code(response: &ErrorResponse) -> Option<&str> { + fn field(response: &DiagnosticResponse, code: u8) -> Option<&str> { response .fields .iter() - .find(|f| f.code == ErrorResponseCode::Code) - .map(|f| f.value.as_str()) - } - - fn error_message(response: &ErrorResponse) -> Option<&str> { - response - .fields - .iter() - .find(|f| f.code == ErrorResponseCode::Message) - .map(|f| f.value.as_str()) + .find(|field| field.code == code) + .and_then(|field| std::str::from_utf8(&field.value).ok()) } #[test] @@ -117,9 +97,9 @@ mod tests { duration: Duration::from_millis(5000), }; let response = handler.error_to_response(err); - assert_eq!(error_code(&response), Some(CODE_IDLE_SESSION_TIMEOUT)); + assert_eq!(field(&response, b'C'), Some(CODE_IDLE_SESSION_TIMEOUT)); assert_eq!( - error_message(&response), + field(&response, b'M'), Some("Connection timed out after 5000 ms") ); } @@ -129,6 +109,23 @@ mod tests { let handler = TestHandler; let err = Error::Unknown; let response = handler.error_to_response(err); - assert_eq!(error_code(&response), Some(CODE_SYSTEM_ERROR)); + assert_eq!(field(&response, b'C'), Some(CODE_SYSTEM_ERROR)); + } + + #[test] + fn internal_protocol_error_does_not_expose_implementation_details() { + let handler = TestHandler; + let err = Error::Protocol(ProtocolError::HeldDataRowOperationMismatch); + let response = handler.error_to_response(err); + + assert_eq!(field(&response, b'C'), Some(CODE_SYSTEM_ERROR)); + assert_eq!( + field(&response, b'M'), + Some(concat!( + "CipherStash Proxy encountered an internal PostgreSQL protocol error. ", + "For help visit https://github.com/cipherstash/proxy/blob/main/docs/errors.md", + "#protocol-internal-error" + )) + ); } } diff --git a/packages/cipherstash-proxy/src/postgresql/handler.rs b/packages/cipherstash-proxy/src/postgresql/handler.rs deleted file mode 100644 index c1fe8e50e..000000000 --- a/packages/cipherstash-proxy/src/postgresql/handler.rs +++ /dev/null @@ -1,394 +0,0 @@ -use super::backend::Backend; -use super::frontend::Frontend; -use super::protocol::StartupCode; -use crate::connect::ChannelWriter; -use crate::error::ConfigError; -use crate::log::{AUTHENTICATION, PROTOCOL}; -use crate::postgresql::messages::authentication::auth::{AuthenticationMethod, SaslMechanism}; -use crate::postgresql::messages::authentication::sasl::SASLResponse; -use crate::postgresql::messages::authentication::{ - Authentication, PasswordMessage, SASLInitialResponse, -}; -use crate::postgresql::messages::error_response::ErrorResponse; -use crate::postgresql::{protocol, startup}; -use crate::proxy::ZeroKms; -use crate::{ - connect::AsyncStream, - error::{Error, ProtocolError}, - postgresql::context::Context, - tls, -}; -use bytes::BytesMut; -use md5::{Digest, Md5}; -use postgres_protocol::authentication::sasl::{ChannelBinding, ScramSha256}; -use rand::Rng; -use tokio::io::{AsyncRead, AsyncWrite, AsyncWriteExt}; -use tracing::{debug, error, info, warn}; -/// -/// -/// Entry point for handling postgres protocol connections -/// Each inbound client connection is mapped to a database connection -/// Hilarity ensues -/// -/// Startup flow -/// -/// Connect to database with TLS if required -/// First message is either: -/// - SSLRequest -/// - ProtocolVersionNumber -/// - CancelRequest -/// -/// On SSLRequest -/// Send SSLResponse -/// Connect with TLS if configured -/// -/// On TLS Connect -/// Expect message containing ProtocolVersionNumber is sent -/// -/// On CancelRequest -/// Propagate and disconnect -/// -/// On ProtocolVersionNumber -/// Propagate and continue -/// -/// -pub async fn handler(client_stream: AsyncStream, context: Context) -> Result<(), Error> { - let mut client_stream = client_stream; - let client_id = context.client_id; - - // Connect to the database server, using TLS if configured - let stream = AsyncStream::connect(&context.database_socket_address()).await?; - let mut database_stream = startup::with_tls(stream, context.config()).await?; - info!( - msg = "Client connected", - database = context.database_socket_address(), - client_id = client_id, - ); - - loop { - let startup_message = - match startup::read_message(&mut client_stream, context.connection_timeout()).await { - Ok(msg) => msg, - Err(err @ Error::ConnectionTimeout { .. }) => { - send_timeout_error(&mut client_stream, &err).await; - return Err(err); - } - Err(err) => return Err(err), - }; - - match &startup_message.code { - StartupCode::SSLRequest => { - startup::send_ssl_response(&mut client_stream, context.use_tls()).await?; - if let Some(ref tls) = context.tls_config() { - match client_stream { - AsyncStream::Tcp(stream) => { - // The Client is connecting to our Server - let tls_stream = tls::server(stream, tls).await?; - client_stream = AsyncStream::Tls(Box::new(tls_stream)); - } - AsyncStream::Tls(_) => { - unreachable!(); - } - } - } - } - StartupCode::CancelRequest => { - database_stream.write_all(&startup_message.bytes).await?; - return Err(Error::CancelRequest); - } - StartupCode::ProtocolVersionNumber => { - database_stream.write_all(&startup_message.bytes).await?; - break; - } - } - } - - // Proxy -> Client Authentication - // Uses MD5 - // SASL is not supported because I need to RTFM https://datatracker.ietf.org/doc/html/rfc5802 - // - // Proxy -> Send AuthenticationMD5Password - // Client -> Send PasswordMessage - // - { - let salt = generate_md5_password_salt(); - - let username = context.database_username().as_bytes(); - let password = context.database_password(); - - let password = password.as_bytes(); - - let hash = md5_hash(username, password, &salt); - - let message = Authentication::md5_password(salt); - let bytes = BytesMut::try_from(message)?; - client_stream.write_all(&bytes).await?; - - let connection_timeout = context.connection_timeout(); - let (_code, bytes) = - match protocol::read_message(&mut client_stream, client_id, connection_timeout).await { - Ok(result) => result, - Err(err @ Error::ConnectionTimeout { .. }) => { - send_timeout_error(&mut client_stream, &err).await; - return Err(err); - } - Err(err) => return Err(err), - }; - - let password_message = PasswordMessage::try_from(&bytes)?; - - if hash == password_message.password { - let message = Authentication::authentication_ok(); - debug!(target: AUTHENTICATION, msg = "Client AuthenticationOk"); - let bytes = BytesMut::try_from(message)?; - client_stream.write_all(&bytes).await?; - } else { - let message = ProtocolError::ClientAuthenticationFailed.to_string(); - error!(msg = message); - - let message = ErrorResponse::invalid_password(message); - let bytes = BytesMut::try_from(message)?; - client_stream.write_all(&bytes).await?; - } - } - - // Database authentication flow - // 1. Database -> Authentication message (SASL) - // -> Proxy -> Auth Reponse flow with SASL - // - // 2. Proxy -> Auth message to the client Md5, SASL etc - // -> Client -> Auth response - // - - // First message should always be Auth - let auth = protocol::read_auth_message(&mut database_stream, client_id).await?; - - match &auth.method { - AuthenticationMethod::AuthenticationOk => { - debug!(target: AUTHENTICATION, msg = "AuthenticationOk"); - } - AuthenticationMethod::AuthenticationCleartextPassword => { - debug!(target: AUTHENTICATION, msg = "AuthenticationCleartextPassword"); - let password = context.database_password(); - let message = PasswordMessage::new(password); - let bytes = BytesMut::try_from(message)?; - database_stream.write_all(&bytes).await?; - } - AuthenticationMethod::Md5Password { salt } => { - debug!(target: AUTHENTICATION, msg = "Md5Password"); - let username = context.database_username().as_bytes(); - let password = context.database_password(); - let password = password.as_bytes(); - - let hash = md5_hash(username, password, salt); - let message = PasswordMessage::new(hash); - let bytes = BytesMut::try_from(message)?; - database_stream.write_all(&bytes).await?; - } - AuthenticationMethod::Sasl { .. } => { - debug!(target: AUTHENTICATION, msg = "Sasl"); - let mechanism = auth.sasl_mechanism()?; - sanity_check_sasl_mechanism(&mechanism, &client_stream); - - // Toby: I don't think we need to do anything here - // If we are connected via TLS, we can support SCRAM-SHA-256-PLUS - // If we are not connected via TLS, the database won't ask for SCRAM-SHA-256-PLUS - let channel_binding = database_stream.channel_binding(); - let password = context.database_password(); - let password = password.as_bytes(); - scram_sha_256_plus_handler(&mut database_stream, mechanism, password, channel_binding) - .await?; - } - AuthenticationMethod::Other { method_code, .. } => { - debug!(target: AUTHENTICATION, msg = "UnsupportedAuthentication"); - return Err(ProtocolError::UnsupportedAuthentication { - method_code: *method_code, - } - .into()); - } - method => { - debug!(target: AUTHENTICATION, msg = "UnexpectedStartupMessage", authentication_method = ?method); - return Err(ProtocolError::UnexpectedStartupMessage.into()); - } - } - - if context.require_tls() && !client_stream.is_tls() { - let message = ErrorResponse::tls_required(); - let bytes = BytesMut::try_from(message)?; - client_stream.write_all(&bytes).await?; - - error!(msg = "Client must connect with Transport Layer Security (TLS)"); - return Err(ConfigError::TlsRequired.into()); - } - - let (client_reader, client_writer) = client_stream.split(); - let (server_reader, server_writer) = database_stream.split(); - - let channel_writer = ChannelWriter::new(client_writer, client_id); - - let mut frontend = Frontend::new( - client_reader, - channel_writer.sender(), - server_writer, - context.clone(), - ); - let mut backend = Backend::new(channel_writer.sender(), server_reader, context.clone()); - - if context.is_passthrough() { - if context.use_structured_logging() { - warn!(msg = "RUNNING IN PASSTHROUGH MODE"); - warn!(msg = "DATA IS NOT PROTECTED WITH ENCRYPTION"); - } else { - warn!(msg = "========================================"); - warn!(msg = "RUNNING IN PASSTHROUGH MODE"); - warn!(msg = "DATA IS NOT PROTECTED WITH ENCRYPTION"); - warn!(msg = "========================================"); - } - } - - let timeout_sender = channel_writer.sender(); - let channel_writer_task = tokio::spawn(channel_writer.receive()); - - let client_to_server = async { - loop { - let result = frontend.rewrite().await; - // Ensure the connection is terminated if the client closes the connection - // The client ConnectionClosed error is triggered before the terminate message is passed through - if matches!(result, Err(Error::ConnectionClosed)) { - frontend.terminate().await? - } - result?; - } - // Unreachable, but helps the compiler understand the return type - // TODO: extract into a function or something with type - #[allow(unreachable_code)] - Ok::<(), Error>(()) - }; - - let server_to_client = async { - loop { - backend.rewrite().await?; - } - #[allow(unreachable_code)] - Ok::<(), Error>(()) - }; - - // Run frontend and backend tasks - let result = tokio::try_join!(client_to_server, server_to_client); - - if let Err(ref err @ Error::ConnectionTimeout { .. }) = &result { - let error_response = ErrorResponse::connection_timeout(err.to_string()); - if let Ok(bytes) = BytesMut::try_from(error_response) { - let _ = timeout_sender.send(bytes); - } - // Best-effort yield to allow ChannelWriter to flush the error response - // before the connection tears down. Not guaranteed — if the runtime doesn't - // schedule the writer task before teardown, the client may see a connection - // reset instead of the ErrorResponse. - tokio::task::yield_now().await; - } - - // Drop frontend and backend to drop their senders and close the channel - // The async blocks above captured frontend/backend by reference, so they're still alive - drop(frontend); - drop(backend); - - // Wait for channel writer to finish shutdown sequence - // The senders are now dropped, which closes the channel and allows - // the writer task to complete its shutdown - if let Err(err) = channel_writer_task.await { - error!( - client_id, - msg = "Channel writer task panicked", - error = ?err - ); - } - - result?; - Ok(()) -} - -// Keep for debugging -fn sanity_check_sasl_mechanism(mechanism: &SaslMechanism, client_stream: &AsyncStream) { - match mechanism { - SaslMechanism::ScramSha256 => { - if client_stream.is_tls() { - debug!( - PROTOCOL, - msg = "Database requested SCRAM-SHA-256, but Proxy has a TLS connection" - ); - } - } - SaslMechanism::ScramSha256Plus => { - if client_stream.is_tcp() { - debug!( - PROTOCOL, - msg = "Database requested SCRAM-SHA-256-PLUS, but Proxy has a TCP connection" - ); - } - } - } -} - -pub fn md5_hash(username: &[u8], password: &[u8], salt: &[u8; 4]) -> String { - let mut md5 = Md5::new(); - md5.update(password); - md5.update(username); - let output = md5.finalize_reset(); - md5.update(format!("{output:x}")); - md5.update(salt); - format!("md5{:x}", md5.finalize()) -} - -fn generate_md5_password_salt() -> [u8; 4] { - let mut rng = rand::rng(); - let mut bytes = [0u8; 4]; - rng.fill(&mut bytes); - bytes -} - -async fn scram_sha_256_plus_handler( - mut stream: S, - mechanism: SaslMechanism, - password: &[u8], - channel_binding: ChannelBinding, -) -> Result<(), Error> { - let mut scram = ScramSha256::new(password, channel_binding); - let bytes = scram.message().to_vec(); - - let sasl_initial_response = SASLInitialResponse::new(mechanism, bytes); - let bytes = BytesMut::try_from(sasl_initial_response)?; - stream.write_all(&bytes).await?; - - let auth = protocol::read_auth_message(&mut stream, 1).await?; - - let bytes = auth.sasl_continue()?; - scram.update(bytes)?; - - let sasl_response = SASLResponse::new(scram.message().to_vec()); - - let bytes = BytesMut::try_from(sasl_response)?; - stream.write_all(&bytes).await?; - - let auth = protocol::read_auth_message(&mut stream, 1).await?; - let bytes = auth.sasl_final()?; - scram.finish(bytes)?; - - let auth = protocol::read_auth_message(&mut stream, 1).await?; - - if auth.is_ok() { - debug!(target: AUTHENTICATION, msg = "SASL authentication successful"); - Ok(()) - } else { - Err(ProtocolError::AuthenticationFailed.into()) - } -} - -/// Best-effort send of a connection timeout ErrorResponse directly to a client stream. -/// Used for pre-split timeout sites where no ChannelWriter exists yet. -async fn send_timeout_error(stream: &mut S, err: &Error) { - let error_response = ErrorResponse::connection_timeout(err.to_string()); - if let Ok(bytes) = BytesMut::try_from(error_response) { - let _ = stream.write_all(&bytes).await; - } -} diff --git a/packages/cipherstash-proxy/src/postgresql/message_buffer.rs b/packages/cipherstash-proxy/src/postgresql/message_buffer.rs deleted file mode 100644 index 750b79b2c..000000000 --- a/packages/cipherstash-proxy/src/postgresql/message_buffer.rs +++ /dev/null @@ -1,42 +0,0 @@ -use super::messages::data_row::DataRow; - -pub struct MessageBuffer { - // buffer: RwLock>, - buffer: Vec, -} - -impl MessageBuffer { - /// Default number of rows to keep in the buffer. - /// Larger rows will require more memory. - const DEFAULT_RESPONSE_BUFFER_SIZE: usize = 4096; - - pub fn new() -> Self { - Self { - buffer: Vec::with_capacity(Self::DEFAULT_RESPONSE_BUFFER_SIZE), - } - } - - pub fn push(&mut self, row: DataRow) { - self.buffer.push(row); - } - - pub fn drain(&mut self) -> Vec { - self.buffer.drain(..).collect() - } - - pub fn clear(&mut self) { - self.buffer.clear(); - } - - pub fn len(&self) -> usize { - self.buffer.len() - } - - pub fn is_empty(&self) -> bool { - self.buffer.is_empty() - } - - pub fn at_capacity(&self) -> bool { - self.buffer.len() >= Self::DEFAULT_RESPONSE_BUFFER_SIZE - } -} diff --git a/packages/cipherstash-proxy/src/postgresql/messages/authentication/auth.rs b/packages/cipherstash-proxy/src/postgresql/messages/authentication/auth.rs deleted file mode 100644 index c96f9e119..000000000 --- a/packages/cipherstash-proxy/src/postgresql/messages/authentication/auth.rs +++ /dev/null @@ -1,358 +0,0 @@ -use crate::error::{Error, ProtocolError}; -use crate::postgresql::messages::{BackendCode, FrontendCode}; -use crate::postgresql::protocol::BytesMutReadString; -use crate::SIZE_I32; -use bytes::{Buf, BufMut, BytesMut}; - -use std::convert::TryFrom; -use std::ffi::CString; -use std::fmt::{self, Display, Formatter}; -use std::io::{Cursor, Read}; - -const SIZE_NULL_BYTE: usize = 1; - -pub const SCRAM_SHA_256_PLUS: &str = "SCRAM-SHA-256-PLUS"; -pub const SCRAM_SHA_256: &str = "SCRAM-SHA-256"; - -#[derive(Debug, Clone, Copy, PartialEq)] -pub enum SaslMechanism { - ScramSha256, - ScramSha256Plus, -} - -#[derive(Debug, Clone)] -pub struct Authentication { - #[allow(dead_code)] - code: u8, - pub method: AuthenticationMethod, -} - -#[derive(Clone, Debug)] -#[repr(i32)] -pub enum AuthenticationMethod { - AuthenticationOk = 0, - AuthenticationCleartextPassword = 3, - Md5Password { salt: [u8; 4] } = 5, - Sasl { mechanisms: Vec } = 10, - AuthenticationSASLContinue { bytes: Vec } = 11, - AuthenticationSASLFinal { bytes: Vec } = 12, - Other { method_code: i32, bytes: Vec }, -} - -#[derive(Clone, Debug)] -pub struct PasswordMessage { - code: u8, - pub password: String, -} - -impl Authentication { - pub fn is_ok(&self) -> bool { - matches!(self.method, AuthenticationMethod::AuthenticationOk) - } - - pub fn is_sasl(&self) -> bool { - matches!(self.method, AuthenticationMethod::Sasl { .. }) - } - - pub fn is_scram_sha_256_plus(&self) -> bool { - match self.method { - AuthenticationMethod::Sasl { ref mechanisms } => { - mechanisms.contains(&SaslMechanism::ScramSha256Plus) - } - _ => false, - } - } - - /// - /// Returns the first mechanism in the list of mechanisms - /// If the method is not SASL, it will return an error - /// If the method is SASL and there are no mechanisms, it will return an error - /// - This should never happen as the server should always return at least one mechanism - /// - If it does, it is a protocol error and the message parse should already have returned an error - /// - This is a safety check to ensure that the server is behaving as expected - /// - pub fn sasl_mechanism(&self) -> Result { - let mechanism = match self.method { - AuthenticationMethod::Sasl { ref mechanisms } => mechanisms.first(), - _ => None, - }; - - match mechanism { - Some(m) => Ok(*m), - None => { - Err(ProtocolError::UnexpectedSaslAuthenticationMethod("None".to_string()).into()) - } - } - } - - pub fn sasl_continue(&self) -> Result<&Vec, Error> { - match self.method { - AuthenticationMethod::AuthenticationSASLContinue { ref bytes } => Ok(bytes), - _ => Err(ProtocolError::UnexpectedAuthenticationResponse { - expected: "SASLContinue".into(), - received: (&self.method).into(), - } - .into()), - } - } - - pub fn sasl_final(&self) -> Result<&Vec, Error> { - match self.method { - AuthenticationMethod::AuthenticationSASLFinal { ref bytes } => Ok(bytes), - _ => Err(ProtocolError::UnexpectedAuthenticationResponse { - expected: "SASLFinal".into(), - received: (&self.method).into(), - } - .into()), - } - } - - pub fn md5_password(salt: [u8; 4]) -> Authentication { - Authentication { - code: BackendCode::Authentication.into(), - method: AuthenticationMethod::Md5Password { salt }, - } - } - - pub fn authentication_ok() -> Authentication { - Authentication { - code: BackendCode::Authentication.into(), - method: AuthenticationMethod::AuthenticationOk, - } - } -} - -impl PasswordMessage { - pub fn new(password: String) -> PasswordMessage { - PasswordMessage { - code: FrontendCode::PasswordMessage.into(), - password, - } - } -} - -impl TryFrom<&BytesMut> for Authentication { - type Error = Error; - - fn try_from(bytes: &BytesMut) -> Result { - let mut cursor = Cursor::new(bytes); - let code = cursor.get_u8(); - - if BackendCode::from(code) != BackendCode::Authentication { - return Err(ProtocolError::UnexpectedMessageCode { - expected: BackendCode::Authentication.into(), - received: code as char, - } - .into()); - } - - let len = cursor.get_i32(); // read and progress cursor - let method_code = cursor.get_i32(); - - let method = match method_code { - 0 => AuthenticationMethod::AuthenticationOk, - 5 => { - let mut salt = [0; 4]; - cursor.read_exact(&mut salt)?; - AuthenticationMethod::Md5Password { salt } - } - 10 => { - let mut mechanisms = Vec::new(); - let mut count = SIZE_I32 // message len - + SIZE_I32 // method_code - + SIZE_NULL_BYTE; // terminating null byte; - - while count < (len as usize) { - let m = cursor.read_string()?; - count += m.len() + SIZE_NULL_BYTE; - mechanisms.push(SaslMechanism::try_from(m)?); - } - AuthenticationMethod::Sasl { mechanisms } - } - 11 => { - let mut bytes = Vec::new(); - cursor.read_to_end(&mut bytes)?; - AuthenticationMethod::AuthenticationSASLContinue { bytes } - } - 12 => { - let mut bytes = Vec::new(); - cursor.read_to_end(&mut bytes)?; - AuthenticationMethod::AuthenticationSASLFinal { bytes } - } - _ => { - // Get any remaining bytes from the cursor - let mut bytes = Vec::new(); - cursor.read_to_end(&mut bytes)?; - AuthenticationMethod::Other { method_code, bytes } - } - }; - - Ok(Authentication { code, method }) - } -} - -impl TryFrom for BytesMut { - type Error = Error; - - fn try_from(auth: Authentication) -> Result { - let mut method_bytes = BytesMut::new(); - - let method_code = (&auth.method).into(); - method_bytes.put_i32(method_code); - - match auth.method { - AuthenticationMethod::AuthenticationOk => {} - AuthenticationMethod::AuthenticationCleartextPassword => {} - AuthenticationMethod::Md5Password { salt } => { - method_bytes.put_slice(&salt); - } - AuthenticationMethod::Sasl { mechanisms } => { - for m in mechanisms { - let s = m.to_string(); - let c = CString::new(s)?; - let s = c.as_bytes_with_nul(); - method_bytes.put_slice(s); - } - method_bytes.put_u8(0); // null byte - } - AuthenticationMethod::AuthenticationSASLContinue { bytes, .. } => { - method_bytes.put_slice(&bytes); - } - AuthenticationMethod::AuthenticationSASLFinal { bytes, .. } => { - method_bytes.put_slice(&bytes); - } - AuthenticationMethod::Other { bytes, .. } => { - method_bytes.put_slice(&bytes); - } - } - - let mut bytes = BytesMut::new(); - - let len = SIZE_I32 // message len - + method_bytes.len(); // method_code - - bytes.put_u8(BackendCode::Authentication.into()); - bytes.put_i32(len as i32); - bytes.put_slice(&method_bytes); - - Ok(bytes) - } -} - -impl TryFrom<&BytesMut> for PasswordMessage { - type Error = Error; - - fn try_from(bytes: &BytesMut) -> Result { - let mut cursor = Cursor::new(bytes); - let code = cursor.get_u8(); - - if FrontendCode::from(code) != FrontendCode::PasswordMessage { - return Err(ProtocolError::UnexpectedMessageCode { - expected: FrontendCode::PasswordMessage.into(), - received: code as char, - } - .into()); - } - - let _len = cursor.get_i32(); // read and progress cursor - let password = cursor.read_string()?; - - Ok(PasswordMessage { code, password }) - } -} - -impl TryFrom for BytesMut { - type Error = Error; - - fn try_from(password_message: PasswordMessage) -> Result { - let mut bytes = BytesMut::new(); - - let password = CString::new(password_message.password)?; - let password = password.as_bytes_with_nul(); - - let len = SIZE_I32 // message len - + password.len(); // password - - bytes.put_u8(FrontendCode::PasswordMessage.into()); - bytes.put_i32(len as i32); - bytes.put_slice(password); - - Ok(bytes) - } -} - -impl TryFrom for SaslMechanism { - type Error = Error; - fn try_from(s: String) -> Result { - match s.as_str() { - SCRAM_SHA_256 => Ok(SaslMechanism::ScramSha256), - SCRAM_SHA_256_PLUS => Ok(SaslMechanism::ScramSha256Plus), - s => Err(ProtocolError::UnexpectedSaslAuthenticationMethod(s.to_owned()).into()), - } - } -} - -impl Display for SaslMechanism { - fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result { - let s = match self { - SaslMechanism::ScramSha256 => SCRAM_SHA_256.to_owned(), - SaslMechanism::ScramSha256Plus => SCRAM_SHA_256_PLUS.to_owned(), - }; - write!(f, "{s}") - } -} - -impl From<&AuthenticationMethod> for i32 { - fn from(method: &AuthenticationMethod) -> Self { - match method { - AuthenticationMethod::AuthenticationOk => 0, - AuthenticationMethod::AuthenticationCleartextPassword => 3, - AuthenticationMethod::Md5Password { .. } => 5, - AuthenticationMethod::Sasl { .. } => 10, - AuthenticationMethod::AuthenticationSASLContinue { .. } => 11, - AuthenticationMethod::AuthenticationSASLFinal { .. } => 12, - AuthenticationMethod::Other { method_code, .. } => *method_code, - } - } -} - -#[cfg(test)] -mod tests { - use bytes::BytesMut; - - use crate::{config::LogConfig, log}; - - use super::Authentication; - - fn to_message(s: &[u8]) -> BytesMut { - BytesMut::from(s) - } - - #[test] - pub fn parse_auth_message() { - log::init(LogConfig::default()); - - let bytes = to_message(b"R\0\0\0*\0\0\0\nSCRAM-SHA-256-PLUS\0SCRAM-SHA-256\0\0"); - - let auth = Authentication::try_from(&bytes).unwrap(); - - assert!(matches!( - auth.method, - super::AuthenticationMethod::Sasl { .. } - )); - - let auth_bytes = BytesMut::try_from(auth).unwrap(); - - assert_eq!(bytes, auth_bytes); - } - - #[test] - pub fn is_scram_sha_256_plus() { - log::init(LogConfig::default()); - - let bytes = to_message(b"R\0\0\0*\0\0\0\nSCRAM-SHA-256-PLUS\0SCRAM-SHA-256\0\0"); - let auth = Authentication::try_from(&bytes).unwrap(); - - assert!(auth.is_scram_sha_256_plus()); - } -} diff --git a/packages/cipherstash-proxy/src/postgresql/messages/authentication/mod.rs b/packages/cipherstash-proxy/src/postgresql/messages/authentication/mod.rs deleted file mode 100644 index 40ca7809a..000000000 --- a/packages/cipherstash-proxy/src/postgresql/messages/authentication/mod.rs +++ /dev/null @@ -1,6 +0,0 @@ -pub mod auth; -pub mod sasl; - -pub use auth::Authentication; -pub use auth::PasswordMessage; -pub use sasl::SASLInitialResponse; diff --git a/packages/cipherstash-proxy/src/postgresql/messages/authentication/sasl.rs b/packages/cipherstash-proxy/src/postgresql/messages/authentication/sasl.rs deleted file mode 100644 index b5806619f..000000000 --- a/packages/cipherstash-proxy/src/postgresql/messages/authentication/sasl.rs +++ /dev/null @@ -1,133 +0,0 @@ -use std::{ - ffi::CString, - io::{Cursor, Read}, -}; - -use bytes::{Buf, BufMut, BytesMut}; - -use crate::{ - error::{Error, ProtocolError}, - postgresql::{messages::FrontendCode, protocol::BytesMutReadString}, - SIZE_I32, -}; - -use super::auth::{self, SaslMechanism}; - -#[derive(Clone, Debug)] -pub struct SASLInitialResponse { - #[allow(dead_code)] - code: u8, - pub mechanism: String, - pub response: Vec, -} - -#[derive(Clone, Debug)] -pub struct SASLResponse { - #[allow(dead_code)] - code: u8, - response: Vec, -} - -impl SASLInitialResponse { - pub fn new(mechanism: SaslMechanism, response: Vec) -> Self { - let mechanism = mechanism.to_string(); - - SASLInitialResponse { - code: FrontendCode::SASLInitialResponse.into(), - mechanism, - - response, - } - } - - pub fn is_scram_sha_256(&self) -> bool { - self.mechanism == auth::SCRAM_SHA_256 - } - - pub fn is_scram_sha_256_plus(&self) -> bool { - self.mechanism == auth::SCRAM_SHA_256_PLUS - } -} - -impl SASLResponse { - pub fn new(response: Vec) -> Self { - SASLResponse { - code: FrontendCode::SASLResponse.into(), - response, - } - } -} - -impl TryFrom<&BytesMut> for SASLInitialResponse { - type Error = Error; - - fn try_from(bytes: &BytesMut) -> Result { - let mut cursor = Cursor::new(bytes); - let code = cursor.get_u8(); - - // Note: all password messages use the 'p' code - if code != b'p' { - return Err(ProtocolError::UnexpectedMessageCode { - expected: FrontendCode::SASLInitialResponse.into(), - received: code as char, - } - .into()); - } - let _len = cursor.get_i32(); - let mechanism = cursor.read_string()?; - let _response_len = cursor.get_i32(); - let mut bytes = Vec::new(); - cursor.read_to_end(&mut bytes)?; - - Ok(SASLInitialResponse { - code, - mechanism, - response: bytes, - }) - } -} - -impl TryFrom for BytesMut { - type Error = Error; - - fn try_from(response: SASLInitialResponse) -> Result { - let mut bytes = BytesMut::new(); - - let mechanism = CString::new(response.mechanism)?; - let mechanism = mechanism.as_bytes_with_nul(); - - let response_len = response.response.len(); - - let len = SIZE_I32 // len length - + mechanism.len() - + SIZE_I32 // response_len - + response_len; - - bytes.put_u8(FrontendCode::SASLInitialResponse.into()); - bytes.put_i32(len as i32); - bytes.put_slice(mechanism); - bytes.put_i32(response_len as i32); - bytes.put_slice(&response.response); - - Ok(bytes) - } -} - -impl TryFrom for BytesMut { - type Error = Error; - - fn try_from(response: SASLResponse) -> Result { - let mut bytes = BytesMut::new(); - - let response_len = response.response.len(); - - let len = SIZE_I32 // len length - + response_len; - - bytes.put_u8(FrontendCode::SASLResponse.into()); - bytes.put_i32(len as i32); - bytes.put_slice(&response.response); - - Ok(bytes) - } -} diff --git a/packages/cipherstash-proxy/src/postgresql/messages/close.rs b/packages/cipherstash-proxy/src/postgresql/messages/close.rs deleted file mode 100644 index 1a8fb11d2..000000000 --- a/packages/cipherstash-proxy/src/postgresql/messages/close.rs +++ /dev/null @@ -1,144 +0,0 @@ -use crate::error::{Error, ProtocolError}; -use crate::postgresql::protocol::BytesMutReadString; -use crate::{SIZE_I32, SIZE_U8}; - -use bytes::{Buf, BufMut, BytesMut}; -use std::convert::TryFrom; -use std::ffi::CString; -use std::io::Cursor; - -use super::target::Target; -use super::{FrontendCode, Name}; - -/// -/// Close b'C' (Frontend) message. -/// -/// See: -/// -/// Byte1('C') -/// Identifies the message as a Close command. -/// -/// Int32 -/// Length of message contents in bytes, including self. -/// -/// Byte1 -/// 'S' to close a prepared statement; or 'P' to close a portal. -/// -/// String -/// The name of the prepared statement or portal to close (an empty string selects the unnamed prepared statement or portal). - -#[derive(Debug, Clone)] -pub(crate) struct Close { - pub target: Target, - pub name: Name, -} - -impl TryFrom<&BytesMut> for Close { - type Error = Error; - - fn try_from(bytes: &BytesMut) -> Result { - let mut cursor = Cursor::new(bytes); - let code = cursor.get_u8(); - - if FrontendCode::from(code) != FrontendCode::Close { - return Err(ProtocolError::UnexpectedMessageCode { - expected: FrontendCode::Close.into(), - received: code as char, - } - .into()); - } - - let _len = cursor.get_i32(); // read and progress cursor - let target = cursor.get_u8(); - let target = Target::try_from(target)?; - let name = cursor.read_string()?; - let name = Name::from(name); - - Ok(Close { target, name }) - } -} - -impl TryFrom for BytesMut { - type Error = Error; - - fn try_from(close: Close) -> Result { - let mut bytes = BytesMut::new(); - - let name = CString::new(close.name.as_str())?; - let name = name.as_bytes_with_nul(); - - let len = SIZE_I32 + SIZE_U8 + name.len(); - - bytes.put_u8(FrontendCode::Close.into()); - bytes.put_i32(len as i32); - bytes.put_u8(close.target.into()); - bytes.put_slice(name); - - Ok(bytes) - } -} - -#[cfg(test)] -mod tests { - use super::*; - use crate::{config::LogConfig, log, postgresql::messages::Name}; - use bytes::BytesMut; - use std::convert::TryFrom; - - fn to_message(s: &[u8]) -> BytesMut { - BytesMut::from(s) - } - - #[test] - pub fn test_close_statement() { - log::init(LogConfig::default()); - - // Close unnamed prepared statement: C\0\0\0\x06S\0 - let bytes = to_message(b"C\0\0\0\x06S\0"); - let close = Close::try_from(&bytes).unwrap(); - - assert!(matches!(close.target, Target::Statement)); - assert!(close.name.is_unnamed()); - } - - #[test] - pub fn test_close_portal() { - log::init(LogConfig::default()); - - // Close unnamed portal: C\0\0\0\x06P\0 - let bytes = to_message(b"C\0\0\0\x06P\0"); - let close = Close::try_from(&bytes).unwrap(); - - assert!(matches!(close.target, Target::Portal)); - assert!(close.name.is_unnamed()); - } - - #[test] - pub fn test_close_named_statement() { - log::init(LogConfig::default()); - - // Close named prepared statement "stmt1": C\0\0\0\x0bSstmt1\0 - let bytes = to_message(b"C\0\0\0\x0bSstmt1\0"); - let close = Close::try_from(&bytes).unwrap(); - - assert!(matches!(close.target, Target::Statement)); - assert_eq!(close.name.as_str(), "stmt1"); - assert!(!close.name.is_unnamed()); - } - - #[test] - pub fn test_close_to_bytes() { - log::init(LogConfig::default()); - - let close = Close { - target: Target::Portal, - name: Name::from("portal1"), - }; - - let bytes = BytesMut::try_from(close).unwrap(); - let parsed = Close::try_from(&bytes).unwrap(); - - assert!(matches!(parsed.target, Target::Portal)); - assert_eq!(parsed.name.as_str(), "portal1"); - } -} diff --git a/packages/cipherstash-proxy/src/postgresql/messages/data_row.rs b/packages/cipherstash-proxy/src/postgresql/messages/data_row.rs deleted file mode 100644 index 96d5cb8af..000000000 --- a/packages/cipherstash-proxy/src/postgresql/messages/data_row.rs +++ /dev/null @@ -1,470 +0,0 @@ -use super::{BackendCode, NULL}; -use crate::EqlCiphertext; -use crate::{ - error::{EncryptError, Error, ProtocolError}, - log::DECRYPT, - postgresql::Column, -}; -use bytes::{Buf, BufMut, BytesMut}; -use std::io::Cursor; -use tracing::{debug, error}; - -/// Leading byte of `jsonb`'s binary wire format. PostgreSQL has only ever -/// emitted version 1. -const JSONB_BINARY_VERSION: u8 = 1; - -#[derive(Debug, Clone)] -pub struct DataRow { - pub columns: Vec, -} - -#[derive(Debug, Clone)] -pub struct DataColumn { - bytes: Option, -} - -impl DataRow { - pub fn as_ciphertext( - &mut self, - column_configuration: &Vec>, - ) -> Vec> { - let mut result = vec![]; - for (data_column, column_config) in self.columns.iter_mut().zip(column_configuration) { - let encrypted = column_config - .as_ref() - .filter(|_| data_column.is_not_null()) - .and_then(|config| { - data_column - .to_eql_ciphertext() - .inspect_err(|err| match err { - Error::Encrypt(EncryptError::ColumnIsNull) => { - debug!(target: DECRYPT, msg ="ColumnIsNull", ?config); - // Not an error, as you were - data_column.set_null(); - } - _ => { - let err = EncryptError::ColumnCouldNotBeDeserialised { - table: config.identifier.table.to_owned(), - column: config.identifier.column.to_owned(), - }; - error!(target: DECRYPT, msg = err.to_string()); - } - }) - .ok() - }); - result.push(encrypted); - } - - result - } - - pub fn column_count(&self) -> usize { - self.columns.len() - } - - fn len_of_columns(&self) -> usize { - let column_len_size = size_of::(); // len of column len - - self.columns - .iter() - .map(|col| column_len_size + col.bytes.as_ref().map(|b| b.len()).unwrap_or(0)) - .sum() - } - - pub fn rewrite(&mut self, plaintexts: &[Option]) -> Result<(), Error> { - for (idx, pt) in plaintexts.iter().enumerate() { - if let Some(bytes) = pt { - self.columns[idx].rewrite(bytes); - } - } - Ok(()) - } -} - -impl DataColumn { - pub fn is_not_null(&self) -> bool { - self.bytes.is_some() - } - - pub fn set_null(&mut self) { - self.bytes = None; - } - - pub fn rewrite(&mut self, b: &[u8]) { - if let Some(ref mut bytes) = self.bytes { - bytes.clear(); - bytes.extend_from_slice(b); - } - } -} - -impl TryFrom<&BytesMut> for DataRow { - type Error = Error; - - fn try_from(buf: &BytesMut) -> Result { - let mut cursor = Cursor::new(buf); - - let code = cursor.get_u8(); - - if BackendCode::from(code) != BackendCode::DataRow { - return Err(ProtocolError::UnexpectedMessageCode { - expected: BackendCode::DataRow.into(), - received: code as char, - } - .into()); - } - - let _len = cursor.get_i32(); - - let num_columns = cursor.get_i16(); - - let mut columns = Vec::new(); - for _ in 0..num_columns { - let len = cursor.get_i32(); - - if len == NULL { - columns.push(DataColumn { bytes: None }); - } else { - let len = len as usize; - - let mut bytes = BytesMut::with_capacity(len); - bytes.resize(len, 0); - cursor.copy_to_slice(&mut bytes); - - columns.push(DataColumn { bytes: Some(bytes) }); - } - } - - Ok(DataRow { columns }) - } -} - -impl TryFrom for BytesMut { - type Error = Error; - - fn try_from(data_row: DataRow) -> Result { - let mut bytes = BytesMut::new(); - - let len = size_of::() // len of len - + size_of::() // num columns - + data_row.len_of_columns(); // len data columns - - bytes.put_u8(BackendCode::DataRow.into()); - bytes.put_i32(len as i32); - bytes.put_i16(data_row.columns.len() as i16); - - for col in data_row.columns.into_iter() { - let b = BytesMut::try_from(col)?; - bytes.put_slice(&b); - } - - Ok(bytes) - } -} - -impl TryFrom for BytesMut { - type Error = Error; - - fn try_from(data_column: DataColumn) -> Result { - let mut bytes = BytesMut::new(); - - if let Some(data) = data_column.bytes { - bytes.put_i32(data.len() as i32); - bytes.put_slice(&data); - } else { - bytes.put_i32(NULL); - } - - Ok(bytes) - } -} - -impl DataColumn { - /// Parse this column's bytes into an [`EqlCiphertext`]. - /// - /// EQL v3 column types (`eql_v3_text_eq`, `eql_v3_integer_ord`, …) are - /// DOMAINS over `jsonb`, so a value arrives with jsonb's representation. - /// - /// EQL v2's `eql_v2_encrypted` was a composite type, which is why this - /// used to strip a `("…")` wrapper in text and a 12-byte rowtype header - /// in binary. Neither exists any more — a domain is wire-identical to its - /// base type. - /// - /// text — the JSON object itself, no wrapper and no doubled quotes - /// binary — a 1-byte jsonb version header followed by the JSON text - /// - /// The two are told apart by the leading byte: the version header is - /// `0x01`, and JSON text for an EQL payload always starts with `{`. - /// - /// The JSON is usually a self-describing payload — a scalar `{v,i,c,…}` or - /// a SteVec document `{v,k:"sv",i,h,sv}` — and deserialises directly. The - /// exception is a JSON field access (`eql_v3."->"(…)` / - /// `eql_v3.jsonb_path_query(…)`), whose result is a single - /// `eql_v3_json_entry` (`{v,i,h,s,c,op}`) — one SteVec entry merged with - /// its document envelope. That has a `c`, so it would masquerade as a - /// scalar `Encrypted` payload, but its `c` is an *entry* ciphertext that - /// only decrypts with the entry's selector-derived nonce. So when the - /// payload is a bare entry (see [`is_json_entry`]) it is reshaped into a - /// one-entry SteVec document (see [`json_entry_into_ste_vec_document`]) and - /// the ordinary SteVec decrypt path recovers the field value. - fn to_eql_ciphertext(&self) -> Result { - let Some(bytes) = &self.bytes else { - return Err(EncryptError::ColumnCouldNotBeParsed.into()); - }; - - let json = match bytes.first() { - Some(&JSONB_BINARY_VERSION) => &bytes[1..], - Some(_) => &bytes[..], - None => return Err(EncryptError::ColumnCouldNotBeParsed.into()), - }; - - let mut value: serde_json::Value = - serde_json::from_slice(json).map_err(log_deserialise_error)?; - - if is_json_entry(&value) { - json_entry_into_ste_vec_document(&mut value)?; - } - - serde_json::from_value(value).map_err(log_deserialise_error) - } -} - -/// Whether a decoded EQL payload is a bare `eql_v3_json_entry` — the result of -/// a JSON field access (`eql_v3."->"(…)` / `eql_v3.jsonb_path_query(…)`). -/// -/// A root-level selector `s` is the tell: a scalar `Encrypted` payload has no -/// selector at all, and a SteVec document carries selectors only inside its -/// `sv[]` entries, never at the root. -fn is_json_entry(value: &serde_json::Value) -> bool { - value.get("s").is_some() -} - -/// Reshape a single `eql_v3_json_entry` into a one-entry SteVec document. -/// -/// The entry is `{v,i,h,s,c,op}`: document-envelope fields (`v`, `i`, `h`) -/// alongside one SteVec entry's fields (`s`, `c`, the optional array marker -/// `a`, and the optional ordering term `op`). Move the entry fields under -/// `sv:[{…}]` and tag the object as a SteVec (`k:"sv"`), yielding -/// `{v,k:"sv",i,h,sv:[{s,c,a?,op?}]}` — the shape an [`EqlCiphertext`] SteVec -/// document deserialises from and the decrypt path knows how to open. -fn json_entry_into_ste_vec_document(value: &mut serde_json::Value) -> Result<(), Error> { - use serde_json::Value; - - let object = value - .as_object_mut() - .ok_or(EncryptError::ColumnCouldNotBeParsed)?; - - let mut entry = serde_json::Map::new(); - for key in ["s", "c", "a", "op"] { - if let Some(field) = object.remove(key) { - entry.insert(key.to_owned(), field); - } - } - - object.insert("k".to_owned(), Value::String("sv".to_owned())); - object.insert("sv".to_owned(), Value::Array(vec![Value::Object(entry)])); - - Ok(()) -} - -fn log_deserialise_error(err: serde_json::Error) -> Error { - debug!(target: DECRYPT, error = err.to_string()); - err.into() -} - -#[cfg(test)] -mod tests { - use super::DataRow; - use crate::Identifier; - use crate::{ - config::{LogConfig, LogLevel}, - log, - postgresql::{messages::data_row::DataColumn, Column}, - }; - use bytes::BytesMut; - use cipherstash_client::schema::{ColumnConfig, ColumnType}; - - fn to_message(s: &[u8]) -> BytesMut { - BytesMut::from(s) - } - - fn column_config(column: &str) -> Option { - let identifier = Identifier::new("encrypted", column); - let config = ColumnConfig::build("column".to_string()).casts_as(ColumnType::SmallInt); - let column = Column::new(identifier, config, None, eql_mapper::EqlTermVariant::Full); - Some(column) - } - - fn column_config_with_id(column: &str) -> Vec> { - vec![None, column_config(column)] - } - - // The four `to_ciphertext_*` fixtures below are REAL EQL v3 wire captures - // taken from Postgres -> Proxy `DataRow` messages for the `encrypted` test - // table, via a live encrypt round-trip against ZeroKMS. They exercise - // `DataRow::try_from` + `as_ciphertext` across the binary (jsonb `0x01` - // version header) and text (bare JSON) wire encodings, and NULL columns. - // - // Captured against EQL v3.0.2. The build has since moved to the version - // pinned by `CS_EQL_VERSION` in `mise.toml`, and these still pass — the - // shapes under test (the jsonb version header, the bare-JSON text form, and - // the payload's `i`/`v` fields) have not changed. Regenerate against the - // pinned version, not against 3.0.2, if a future release does change them. - #[test] - pub fn to_ciphertext_with_binary_encoding() { - log::init(LogConfig::with_level(LogLevel::Debug)); - - // `SELECT encrypted_text FROM encrypted WHERE id = $1` (extended/binary): - // the jsonb column arrives as `0x01` + the v3 EqlCiphertextV3 JSON. - let bytes = to_message(b"D\x00\x00\x03\x16\x00\x01\x00\x00\x03\x0c\x01{\"c\": \"mBbL3gJuL?E})+>NeOq5<7N279rs9aRhBwjz3>wOdg{d64myql`6cXIurM_?B|pR<+M8(SeOLoLt~axenSv%=hCOb&m`FC5F;fS-ykq76u4Qgxa(QrcWn^D;Wq5SN5EJ90LtnW_NroxKJj=JLK>\", \"i\": {\"c\": \"encrypted_text\", \"t\": \"encrypted\"}, \"v\": 3, \"bf\": [1512, 1681, 836, 288, 1837, 1131, 415, 1430, 60, 812, 1990, 1211, 1368, 343, 1473, 1980, 598, 1549, 457, 1389, 1557, 941, 494, 1009, 1604, 1033, 2046, 222, 2012, 671, 7, 1525, 265, 901, 743, 543, 1771, 1149, 890, 755, 1974, 1960, 387, 1947, 1298, 130, 1758, 1060, 268, 844, 1375, 746, 1251, 2040], \"hm\": \"96aeaf9852416229d6b33ceb018d9abc90d70cbe7632539d69ef1462c9aa86a0\", \"op\": \"00bf0281ccb68cc6fe496bb1c8277e3484f6392517d5b8425536af7ec00ad7cc40e17e6336568ac4ed98dd659f7581f8a113fe5669b89833d9dd8eadc587a8950b6bd94f872e7f4205a6859e071df47134d3cccf1e53295417\"}"); - let mut data_row = DataRow::try_from(&bytes).unwrap(); - - let column_config = vec![column_config("encrypted_text")]; - let encrypted = data_row.as_ciphertext(&column_config); - - assert_eq!(encrypted.len(), 1); - assert!(encrypted[0].is_some()); - assert_eq!( - &column_config[0].as_ref().unwrap().identifier, - encrypted[0].as_ref().unwrap().identifier() - ); - } - - #[test] - pub fn to_ciphertext_with_binary_encoding_and_null() { - log::init(LogConfig::with_level(LogLevel::Debug)); - - // `SELECT encrypted_text, encrypted_bool FROM encrypted WHERE id = $1` - // (binary), encrypted_text set, encrypted_bool NULL. - let bytes = to_message(b"D\x00\x00\x03\x1a\x00\x02\x00\x00\x03\x0c\x01{\"c\": \"mBbL3gJuL?E})+>NeOq5<7N279rs9aRhBwjz3>wOdg{d64myql`6cXIurM_?B|pR<+M8(SeOLoLt~axenSv%=hCOb&m`FC5F;fS-ykq76u4Qgxa(QrcWn^D;Wq5SN5EJ90LtnW_NroxKJj=JLK>\", \"i\": {\"c\": \"encrypted_text\", \"t\": \"encrypted\"}, \"v\": 3, \"bf\": [1512, 1681, 836, 288, 1837, 1131, 415, 1430, 60, 812, 1990, 1211, 1368, 343, 1473, 1980, 598, 1549, 457, 1389, 1557, 941, 494, 1009, 1604, 1033, 2046, 222, 2012, 671, 7, 1525, 265, 901, 743, 543, 1771, 1149, 890, 755, 1974, 1960, 387, 1947, 1298, 130, 1758, 1060, 268, 844, 1375, 746, 1251, 2040], \"hm\": \"96aeaf9852416229d6b33ceb018d9abc90d70cbe7632539d69ef1462c9aa86a0\", \"op\": \"00bf0281ccb68cc6fe496bb1c8277e3484f6392517d5b8425536af7ec00ad7cc40e17e6336568ac4ed98dd659f7581f8a113fe5669b89833d9dd8eadc587a8950b6bd94f872e7f4205a6859e071df47134d3cccf1e53295417\"}\xff\xff\xff\xff"); - let mut data_row = DataRow::try_from(&bytes).unwrap(); - - let column_config = vec![ - column_config("encrypted_text"), - column_config("encrypted_bool"), - ]; - let encrypted = data_row.as_ciphertext(&column_config); - - assert_eq!(encrypted.len(), 2); - assert!(encrypted[0].is_some()); - assert!(encrypted[1].is_none()); - } - - #[test] - pub fn to_ciphertext_with_text_encoding() { - log::init(LogConfig::with_level(LogLevel::Debug)); - - // `SELECT encrypted_jsonb FROM encrypted WHERE id = 2` (simple/text): the - // jsonb column arrives as bare JSON text, no version header. - let bytes = to_message(b"D\x00\x00\x027\x00\x01\x00\x00\x02-{\"h\": \"l*AC8+7wO)sD**%APm>F3Bc9FAg#FNCmyISKh%bW{NbL}o`gZpBwFD}ye0IoZJ}<8La$|RV{&NeOq5<7N279rs9aRhBwjz3>wOdg{d64myql`6cXIurM_?B|pR<+M8(SeOLoLt~axenSv%=hCOb&m`FC5F;fS-ykq76u4Qgxa(QrcWn^D;Wq5SN5EJ90LtnW_NroxKJj=JLK>\", \"i\": {\"c\": \"encrypted_text\", \"t\": \"encrypted\"}, \"v\": 3, \"bf\": [1512, 1681, 836, 288, 1837, 1131, 415, 1430, 60, 812, 1990, 1211, 1368, 343, 1473, 1980, 598, 1549, 457, 1389, 1557, 941, 494, 1009, 1604, 1033, 2046, 222, 2012, 671, 7, 1525, 265, 901, 743, 543, 1771, 1149, 890, 755, 1974, 1960, 387, 1947, 1298, 130, 1758, 1060, 268, 844, 1375, 746, 1251, 2040], \"hm\": \"96aeaf9852416229d6b33ceb018d9abc90d70cbe7632539d69ef1462c9aa86a0\", \"op\": \"00bf0281ccb68cc6fe496bb1c8277e3484f6392517d5b8425536af7ec00ad7cc40e17e6336568ac4ed98dd659f7581f8a113fe5669b89833d9dd8eadc587a8950b6bd94f872e7f4205a6859e071df47134d3cccf1e53295417\"}\xff\xff\xff\xff"); - let mut data_row = DataRow::try_from(&bytes).unwrap(); - - let column_config = vec![ - column_config("encrypted_text"), - column_config("encrypted_bool"), - ]; - let encrypted = data_row.as_ciphertext(&column_config); - - assert_eq!(encrypted.len(), 2); - assert!(encrypted[0].is_some()); - assert!(encrypted[1].is_none()); - } - - #[test] - pub fn parse_data_row() { - log::init(LogConfig::with_level(LogLevel::Debug)); - - let messages = vec![ - to_message(b"D\0\0\0\x0e\0\x01\0\0\0\x04\0\0\x1e\xa2"), - // SELECT encrypted_jsonb FROM encrypted LIMIT 1 - to_message(b"D\0\0\x03\xba\0\x01\0\0\x03\xb0(\"{\"\"b\"\": null, \"\"c\"\": \"\"mBbLR(BvRN1BF^PAFs!B^`U;mA>uOUiFLgDpZXhU#s#%c4wyi&Z7`(d0IxUty-cI#Yp%o~QFF39^sRf>4*EG{zlk;}ArEQ}NQHa9@;T73aPOSTpuh\"\", \"\"i\"\": {\"\"c\"\": \"\"encrypted_jsonb\"\", \"\"t\"\": \"\"encrypted\"\"}, \"\"m\"\": null, \"\"o\"\": null, \"\"s\"\": null, \"\"u\"\": null, \"\"v\"\": 1, \"\"sv\"\": [{\"\"b\"\": \"\"8067db44a848ab32c3056a3dbe4edf16\"\", \"\"c\"\": \"\"mBbLR(BvRN1BF^PAFs!B^`U;mA>uOUiFLgDpZXhU#s#%c4wyi&Z7`(d0IxUty-cI#Yp%o~QFF39^sRf>4*EG{zlk;}ArEQ}NQHa9@;T73aPOSTpuh\"\", \"\"m\"\": null, \"\"o\"\": null, \"\"s\"\": \"\"9493d6010fe7845d52149b697729c745\"\", \"\"u\"\": null, \"\"sv\"\": null, \"\"ocf\"\": null, \"\"ocv\"\": null}, {\"\"b\"\": null, \"\"c\"\": \"\"mBbLR(BvRN1BF^PAFs!B^`U;m8QkTKr|h>Q`^NbW(CC|>SD}UM=o%mz(Fw#LQFF39^sRf>4*EG{zlk;}ArEQ}NQHa9@;T73aPOSTpuh\"\", \"\"m\"\": null, \"\"o\"\": null, \"\"s\"\": \"\"b1f0e4bb3855bc33936ef1fddf532765\"\", \"\"u\"\": null, \"\"sv\"\": null, \"\"ocf\"\": null, \"\"ocv\"\": \"\"fbc7a11fc81f2a31c904c5b05572b054824e3b5f5ece78f1b711f93175f0a4a9726157cea247e107\"\"}], \"\"ocf\"\": null, \"\"ocv\"\": null}\")"), - ]; - - for bytes in messages { - let expected = bytes.clone(); - - let data_row = DataRow::try_from(&bytes).unwrap(); - - let bytes = BytesMut::try_from(data_row).unwrap(); - assert_eq!(bytes, expected); - } - } - - #[test] - pub fn parse_data_row_with_columns() { - let bytes = to_message( - b"D\0\0\09\0\x03\0\0\0\x08blahvtha\0\0\0\x0242\0\0\0\x1d2023-12-16 01:52:25.031985+00", - ); - - let data_row = DataRow::try_from(&bytes).unwrap(); - - let data_col = data_row.columns.first().unwrap(); - - let buf: &[u8] = data_col.bytes.as_ref().unwrap(); - let value = String::from_utf8_lossy(buf).to_string(); - assert_eq!(value, "blahvtha"); - } - - #[test] - pub fn parse_data_row_with_null_column() { - let bytes = to_message(b"D\0\0\0\n\0\x01\xff\xff\xff\xff"); - - let data_row = DataRow::try_from(&bytes).unwrap(); - - let data_col = data_row.columns.first().unwrap(); - - assert_eq!(data_col.bytes, None); - } - - #[test] - pub fn data_row_column_len() { - let column = DataColumn { bytes: None }; - let data_row = DataRow { - columns: vec![column], - }; - assert_eq!(data_row.len_of_columns(), 4); - - let data = BytesMut::from(""); - let column = DataColumn { bytes: Some(data) }; - let data_row = DataRow { - columns: vec![column], - }; - assert_eq!(data_row.len_of_columns(), 4); - - let data = BytesMut::from("blah"); - let column = DataColumn { bytes: Some(data) }; - let data_row = DataRow { - columns: vec![column], - }; - assert_eq!(data_row.len_of_columns(), 8); - - let mut columns = Vec::new(); - for _ in 1..5 { - let data = BytesMut::from("blah"); - let column = DataColumn { bytes: Some(data) }; - columns.push(column); - } - let data_row = DataRow { columns }; - assert_eq!(data_row.len_of_columns(), 32); - } -} diff --git a/packages/cipherstash-proxy/src/postgresql/messages/describe.rs b/packages/cipherstash-proxy/src/postgresql/messages/describe.rs deleted file mode 100644 index 6f0bbd8e3..000000000 --- a/packages/cipherstash-proxy/src/postgresql/messages/describe.rs +++ /dev/null @@ -1,79 +0,0 @@ -use crate::error::{Error, ProtocolError}; -use crate::postgresql::protocol::BytesMutReadString; -use crate::{SIZE_I32, SIZE_U8}; - -use bytes::{Buf, BufMut, BytesMut}; -use std::convert::TryFrom; -use std::ffi::CString; -use std::io::Cursor; - -use super::target::Target; -use super::{FrontendCode, Name}; - -/// -/// Describe b'D' (Frontend) message. -/// -/// See: -/// -/// Byte1('D') -/// Identifies the message as a Describe command. -/// -/// Int32 -/// Length of message contents in bytes, including self. -/// -/// Byte1 -/// 'S' to describe a prepared statement; or 'P' to describe a portal. -/// -/// String -/// The name of the prepared statement or portal to describe (an empty string selects the unnamed prepared statement or portal). - -#[derive(Debug, Clone)] -pub struct Describe { - pub target: Target, - pub name: Name, -} - -impl TryFrom<&BytesMut> for Describe { - type Error = Error; - - fn try_from(bytes: &BytesMut) -> Result { - let mut cursor = Cursor::new(bytes); - let code = cursor.get_u8(); - - if FrontendCode::from(code) != FrontendCode::Describe { - return Err(ProtocolError::UnexpectedMessageCode { - expected: FrontendCode::Describe.into(), - received: code as char, - } - .into()); - } - - let _len = cursor.get_i32(); // read and progress cursor - let target = cursor.get_u8(); - let target = Target::try_from(target)?; - let name = cursor.read_string()?; - let name = Name::from(name); - - Ok(Describe { target, name }) - } -} - -impl TryFrom for BytesMut { - type Error = Error; - - fn try_from(describe: Describe) -> Result { - let mut bytes = BytesMut::new(); - - let name = CString::new(describe.name.as_str())?; - let name = name.as_bytes_with_nul(); - - let len = SIZE_I32 + SIZE_U8 + name.len(); - - bytes.put_u8(FrontendCode::Describe.into()); - bytes.put_i32(len as i32); - bytes.put_u8(describe.target.into()); - bytes.put_slice(name); - - Ok(bytes) - } -} diff --git a/packages/cipherstash-proxy/src/postgresql/messages/error_response.rs b/packages/cipherstash-proxy/src/postgresql/messages/error_response.rs deleted file mode 100644 index 869edfb5d..000000000 --- a/packages/cipherstash-proxy/src/postgresql/messages/error_response.rs +++ /dev/null @@ -1,566 +0,0 @@ -use super::BackendCode; -use crate::error::{Error, ProtocolError}; -use crate::postgresql::protocol::BytesMutReadString; -use crate::SIZE_I32; -use bytes::{Buf, BufMut, BytesMut}; -use core::fmt; -use regex::Regex; -use std::io::Cursor; -use std::sync::LazyLock; -use std::{convert::TryFrom, ffi::CString}; -/// -/// Postgres Error Codes -/// https://www.postgresql.org/docs/current/errcodes-appendix.html -pub const CODE_UNDEFINED_COLUMN: &str = "42703"; -pub const CODE_INVALID_PASSWORD: &str = "28P01"; -pub const CODE_RAISE_EXCEPTION: &str = "P0001"; -pub const CODE_SYNTAX_ERROR: &str = "42601"; -pub const CODE_INVALID_TEXT_REPRESENTATION: &str = "22P02"; -pub const CODE_IDLE_SESSION_TIMEOUT: &str = "57P05"; -pub const CODE_SYSTEM_ERROR: &str = "58000"; - -/// -/// ErrorResponse (B) -/// https://www.postgresql.org/docs/current/protocol-message-formats.html#PROTOCOL-MESSAGE-FORMATS-ERRORRESPONSE -/// -#[derive(Debug, Clone)] -pub struct ErrorResponse { - pub fields: Vec, -} - -#[derive(Debug, Clone)] -pub struct Field { - pub code: ErrorResponseCode, - pub value: String, -} - -/// ErrorResponseCodes -/// https://www.postgresql.org/docs/current/protocol-error-fields.html -#[derive(Debug, Clone, PartialEq)] -pub enum ErrorResponseCode { - Severity, - SeverityLegacy, - Code, - Message, - Detail, - Hint, - Position, - InternalPosition, - InternalQuery, - Where, - Schema, - Table, - Column, - DataType, - Constraint, - File, - Line, - Routine, - Unknown(char), -} - -impl ErrorResponse { - /// Create a FATAL error response for connection timeout. - /// - /// Uses PostgreSQL error code 57P05 (idle_session_timeout). While this code - /// is technically for idle session timeouts, it is the closest match for a - /// proxy-enforced connection timeout. The alternative 08006 (connection_failure) - /// implies a network-level failure, which is misleading — the proxy is - /// deliberately terminating a connection that exceeded its time limit. - pub fn connection_timeout(message: String) -> Self { - Self { - fields: vec![ - Field { - code: ErrorResponseCode::Severity, - value: "FATAL".to_string(), - }, - Field { - code: ErrorResponseCode::SeverityLegacy, - value: "FATAL".to_string(), - }, - Field { - code: ErrorResponseCode::Code, - value: CODE_IDLE_SESSION_TIMEOUT.to_string(), - }, - Field { - code: ErrorResponseCode::Message, - value: message, - }, - ], - } - } - - /// Whether this error carries FATAL severity — the client abandons the - /// connection on receipt. - pub fn is_fatal(&self) -> bool { - self.fields - .iter() - .any(|field| field.code == ErrorResponseCode::Severity && field.value == "FATAL") - } - - pub fn invalid_password(message: String) -> Self { - Self { - fields: vec![ - Field { - code: ErrorResponseCode::Severity, - value: "FATAL".to_string(), - }, - Field { - code: ErrorResponseCode::SeverityLegacy, - value: "FATAL".to_string(), - }, - Field { - code: ErrorResponseCode::Code, - value: CODE_INVALID_PASSWORD.to_string(), - }, - Field { - code: ErrorResponseCode::Message, - value: message, - }, - ], - } - } - - /// - /// SQL Parse error as PostgreSQL error - /// Code: 42601 syntax_errpr - /// - /// As EncryptError is an enum, this can be passed a different error variation. - /// - pub fn invalid_sql_statement(message: String) -> Self { - let line = extract_line_from_parse_error(&message); - let position: Option = extract_position_from_parse_error(&message); - - let mut fields = vec![ - Field { - code: ErrorResponseCode::Severity, - value: "ERROR".to_string(), - }, - Field { - code: ErrorResponseCode::SeverityLegacy, - value: "ERROR".to_string(), - }, - Field { - code: ErrorResponseCode::Code, - value: CODE_SYNTAX_ERROR.to_string(), - }, - Field { - code: ErrorResponseCode::Message, - value: message, - }, - ]; - - if let Some(line) = line { - fields.push(Field { - code: ErrorResponseCode::Line, - value: line.to_string(), - }); - } - if let Some(position) = position { - fields.push(Field { - code: ErrorResponseCode::Position, - value: position.to_string(), - }); - } - - Self { fields } - } - - /// - /// Invalid parameter as PostgreSQL error - /// Code: 22P02 invalid_text_representation - /// - /// As EncryptError is an enum, this can be passed a different error variation. - /// - pub fn invalid_parameter(message: String, table: &str, column: &str) -> Self { - Self { - fields: vec![ - Field { - code: ErrorResponseCode::Severity, - value: "ERROR".to_string(), - }, - Field { - code: ErrorResponseCode::SeverityLegacy, - value: "ERROR".to_string(), - }, - Field { - code: ErrorResponseCode::Code, - value: CODE_INVALID_TEXT_REPRESENTATION.to_string(), - }, - Field { - code: ErrorResponseCode::Message, - value: message, - }, - // TODO: make this work more good - // URL is curently in message, so this looks like a bug atm - // Field { - // code: ErrorResponseCode::Detail, - // value: ERROR_DOC_ENCRYPT_INVALID_PARAMETER_URL.to_string(), - // }, - Field { - code: ErrorResponseCode::Table, - value: table.to_string(), - }, - Field { - code: ErrorResponseCode::Column, - value: column.to_string(), - }, - ], - } - } - - /// - /// Unknown encrypted column as PostgreSQL error - /// Code: 42703 undefined_column - /// - pub fn unknown_column(message: String, table: &str, column: &str) -> Self { - Self { - fields: vec![ - Field { - code: ErrorResponseCode::Severity, - value: "ERROR".to_string(), - }, - Field { - code: ErrorResponseCode::SeverityLegacy, - value: "ERROR".to_string(), - }, - Field { - code: ErrorResponseCode::Code, - value: CODE_UNDEFINED_COLUMN.to_string(), - }, - Field { - code: ErrorResponseCode::Message, - value: message, - }, - // TODO: make this work more good - // URL is curently in message, so this looks like a bug atm - // Field { - // code: ErrorResponseCode::Detail, - // value: ERROR_DOC_ENCRYPT_UNKNOWN_COLUMN_URL.to_string(), - // }, - Field { - code: ErrorResponseCode::Table, - value: table.to_string(), - }, - Field { - code: ErrorResponseCode::Column, - value: column.to_string(), - }, - Field { - code: ErrorResponseCode::Routine, - value: "cipherstash-proxy".to_string(), - }, - ], - } - } - - pub fn system_error(message: String) -> Self { - Self { - fields: vec![ - Field { - code: ErrorResponseCode::Severity, - value: "FATAL".to_string(), - }, - Field { - code: ErrorResponseCode::SeverityLegacy, - value: "FATAL".to_string(), - }, - Field { - code: ErrorResponseCode::Code, - value: CODE_SYSTEM_ERROR.to_string(), - }, - Field { - code: ErrorResponseCode::Message, - value: message, - }, - ], - } - } - - pub fn tls_required() -> Self { - Self { - fields: vec![ - Field { - code: ErrorResponseCode::Severity, - value: "FATAL".to_string(), - }, - Field { - code: ErrorResponseCode::SeverityLegacy, - value: "FATAL".to_string(), - }, - Field { - code: ErrorResponseCode::Code, - value: "08001".to_string(), - }, - Field { - code: ErrorResponseCode::Message, - value: "Transport Layer Security (TLS) connection is required".to_string(), - }, - ], - } - } -} - -/// -/// Extracts line (if present) from a SQL Parser error message -/// -fn extract_line_from_parse_error(error_message: &str) -> Option { - static RE: LazyLock = LazyLock::new(|| Regex::new(r"\s*Line:\s*(\d+)").unwrap()); - RE.captures(error_message) - .and_then(|c| c.get(1)?.as_str().parse::().ok()) -} - -/// -/// Extracts position (if present) from a SQL Parser error message -/// Column in the error message is the "text" column, not a reference to a database column -/// -fn extract_position_from_parse_error(error_message: &str) -> Option { - static RE: LazyLock = LazyLock::new(|| Regex::new(r"\s*Column:\s*(\d+)").unwrap()); - - RE.captures(error_message) - .and_then(|c| c.get(1)?.as_str().parse::().ok()) -} - -impl TryFrom<&BytesMut> for ErrorResponse { - type Error = Error; - - fn try_from(buf: &BytesMut) -> Result { - let mut cursor = Cursor::new(buf); - let code = cursor.get_u8(); - - if BackendCode::from(code) != BackendCode::ErrorResponse { - return Err(ProtocolError::UnexpectedMessageCode { - expected: BackendCode::ErrorResponse.into(), - received: code as char, - } - .into()); - } - - let _len = cursor.get_i32(); - - // The message body consists of one or more identified fields, followed by a zero byte as a terminator. - let mut fields = Vec::new(); - - loop { - let code = cursor.get_u8(); - - // zero byte is terminator - if code == 0 { - break; - } - - let value = cursor.read_string()?; - let field = Field { - code: code.into(), - value, - }; - fields.push(field); - } - - Ok(ErrorResponse { fields }) - } -} - -impl TryFrom for BytesMut { - type Error = Error; - - fn try_from(error_response: ErrorResponse) -> Result { - let mut field_bytes = BytesMut::new(); - - for field in error_response.fields { - let value = CString::new(field.value)?; - let value = value.as_bytes_with_nul(); - - field_bytes.put_u8(field.code.into()); - field_bytes.put_slice(value); - } - field_bytes.put_u8(0); // field terminator - - let mut bytes = BytesMut::new(); - - let len = SIZE_I32 + field_bytes.len(); // len + fields - - bytes.put_u8(BackendCode::ErrorResponse.into()); - bytes.put_i32(len as i32); - bytes.put_slice(&field_bytes); - - Ok(bytes) - } -} - -impl fmt::Display for ErrorResponse { - fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { - for field in self.fields.iter() { - let msg = match field.code { - ErrorResponseCode::Severity => "Severity", - ErrorResponseCode::SeverityLegacy => continue, // skipped, always appears with `S` in versions we support - ErrorResponseCode::Code => "Code", - ErrorResponseCode::Message => "Message", - ErrorResponseCode::Detail => "Detail", - ErrorResponseCode::Hint => "Hint", - ErrorResponseCode::Position => "Position", - ErrorResponseCode::InternalPosition => "Internal Position", - ErrorResponseCode::InternalQuery => "Internal Query", - ErrorResponseCode::Where => "Where", - ErrorResponseCode::Schema => "Schema", - ErrorResponseCode::Table => "Table", - ErrorResponseCode::Column => "Column", - ErrorResponseCode::DataType => "Data Type", - ErrorResponseCode::Constraint => "Constraint", - ErrorResponseCode::File => "File", - ErrorResponseCode::Line => "Line", - ErrorResponseCode::Routine => "Routine", - ErrorResponseCode::Unknown(_) => "Unknown", - }; - write!(f, "{} ({}): {} ", msg, char::from(&field.code), field.value)?; - } - - Ok(()) - } -} - -impl From for u8 { - fn from(code: ErrorResponseCode) -> Self { - match code { - ErrorResponseCode::Severity => b'S', - ErrorResponseCode::SeverityLegacy => b'V', - ErrorResponseCode::Code => b'C', - ErrorResponseCode::Message => b'M', - ErrorResponseCode::Detail => b'D', - ErrorResponseCode::Hint => b'H', - ErrorResponseCode::Position => b'P', - ErrorResponseCode::InternalPosition => b'p', - ErrorResponseCode::InternalQuery => b'q', - ErrorResponseCode::Where => b'W', - ErrorResponseCode::Schema => b's', - ErrorResponseCode::Table => b't', - ErrorResponseCode::Column => b'c', - ErrorResponseCode::DataType => b'd', - ErrorResponseCode::Constraint => b'n', - ErrorResponseCode::File => b'F', - ErrorResponseCode::Line => b'L', - ErrorResponseCode::Routine => b'R', - ErrorResponseCode::Unknown(c) => c as u8, - } - } -} - -impl From<&ErrorResponseCode> for char { - fn from(code: &ErrorResponseCode) -> Self { - match code { - ErrorResponseCode::Severity => 'S', - ErrorResponseCode::SeverityLegacy => 'V', - ErrorResponseCode::Code => 'C', - ErrorResponseCode::Message => 'M', - ErrorResponseCode::Detail => 'D', - ErrorResponseCode::Hint => 'H', - ErrorResponseCode::Position => 'P', - ErrorResponseCode::InternalPosition => 'p', - ErrorResponseCode::InternalQuery => 'q', - ErrorResponseCode::Where => 'W', - ErrorResponseCode::Schema => 's', - ErrorResponseCode::Table => 't', - ErrorResponseCode::Column => 'c', - ErrorResponseCode::DataType => 'd', - ErrorResponseCode::Constraint => 'n', - ErrorResponseCode::File => 'F', - ErrorResponseCode::Line => 'L', - ErrorResponseCode::Routine => 'R', - ErrorResponseCode::Unknown(c) => c.to_owned(), - } - } -} - -impl From for ErrorResponseCode { - fn from(byte: u8) -> Self { - match byte { - b'S' => ErrorResponseCode::Severity, - b'V' => ErrorResponseCode::SeverityLegacy, - b'C' => ErrorResponseCode::Code, - b'M' => ErrorResponseCode::Message, - b'D' => ErrorResponseCode::Detail, - b'H' => ErrorResponseCode::Hint, - b'P' => ErrorResponseCode::Position, - b'p' => ErrorResponseCode::InternalPosition, - b'q' => ErrorResponseCode::InternalQuery, - b'W' => ErrorResponseCode::Where, - b's' => ErrorResponseCode::Schema, - b't' => ErrorResponseCode::Table, - b'c' => ErrorResponseCode::Column, - b'd' => ErrorResponseCode::DataType, - b'n' => ErrorResponseCode::Constraint, - b'F' => ErrorResponseCode::File, - b'L' => ErrorResponseCode::Line, - b'R' => ErrorResponseCode::Routine, - c => ErrorResponseCode::Unknown(c as char), - } - } -} - -#[cfg(test)] -mod tests { - use super::ErrorResponseCode; - use crate::postgresql::messages::error_response::ErrorResponse; - use bytes::BytesMut; - - fn to_message(s: &[u8]) -> BytesMut { - BytesMut::from(s) - } - - #[test] - pub fn parse_error_response_message() { - let message = to_message(b"E\0\0\0kSERROR\0VERROR\0C26000\0Mprepared statement \"a37\" does not exist\0Fprepare.c\0L454\0RFetchPreparedStatement\0\0Z\0\0\0\x05I"); - - let error_response = ErrorResponse::try_from(&message).unwrap(); - assert_eq!(error_response.fields.len(), 7); - - // let next = cursor.get_u8() as char; - // assert_eq!(next, 'Z'); - - let bytes = BytesMut::try_from(error_response).unwrap(); - let message = to_message(b"E\0\0\0kSERROR\0VERROR\0C26000\0Mprepared statement \"a37\" does not exist\0Fprepare.c\0L454\0RFetchPreparedStatement\0\0"); - assert_eq!(bytes, message); - } - - #[test] - pub fn sql_parse_error_response() { - let response = ErrorResponse::invalid_sql_statement( - "sql syntax error in blah vtha Line: 1, Column: 2".to_string(), - ); - - let line = response - .fields - .iter() - .find(|f| f.code == ErrorResponseCode::Line) - .unwrap(); - - assert_eq!(line.value, "1".to_string()); - - let position = response - .fields - .iter() - .find(|f| f.code == ErrorResponseCode::Position) - .unwrap(); - - assert_eq!(position.value, "2".to_string()); - - let response = ErrorResponse::invalid_sql_statement( - "sql syntax error in blah vtha Column: 2".to_string(), - ); - - let line = response - .fields - .iter() - .find(|f| f.code == ErrorResponseCode::Line); - - assert!(line.is_none()); - - let position = response - .fields - .iter() - .find(|f| f.code == ErrorResponseCode::Position) - .unwrap(); - - assert_eq!(position.value, "2".to_string()); - } -} diff --git a/packages/cipherstash-proxy/src/postgresql/messages/execute.rs b/packages/cipherstash-proxy/src/postgresql/messages/execute.rs deleted file mode 100644 index f5e17ae29..000000000 --- a/packages/cipherstash-proxy/src/postgresql/messages/execute.rs +++ /dev/null @@ -1,37 +0,0 @@ -use super::{FrontendCode, Name}; -use crate::error::{Error, ProtocolError}; -use crate::postgresql::protocol::BytesMutReadString; -use bytes::{Buf, BytesMut}; -use std::convert::TryFrom; -use std::io::Cursor; - -#[derive(Debug, Clone)] -pub(crate) struct Execute { - pub portal: Name, - pub max_rows: i32, -} - -impl TryFrom<&BytesMut> for Execute { - type Error = Error; - - fn try_from(bytes: &BytesMut) -> Result { - let mut cursor = Cursor::new(bytes); - let code = cursor.get_u8(); - - if FrontendCode::from(code) != FrontendCode::Execute { - return Err(ProtocolError::UnexpectedMessageCode { - expected: FrontendCode::Execute.into(), - received: code as char, - } - .into()); - } - - let _len = cursor.get_i32(); // read and progress cursor - - let portal = cursor.read_string()?; - let portal = Name::from(portal); - let max_rows = cursor.get_i32(); - - Ok(Execute { portal, max_rows }) - } -} diff --git a/packages/cipherstash-proxy/src/postgresql/messages/mod.rs b/packages/cipherstash-proxy/src/postgresql/messages/mod.rs deleted file mode 100644 index b271c001c..000000000 --- a/packages/cipherstash-proxy/src/postgresql/messages/mod.rs +++ /dev/null @@ -1,304 +0,0 @@ -use std::fmt; - -use bytes::BytesMut; - -pub mod authentication; -pub mod bind; -pub mod close; -pub mod data_row; -pub mod describe; -pub mod error_response; -pub mod execute; -pub mod name; -pub mod param_description; -pub mod parse; -pub mod query; -pub mod ready_for_query; -pub mod row_description; -pub mod target; -pub mod terminate; - -// Re-export commonly used types -pub use name::Name; -pub use target::Target; - -pub const NULL: i32 = -1; - -/// PostgreSQL's "unspecified type, infer it" param OID, used in `Parse` and -/// when a param's type is not known to the proxy. -pub const UNSPECIFIED_TYPE_OID: i32 = 0; - -#[derive(Clone, Copy, Debug, PartialEq)] -pub enum FrontendCode { - Bind, - Close, - Describe, - Execute, - Flush, - Parse, - PasswordMessage, - Query, - SASLInitialResponse, - SASLResponse, - Sync, - Terminate, - Unknown(char), -} - -#[derive(Clone, Copy, Debug, PartialEq)] -pub enum BackendCode { - Authentication, - BindComplete, - BackendKeyData, - CloseComplete, - CommandComplete, - CopyBothResponse, - CopyInResponse, - CopyOutResponse, - DataRow, - EmptyQueryResponse, - ErrorResponse, - NoData, - NoticeResponse, - NotificationResponse, - ParameterDescription, - ParameterStatus, - ParseComplete, - PortalSuspended, - ReadyForQuery, - RowDescription, - Unknown(char), -} - -impl From for FrontendCode { - fn from(code: u8) -> Self { - (code as char).into() - } -} - -impl From for FrontendCode { - fn from(code: char) -> Self { - match code { - 'B' => FrontendCode::Bind, - 'C' => FrontendCode::Close, - 'D' => FrontendCode::Describe, - 'E' => FrontendCode::Execute, - 'H' => FrontendCode::Flush, - 'p' => FrontendCode::PasswordMessage, - 'P' => FrontendCode::Parse, - 'Q' => FrontendCode::Query, - #[allow(unreachable_patterns)] - 'p' => FrontendCode::SASLInitialResponse, // Uses same char, here for completeness - #[allow(unreachable_patterns)] - 'p' => FrontendCode::SASLResponse, // Uses same char, here for completeness - 'S' => FrontendCode::Sync, - 'X' => FrontendCode::Terminate, - _ => FrontendCode::Unknown(code), - } - } -} - -impl From for u8 { - fn from(code: FrontendCode) -> Self { - match code { - FrontendCode::Bind => b'B', - FrontendCode::Close => b'C', - FrontendCode::Describe => b'D', - FrontendCode::Execute => b'E', - FrontendCode::Flush => b'F', - FrontendCode::Parse => b'P', - FrontendCode::PasswordMessage => b'p', - FrontendCode::Query => b'Q', - FrontendCode::SASLInitialResponse => b'p', - FrontendCode::SASLResponse => b'p', - FrontendCode::Sync => b'S', - FrontendCode::Terminate => b'X', - FrontendCode::Unknown(c) => c as u8, - } - } -} - -impl From for char { - fn from(code: FrontendCode) -> Self { - match code { - FrontendCode::Bind => 'B', - FrontendCode::Close => 'C', - FrontendCode::Describe => 'D', - FrontendCode::Execute => 'E', - FrontendCode::Flush => 'F', - FrontendCode::Parse => 'P', - FrontendCode::PasswordMessage => 'p', - FrontendCode::Query => 'Q', - FrontendCode::SASLInitialResponse => 'p', - FrontendCode::SASLResponse => 'p', - FrontendCode::Sync => 'S', - FrontendCode::Terminate => 'X', - FrontendCode::Unknown(c) => c, - } - } -} - -impl From for BackendCode { - fn from(code: u8) -> Self { - match code as char { - 'R' => BackendCode::Authentication, - 'K' => BackendCode::BackendKeyData, - '2' => BackendCode::BindComplete, - '3' => BackendCode::CloseComplete, - 'C' => BackendCode::CommandComplete, - 'W' => BackendCode::CopyBothResponse, - 'G' => BackendCode::CopyInResponse, - 'H' => BackendCode::CopyOutResponse, - 'D' => BackendCode::DataRow, - 'I' => BackendCode::EmptyQueryResponse, - 'E' => BackendCode::ErrorResponse, - 'n' => BackendCode::NoData, - 'N' => BackendCode::NoticeResponse, - 'A' => BackendCode::NotificationResponse, - 't' => BackendCode::ParameterDescription, - 'S' => BackendCode::ParameterStatus, - '1' => BackendCode::ParseComplete, - 's' => BackendCode::PortalSuspended, - 'Z' => BackendCode::ReadyForQuery, - 'T' => BackendCode::RowDescription, - _ => BackendCode::Unknown(code as char), - } - } -} - -impl From for u8 { - fn from(code: BackendCode) -> Self { - match code { - BackendCode::Authentication => b'R', - BackendCode::BackendKeyData => b'K', - BackendCode::BindComplete => b'2', - BackendCode::CloseComplete => b'3', - BackendCode::CommandComplete => b'C', - BackendCode::CopyBothResponse => b'W', - BackendCode::CopyInResponse => b'G', - BackendCode::CopyOutResponse => b'H', - BackendCode::DataRow => b'D', - BackendCode::EmptyQueryResponse => b'I', - BackendCode::ErrorResponse => b'E', - BackendCode::NoData => b'n', - BackendCode::NoticeResponse => b'N', - BackendCode::NotificationResponse => b'A', - BackendCode::ParameterDescription => b't', - BackendCode::ParameterStatus => b'S', - BackendCode::ParseComplete => b'1', - BackendCode::PortalSuspended => b's', - BackendCode::ReadyForQuery => b'Z', - BackendCode::RowDescription => b'T', - BackendCode::Unknown(c) => c as u8, - } - } -} - -impl From for char { - fn from(code: BackendCode) -> Self { - match code { - BackendCode::Authentication => 'R', - BackendCode::BackendKeyData => 'K', - BackendCode::BindComplete => '2', - BackendCode::CloseComplete => '3', - BackendCode::CommandComplete => 'C', - BackendCode::CopyBothResponse => 'W', - BackendCode::CopyInResponse => 'G', - BackendCode::CopyOutResponse => 'H', - BackendCode::DataRow => 'D', - BackendCode::EmptyQueryResponse => 'I', - BackendCode::ErrorResponse => 'E', - BackendCode::NoData => 'n', - BackendCode::NoticeResponse => 'N', - BackendCode::NotificationResponse => 'A', - BackendCode::ParameterDescription => 't', - BackendCode::ParameterStatus => 'S', - BackendCode::ParseComplete => '1', - BackendCode::PortalSuspended => 's', - BackendCode::ReadyForQuery => 'Z', - BackendCode::RowDescription => 'T', - BackendCode::Unknown(c) => c, - } - } -} - -impl fmt::Display for BackendCode { - fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - match self { - BackendCode::Authentication => write!(f, "BackendCode::Authentication"), - BackendCode::BackendKeyData => write!(f, "BackendCode::BackendKeyData"), - BackendCode::BindComplete => write!(f, "BackendCode::BindComplete"), - BackendCode::CloseComplete => write!(f, "BackendCode::CloseComplete"), - BackendCode::CommandComplete => write!(f, "BackendCode::CommandComplete"), - BackendCode::CopyBothResponse => write!(f, "BackendCode::CopyBothResponse"), - BackendCode::CopyInResponse => write!(f, "BackendCode::CopyInResponse"), - BackendCode::CopyOutResponse => write!(f, "BackendCode::CopyOutResponse"), - BackendCode::DataRow => write!(f, "BackendCode::DataRow"), - BackendCode::EmptyQueryResponse => write!(f, "BackendCode::EmptyQueryResponse"), - BackendCode::ErrorResponse => write!(f, "BackendCode::ErrorResponse"), - BackendCode::NoData => write!(f, "BackendCode::NoData"), - BackendCode::NoticeResponse => write!(f, "BackendCode::NoticeResponse"), - BackendCode::NotificationResponse => write!(f, "BackendCode::NotificationResponse"), - BackendCode::ParameterDescription => write!(f, "BackendCode::ParameterDescription"), - BackendCode::ParameterStatus => write!(f, "BackendCode::ParameterStatus"), - BackendCode::ParseComplete => write!(f, "BackendCode::ParseComplete"), - BackendCode::PortalSuspended => write!(f, "BackendCode::PortalSuspended"), - BackendCode::ReadyForQuery => write!(f, "BackendCode::ReadyForQuery"), - BackendCode::RowDescription => write!(f, "BackendCode::RowDescription"), - BackendCode::Unknown(c) => write!(f, "BackendCode::Unknown('{}')", c), - } - } -} - -impl fmt::Display for FrontendCode { - fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - match self { - FrontendCode::Bind => write!(f, "FrontendCode::Bind"), - FrontendCode::Close => write!(f, "FrontendCode::Close"), - FrontendCode::Describe => write!(f, "FrontendCode::Describe"), - FrontendCode::Execute => write!(f, "FrontendCode::Execute"), - FrontendCode::Flush => write!(f, "FrontendCode::Flush"), - FrontendCode::Parse => write!(f, "FrontendCode::Parse"), - FrontendCode::PasswordMessage => write!(f, "FrontendCode::PasswordMessage"), - FrontendCode::Query => write!(f, "FrontendCode::Query"), - FrontendCode::SASLInitialResponse => write!(f, "FrontendCode::SASLInitialResponse"), - FrontendCode::SASLResponse => write!(f, "FrontendCode::SASLResponse"), - FrontendCode::Sync => write!(f, "FrontendCode::Sync"), - FrontendCode::Terminate => write!(f, "FrontendCode::Terminate"), - FrontendCode::Unknown(c) => write!(f, "FrontendCode::Unknown('{}')", c), - } - } -} - -/// -/// Peaks at the first byte char. -/// Assumes that a leading `{` may be a JSON value -/// The Plaintext Payload is always a JSON object so this is a pretty naive approach -/// We are not worried about an exhaustive check here -/// -pub fn maybe_json(bytes: &BytesMut) -> bool { - if bytes.is_empty() { - return false; - } - - let b = bytes.as_ref()[0]; - b == b'{' -} - -/// -/// Postgres binary json is regular json with a leading header byte -/// The header byte is always 1 -/// -pub fn maybe_jsonb(bytes: &BytesMut) -> bool { - // Empty JSONB is at least 3 bytes - // `1{}`` - if bytes.len() <= 3 { - return false; - } - - let b = bytes.as_ref(); - - let header = b[0]; - let first = b[1]; - header == 1 && first == b'{' -} diff --git a/packages/cipherstash-proxy/src/postgresql/messages/name.rs b/packages/cipherstash-proxy/src/postgresql/messages/name.rs deleted file mode 100644 index f4f1b9acc..000000000 --- a/packages/cipherstash-proxy/src/postgresql/messages/name.rs +++ /dev/null @@ -1,50 +0,0 @@ -#[derive(Debug, Clone, Hash, Eq, PartialEq)] -pub enum Name { - Named(String), - Unnamed, -} - -impl Name { - pub fn unnamed() -> Name { - Name::Unnamed - } - - pub fn is_unnamed(&self) -> bool { - matches!(self, Name::Unnamed) - } - - pub fn as_str(&self) -> &str { - match self { - Name::Named(s) => s, - Name::Unnamed => "", - } - } -} - -impl std::ops::Deref for Name { - type Target = str; - - fn deref(&self) -> &str { - self.as_str() - } -} - -impl From for Name { - fn from(s: String) -> Self { - if s.is_empty() { - Name::Unnamed - } else { - Name::Named(s) - } - } -} - -impl From<&str> for Name { - fn from(s: &str) -> Self { - if s.is_empty() { - Name::Unnamed - } else { - Name::Named(s.to_string()) - } - } -} diff --git a/packages/cipherstash-proxy/src/postgresql/messages/param_description.rs b/packages/cipherstash-proxy/src/postgresql/messages/param_description.rs deleted file mode 100644 index 6a9b4c19b..000000000 --- a/packages/cipherstash-proxy/src/postgresql/messages/param_description.rs +++ /dev/null @@ -1,197 +0,0 @@ -use super::BackendCode; -use crate::{ - error::{Error, ProtocolError}, - log::MAPPER, - SIZE_I16, SIZE_I32, -}; -use bytes::{Buf, BufMut, BytesMut}; -use postgres_types::Type; -use std::io::Cursor; -use tracing::debug; - -/// -/// Describe b't' (Backend) message. -/// -/// See: -/// -/// Byte1('t') -/// Identifies the message as a parameter description. -/// -/// Int32 -/// Length of message contents in bytes, including self. -/// -/// Int16 -/// The number of parameters used by the statement (can be zero). -/// -/// For each parameter: -/// Int32 -/// Specifies the object ID of the parameter data type. -/// - -#[derive(Debug)] -pub struct ParamDescription { - pub types: Vec, - dirty: bool, -} - -impl ParamDescription { - pub fn map_types(&mut self, mapped_types: &[Option]) { - debug!(target: MAPPER, ?mapped_types); - - for (idx, t) in mapped_types.iter().enumerate() { - if let Some(t) = t { - self.types[idx] = t.oid() as i32; - self.dirty = true; - } - } - } - - /// Replaces the described params wholesale. - /// - /// PostgreSQL describes the params of the *rewritten* statement, but the - /// client must be told about the params it wrote — a rewrite that fuses two - /// params into one would otherwise describe too few, and the client would - /// bind the wrong number of values. - pub fn set_types(&mut self, types: Vec) { - debug!(target: MAPPER, ?types); - - if types != self.types { - self.types = types; - self.dirty = true; - } - } - - pub fn requires_rewrite(&self) -> bool { - self.dirty - } -} - -impl TryFrom<&BytesMut> for ParamDescription { - type Error = Error; - - fn try_from(bytes: &BytesMut) -> Result { - let mut cursor = Cursor::new(bytes); - - let code = cursor.get_u8(); - - if BackendCode::from(code) != BackendCode::ParameterDescription { - return Err(ProtocolError::UnexpectedMessageCode { - expected: BackendCode::ParameterDescription.into(), - received: code as char, - } - .into()); - } - - let _len = cursor.get_i32(); // move the cursor - let count = cursor.get_i16() as usize; - - let mut types = vec![]; - for _idx in 0..count { - let type_oid = cursor.get_i32(); - types.push(type_oid) - } - - Ok(ParamDescription { - types, - dirty: false, - }) - } -} - -impl TryFrom for BytesMut { - type Error = Error; - - fn try_from(parameter_description: ParamDescription) -> Result { - let mut bytes = BytesMut::new(); - - let count = parameter_description.types.len(); - let size_of_types = count * SIZE_I32; - - let len = SIZE_I32 + SIZE_I16 + size_of_types; - - bytes.put_u8(BackendCode::ParameterDescription.into()); - bytes.put_i32(len as i32); - bytes.put_i16(count as i16); - - for type_oid in parameter_description.types.into_iter() { - bytes.put_i32(type_oid); - } - - Ok(bytes) - } -} - -#[cfg(test)] -mod tests { - - use bytes::BytesMut; - use tracing::info; - - use crate::{config::LogConfig, log}; - - use super::ParamDescription; - - fn to_message(s: &[u8]) -> BytesMut { - BytesMut::from(s) - } - - #[test] - pub fn map_parameter_types() { - log::init(LogConfig::default()); - - let mut pd = ParamDescription { - types: vec![ - postgres_types::Type::TEXT.oid() as i32, - postgres_types::Type::INT4.oid() as i32, - postgres_types::Type::INT8.oid() as i32, - ], - dirty: false, - }; - - // No types to map, should not rewrite - let mapped_types = vec![None, None, None]; - pd.map_types(&mapped_types); - assert!(!pd.requires_rewrite()); - - let mapped_types = vec![ - Some(postgres_types::Type::TEXT), - None, - Some(postgres_types::Type::TEXT), - ]; - pd.map_types(&mapped_types); - assert!(pd.requires_rewrite()); - - let expected = vec![ - postgres_types::Type::TEXT.oid() as i32, - postgres_types::Type::INT4.oid() as i32, - postgres_types::Type::TEXT.oid() as i32, - ]; - - assert_eq!(pd.types, expected); - } - - #[test] - pub fn parse_parameter_description() { - log::init(LogConfig::default()); - let bytes = to_message(b"t\0\0\0\x0e\0\x02\0\0\0\x14\0\0\x0e\xda"); - - let expected = bytes.clone(); - - let description = ParamDescription::try_from(&bytes).unwrap(); - - info!("{:?}", description); - - assert_eq!(description.types.len(), 2); - assert_eq!( - description.types[0], - postgres_types::Type::INT8.oid() as i32 - ); - assert_eq!( - description.types[1], - postgres_types::Type::JSONB.oid() as i32 - ); - - let bytes = BytesMut::try_from(description).unwrap(); - assert_eq!(bytes, expected); - } -} diff --git a/packages/cipherstash-proxy/src/postgresql/messages/parse.rs b/packages/cipherstash-proxy/src/postgresql/messages/parse.rs deleted file mode 100644 index 8f7c9666d..000000000 --- a/packages/cipherstash-proxy/src/postgresql/messages/parse.rs +++ /dev/null @@ -1,259 +0,0 @@ -use super::{FrontendCode, Name, UNSPECIFIED_TYPE_OID}; -use crate::{ - error::{Error, ProtocolError}, - postgresql::{context::statement::OutputParam, protocol::BytesMutReadString}, - SIZE_I16, SIZE_I32, -}; -use bytes::{Buf, BufMut, BytesMut}; -use eql_mapper::EqlTermVariant; -use postgres_types::Type; -use std::{ffi::CString, io::Cursor}; - -#[derive(Debug, Clone)] -pub struct Parse { - pub code: char, - pub name: Name, - pub statement: String, - pub num_params: i16, - pub param_types: Vec, - dirty: bool, -} - -impl Parse { - pub fn requires_rewrite(&self) -> bool { - self.dirty - } - - /// Rewrites the declared param types to describe the params of the - /// *rewritten* statement. - /// - /// EQL v3 encrypted columns are JSONB-backed domain types (e.g. - /// `eql_v3_text_search`). JSONB is declared rather than the domain itself to - /// avoid loading each domain's OID — PostgreSQL coerces JSONB to the domain - /// if it passes the CHECK constraint. - /// - /// The client declares types for the params it wrote; the rewrite may have - /// dropped or fused some of those, so each declaration is carried across to - /// the output param that consumes it. An output param that carries an - /// encrypted value is declared JSONB regardless — that is the wire type of - /// every EQL payload, whatever the client thought it was binding. - /// - /// A JSON *selector* is the exception. It is passed to the rewritten - /// function as bare encrypted text — `eql_v3."->"(json, text)`, - /// `eql_v3.jsonb_path_exists(json, text)` — not as a jsonb query payload, so - /// declaring JSONB leaves PostgreSQL looking for an overload that does not - /// exist: - /// - /// ```text - /// ERROR: function eql_v3.jsonb_path_exists(eql_v3_json_search, jsonb) does not exist - /// ``` - /// - /// A client that declares no types at all (the common case — it lets the - /// server infer them) is left alone: every output param is referenced by the - /// rewritten SQL, so PostgreSQL can always infer them. That is why this only - /// bites clients that send their own Parse OIDs, such as pgx in - /// `cache_describe` mode. - pub fn rewrite_param_types(&mut self, output_params: &[OutputParam]) { - if self.param_types.is_empty() { - return; - } - - let param_types = output_params - .iter() - .map(|output| match &output.column { - Some(column) => match column.eql_term { - EqlTermVariant::JsonAccessor | EqlTermVariant::JsonPath => { - Type::TEXT.oid() as i32 - } - _ => Type::JSONB.oid() as i32, - }, - None => self - .param_types - .get(output.source.primary_input()) - .copied() - .unwrap_or(UNSPECIFIED_TYPE_OID), - }) - .collect::>(); - - if param_types != self.param_types { - self.num_params = param_types.len() as i16; - self.param_types = param_types; - self.dirty = true; - } - } - - pub fn rewrite_statement(&mut self, statement: String) { - self.statement = statement; - self.dirty = true; - } -} - -impl TryFrom<&BytesMut> for Parse { - type Error = Error; - - fn try_from(buf: &BytesMut) -> Result { - let mut cursor = Cursor::new(buf); - let code = cursor.get_u8() as char; - - if FrontendCode::from(code) != FrontendCode::Parse { - return Err(ProtocolError::UnexpectedMessageCode { - expected: FrontendCode::Parse.into(), - received: code, - } - .into()); - } - - let _len = cursor.get_i32(); - let name = cursor.read_string()?; - let name = Name::from(name); - - let statement = cursor.read_string()?; - let num_params = cursor.get_i16(); - let mut param_types = Vec::new(); - - for _ in 0..num_params { - param_types.push(cursor.get_i32()); - } - - Ok(Parse { - code, - name, - statement, - num_params, - param_types, - dirty: false, - }) - } -} - -impl TryFrom for BytesMut { - type Error = Error; - - fn try_from(parse: Parse) -> Result { - let mut bytes = BytesMut::new(); - - let name = CString::new(parse.name.as_str())?; - let name = name.as_bytes_with_nul(); - - let statement = CString::new(parse.statement)?; - let statement = statement.as_bytes_with_nul(); - - let len = SIZE_I32 // len - + name.len() - + statement.len() - + SIZE_I16 // num_params - + SIZE_I32 * parse.param_types.len(); - - bytes.put_u8(FrontendCode::Parse.into()); - bytes.put_i32(len as i32); - bytes.put_slice(name); - bytes.put_slice(statement); - bytes.put_i16(parse.num_params); - for param in parse.param_types { - bytes.put_i32(param); - } - - Ok(bytes) - } -} - -#[cfg(test)] -mod tests { - use crate::{ - config::LogConfig, - log, - postgresql::{ - context::statement::{OutputParam, OutputParamSource}, - messages::parse::Parse, - Column, - }, - Identifier, - }; - use bytes::BytesMut; - use cipherstash_client::schema::{ColumnConfig, ColumnType}; - - fn to_message(s: &[u8]) -> BytesMut { - BytesMut::from(s) - } - - #[test] - pub fn test_parse() { - log::init(LogConfig::default()); - let bytes = to_message( - b"P\0\0\0J\0INSERT INTO encrypted (id, encrypted_int2) VALUES ($1, $2)\0\0\x02\0\0\0\x15\0\0\0\x15" - ); - - let expected = bytes.clone(); - - let parse = Parse::try_from(&bytes).unwrap(); - - let bytes = BytesMut::try_from(parse).unwrap(); - assert_eq!(bytes, expected); - } - - #[test] - pub fn test_parse_rewrite_param_types() { - log::init(LogConfig::default()); - let bytes = to_message( - b"P\0\0\0J\0INSERT INTO encrypted (id, encrypted_int2) VALUES ($1, $2)\0\0\x02\0\0\0\x15\0\0\0\x15" - ); - - let mut parse = Parse::try_from(&bytes).unwrap(); - - let identifier = Identifier::new("table", "column"); - - let config = ColumnConfig::build("column".to_string()).casts_as(ColumnType::SmallInt); - - let column = Column::new(identifier, config, None, eql_mapper::EqlTermVariant::Full); - let output_params = vec![ - OutputParam { - column: None, - source: OutputParamSource::Input(0), - query_operand: false, - }, - OutputParam { - column: Some(column), - source: OutputParamSource::Input(1), - query_operand: false, - }, - ]; - - parse.rewrite_param_types(&output_params); - assert!(parse.requires_rewrite()); - assert_eq!( - parse.param_types, - vec![ - postgres_types::Type::INT2.oid() as i32, - postgres_types::Type::JSONB.oid() as i32 - ] - ); - } - - /// A rewrite that fuses two params into one must leave the client's - /// declaration for the surviving param, not the one it happened to sit at. - #[test] - pub fn test_parse_rewrite_param_types_after_fusion() { - log::init(LogConfig::default()); - let bytes = to_message( - b"P\0\0\0J\0INSERT INTO encrypted (id, encrypted_int2) VALUES ($1, $2)\0\0\x02\0\0\0\x15\0\0\0\x15" - ); - - let mut parse = Parse::try_from(&bytes).unwrap(); - - // Two input params collapse to a single native output param sourced - // from input 1. - let output_params = vec![OutputParam { - column: None, - source: OutputParamSource::Input(1), - query_operand: false, - }]; - - parse.rewrite_param_types(&output_params); - assert!(parse.requires_rewrite()); - assert_eq!(parse.num_params, 1); - assert_eq!( - parse.param_types, - vec![postgres_types::Type::INT2.oid() as i32] - ); - } -} diff --git a/packages/cipherstash-proxy/src/postgresql/messages/query.rs b/packages/cipherstash-proxy/src/postgresql/messages/query.rs deleted file mode 100644 index 88a5cd570..000000000 --- a/packages/cipherstash-proxy/src/postgresql/messages/query.rs +++ /dev/null @@ -1,79 +0,0 @@ -use crate::error::{Error, ProtocolError}; -use crate::postgresql::protocol::BytesMutReadString; -use crate::SIZE_I32; - -use bytes::{Buf, BufMut, BytesMut}; -use std::convert::TryFrom; -use std::ffi::CString; -use std::io::Cursor; - -use super::FrontendCode; - -#[derive(Debug, Clone)] -pub struct Query { - pub statement: String, - // Used to mark that a Query message requires rewrite - dirty: bool, -} - -impl Query { - pub fn new(statement: String) -> Self { - Self { - statement, - dirty: false, - } - } - - pub fn requires_rewrite(&self) -> bool { - self.dirty - } - - pub fn rewrite(&mut self, statement: String) { - self.statement = statement; - self.dirty = true; - } -} - -impl TryFrom<&BytesMut> for Query { - type Error = Error; - - fn try_from(bytes: &BytesMut) -> Result { - let mut cursor = Cursor::new(bytes); - let code = cursor.get_u8(); - - if FrontendCode::from(code) != FrontendCode::Query { - return Err(ProtocolError::UnexpectedMessageCode { - expected: FrontendCode::Query.into(), - received: code as char, - } - .into()); - } - - let _len = cursor.get_i32(); // read and progress cursor - let query = cursor.read_string()?; - - Ok(Query { - statement: query, - dirty: false, - }) - } -} - -impl TryFrom for BytesMut { - type Error = Error; - - fn try_from(query: Query) -> Result { - let mut bytes = BytesMut::new(); - - let statement = CString::new(query.statement).map_err(|_| ProtocolError::UnexpectedNull)?; - let statement_bytes = statement.as_bytes_with_nul(); - - let len = SIZE_I32 + statement_bytes.len(); // len of query - - bytes.put_u8(FrontendCode::Query.into()); - bytes.put_i32(len as i32); - bytes.put_slice(statement_bytes); - - Ok(bytes) - } -} diff --git a/packages/cipherstash-proxy/src/postgresql/messages/ready_for_query.rs b/packages/cipherstash-proxy/src/postgresql/messages/ready_for_query.rs deleted file mode 100644 index 324a03d25..000000000 --- a/packages/cipherstash-proxy/src/postgresql/messages/ready_for_query.rs +++ /dev/null @@ -1,21 +0,0 @@ -use crate::{postgresql::messages::BackendCode, SIZE_I32, SIZE_U8}; -use bytes::{BufMut, BytesMut}; - -/// Bind (Z) message. -/// See: -#[derive(Clone, Debug)] -pub struct ReadyForQuery; - -impl From for BytesMut { - fn from(_: ReadyForQuery) -> BytesMut { - let mut bytes = BytesMut::new(); - - let len = SIZE_I32 + SIZE_U8; - - bytes.put_u8(BackendCode::ReadyForQuery.into()); - bytes.put_i32(len as i32); - bytes.put_u8(b'I'); - - bytes - } -} diff --git a/packages/cipherstash-proxy/src/postgresql/messages/row_description.rs b/packages/cipherstash-proxy/src/postgresql/messages/row_description.rs deleted file mode 100644 index 7fd8aea5e..000000000 --- a/packages/cipherstash-proxy/src/postgresql/messages/row_description.rs +++ /dev/null @@ -1,244 +0,0 @@ -use std::{ffi::CString, io::Cursor}; - -use bytes::{Buf, BufMut, BytesMut}; -use postgres_types::Type; - -use crate::{ - error::{Error, ProtocolError}, - postgresql::{format_code::FormatCode, protocol::BytesMutReadString}, - SIZE_I16, SIZE_I32, -}; - -use super::BackendCode; - -#[derive(Debug)] -pub struct RowDescription { - pub fields: Vec, -} - -#[derive(Debug)] -pub struct RowDescriptionField { - pub name: String, - pub table_oid: i32, - pub table_column: i16, - pub type_oid: i32, - pub type_size: i16, - pub type_modifier: i32, - pub format_code: FormatCode, - dirty: bool, -} - -impl RowDescription { - pub fn requires_rewrite(&self) -> bool { - self.fields.iter().any(|f| f.requires_rewrite()) - } - - pub fn map_types(&mut self, projection_types: &[Option]) { - self.fields - .iter_mut() - .zip(projection_types.iter()) - .for_each(|(field, t)| { - if let Some(t) = t { - field.rewrite_type_oid(t.clone()); - } - }); - } -} - -impl RowDescriptionField { - pub fn rewrite_type_oid(&mut self, postgres_type: postgres_types::Type) { - self.type_oid = postgres_type.oid() as i32; - self.dirty = true; - } - - pub fn requires_rewrite(&self) -> bool { - self.dirty - } -} - -impl TryFrom<&BytesMut> for RowDescription { - type Error = Error; - - fn try_from(bytes: &BytesMut) -> Result { - let mut cursor = Cursor::new(bytes); - - let code = cursor.get_u8(); - - if BackendCode::from(code) != BackendCode::RowDescription { - return Err(ProtocolError::UnexpectedMessageCode { - expected: BackendCode::RowDescription.into(), - received: code as char, - } - .into()); - } - - let _len = cursor.get_i32(); // move the cursor - let num_fields = cursor.get_i16() as usize; - - let fields = std::iter::repeat_with(|| RowDescriptionField::try_from(&mut cursor)) - .take(num_fields) - .collect::>()?; - - Ok(RowDescription { fields }) - } -} - -impl TryFrom for BytesMut { - type Error = Error; - - fn try_from(row_description: RowDescription) -> Result { - let mut bytes = BytesMut::new(); - - // Convert each field to bytes - let fields = row_description - .fields - .into_iter() - .map(BytesMut::try_from) - .collect::, _>>()?; - - let field_count = fields.len(); - let field_size = fields.iter().map(|x| x.len()).sum::(); - - let len = SIZE_I32 + SIZE_I16 + field_size; - - bytes.put_u8(BackendCode::RowDescription.into()); - bytes.put_i32(len as i32); - bytes.put_i16(field_count as i16); - - for field in fields.into_iter() { - bytes.put_slice(&field); - } - - Ok(bytes) - } -} - -// impl TryFrom<&BytesMut> for RowDescriptionField { -impl TryFrom<&mut Cursor<&BytesMut>> for RowDescriptionField { - type Error = Error; - - fn try_from(cursor: &mut Cursor<&BytesMut>) -> Result { - let name = cursor.read_string()?; - - let table_oid = cursor.get_i32(); - let table_column = cursor.get_i16(); - let type_oid = cursor.get_i32(); - - let type_size = cursor.get_i16(); - let type_modifier = cursor.get_i32(); - let format_code = cursor.get_i16().into(); - - Ok(Self { - name, - table_oid, - table_column, - type_oid, - type_size, - type_modifier, - format_code, - dirty: false, - }) - } -} - -impl TryFrom for BytesMut { - type Error = Error; - - fn try_from(field: RowDescriptionField) -> Result { - let mut bytes = BytesMut::new(); - - let name = CString::new(field.name)?; - let name = name.as_bytes_with_nul(); - - bytes.put_slice(name); - bytes.put_i32(field.table_oid); - bytes.put_i16(field.table_column); - bytes.put_i32(field.type_oid); - bytes.put_i16(field.type_size); - bytes.put_i32(field.type_modifier); - bytes.put_i16(field.format_code.into()); - - Ok(bytes) - } -} - -#[cfg(test)] -mod tests { - - use crate::{config::LogConfig, log, postgresql::messages::row_description::RowDescription}; - use bytes::BytesMut; - use tracing::info; - - fn to_message(s: &[u8]) -> BytesMut { - BytesMut::from(s) - } - - #[test] - pub fn map_projection_types() { - log::init(LogConfig::default()); - - // let mut pd = RowDescription { - // types: vec![ - // postgres_types::Type::TEXT, - // postgres_types::Type::INT4, - // postgres_types::Type::INT8, - // ], - // }; - - // let mapped_types = vec![ - // Some(postgres_types::Type::TEXT), - // None, - // Some(postgres_types::Type::TEXT), - // ]; - - // pd.map_types(&mapped_types); - - // let expected = vec![ - // postgres_types::Type::TEXT, - // postgres_types::Type::INT4, - // postgres_types::Type::TEXT, - // ]; - - // assert_eq!(pd.types, expected); - } - - #[test] - pub fn parse_row_description() { - log::init(LogConfig::default()); - let bytes = to_message( - b"T\0\0\0!\0\x01TimeZone\0\0\0\0\0\0\0\0\0\0\x19\xff\xff\xff\xff\xff\xff\0\0", - ); - - let expected = bytes.clone(); - - let row_description = RowDescription::try_from(&bytes).unwrap(); - - info!("{:?}", row_description); - - assert_eq!(row_description.fields.len(), 1); - assert_eq!(row_description.fields[0].name, "TimeZone"); - - let bytes = BytesMut::try_from(row_description).unwrap(); - assert_eq!(bytes, expected); - } - - #[test] - pub fn parse_row_description_with_many_fields() { - log::init(LogConfig::default()); - let bytes = to_message( - b"T\0\0\0J\0\x03id\0\0\0h,\0\x01\0\0\0\x14\0\x08\xff\xff\xff\xff\0\0name\0\0\0h,\0\x02\0\0\0\x19\xff\xff\xff\xff\xff\xff\0\0email\0\0\0h,\0\x03\0\0\x0e\xda\xff\xff\xff\xff\xff\xff\0\0" - ); - - let expected = bytes.clone(); - - let row_description = RowDescription::try_from(&bytes).unwrap(); - - assert_eq!(row_description.fields.len(), 3); - assert_eq!(row_description.fields[0].name, "id"); - assert_eq!(row_description.fields[1].name, "name"); - assert_eq!(row_description.fields[2].name, "email"); - - let bytes = BytesMut::try_from(row_description).unwrap(); - assert_eq!(bytes, expected); - } -} diff --git a/packages/cipherstash-proxy/src/postgresql/messages/target.rs b/packages/cipherstash-proxy/src/postgresql/messages/target.rs deleted file mode 100644 index cf9bfacd2..000000000 --- a/packages/cipherstash-proxy/src/postgresql/messages/target.rs +++ /dev/null @@ -1,46 +0,0 @@ -use crate::error::{Error, ProtocolError}; -use std::convert::TryFrom; - -/// -/// The target of describe or close messages. -/// -/// Valid values are PreparedStatement or Portal -/// -/// A Portal is a parsed statement PLUS any bound parameters -/// Describe with `Target::Portal` returns the RowDescription describing the result set. -/// The assumption is that the parameters are already bound to the portal, so the Describe message is not required to include any parameter information. -/// -/// Calls to Execute are made on a Portal (not a prepared statement) as execute requires any bound parameters -/// -/// A Statement is the parsed statement -/// Describe with `Target::Statement` returns a ParameterDescription followed by the RowDescription. -/// -/// -/// See https://www.postgresql.org/docs/current/protocol-flow.html#PROTOCOL-FLOW-EXT-QUERY -/// -#[derive(Debug, Clone)] -pub enum Target { - Portal, - Statement, -} - -impl TryFrom for Target { - type Error = Error; - - fn try_from(t: u8) -> Result { - match t as char { - 'S' => Ok(Target::Statement), - 'P' => Ok(Target::Portal), - t => Err(ProtocolError::UnexpectedDescribeTarget(t).into()), - } - } -} - -impl From for u8 { - fn from(target: Target) -> u8 { - match target { - Target::Statement => b'S', - Target::Portal => b'P', - } - } -} diff --git a/packages/cipherstash-proxy/src/postgresql/messages/terminate.rs b/packages/cipherstash-proxy/src/postgresql/messages/terminate.rs deleted file mode 100644 index 0024c10ef..000000000 --- a/packages/cipherstash-proxy/src/postgresql/messages/terminate.rs +++ /dev/null @@ -1,15 +0,0 @@ -use super::FrontendCode; -use bytes::{BufMut, BytesMut}; - -pub struct Terminate; - -impl Terminate { - pub fn message() -> BytesMut { - let mut bytes = BytesMut::new(); - - bytes.put_u8(FrontendCode::Terminate.into()); - bytes.put_i32(4); - - bytes - } -} diff --git a/packages/cipherstash-proxy/src/postgresql/backend.rs b/packages/cipherstash-proxy/src/postgresql/middleware/backend.rs similarity index 53% rename from packages/cipherstash-proxy/src/postgresql/backend.rs rename to packages/cipherstash-proxy/src/postgresql/middleware/backend.rs index d22730092..f0637a41d 100644 --- a/packages/cipherstash-proxy/src/postgresql/backend.rs +++ b/packages/cipherstash-proxy/src/postgresql/middleware/backend.rs @@ -1,30 +1,25 @@ -use super::context::Context; -use super::data::to_sql; -use super::error_handler::PostgreSqlErrorHandler; -use super::message_buffer::MessageBuffer; -use super::messages::error_response::ErrorResponse; -use super::messages::row_description::RowDescription; -use super::messages::{BackendCode, UNSPECIFIED_TYPE_OID}; -use super::Column; -use crate::connect::Sender; -use crate::error::{EncryptError, Error}; +use super::super::context::Context; +use super::super::data::to_sql; +use super::super::error_handler::PostgreSqlErrorHandler; +use super::super::rewrite::UNSPECIFIED_TYPE_OID; +use super::super::Column; +use crate::error::{EncryptError, Error, ProtocolError}; use crate::log::{CONTEXT, DEVELOPMENT, MAPPER, PROTOCOL}; use crate::postgresql::context::Portal; -use crate::postgresql::messages::data_row::DataRow; -use crate::postgresql::messages::param_description::ParamDescription; -use crate::postgresql::protocol::{self}; +use crate::postgresql::rewrite::data_row; use crate::prometheus::{ - CLIENTS_BYTES_SENT_TOTAL, DECRYPTED_VALUES_TOTAL, DECRYPTION_DURATION_SECONDS, - DECRYPTION_ERROR_TOTAL, DECRYPTION_REQUESTS_TOTAL, ROWS_ENCRYPTED_TOTAL, - ROWS_PASSTHROUGH_TOTAL, ROWS_TOTAL, SERVER_BYTES_RECEIVED_TOTAL, + DECRYPTED_VALUES_TOTAL, DECRYPTION_DURATION_SECONDS, DECRYPTION_ERROR_TOTAL, + DECRYPTION_REQUESTS_TOTAL, ROWS_ENCRYPTED_TOTAL, ROWS_PASSTHROUGH_TOTAL, ROWS_TOTAL, }; use crate::proxy::EncryptionService; use crate::EqlCiphertext; -use bytes::BytesMut; use metrics::{counter, histogram}; +use pg_proto::{ + AttributedBackendMessages, BackendBatchOutput, BackendMessage, BackendMiddlewareOutput, + OperationId, +}; use std::time::Instant; -use tokio::io::AsyncRead; -use tracing::{debug, error, info, warn}; +use tracing::{debug, error, info}; /// The PostgreSQL proxy backend that handles server-to-client message processing. /// @@ -70,26 +65,14 @@ use tracing::{debug, error, info, warn}; /// - `RowDescription`: Result column metadata (modified for encrypted columns) /// - `ParameterDescription`: Parameter metadata (modified for encrypted parameters) /// - `ReadyForQuery`: Session ready state (triggers schema reload if needed) -pub struct Backend -where - R: AsyncRead + Unpin, - S: EncryptionService, -{ - /// Sender for outgoing messages to client - client_sender: Sender, - /// Reader for incoming messages from server - server_reader: R, +pub struct Backend { /// Session context with portal and statement metadata context: Context, - /// Buffer for batching DataRow messages before decryption - buffer: MessageBuffer, } -impl Backend -where - R: AsyncRead + Unpin, - S: EncryptionService, -{ +impl Backend { + const RESPONSE_BUFFER_SIZE: usize = 4096; + /// Creates a new Backend instance. /// /// # Arguments @@ -98,162 +81,101 @@ where /// * `server_reader` - Stream for reading messages from the PostgreSQL server /// * `encrypt` - Encryption service for handling column decryption /// * `context` - Session context shared with the frontend - pub fn new(client_sender: Sender, server_reader: R, context: Context) -> Self { - let buffer = MessageBuffer::new(); - Backend { - client_sender, - server_reader, - context, - buffer, - } + pub fn new(context: Context) -> Self { + Backend { context } } - /// Main message processing loop for handling server messages. - /// - /// Reads messages from the PostgreSQL server, processes them based on message type, - /// performs decryption for encrypted result data, and forwards messages to the client. - /// - /// # PostgreSQL Protocol Phases - /// - /// ## Execute Phase - /// Execute operations produce a stream of DataRow messages followed by exactly one of: - /// - `CommandComplete` - Successful completion - /// - `EmptyQueryResponse` - Empty query completed - /// - `ErrorResponse` - Error occurred - /// - `PortalSuspended` - Portal execution suspended (LIMIT reached) - /// - /// ## Describe Phase - /// Describe operations return metadata about statements or portals: - /// - `ParameterDescription` - Parameter metadata (for statements) - /// - `RowDescription` - Result column metadata - /// - `NoData` - No result columns - /// - /// # Message Processing Flow - /// - /// 1. **Read Message**: Read and parse PostgreSQL wire protocol message - /// 2. **Check Passthrough**: Skip processing if encryption is disabled - /// 3. **Handle by Type**: Route to appropriate handler based on message code - /// 4. **Buffer Management**: Buffer DataRows, flush on completion/errors - /// 5. **Forward**: Send processed message to PostgreSQL client - /// - /// # Buffering Behavior - /// - /// DataRow messages are buffered for batch decryption to improve performance. - /// The buffer is automatically flushed when: - /// - Buffer reaches capacity - /// - Execute phase completes (CommandComplete, ErrorResponse, etc.) - /// - Non-DataRow message is encountered - /// - /// # Returns - /// - /// Returns `Ok(())` on successful message processing, or an `Error` if a fatal - /// error occurs that should terminate the connection. - pub async fn rewrite(&mut self) -> Result<(), Error> { - let read_start = Instant::now(); - let (code, mut bytes) = protocol::read_message( - &mut self.server_reader, - self.context.client_id, - self.context.connection_timeout(), - ) - .await?; - let read_duration = read_start.elapsed(); - self.context.record_execute_server_timing(read_duration); + pub async fn intercept( + &mut self, + operation: Option, + message: BackendMessage, + ) -> Result { + self.intercept_backend(operation, message).await + } - let sent: u64 = bytes.len() as u64; - counter!(SERVER_BYTES_RECEIVED_TOTAL).increment(sent); + pub async fn flush_held( + &mut self, + held: AttributedBackendMessages<'_>, + ) -> Result { + self.decrypt_held(held).await + } - // Log slow database responses (configurable threshold, default 100ms) - if read_duration > self.context.slow_db_response_min_duration() { - warn!( - client_id = self.context.client_id, - msg = "Slow database response", - duration_ms = read_duration.as_millis(), - message_code = ?code, - ); - } + async fn intercept_backend( + &mut self, + operation: Option, + protocol_message: BackendMessage, + ) -> Result { + let mut outbound_message = protocol_message.clone(); if self.context.is_passthrough() { debug!(target: DEVELOPMENT, client_id = self.context.client_id, msg = "Passthrough enabled" ); - self.write_with_flush(bytes).await?; - - // The frontend starts a session and enqueues an execute for every - // statement (start_session / set_execute), regardless of whether - // the statement is mapped. Those per-connection queues are only - // drained by complete_execution()/finish_session(), which are - // normally called when an execute terminates (below). Because the - // passthrough path returns early, we must drain them here too — - // otherwise the execute and session_metrics queues grow by one - // entry per statement and never shrink, leaking memory until the - // process is OOM-killed. See BUG-300. - match code.into() { - BackendCode::CommandComplete - | BackendCode::EmptyQueryResponse - | BackendCode::PortalSuspended - | BackendCode::ErrorResponse => { - self.context.complete_execution(); - self.context.finish_session(); + // CipherStash metadata is operation-keyed even in passthrough mode, + // and must be released when pg-proto identifies its terminal response. + match protocol_message { + BackendMessage::CommandComplete(_) + | BackendMessage::EmptyQueryResponse + | BackendMessage::PortalSuspended + | BackendMessage::ErrorResponse(_) => { + if let Some(operation) = operation { + let session = self.context.complete_execution(operation); + self.context.finish_session(session); + } + } + BackendMessage::RowDescription(_) | BackendMessage::NoData => { + if let Some(operation) = operation { + self.context.complete_describe(operation); + } + } + BackendMessage::ReadyForQuery(status) => { + self.context.set_transaction_status(status); } _ => {} } - return Ok(()); + return Ok(BackendMiddlewareOutput::Forward(outbound_message)); } let keyset_id = self.context.keyset_identifier(); debug!(target: CONTEXT, client_id = ?self.context.client_id, ?keyset_id); - match code.into() { - BackendCode::DataRow => { + match protocol_message { + BackendMessage::DataRow(_) => { // Encrypted DataRows are added to the buffer and we return early // Otherwise, continue and write - if self.data_row_handler(&bytes).await? { - return Ok(()); + if self.data_row_handler(operation).await? { + return Ok(BackendMiddlewareOutput::Hold); } } // Execute phase is always terminated by the appearance of exactly one of these messages: // CommandComplete, EmptyQueryResponse (if the portal was created from an empty query string), ErrorResponse, or PortalSuspended. - BackendCode::CommandComplete - | BackendCode::EmptyQueryResponse - | BackendCode::PortalSuspended => { + BackendMessage::CommandComplete(_) + | BackendMessage::EmptyQueryResponse + | BackendMessage::PortalSuspended => { debug!(target: PROTOCOL, client_id = self.context.client_id, msg = "CommandComplete | EmptyQueryResponse | PortalSuspended"); - match self.flush().await { - Ok(_) => (), - Err(err) => { - warn!(client_id = self.client_id(), error = err.to_string()); - self.send_error_response(err)?; - } + if let Some(operation) = operation { + let session = self.context.complete_execution(operation); + self.context.finish_session(session); } - - self.context.complete_execution(); - self.context.finish_session(); } - BackendCode::ErrorResponse => { - if let Some(b) = self.error_response_handler(&bytes)? { - bytes = b - } + BackendMessage::ErrorResponse(ref response) => { + self.error_response_handler(response); - match self.flush().await { - Ok(_) => (), - Err(err) => { - warn!(client_id = self.client_id(), error = err.to_string()); - self.send_error_response(err)?; - } + if let Some(operation) = operation { + let session = self.context.complete_execution(operation); + self.context.finish_session(session); } - - self.context.complete_execution(); - self.context.finish_session(); } // Describe with Target:Statement // Returns a ParameterDescription followed by RowDescription // The Describe is complete after the RowDescription - BackendCode::ParameterDescription => { - if let Some(b) = self.parameter_description_handler(&bytes).await? { - bytes = b + BackendMessage::ParameterDescription(types) => { + if let Some(message) = self.parameter_description_handler(operation, types).await? { + outbound_message = message; } } // Describe with Target:Statement or Target::Portal @@ -261,21 +183,26 @@ where // Target::Portal returns a RowDescription // If no rows are returned, NoData is returned instead of a RowDescription // Complete the Describe - BackendCode::RowDescription => { - if let Some(b) = self.row_description_handler(&bytes).await? { - bytes = b + BackendMessage::RowDescription(description) => { + if let Some(message) = self.row_description_handler(operation, description).await? { + outbound_message = message; + } + if let Some(operation) = operation { + self.context.complete_describe(operation); } - self.context.complete_describe(); } // Describe with Target:Statement or Target::Portal // If the statement returns no rows, NoData is returned instead of a RowDescription - BackendCode::NoData => { - self.context.complete_describe(); + BackendMessage::NoData => { + if let Some(operation) = operation { + self.context.complete_describe(operation); + } } // Reload for SompleQuery flow // Reload is potentially triggered by a FrontEnd Sync message. // However, the SimpleQuery flow does not use Sync so we check here as well - BackendCode::ReadyForQuery => { + BackendMessage::ReadyForQuery(status) => { + self.context.set_transaction_status(status); debug!(target: PROTOCOL, client_id = self.context.client_id, msg = "ReadyForQuery" @@ -285,18 +212,16 @@ where } } - code => { + _ => { debug!(target: PROTOCOL, client_id = self.context.client_id, msg = "Passthrough", - ?code, + message = ?protocol_message, ); } } - self.write_with_flush(bytes).await?; - - Ok(()) + Ok(BackendMiddlewareOutput::Forward(outbound_message)) } /// Handles PostgreSQL ErrorResponse messages from the server. @@ -335,11 +260,9 @@ where /// /// Always returns `Some(bytes)` containing the original error response /// to forward to the client unchanged. - fn error_response_handler(&mut self, bytes: &BytesMut) -> Result, Error> { - let error_response = ErrorResponse::try_from(bytes)?; - error!(msg = "PostgreSQL Error", error = ?error_response); + fn error_response_handler(&mut self, response: &pg_proto::DiagnosticResponse) { + error!(msg = "PostgreSQL Error", fields = ?response.fields); info!(msg = "PostgreSQL Errors originate in the database"); - Ok(Some(bytes.to_owned())) } /// @@ -349,49 +272,6 @@ where /// - when the buffer is full /// - when any other message type is written /// - async fn buffer(&mut self, data_row: DataRow) -> Result<(), Error> { - self.buffer.push(data_row); - if self.buffer.at_capacity() { - debug!(target: DEVELOPMENT, client_id = self.context.client_id, msg = "Flush message buffer"); - self.flush().await?; - } - Ok(()) - } - - /// - /// Write a message to the client - /// Flushes all messages in the buffer before writing the message - /// - pub async fn write_with_flush(&mut self, bytes: BytesMut) -> Result<(), Error> { - debug!(target: DEVELOPMENT, client_id = self.context.client_id, msg = "Write"); - - match self.flush().await { - Ok(_) => (), - Err(err) => { - warn!(client_id = self.client_id(), error = err.to_string()); - self.send_error_response(err)?; - } - } - - self.write(bytes).await?; - Ok(()) - } - - /// - /// Write a message to the client - /// - pub async fn write(&mut self, bytes: BytesMut) -> Result<(), Error> { - let sent: u64 = bytes.len() as u64; - counter!(CLIENTS_BYTES_SENT_TOTAL).increment(sent); - - let start = Instant::now(); - self.client_sender.send(bytes)?; - let duration = start.elapsed(); - self.context.add_client_write_duration_for_execute(duration); - - Ok(()) - } - /// Flushes all buffered DataRow messages by performing batch decryption. /// /// This is the core decryption logic that processes buffered DataRow messages, @@ -431,32 +311,40 @@ where /// appropriate error responses and recorded in metrics. The error mapping /// implemented in the encryption service ensures proper keyset ID context /// is preserved in error messages. - async fn flush(&mut self) -> Result<(), Error> { - if self.buffer.is_empty() { - debug!(target: MAPPER, client_id = self.context.client_id, msg = "Empty buffer"); + async fn decrypt_held( + &mut self, + held: AttributedBackendMessages<'_>, + ) -> Result { + let mut operation = None; + let mut rows = Vec::with_capacity(held.iter().len()); + for (row_operation, message) in held.iter() { + let row_operation = row_operation.ok_or(ProtocolError::HeldDataRowMissingOperation)?; + if operation + .replace(row_operation) + .is_some_and(|current| current != row_operation) + { + return Err(ProtocolError::HeldDataRowOperationMismatch.into()); + } + let BackendMessage::DataRow(row) = message else { + return Err(ProtocolError::HeldBackendMessageNotDataRow.into()); + }; + rows.push(row.clone()); } - let portal = self.context.get_portal_from_execute(); + let portal = + operation.and_then(|operation| self.context.get_portal_from_execute(operation)); let portal = match portal.as_deref() { Some(Portal::Encrypted { .. }) => portal.unwrap(), _ => { debug!(target: MAPPER, client_id = self.context.client_id, msg = "Passthrough portal"); - if !self.buffer.is_empty() { - error!( - client_id = self.context.client_id, - msg = "Buffer is not empty" - ); - } - return Ok(()); + return Err(ProtocolError::HeldDataRowsNotEncrypted.into()); } }; - - let mut rows: Vec = self.buffer.drain().into_iter().collect(); debug!(target: DEVELOPMENT, client_id = self.context.client_id, rows = rows.len()); let result_column_count = match rows.first() { - Some(row) => row.column_count(), - None => return Ok(()), + Some(row) => row.columns.len(), + None => return Ok(BackendBatchOutput::ReplaceOneToOne(Vec::new())), }; // Result Column Format Codes are passed with the Bind message @@ -470,7 +358,7 @@ where // Each row is converted into Vec> let ciphertexts: Vec> = rows .iter_mut() - .flat_map(|row| row.as_ciphertext(projection_columns)) + .flat_map(|row| data_row::as_ciphertext(row, projection_columns)) .collect::>(); let start = Instant::now(); @@ -492,7 +380,10 @@ where let duration = Instant::now().duration_since(start); // Always record for slow-statement diagnostics - self.context.add_decrypt_duration_for_execute(duration); + if let Some(operation) = operation { + self.context + .add_decrypt_duration_for_execute(operation, duration); + } // Prometheus metrics remain gated if self.context.prometheus_enabled() { @@ -511,6 +402,7 @@ where // Stitch Plaintext back into Rows encoded with the appropriate Format Code // Each chunk is written to the client + let mut messages = Vec::with_capacity(held.iter().len()); for (chunk, mut row) in rows { let data = chunk .iter() @@ -521,13 +413,11 @@ where }) .collect::, _>>()?; - row.rewrite(&data)?; + data_row::rewrite(&mut row, &data)?; - let bytes = BytesMut::try_from(row)?; - self.write(bytes).await?; + messages.push(BackendMessage::DataRow(row)); } - - Ok(()) + Ok(BackendBatchOutput::ReplaceOneToOne(messages)) } fn check_column_config( @@ -565,13 +455,14 @@ where async fn parameter_description_handler( &self, - bytes: &BytesMut, - ) -> Result, Error> { - let mut description = ParamDescription::try_from(bytes)?; - + operation: Option, + description: Vec, + ) -> Result, Error> { debug!(target: PROTOCOL, client_id = self.context.client_id, ParamDescription = ?description); - if let Some(statement) = self.context.get_statement_from_describe() { + if let Some(statement) = + operation.and_then(|operation| self.context.get_statement_from_describe(operation)) + { // Describe the params the CLIENT wrote, not the ones PostgreSQL was // sent. A rewrite may have fused or dropped params, in which case // the server's description is both shorter than and shifted from @@ -592,23 +483,25 @@ where .output_params .iter() .position(|output| output.source.primary_input() == idx) - .and_then(|output_idx| description.types.get(output_idx).copied()) + .and_then(|output_idx| description.get(output_idx).copied()) + .map(|oid| oid as i32) .unwrap_or(UNSPECIFIED_TYPE_OID), }) .collect::>(); debug!(target: MAPPER, client_id = self.context.client_id, param_types = ?param_types); - description.set_types(param_types); - } - - if description.requires_rewrite() { - let bytes = BytesMut::try_from(description)?; - debug!(target: MAPPER, client_id = self.context.client_id, msg = "Rewrite ParamDescription", bytes = ?bytes); - Ok(Some(bytes)) - } else { - Ok(None) + let rewritten = param_types + .into_iter() + .map(|oid| oid as u32) + .collect::>(); + if rewritten != description { + let message = BackendMessage::ParameterDescription(rewritten); + debug!(target: MAPPER, client_id = self.context.client_id, msg = "Rewrite ParamDescription", ?message); + return Ok(Some(message)); + } } + Ok(None) } /// @@ -620,13 +513,14 @@ where /// async fn row_description_handler( &mut self, - bytes: &BytesMut, - ) -> Result, Error> { - let mut description = RowDescription::try_from(bytes)?; - + operation: Option, + mut description: pg_proto::RowDescription, + ) -> Result, Error> { debug!(target: PROTOCOL, client_id = self.context.client_id, RowDescription = ?description); - if let Some(statement) = self.context.get_statement_for_row_decription() { + if let Some(statement) = + operation.and_then(|operation| self.context.get_statement_for_operation(operation)) + { let projection_types = statement .projection_columns .iter() @@ -635,16 +529,21 @@ where debug!(target: MAPPER, client_id = self.context.client_id, projection_types = ?projection_types); - description.map_types(&projection_types); - } - - if description.requires_rewrite() { - let bytes = BytesMut::try_from(description)?; - debug!(target: MAPPER, client_id = self.context.client_id, msg = "Rewrite RowDescription", bytes = ?bytes); - Ok(Some(bytes)) - } else { - Ok(None) + let mut rewritten = false; + for (field, postgres_type) in description.fields.iter_mut().zip(projection_types) { + if let Some(postgres_type) = postgres_type { + let oid = postgres_type.oid(); + rewritten |= field.type_oid != oid; + field.type_oid = oid; + } + } + if rewritten { + let message = BackendMessage::RowDescription(description); + debug!(target: MAPPER, client_id = self.context.client_id, msg = "Rewrite RowDescription", ?message); + return Ok(Some(message)); + } } + Ok(None) } /// Handles PostgreSQL DataRow messages containing query result data. @@ -681,15 +580,15 @@ where /// /// Records metrics for both encrypted and passthrough row processing to /// track proxy performance and encryption usage patterns. - async fn data_row_handler(&mut self, bytes: &BytesMut) -> Result { + async fn data_row_handler(&mut self, operation: Option) -> Result { counter!(ROWS_TOTAL).increment(1); - match self.context.get_portal_from_execute().as_deref() { + match operation + .and_then(|operation| self.context.get_portal_from_execute(operation)) + .as_deref() + { Some(Portal::Encrypted { .. }) => { debug!(target: MAPPER, client_id = self.context.client_id, msg = "Encrypted"); - let data_row = DataRow::try_from(bytes)?; - self.buffer(data_row).await?; - counter!(ROWS_ENCRYPTED_TOTAL).increment(1); Ok(true) } @@ -703,54 +602,19 @@ where } /// Implementation of PostgreSQL error handling for the Backend component. -impl PostgreSqlErrorHandler for Backend -where - R: AsyncRead + Unpin, - S: EncryptionService, -{ - fn client_sender(&mut self) -> &mut Sender { - &mut self.client_sender - } - +impl PostgreSqlErrorHandler for Backend { fn client_id(&self) -> i32 { self.context.client_id } - - /// Backend-specific error response handling. - /// - /// Unlike the frontend, the backend doesn't need to set an error state - /// since errors during result processing should immediately terminate - /// the current query execution. - fn send_error_response(&mut self, err: Error) -> Result<(), Error> { - let error_response = self.error_to_response(err); - // Ensure any buffered data is cleared before sending error - self.buffer.clear(); - - let message = BytesMut::try_from(error_response)?; - - debug!( - target: "PROTOCOL", - client_id = self.context.client_id, - msg = "backend_send_error_response", - ?message, - ); - - self.client_sender.send(message)?; - - Ok(()) - } } #[cfg(test)] mod tests { use super::*; - use crate::config::{LogConfig, TandemConfig}; - use crate::log; + use crate::config::TandemConfig; use crate::postgresql::context::KeysetIdentifier; - use crate::postgresql::messages::Name; use crate::proxy::{EncryptConfig, EncryptionService}; use eql_mapper::Schema; - use std::io::Cursor; use std::sync::Arc; use tokio::sync::mpsc; @@ -793,116 +657,4 @@ mod tests { reload_sender, ) } - - /// Encodes a backend message as wire bytes (one per execute-terminating code). - type MessageEncoder = fn() -> BytesMut; - - /// Frame a backend message on the wire: 1-byte code + Int32 length - /// (body length + 4) + body. Sufficient for the passthrough path, which - /// matches on the code only and does not parse the body. - fn backend_message(code: u8, body: &[u8]) -> BytesMut { - let len = (body.len() + 4) as i32; - - let mut bytes = BytesMut::new(); - bytes.extend_from_slice(&[code]); - bytes.extend_from_slice(&len.to_be_bytes()); - bytes.extend_from_slice(body); - bytes - } - - /// `'C'` CommandComplete, carrying a command tag. - fn command_complete_bytes() -> BytesMut { - backend_message(b'C', b"SELECT 1\0") - } - - /// `'I'` EmptyQueryResponse — no body. - fn empty_query_response_bytes() -> BytesMut { - backend_message(b'I', b"") - } - - /// `'s'` PortalSuspended — no body. - fn portal_suspended_bytes() -> BytesMut { - backend_message(b's', b"") - } - - /// `'E'` ErrorResponse — a sequence of (field-type, C-string) pairs - /// terminated by a zero byte. Content is irrelevant here: the passthrough - /// path forwards the bytes and matches on the code without parsing them. - fn error_response_bytes() -> BytesMut { - backend_message(b'E', b"SERROR\0CXX000\0Mboom\0\0") - } - - /// Regression test for BUG-300 (passthrough memory leak). - /// - /// The frontend enqueues a session + execute for *every* statement. Those - /// per-connection `execute` / `session_metrics` queues are only drained by - /// `complete_execution()` / `finish_session()`. Before the fix, the - /// passthrough branch in `rewrite()` returned early without calling these, - /// so the queues grew by one entry per statement and leaked until OOM. - /// - /// This drives `Backend::rewrite()` through the passthrough branch and - /// asserts both queues stay empty across many statements — once for *each* - /// execute-terminating message code the fix drains on, so dropping any arm - /// of that match is caught. It fails against the pre-fix backend (which - /// never drained in passthrough) — i.e. it actually guards the bug. - #[tokio::test] - async fn passthrough_drains_queues_on_execute_terminating_message() { - log::init(LogConfig::default()); - - const STATEMENTS: usize = 256; - - // Every code that terminates the execute phase must drain the queues. - let cases: [(&str, MessageEncoder); 4] = [ - ("CommandComplete", command_complete_bytes), - ("EmptyQueryResponse", empty_query_response_bytes), - ("PortalSuspended", portal_suspended_bytes), - ("ErrorResponse", error_response_bytes), - ]; - - for (label, encode) in cases { - let context = passthrough_context(); - assert!( - context.is_passthrough(), - "test context must be in passthrough mode" - ); - - // A stream of identical terminating messages — one per statement — - // that the backend reads from the "server". - let message = encode(); - let mut server_bytes = BytesMut::new(); - for _ in 0..STATEMENTS { - server_bytes.extend_from_slice(&message); - } - - // Keep the client receiver alive so write_with_flush succeeds. - let (client_sender, _client_receiver) = mpsc::unbounded_channel(); - let reader = Cursor::new(server_bytes.to_vec()); - let mut backend = Backend::new(client_sender, reader, context); - - for i in 0..STATEMENTS { - // Frontend: enqueue a session + execute for the statement. - let session_id = backend.context.start_session(); - backend - .context - .set_execute(Name::unnamed(), Some(session_id)); - - // Backend: process the terminating message via the passthrough - // path, which must drain the queues. - backend.rewrite().await.unwrap(); - - // The queues must be drained every iteration — not grow by one - // per statement (the BUG-300 leak). - assert_eq!( - backend.context.execute_queue_len(), - 0, - "{label}: execute queue not drained at statement {i}" - ); - assert_eq!( - backend.context.session_metrics_queue_len(), - 0, - "{label}: session_metrics queue not drained at statement {i}" - ); - } - } - } } diff --git a/packages/cipherstash-proxy/src/postgresql/frontend.rs b/packages/cipherstash-proxy/src/postgresql/middleware/frontend.rs similarity index 72% rename from packages/cipherstash-proxy/src/postgresql/frontend.rs rename to packages/cipherstash-proxy/src/postgresql/middleware/frontend.rs index 0e5436428..98abe046e 100644 --- a/packages/cipherstash-proxy/src/postgresql/frontend.rs +++ b/packages/cipherstash-proxy/src/postgresql/middleware/frontend.rs @@ -1,15 +1,8 @@ -use super::context::phase_timing::PhaseTimer; -use super::context::{Context, SessionId, Statement}; -use super::error_handler::PostgreSqlErrorHandler; -use super::messages::bind::Bind; -use super::messages::describe::Describe; -use super::messages::execute::Execute; -use super::messages::parse::Parse; -use super::messages::query::Query; -use super::messages::FrontendCode as Code; -use super::parser::SqlParser; -use super::protocol::{self}; -use crate::connect::Sender; +use super::super::context::phase_timing::PhaseTimer; +use super::super::context::{Context, SessionId, Statement}; +use super::super::error_handler::PostgreSqlErrorHandler; +use super::super::parser::SqlParser; +use super::super::rewrite::bind::Bind; use crate::error::{EncryptError, Error, MappingError}; use crate::log::{MAPPER, PROTOCOL}; use crate::postgresql::context::column::Column; @@ -21,32 +14,30 @@ use crate::postgresql::context::Portal; use crate::postgresql::data::{ compose_json_selector_path, json_value_selector_plaintext, literal_from_sql, literal_json_value, }; -use crate::postgresql::messages::close::Close; -use crate::postgresql::messages::error_response::ErrorResponseCode; -use crate::postgresql::messages::ready_for_query::ReadyForQuery; -use crate::postgresql::messages::terminate::Terminate; -use crate::postgresql::messages::{Name, Target}; +use crate::postgresql::rewrite::Name; +use crate::postgresql::rewrite::UNSPECIFIED_TYPE_OID; use crate::prometheus::{ - CLIENTS_BYTES_RECEIVED_TOTAL, ENCRYPTED_VALUES_TOTAL, ENCRYPTION_DURATION_SECONDS, - ENCRYPTION_ERROR_TOTAL, ENCRYPTION_REQUESTS_TOTAL, SERVER_BYTES_SENT_TOTAL, - STATEMENTS_ENCRYPTED_TOTAL, STATEMENTS_PASSTHROUGH_MAPPING_DISABLED_TOTAL, - STATEMENTS_PASSTHROUGH_TOTAL, STATEMENTS_UNMAPPABLE_TOTAL, + ENCRYPTED_VALUES_TOTAL, ENCRYPTION_DURATION_SECONDS, ENCRYPTION_ERROR_TOTAL, + ENCRYPTION_REQUESTS_TOTAL, STATEMENTS_ENCRYPTED_TOTAL, + STATEMENTS_PASSTHROUGH_MAPPING_DISABLED_TOTAL, STATEMENTS_PASSTHROUGH_TOTAL, + STATEMENTS_UNMAPPABLE_TOTAL, }; use crate::proxy::EncryptionService; use crate::{EqlOutput, EqlQueryPayload}; -use bytes::BytesMut; use cipherstash_client::encryption::Plaintext; use eql_mapper::{self, EqlMapperError, EqlTermVariant, JsonSelectorSegment, TypeCheckedStatement}; use metrics::{counter, histogram}; -use pg_escape::quote_literal; +use pg_proto::{ + BackendMessage, Close, Describe, DescribeTarget, Execute, FrontendMessage, + FrontendMiddlewareOutput, Parse, +}; use serde::Serialize; use sqltk::parser::ast::{self, Value}; use sqltk::NodeKey; use std::collections::HashMap; use std::sync::Arc; use std::time::Instant; -use tokio::io::{AsyncRead, AsyncWrite, AsyncWriteExt}; -use tracing::{debug, error, info, warn}; +use tracing::{debug, info, warn}; /// The PostgreSQL proxy frontend that handles client-to-server message processing. /// @@ -91,43 +82,12 @@ use tracing::{debug, error, info, warn}; /// Encryption and mapping errors are converted to appropriate PostgreSQL error responses /// and sent back to the client. The frontend maintains error state to properly handle /// the PostgreSQL extended query error recovery protocol. -pub struct Frontend -where - R: AsyncRead + Unpin, - W: AsyncWrite + Unpin, - S: EncryptionService, -{ - /// Reader for incoming client messages - client_reader: R, - /// Sender for outgoing messages to client - client_sender: Sender, - /// Writer for forwarding messages to server - server_writer: W, +pub struct Frontend { /// Session context tracking statements, portals, and keyset IDs context: Context, - /// Error state flag for extended query protocol error handling - error_state: Option, -} - -/// How a frontend failure was delivered, which determines how the batch's -/// Sync must be answered. -#[derive(Debug)] -enum ErrorState { - /// An exception statement was forwarded to the server in place of the - /// failed message. The server produces the ErrorResponse and ReadyForQuery - /// for this batch, so the client's Sync is discarded. - ExceptionInjected, - /// A FATAL error was written directly to the client; nothing was forwarded - /// to the server, so the proxy answers the Sync with its own ReadyForQuery. - ErrorResponseSent, } -impl Frontend -where - R: AsyncRead + Unpin, - W: AsyncWrite + Unpin, - S: EncryptionService, -{ +impl Frontend { /// Creates a new Frontend instance. /// /// # Arguments @@ -136,253 +96,161 @@ where /// * `client_sender` - Channel sender for sending messages back to client /// * `server_writer` - Stream for writing messages to the PostgreSQL server /// * `context` - Session context for tracking statements and portals with service access - pub fn new( - client_reader: R, - client_sender: Sender, - server_writer: W, - context: Context, - ) -> Self { - Frontend { - client_reader, - client_sender, - server_writer, - context, - error_state: None, - } + pub fn new(context: Context) -> Self { + Frontend { context } } - /// Main message processing loop for handling client messages. - /// - /// Reads a message from the client, processes it based on the PostgreSQL message type, - /// performs any necessary encryption/transformation, and forwards it to the server. - /// - /// # Message Processing Flow - /// - /// 1. **Read Message**: Read and parse the PostgreSQL wire protocol message - /// 2. **Check Mapping**: Skip processing if mapping is disabled - /// 3. **Handle by Type**: Route to appropriate handler based on message type - /// 4. **Error Recovery**: Handle extended query protocol error states - /// 5. **Forward**: Send processed message to PostgreSQL server - /// - /// # Error States - /// - /// When an error occurs during extended query processing, the frontend enters - /// error state and discards messages until a Sync message is received, following - /// the PostgreSQL protocol specification. - /// - /// # Returns - /// - /// Returns `Ok(())` on successful message processing, or an `Error` if a fatal - /// error occurs that should terminate the connection. - pub async fn rewrite(&mut self) -> Result<(), Error> { - let (code, mut bytes) = protocol::read_message( - &mut self.client_reader, - self.context.client_id, - self.context.connection_timeout(), - ) - .await?; - - let sent: u64 = bytes.len() as u64; - counter!(CLIENTS_BYTES_RECEIVED_TOTAL).increment(sent); - + pub async fn intercept( + &mut self, + operation: pg_proto::OperationId, + protocol_message: FrontendMessage, + ) -> Result { if self.context.mapping_disabled() { - self.write_to_server(bytes).await?; - return Ok(()); + return Ok(FrontendMiddlewareOutput::Forward(protocol_message)); } - let code = Code::from(code); - - // When an error is detected while processing any extended-query message, the backend issues ErrorResponse, then reads and discards messages until a Sync is reached, - // https://www.postgresql.org/docs/current/protocol-flow.html#PROTOCOL-FLOW-EXT-QUERY - if self.error_state.is_some() { - warn!(target: PROTOCOL, - client_id = self.context.client_id, - error_state = ?self.error_state, - ?code, - ); - if code != Code::Sync { - return Ok(()); - } - } + let mut outbound_message = protocol_message.clone(); - match code { - Code::Query => { - match self.query_handler(&bytes).await { - Ok(Some(mapped)) => bytes = mapped, + match protocol_message { + FrontendMessage::Query(query) => { + match self.query_handler(operation, query).await { + Ok(Some(mapped)) => outbound_message = mapped, // No mapping needed, don't change the bytes Ok(None) => (), Err(err) => { + let session_id = self.context.latest_session_id(); + self.context.finish_session(session_id); warn!( client_id = self.context.client_id, msg = "Query Handler Error", error = ?err.to_string(), ); - // Replace the failed Query with an exception statement. - // The server's ErrorResponse and ReadyForQuery answer - // the client's Query in order. See handle_statement_error. - match self.handle_statement_error(err)? { - Some(exception) => bytes = exception, - // FATAL error written directly to the client; the - // simple protocol still expects a ReadyForQuery. - None => { - self.send_ready_for_query()?; - return Ok(()); - } - } + return Ok(FrontendMiddlewareOutput::Respond { + request: outbound_message, + responses: vec![ + self.error_response(err), + BackendMessage::ReadyForQuery(self.context.transaction_status()), + ], + }); } } } - Code::Describe => { - self.describe_handler(&bytes).await?; + FrontendMessage::Describe(describe) => { + self.describe_handler(operation, describe).await?; } - Code::Execute => { - self.execute_handler(&bytes).await?; + FrontendMessage::Execute(execute) => { + self.execute_handler(operation, execute).await?; } - Code::Parse => { - match self.parse_handler(&bytes).await { - Ok(Some(mapped)) => bytes = mapped, + FrontendMessage::Parse(parse) => { + let statement = parse.statement.clone(); + match self.parse_handler(parse).await { + Ok(Some(mapped)) => outbound_message = mapped, // No mapping needed, don't change the bytes Ok(None) => (), Err(err) => { + self.context.close_statement(&statement); warn!( client_id = self.context.client_id, msg = "Parse Handler Error", error = ?err.to_string(), ); - // Replace the failed Parse with an exception statement - // and drop the rest of the batch until Sync. The - // server's ErrorResponse and ReadyForQuery answer the - // client's batch in order. See handle_statement_error. - match self.handle_statement_error(err)? { - Some(exception) => { - self.error_state = Some(ErrorState::ExceptionInjected); - bytes = exception; - } - None => { - self.error_state = Some(ErrorState::ErrorResponseSent); - return Ok(()); - } - } + return Ok(FrontendMiddlewareOutput::Respond { + request: outbound_message, + responses: vec![self.error_response(err)], + }); } } } - Code::Bind => { - match self.bind_handler(&bytes).await { - Ok(Some(mapped)) => bytes = mapped, + FrontendMessage::Bind(bind) => { + match self.bind_handler(Bind::try_from(bind)?).await { + Ok(Some(mapped)) => outbound_message = mapped, // No mapping needed, don't change the bytes Ok(None) => (), - Err(err) => { - warn!(target: PROTOCOL, - client_id = self.context.client_id, - msg = "Bind Handler Error", - error = ?err.to_string(), - ); - // Replace the failed Bind with an exception statement - // and drop the rest of the batch until Sync. The - // server's ErrorResponse and ReadyForQuery answer the - // client's batch in order. See handle_statement_error. - match self.handle_statement_error(err)? { - Some(exception) => { - self.error_state = Some(ErrorState::ExceptionInjected); - bytes = exception; - } - None => { - self.error_state = Some(ErrorState::ErrorResponseSent); - return Ok(()); - } + Err(err) => match err { + Error::Mapping(MappingError::InvalidParameter(_)) => { + warn!(target: PROTOCOL, + client_id = self.context.client_id, + msg = "EncryptError::InvalidParameter", + ); + return Ok(FrontendMiddlewareOutput::Respond { + request: outbound_message, + responses: vec![self.error_response(err)], + }); } - } + Error::Encrypt(EncryptError::UnknownKeysetIdentifier { .. }) => { + warn!(target: PROTOCOL, + client_id = self.context.client_id, + msg = "EncryptError::UnknownKeysetIdentifier", + ); + return Ok(FrontendMiddlewareOutput::Respond { + request: outbound_message, + responses: vec![self.error_response(err)], + }); + } + _ => { + warn!(target: PROTOCOL, + client_id = self.context.client_id, + msg = "Bind Error", + err = err.to_string() + ); + return Ok(FrontendMiddlewareOutput::Respond { + request: outbound_message, + responses: vec![self.error_response(err)], + }); + } + }, } } - Code::Sync => { + FrontendMessage::Sync => { debug!(target: PROTOCOL, client_id = self.context.client_id, - ?code, + message = ?protocol_message, ); self.context.reload_schema_if_changed().await; - - match self.error_state.take() { - Some(ErrorState::ExceptionInjected) => { - // The exception statement injected on failure already - // elicits the ErrorResponse and ReadyForQuery for this - // batch from the server. Forwarding the Sync would - // produce a second ReadyForQuery and desync the client. - debug!(target: PROTOCOL, - client_id = self.context.client_id, - msg = "Discard Sync for failed batch", - ); - return Ok(()); - } - Some(ErrorState::ErrorResponseSent) => { - // Nothing was forwarded to the server for this batch; - // answer the Sync directly. - self.send_ready_for_query()?; - return Ok(()); - } - None => {} - } } - Code::Close => { - self.close_handler(&bytes).await?; + FrontendMessage::Close(close) => { + self.close_handler(close).await?; } - code => { + _ => { debug!(target: PROTOCOL, client_id = self.context.client_id, msg = "Passthrough", - ?code, + message = ?protocol_message, ); } } - self.write_to_server(bytes).await?; - Ok(()) + Ok(FrontendMiddlewareOutput::Forward(outbound_message)) } - pub async fn write_to_server(&mut self, bytes: BytesMut) -> Result<(), Error> { - debug!(target: PROTOCOL, msg = "Write to server", ?bytes); - let sent: u64 = bytes.len() as u64; - counter!(SERVER_BYTES_SENT_TOTAL).increment(sent); - - let start = Instant::now(); - self.server_writer.write_all(&bytes).await?; - let duration = start.elapsed(); - if let Some(session_id) = self.context.latest_session_id() { - self.context.add_server_write_duration(session_id, duration); - } - - Ok(()) - } - - pub async fn terminate(&mut self) -> Result<(), Error> { - debug!(target: PROTOCOL, msg = "Terminate server connection"); - let bytes = Terminate::message(); - self.write_to_server(bytes).await?; - Ok(()) - } - - async fn describe_handler(&mut self, bytes: &BytesMut) -> Result<(), Error> { - let describe = Describe::try_from(bytes)?; + async fn describe_handler( + &mut self, + operation: pg_proto::OperationId, + describe: Describe, + ) -> Result<(), Error> { debug!(target: PROTOCOL, client_id = self.context.client_id, ?describe); - self.context.set_describe(describe); + self.context.set_describe(operation, describe); Ok(()) } - async fn close_handler(&mut self, bytes: &BytesMut) -> Result<(), Error> { - let close = Close::try_from(bytes)?; + async fn close_handler(&mut self, close: Close) -> Result<(), Error> { debug!(target: PROTOCOL, client_id = self.context.client_id, ?close); match close.target { - Target::Portal => self.context.close_portal(&close.name), - Target::Statement => self.context.close_statement_and_portal(&close.name), + DescribeTarget::Portal => self.context.close_portal(&close.name), + DescribeTarget::Statement => self.context.close_statement_and_portal(&close.name), } Ok(()) } - async fn execute_handler(&mut self, bytes: &BytesMut) -> Result<(), Error> { - let execute = Execute::try_from(bytes)?; + async fn execute_handler( + &mut self, + operation: pg_proto::OperationId, + execute: Execute, + ) -> Result<(), Error> { debug!(target: PROTOCOL, client_id = self.context.client_id, ?execute); self.context - .set_execute_for_portal(execute.portal.to_owned()); + .set_execute_for_portal(operation, execute.portal.to_owned()); Ok(()) } @@ -418,7 +286,11 @@ where /// - `Ok(Some(bytes))` - Transformed query that should replace the original /// - `Ok(None)` - No transformation needed, forward original query /// - `Err(error)` - Processing failed, error should be sent to client - async fn query_handler(&mut self, bytes: &BytesMut) -> Result, Error> { + async fn query_handler( + &mut self, + operation: pg_proto::OperationId, + query: bytes::Bytes, + ) -> Result, Error> { let handler_start = Instant::now(); let session_id = self.context.start_session(); @@ -429,10 +301,9 @@ where let parse_timer = PhaseTimer::start(); - let mut query = Query::try_from(bytes)?; - // Simple Query may contain many statements - let parsed_statements = SqlParser::parse_statements(&query.statement)?; + let query_text = String::from_utf8_lossy(&query).into_owned(); + let parsed_statements = SqlParser::parse_statements(&query_text)?; let mut transformed_statements = vec![]; debug!(target: MAPPER, @@ -565,11 +436,12 @@ where // Set query fingerprint self.context.update_statement_metadata(session_id, |m| { - m.set_query_fingerprint(&query.statement); + m.set_query_fingerprint(&query_text); }); - self.context.add_portal(Name::unnamed(), portal); - self.context.set_execute(Name::unnamed(), Some(session_id)); + self.context.add_portal(Name::new(), portal); + self.context + .set_execute(operation, Name::new(), Some(session_id)); if encrypted { let transformed_statement = transformed_statements @@ -578,9 +450,7 @@ where .collect::>() .join(";"); - query.rewrite(transformed_statement.to_string()); - - let bytes = BytesMut::try_from(query)?; + let message = FrontendMessage::Query(bytes::Bytes::from(transformed_statement.clone())); let handler_duration = handler_start.elapsed(); debug!( target: MAPPER, @@ -588,7 +458,7 @@ where msg = "Rewrite Query", transformed_statement = transformed_statement.to_string(), duration_ms = handler_duration.as_millis(), - bytes = ?bytes, + ?message, ); if handler_duration.as_millis() > 100 { warn!( @@ -597,7 +467,7 @@ where duration_ms = handler_duration.as_millis(), ); } - Ok(Some(bytes)) + Ok(Some(message)) } else { let handler_duration = handler_start.elapsed(); if handler_duration.as_millis() > 50 { @@ -783,7 +653,11 @@ where /// - `Ok(Some(bytes))` - Modified Parse message with transformed SQL/parameters /// - `Ok(None)` - No transformation needed, forward original message /// - `Err(error)` - Processing failed, error should be sent to client - async fn parse_handler(&mut self, bytes: &BytesMut) -> Result, Error> { + async fn parse_handler( + &mut self, + mut message: Parse, + ) -> Result, Error> { + let original_query = message.query.clone(); let session_id = self.context.start_session(); // Set protocol type @@ -793,8 +667,6 @@ where let parse_timer = PhaseTimer::start(); - let mut message = Parse::try_from(bytes)?; - debug!( target: PROTOCOL, client_id = self.context.client_id, @@ -821,11 +693,21 @@ where // BEFORE the new session is recorded — the other way around wipes the // mapping that was just written and every Bind falls back to the // latest-session guess. - self.context.close_statement(&message.name); + self.context.close_statement(&message.statement); self.context - .set_statement_session(message.name.to_owned(), session_id); + .set_statement_session(message.statement.to_owned(), session_id); + + let mut statement_text = String::from_utf8_lossy(&message.query).into_owned(); + let statement = SqlParser::parse_statement(&statement_text)?; - let statement = SqlParser::parse_statement(&message.statement)?; + // Record diagnostics before any passthrough path can return early. + // Statements that do not require EQL type checking (for example, + // plaintext INSERTs and SELECT pg_sleep(...)) still need their real + // statement type and fingerprint in metrics. + self.context.update_statement_metadata(session_id, |m| { + m.statement_type = Some(StatementType::from_statement(&statement)); + m.set_query_fingerprint(&statement_text); + }); if let Some(mapping_disabled) = self.context.maybe_set_unsafe_disable_mapping(&statement) { warn!( @@ -865,7 +747,8 @@ where // Capture the parse message param_types // These override the underlying column type - let param_types = message.param_types.clone(); + let client_param_types = message.parameter_types.clone(); + let param_types = client_param_types.iter().map(|oid| *oid as i32).collect(); let mut parse_duration_recorded = false; @@ -900,15 +783,17 @@ where statement.output_params = output_params_from_plan(&transformed_statement.params, output_columns); - message.rewrite_statement(transformed_statement.statement.to_string()); + statement_text = transformed_statement.statement.to_string(); + message.query = bytes::Bytes::copy_from_slice(statement_text.as_bytes()); } } counter!(STATEMENTS_ENCRYPTED_TOTAL).increment(1); - message.rewrite_param_types(&statement.output_params); + message.parameter_types = + rewrite_parse_param_types(&client_param_types, &statement.output_params); self.context - .add_statement(message.name.to_owned(), statement); + .add_statement(message.statement.to_owned(), statement); } _ => { debug!(target: MAPPER, @@ -925,21 +810,15 @@ where .record_parse_duration(session_id, parse_timer.elapsed()); } - // Set statement type and fingerprint - self.context.update_statement_metadata(session_id, |m| { - m.statement_type = Some(StatementType::from_statement(&statement)); - m.set_query_fingerprint(&message.statement); - }); - - if message.requires_rewrite() { - let bytes = BytesMut::try_from(message)?; + if message.query != original_query || message.parameter_types != client_param_types { + let message = FrontendMessage::Parse(message); debug!(target: MAPPER, client_id = self.context.client_id, msg = "Rewrite Parse", - bytes = ?bytes); + ?message); - Ok(Some(bytes)) + Ok(Some(message)) } else { Ok(None) } @@ -1084,7 +963,7 @@ where /// - `Ok(Some(bytes))` - Modified Bind message with encrypted parameter values /// - `Ok(None)` - No parameter encryption needed, forward original message /// - `Err(error)` - Processing failed, error should be sent to client - async fn bind_handler(&mut self, bytes: &BytesMut) -> Result, Error> { + async fn bind_handler(&mut self, mut bind: Bind) -> Result, Error> { if self.context.unsafe_disable_mapping() { warn!(msg = "Encrypted statement mapping is not enabled"); counter!(STATEMENTS_PASSTHROUGH_MAPPING_DISABLED_TOTAL).increment(1); @@ -1092,8 +971,6 @@ where return Ok(None); } - let mut bind = Bind::try_from(bytes)?; - let session_id = self .context .get_statement_session_or_latest(&bind.prepared_statement); @@ -1129,15 +1006,15 @@ where self.context.add_portal(bind.portal.to_owned(), portal); if bind.requires_rewrite() { - let bytes = BytesMut::try_from(bind)?; + let message = FrontendMessage::from(bind); debug!( target: MAPPER, client_id = self.context.client_id, msg = "Rewrite Bind", - bytes = ?bytes + ?message ); - Ok(Some(bytes)) + Ok(Some(message)) } else { Ok(None) } @@ -1260,108 +1137,28 @@ where } } } +} - /// - /// Send a ReadyForQuery to the client, answering a Sync (or simple Query) - /// for a batch of which nothing reached the server. - /// - fn send_ready_for_query(&mut self) -> Result<(), Error> { - let message = BytesMut::from(ReadyForQuery); - - debug!(target: PROTOCOL, - client_id = self.context.client_id, - msg = "send_ready_for_query", - ?message, - ); - - self.client_sender.send(message)?; - - Ok(()) +fn rewrite_parse_param_types(client_types: &[u32], output_params: &[OutputParam]) -> Vec { + if client_types.is_empty() { + return Vec::new(); } - /// Converts a frontend failure into a `DO ... RAISE EXCEPTION` statement - /// that is sent to the server IN PLACE of the failed message. - /// - /// The proxy must not reply to the client directly: responses to earlier - /// pipelined requests (for example the Close + Sync issued when a client - /// drops a cached prepared statement) may still be in flight from the - /// server, and a direct reply would overtake them — every subsequent - /// response then answers the wrong request and the client observes a - /// protocol desync (CIP-3678). Sending the error THROUGH the server keeps - /// the response stream ordered: the injected statement is a request like - /// any other, so its ErrorResponse and ReadyForQuery arrive exactly where - /// the client expects them, and the transaction state the client observes - /// (aborted on error) matches what PostgreSQL itself would produce. - /// - /// The ErrorResponse fields the proxy would have produced are carried as - /// `RAISE ... USING` options, so the client still receives the proxy's - /// message, SQLSTATE and hints. `RAISE` has no `POSITION` option; parse - /// errors already embed line/column in the message. - /// - /// The exception body is a plain single-quoted literal (no dollar quoting) - /// built with `quote_literal` at both nesting levels, so error text - /// containing quotes, `%` or `$$` cannot escape the statement. - /// - /// Returns `None` for FATAL errors, which are written directly to the - /// client instead: `RAISE` cannot produce a FATAL response, and ordering - /// no longer matters because the client abandons the connection — and any - /// pipelined requests on it — as soon as it reads a FATAL error. - fn handle_statement_error(&mut self, err: Error) -> Result, Error> { - error!(client_id = self.context.client_id, msg = err.to_string(), error = ?err); - - let error_response = self.error_to_response(err); - - if error_response.is_fatal() { - let message = BytesMut::try_from(error_response)?; - - debug!(target: PROTOCOL, - client_id = self.context.client_id, - msg = "send_error_response", - ?message, - ); - - self.client_sender.send(message)?; - return Ok(None); - } - - let options = error_response - .fields - .iter() - .filter_map(|field| { - let option = match field.code { - ErrorResponseCode::Message => "MESSAGE", - ErrorResponseCode::Code => "ERRCODE", - ErrorResponseCode::Detail => "DETAIL", - ErrorResponseCode::Hint => "HINT", - ErrorResponseCode::Schema => "SCHEMA", - ErrorResponseCode::Table => "TABLE", - ErrorResponseCode::Column => "COLUMN", - ErrorResponseCode::DataType => "DATATYPE", - ErrorResponseCode::Constraint => "CONSTRAINT", - // Severity is implied by RAISE EXCEPTION; the remaining - // fields have no RAISE equivalent. - _ => return None, - }; - Some(format!("{option} = {}", quote_literal(&field.value))) - }) - .collect::>() - .join(", "); - - let body = format!("BEGIN RAISE EXCEPTION USING {options}; END;"); - let content = format!("DO {};", quote_literal(&body)); - - debug!( - target: MAPPER, - client_id = self.context.client_id, - msg = "Frontend exception", - content = content.as_str(), - ); - - let query = Query::new(content); - let bytes = BytesMut::try_from(query)?; - - Ok(Some(bytes)) - } + output_params + .iter() + .map(|output| match &output.column { + Some(column) => match column.eql_term { + EqlTermVariant::JsonAccessor | EqlTermVariant::JsonPath => { + postgres_types::Type::TEXT.oid() + } + _ => postgres_types::Type::JSONB.oid(), + }, + None => client_types + .get(output.source.primary_input()) + .copied() + .unwrap_or(UNSPECIFIED_TYPE_OID as u32), + }) + .collect() } /// Projects a stored payload into its query operand when the value is bound in @@ -1528,29 +1325,16 @@ where } /// Implementation of PostgreSQL error handling for the Frontend component. -impl PostgreSqlErrorHandler for Frontend -where - R: AsyncRead + Unpin, - W: AsyncWrite + Unpin, - S: EncryptionService, -{ - fn client_sender(&mut self) -> &mut Sender { - &mut self.client_sender - } - +impl PostgreSqlErrorHandler for Frontend { fn client_id(&self) -> i32 { self.context.client_id } +} - /// Write an ErrorResponse directly to the client. - /// - /// Statement failures must NOT use this — they go through the server via - /// [`Frontend::handle_statement_error`] so the error is sequenced after any - /// in-flight responses (CIP-3678). A direct write is only safe when the - /// connection is being abandoned. - fn send_error_response(&mut self, err: Error) -> Result<(), Error> { +impl Frontend { + fn error_response(&self, err: Error) -> BackendMessage { let error_response = self.error_to_response(err); - let message = BytesMut::try_from(error_response)?; + let message = BackendMessage::ErrorResponse(error_response); debug!(target: PROTOCOL, client_id = self.context.client_id, @@ -1558,10 +1342,6 @@ where ?message, ); - self.client_sender.send(message)?; - // Frontend-specific: set error state for extended query protocol - self.error_state = Some(ErrorState::ErrorResponseSent); - - Ok(()) + message } } diff --git a/packages/cipherstash-proxy/src/postgresql/middleware/mod.rs b/packages/cipherstash-proxy/src/postgresql/middleware/mod.rs new file mode 100644 index 000000000..2b8c4e3f7 --- /dev/null +++ b/packages/cipherstash-proxy/src/postgresql/middleware/mod.rs @@ -0,0 +1,85 @@ +mod backend; +mod frontend; + +use super::Context; +use crate::{error::Error, proxy::EncryptionService}; +use backend::Backend; +use frontend::Frontend; +use pg_proto::{ + AttributedBackendMessages, BackendBatchOutput, BackendFlushReason, BackendMessage, + BackendMiddlewareOutput, FrontendMessage, FrontendMiddlewareOutput, IntermediaryMiddleware, + IntermediaryMiddlewareFactory, OperationId, +}; + +pub struct CipherStashMiddleware { + frontend: Frontend, + backend: Backend, +} + +#[derive(Clone)] +pub struct CipherStashMiddlewareFactory(pub Context); + +impl IntermediaryMiddlewareFactory + for CipherStashMiddlewareFactory +where + S: EncryptionService + Clone, + ServerContext: Sync, + ClientContext: Sync, +{ + type Handler = CipherStashMiddleware; + fn create(&self, _: &ServerContext, _: &ClientContext) -> Self::Handler { + CipherStashMiddleware::new(self.0.clone()) + } +} + +impl CipherStashMiddleware { + pub fn new(context: Context) -> Self { + Self { + frontend: Frontend::new(context.clone()), + backend: Backend::new(context), + } + } +} + +impl IntermediaryMiddleware<(), ServerContext, ClientContext> + for CipherStashMiddleware +where + S: EncryptionService + Clone, + ServerContext: Sync, + ClientContext: Sync, +{ + type Error = Error; + + async fn frontend_operation( + &mut self, + _: &ServerContext, + _: &ClientContext, + _: &mut (), + operation: OperationId, + message: FrontendMessage, + ) -> Result { + self.frontend.intercept(operation, message).await + } + + async fn backend_operation( + &mut self, + _: &ServerContext, + _: &ClientContext, + _: &mut (), + operation: Option, + message: BackendMessage, + ) -> Result { + self.backend.intercept(operation, message).await + } + + async fn flush_backend_operations( + &mut self, + _: &ServerContext, + _: &ClientContext, + _: &mut (), + held: AttributedBackendMessages<'_>, + _: BackendFlushReason, + ) -> Result { + self.backend.flush_held(held).await + } +} diff --git a/packages/cipherstash-proxy/src/postgresql/mod.rs b/packages/cipherstash-proxy/src/postgresql/mod.rs index 71c8df249..a3995b977 100644 --- a/packages/cipherstash-proxy/src/postgresql/mod.rs +++ b/packages/cipherstash-proxy/src/postgresql/mod.rs @@ -1,28 +1,15 @@ -mod backend; mod column_mapper; mod context; mod data; +mod diagnostics; +mod driver; mod error_handler; mod format_code; -mod frontend; -mod handler; -mod message_buffer; -mod messages; +mod middleware; mod parser; -mod protocol; -mod startup; +mod rewrite; pub use context::column::Column; pub use context::Context; pub use context::KeysetIdentifier; -pub use handler::handler; - -pub const PROTOCOL_VERSION_NUMBER: i32 = 196608; - -pub const SSL_REQUEST: i32 = 80877103; - -pub const CANCEL_REQUEST: i32 = 80877102; - -pub const SSL_RESPONSE_NO: u8 = b'N'; - -pub const SSL_RESPONSE_YES: u8 = b'S'; +pub use driver::handler; diff --git a/packages/cipherstash-proxy/src/postgresql/protocol.rs b/packages/cipherstash-proxy/src/postgresql/protocol.rs deleted file mode 100644 index 9c1afb79b..000000000 --- a/packages/cipherstash-proxy/src/postgresql/protocol.rs +++ /dev/null @@ -1,163 +0,0 @@ -use super::{messages::authentication::Authentication, CANCEL_REQUEST, SSL_REQUEST}; -use crate::{ - error::{Error, ProtocolError}, - log::PROTOCOL, - postgresql::PROTOCOL_VERSION_NUMBER, - SIZE_I32, SIZE_U8, -}; -use bytes::{BufMut, BytesMut}; -use std::{ - io::{BufRead, Cursor}, - time::Duration, -}; -use tokio::{ - io::{AsyncRead, AsyncReadExt}, - time::timeout, -}; -use tracing::{debug, error}; - -type Code = u8; - -#[derive(Clone, Debug, PartialEq)] -pub enum StartupCode { - ProtocolVersionNumber, - CancelRequest, - SSLRequest, -} - -#[derive(Clone, Debug)] -pub struct StartupMessage { - pub code: StartupCode, - pub bytes: BytesMut, -} - -#[derive(Clone, Debug)] -pub struct Message { - pub code: u8, - pub bytes: BytesMut, -} - -impl From for StartupCode { - fn from(code: i32) -> Self { - match code { - PROTOCOL_VERSION_NUMBER => StartupCode::ProtocolVersionNumber, - SSL_REQUEST => StartupCode::SSLRequest, - CANCEL_REQUEST => StartupCode::CancelRequest, - _ => panic!("Unexpected startup code {code}"), - } - } -} - -pub trait BytesMutReadString { - fn read_string(&mut self) -> Result; -} - -impl BytesMutReadString for Cursor<&BytesMut> { - /// Should only be used when reading strings from the message protocol. - /// Can be used to read multiple strings from the same message which are separated by the null byte - fn read_string(&mut self) -> Result { - let mut buf = Vec::with_capacity(512); - match self.read_until(b'\0', &mut buf) { - Ok(_) => Ok(String::from_utf8_lossy(&buf[..buf.len() - 1]).to_string()), - Err(err) => Err(err.into()), - } - } -} - -/// -/// Reads an Auth Message from Stream -/// -/// Does not use the default connection timeout as the auth message is expected to be sent immediately -/// 10 seconds is a reasonable timeout for the auth message -/// -/// -pub async fn read_auth_message( - mut stream: S, - client_id: i32, -) -> Result { - let connection_timeout = Duration::from_millis(1000 * 10); - let (_code, bytes) = - read_message_with_timeout(&mut stream, client_id, connection_timeout).await?; - Authentication::try_from(&bytes) -} - -/// -/// Reads a Postgres message from client with an optional timeout -/// -/// Timeout values are in config -/// -/// -pub async fn read_message( - mut stream: S, - client_id: i32, - connection_timeout: Option, -) -> Result<(Code, BytesMut), Error> { - match connection_timeout { - Some(duration) => read_message_with_timeout(stream, client_id, duration).await, - None => read(&mut stream, client_id).await, - } -} - -/// -/// Reads a Postgres message from client with a timeout -/// -/// Timeout values are in config -/// -/// -async fn read_message_with_timeout( - mut stream: S, - client_id: i32, - duration: Duration, -) -> Result<(Code, BytesMut), Error> { - timeout(duration, read(&mut stream, client_id)) - .await - .map_err(|_| Error::ConnectionTimeout { duration })? -} - -/// -/// Reads a Postgres message from client -/// -/// The SSLRequest/Response sequence requires the Backend to inspect the first byte of the message -/// Byte is then passed as `code` to this function to preserve the message structure -/// -/// -async fn read( - mut stream: S, - client_id: i32, -) -> Result<(Code, BytesMut), Error> { - let code = stream.read_u8().await?; - let len = stream.read_i32().await?; - - // Detect unexpected message len and avoid panic on read_exact - // Len must be at least 4 bytes (4 bytes for len/i32) - if (len as usize) < SIZE_I32 { - error!( - msg = "Unexpected PostgreSQL message length", - code = code, - len = len - ); - return Err(ProtocolError::UnexpectedMessageLength { - code, - len: len as usize, - } - .into()); - } - - let capacity = len as usize + SIZE_U8; //len plus len of code - let mut bytes = BytesMut::with_capacity(capacity); - - bytes.put_u8(code); - bytes.put_i32(len); - - let slice_start = bytes.len(); - - // Capacity and len are not the same!! - // resize populates the buffer with 0s - bytes.resize(capacity, 0); - - stream.read_exact(&mut bytes[slice_start..]).await?; - - debug!(target: PROTOCOL, client_id, code = ?(code as char), ?bytes); - - Ok((code, bytes)) -} diff --git a/packages/cipherstash-proxy/src/postgresql/messages/bind.rs b/packages/cipherstash-proxy/src/postgresql/rewrite/bind.rs similarity index 77% rename from packages/cipherstash-proxy/src/postgresql/messages/bind.rs rename to packages/cipherstash-proxy/src/postgresql/rewrite/bind.rs index b3f382f5a..ce151567c 100644 --- a/packages/cipherstash-proxy/src/postgresql/messages/bind.rs +++ b/packages/cipherstash-proxy/src/postgresql/rewrite/bind.rs @@ -1,3 +1,4 @@ +//! CipherStash Bind parameter rewriting. use super::{maybe_json, maybe_jsonb, Name, NULL}; use crate::error::{Error, MappingError, ProtocolError}; use crate::log::MAPPER; @@ -10,29 +11,23 @@ use crate::postgresql::data::{ json_value_selector_plaintext, }; use crate::postgresql::format_code::FormatCode; -use crate::postgresql::protocol::BytesMutReadString; use crate::{EqlOutput, EqlQueryPayload}; -use crate::{SIZE_I16, SIZE_I32}; -use bytes::{Buf, BufMut, BytesMut}; +use bytes::{BufMut, BytesMut}; use cipherstash_client::encryption::Plaintext; +use pg_proto::{Bind as PgBind, FrontendMessage}; use postgres_types::Type; +use std::convert::TryFrom; use std::fmt::{self, Display, Formatter}; -use std::io::Cursor; -use std::{convert::TryFrom, ffi::CString}; use tracing::debug; /// Bind (B) message. /// See: #[derive(Clone, Debug)] pub struct Bind { - pub code: char, pub portal: Name, pub prepared_statement: Name, - pub num_param_format_codes: i16, pub param_format_codes: Vec, - pub num_param_values: i16, pub param_values: Vec, - pub num_result_column_format_codes: i16, pub result_columns_format_codes: Vec, /// Set when the param list was rebuilt because the rewrite reshaped the /// params. The message must then be re-sent even if no individual param was @@ -44,6 +39,7 @@ pub struct Bind { pub struct BindParam { pub format_code: FormatCode, pub bytes: BytesMut, + null: bool, dirty: bool, } @@ -250,8 +246,6 @@ impl Bind { } self.param_format_codes = param_values.iter().map(|param| param.format_code).collect(); - self.num_param_format_codes = self.param_format_codes.len() as i16; - self.num_param_values = param_values.len() as i16; self.param_values = param_values; self.reshaped = true; @@ -306,6 +300,55 @@ impl Bind { } } +impl TryFrom for Bind { + type Error = Error; + + fn try_from(bind: PgBind) -> Result { + let portal = bind.portal; + let prepared_statement = bind.statement; + let param_format_codes = bind + .parameter_formats + .iter() + .copied() + .map(FormatCode::from) + .collect::>(); + let num_param_values = bind.parameters.len(); + let mut param_values = Vec::with_capacity(num_param_values); + for (idx, parameter) in bind.parameters.into_iter().enumerate() { + let format_code = match param_format_codes.len() { + 0 => FormatCode::Text, + 1 => param_format_codes[0], + len if len == num_param_values => param_format_codes[idx], + _ => { + return Err(ProtocolError::ParameterFormatCodesMismatch { + expected: num_param_values, + received: param_format_codes.len(), + } + .into()); + } + }; + match parameter { + None => param_values.push(BindParam::null_with_format(format_code)), + Some(bytes) => { + param_values.push(BindParam::new(format_code, BytesMut::from(&bytes[..]))); + } + } + } + Ok(Self { + portal, + prepared_statement, + param_format_codes, + param_values, + result_columns_format_codes: bind + .result_formats + .into_iter() + .map(FormatCode::from) + .collect(), + reshaped: false, + }) + } +} + /// /// Param type is either provided with Parse message or the column type /// Column type is the cast of the encrypted column @@ -322,6 +365,7 @@ impl BindParam { Self { format_code, bytes, + null: false, dirty: false, } } @@ -330,6 +374,16 @@ impl BindParam { Self { format_code: FormatCode::Text, bytes: BytesMut::new(), + null: true, + dirty: false, + } + } + + fn null_with_format(format_code: FormatCode) -> Self { + Self { + format_code, + bytes: BytesMut::new(), + null: true, dirty: false, } } @@ -357,6 +411,7 @@ impl BindParam { pub fn rewrite(&mut self, bytes: &[u8]) { self.bytes.clear(); + self.null = false; if self.is_binary() { self.bytes.put_u8(1); @@ -376,6 +431,7 @@ impl BindParam { /// stop `->` from matching any stored entry. pub fn rewrite_text(&mut self, bytes: Vec) { self.bytes.clear(); + self.null = false; self.bytes.extend_from_slice(&bytes); self.dirty = true; } @@ -392,6 +448,7 @@ impl BindParam { } self.bytes.clear(); + self.null = true; self.dirty = true; } @@ -415,7 +472,7 @@ impl BindParam { } pub fn is_null(&self) -> bool { - self.bytes.is_empty() + self.null } pub fn is_text(&self) -> bool { @@ -434,147 +491,26 @@ impl Display for BindParam { } } -impl TryFrom<&BytesMut> for Bind { - type Error = Error; - - fn try_from(buf: &BytesMut) -> Result { - let mut cursor = Cursor::new(buf); - let code = cursor.get_u8() as char; - let _len = cursor.get_i32(); - - let portal = cursor.read_string()?; - let portal = Name::from(portal); - - let prepared_statement = cursor.read_string()?; - let prepared_statement = Name::from(prepared_statement); - - let num_param_format_codes = cursor.get_i16(); - let mut param_format_codes = Vec::new(); - - for _ in 0..num_param_format_codes { - param_format_codes.push(cursor.get_i16().into()); - } - - let num_param_values = cursor.get_i16(); - let mut param_values = Vec::new(); - - for idx in 0..num_param_values as usize { - let param_len = cursor.get_i32(); - - let format_code = match num_param_format_codes { - 0 => FormatCode::Text, - 1 => param_format_codes[0], - _ => param_format_codes[idx], - }; - - // NULL parameters have a length of -1 and no bytes - match param_len { - NULL => { - param_values.push(BindParam::null()); - } - _ => { - let mut bytes = BytesMut::with_capacity(param_len as usize); - bytes.resize(param_len as usize, b'0'); - cursor.copy_to_slice(&mut bytes); - param_values.push(BindParam::new(format_code, bytes)); - } - } - } - - let num_result_column_format_codes = cursor.get_i16(); - let mut result_columns_format_codes = Vec::new(); - - for _ in 0..num_result_column_format_codes { - result_columns_format_codes.push(cursor.get_i16().into()); - } - - Ok(Bind { - code, - portal, - prepared_statement, - num_param_format_codes, - param_format_codes, - num_param_values, - param_values, - num_result_column_format_codes, - result_columns_format_codes, - reshaped: false, +impl From for FrontendMessage { + fn from(bind: Bind) -> Self { + Self::Bind(PgBind { + portal: bind.portal, + statement: bind.prepared_statement, + parameter_formats: bind.param_format_codes.into_iter().map(i16::from).collect(), + parameters: bind + .param_values + .into_iter() + .map(|param| (!param.null).then(|| param.bytes.freeze())) + .collect(), + result_formats: bind + .result_columns_format_codes + .into_iter() + .map(i16::from) + .collect(), }) } } -impl TryFrom for BytesMut { - type Error = Error; - - fn try_from(bind: Bind) -> Result { - let mut bytes = BytesMut::new(); - - let portal_binding = CString::new(&*bind.portal)?; - let portal = portal_binding.as_bytes_with_nul(); - - let prepared_statement_binding = CString::new(&*bind.prepared_statement)?; - let prepared_statement = prepared_statement_binding.as_bytes_with_nul(); - - if bind.num_param_format_codes != bind.param_format_codes.len() as i16 { - let err = ProtocolError::ParameterFormatCodesMismatch { - expected: bind.num_param_format_codes as usize, - received: bind.param_format_codes.len(), - }; - return Err(err.into()); - } - - if bind.num_result_column_format_codes != bind.result_columns_format_codes.len() as i16 { - let err = ProtocolError::ParameterResultFormatCodesMismatch { - expected: bind.num_result_column_format_codes as usize, - received: bind.result_columns_format_codes.len(), - }; - return Err(err.into()); - } - - // sum of param byte_lens (the *actual* byte lengths of the parameters) - let param_byte_len = &bind - .param_values - .iter() - .fold(0, |acc, param| acc + SIZE_I32 + param.byte_len()); - - let len = SIZE_I32 // self/len of len - + portal.len() - + prepared_statement.len() - + SIZE_I16 // num_param_format_codes - + SIZE_I16 * bind.num_param_format_codes as usize // num_param_format_codes - + SIZE_I16 // num_param_values - + param_byte_len // parameter bytes - + SIZE_I16 // num_result_column_format_codes - + SIZE_I16 * bind.num_result_column_format_codes as usize; - - bytes.put_u8(bind.code as u8); - bytes.put_i32(len as i32); - bytes.put_slice(portal); - bytes.put_slice(prepared_statement); - bytes.put_i16(bind.num_param_format_codes); - for param_format_code in bind.param_format_codes { - bytes.put_i16(param_format_code.into()); - } - - let num_param_values = bind.param_values.len() as i16; - bytes.put_i16(num_param_values); - - for p in bind.param_values { - // len is not the same as byte_len - // A NULL param len is -1 - bytes.put_i32(p.len()); - bytes.put_slice(&p.bytes); - } - - bytes.put_i16(bind.num_result_column_format_codes); - for result_column_format_code in bind.result_columns_format_codes { - bytes.put_i16(result_column_format_code.into()); - } - - Ok(bytes) - } -} - #[cfg(test)] mod tests { use super::{BindParam, JsonSelectorPath, JsonSelectorStep, OutputParam, OutputParamSource}; @@ -582,33 +518,34 @@ mod tests { config::LogConfig, log, postgresql::{ - context::column::Column, format_code::FormatCode, messages::bind::Bind, messages::Name, + context::column::Column, format_code::FormatCode, rewrite::bind::Bind, rewrite::Name, }, Identifier, }; - use bytes::BytesMut; + use bytes::{Bytes, BytesMut}; use cipherstash_client::schema::{ColumnConfig, ColumnMode, ColumnType}; use eql_mapper::EqlTermVariant; - - fn to_message(s: &[u8]) -> BytesMut { - BytesMut::from(s) - } + use pg_proto::{Bind as PgBind, FrontendMessage}; #[test] pub fn parse_bind() { log::init(LogConfig::default()); - let bytes = - to_message(b"B\0\0\0\x18\0\0\0\x01\0\x01\0\x01\0\0\0\x04.\xbe\x8a\xd4\0\x01\0\x01"); - - let expected = bytes.clone(); - - let bind = Bind::try_from(&bytes).unwrap(); + let expected = PgBind { + portal: Bytes::new(), + statement: Bytes::new(), + parameter_formats: vec![1], + parameters: vec![Some(Bytes::from_static(b".\xbe\x8a\xd4"))], + result_formats: vec![1], + }; + let bind = Bind::try_from(expected.clone()).unwrap(); assert_eq!(bind.param_values.len(), 1); assert_eq!(bind.result_columns_format_codes[0], FormatCode::Binary); - let bytes = BytesMut::try_from(bind).unwrap(); - assert_eq!(bytes, expected); + let FrontendMessage::Bind(actual) = bind.into() else { + panic!("expected Bind") + }; + assert_eq!(actual, expected); } #[test] @@ -616,17 +553,46 @@ mod tests { log::init(LogConfig::default()); // Bind message from statement INSERT INTO encrypted (id, plaintext, plaintext_date, encrypted_text) VALUES ($1, $2, $3, $4) - let bytes = - to_message(b"B\0\0\0N\0s0\0\0\x04\0\x01\0\x01\0\x01\0\x01\0\x04\0\0\0\x084\xd8\x1d@\x83U\x0em\0\0\0\tplaintext\xff\xff\xff\xff\0\0\0\x15hello@cipherstash.com\0\x01\0\x01"); + let expected = PgBind { + portal: Bytes::new(), + statement: Bytes::from_static(b"s0"), + parameter_formats: vec![1, 1, 1, 1], + parameters: vec![ + Some(Bytes::from_static(b"4\xd8\x1d@\x83U\x0em")), + Some(Bytes::from_static(b"plaintext")), + None, + Some(Bytes::from_static(b"hello@cipherstash.com")), + ], + result_formats: vec![1], + }; + let bind = Bind::try_from(expected.clone()).unwrap(); - let expected = bytes.clone(); + assert_eq!(bind.param_values.len(), 4); - let bind = Bind::try_from(&bytes).unwrap(); + let FrontendMessage::Bind(actual) = bind.into() else { + panic!("expected Bind") + }; + assert_eq!(actual, expected); + } - assert_eq!(bind.param_values.len(), 4); + #[test] + pub fn preserves_empty_and_null_params_distinctly() { + let expected = PgBind { + portal: Bytes::new(), + statement: Bytes::new(), + parameter_formats: vec![], + parameters: vec![Some(Bytes::new()), None], + result_formats: vec![], + }; + let bind = Bind::try_from(expected.clone()).unwrap(); - let bytes = BytesMut::try_from(bind).unwrap(); - assert_eq!(bytes, expected); + assert!(!bind.param_values[0].is_null()); + assert_eq!(bind.param_values[0].byte_len(), 0); + assert!(bind.param_values[1].is_null()); + let FrontendMessage::Bind(actual) = bind.into() else { + panic!("expected Bind") + }; + assert_eq!(actual, expected); } #[test] @@ -662,14 +628,10 @@ mod tests { fn bind_with(param_values: Vec) -> Bind { Bind { - code: 'B', - portal: Name::unnamed(), - prepared_statement: Name::unnamed(), - num_param_format_codes: param_values.len() as i16, + portal: Name::new(), + prepared_statement: Name::new(), param_format_codes: param_values.iter().map(|p| p.format_code).collect(), - num_param_values: param_values.len() as i16, param_values, - num_result_column_format_codes: 0, result_columns_format_codes: vec![], reshaped: false, } diff --git a/packages/cipherstash-proxy/src/postgresql/rewrite/data_row.rs b/packages/cipherstash-proxy/src/postgresql/rewrite/data_row.rs new file mode 100644 index 000000000..b06c13774 --- /dev/null +++ b/packages/cipherstash-proxy/src/postgresql/rewrite/data_row.rs @@ -0,0 +1,308 @@ +//! CipherStash DataRow rewriting. +use crate::EqlCiphertext; +use crate::{ + error::{EncryptError, Error, ProtocolError}, + log::DECRYPT, + postgresql::Column, +}; +use bytes::{Bytes, BytesMut}; +use pg_proto::DataRow; +use tracing::{debug, error}; + +/// Leading byte of `jsonb`'s binary wire format. PostgreSQL has only ever +/// emitted version 1. +const JSONB_BINARY_VERSION: u8 = 1; + +pub fn as_ciphertext( + row: &mut DataRow, + column_configuration: &[Option], +) -> Vec> { + let mut result = vec![]; + for (data_column, column_config) in row.columns.iter_mut().zip(column_configuration) { + let encrypted = column_config + .as_ref() + .filter(|_| data_column.is_some()) + .and_then(|config| { + to_eql_ciphertext(data_column) + .inspect_err(|err| match err { + Error::Encrypt(EncryptError::ColumnIsNull) => { + debug!(target: DECRYPT, msg ="ColumnIsNull", ?config); + // Not an error, as you were + *data_column = None; + } + _ => { + let err = EncryptError::ColumnCouldNotBeDeserialised { + table: config.identifier.table.to_owned(), + column: config.identifier.column.to_owned(), + }; + error!(target: DECRYPT, msg = err.to_string()); + } + }) + .ok() + }); + result.push(encrypted); + } + + result +} + +pub fn rewrite(row: &mut DataRow, plaintexts: &[Option]) -> Result<(), Error> { + if plaintexts.len() > row.columns.len() { + return Err(ProtocolError::DataRowColumnCountMismatch { + expected: plaintexts.len(), + received: row.columns.len(), + } + .into()); + } + for (idx, pt) in plaintexts.iter().enumerate() { + if let Some(bytes) = pt { + if let Some(column @ Some(_)) = row.columns.get_mut(idx) { + *column = Some(Bytes::copy_from_slice(bytes)); + } + } + } + Ok(()) +} + +/// Parse this column's bytes into an [`EqlCiphertext`]. +/// +/// EQL v3 column types (`eql_v3_text_eq`, `eql_v3_integer_ord`, …) are +/// DOMAINS over `jsonb`, so a value arrives with jsonb's representation. +/// +/// EQL v2's `eql_v2_encrypted` was a composite type, which is why this +/// used to strip a `("…")` wrapper in text and a 12-byte rowtype header +/// in binary. Neither exists any more — a domain is wire-identical to its +/// base type. +/// +/// text — the JSON object itself, no wrapper and no doubled quotes +/// binary — a 1-byte jsonb version header followed by the JSON text +/// +/// The two are told apart by the leading byte: the version header is +/// `0x01`, and JSON text for an EQL payload always starts with `{`. +/// +/// The JSON is usually a self-describing payload — a scalar `{v,i,c,…}` or +/// a SteVec document `{v,k:"sv",i,h,sv}` — and deserialises directly. The +/// exception is a JSON field access (`eql_v3."->"(…)` / +/// `eql_v3.jsonb_path_query(…)`), whose result is a single +/// `eql_v3_json_entry` (`{v,i,h,s,c,op}`) — one SteVec entry merged with +/// its document envelope. That has a `c`, so it would masquerade as a +/// scalar `Encrypted` payload, but its `c` is an *entry* ciphertext that +/// only decrypts with the entry's selector-derived nonce. So when the +/// payload is a bare entry (see [`is_json_entry`]) it is reshaped into a +/// one-entry SteVec document (see [`json_entry_into_ste_vec_document`]) and +/// the ordinary SteVec decrypt path recovers the field value. +fn to_eql_ciphertext(column: &Option) -> Result { + let Some(bytes) = column else { + return Err(EncryptError::ColumnCouldNotBeParsed.into()); + }; + + let json = match bytes.first() { + Some(&JSONB_BINARY_VERSION) => &bytes[1..], + Some(_) => &bytes[..], + None => return Err(EncryptError::ColumnCouldNotBeParsed.into()), + }; + + let mut value: serde_json::Value = + serde_json::from_slice(json).map_err(log_deserialise_error)?; + + if is_json_entry(&value) { + json_entry_into_ste_vec_document(&mut value)?; + } + + serde_json::from_value(value).map_err(log_deserialise_error) +} + +/// Whether a decoded EQL payload is a bare `eql_v3_json_entry` — the result of +/// a JSON field access (`eql_v3."->"(…)` / `eql_v3.jsonb_path_query(…)`). +/// +/// A root-level selector `s` is the tell: a scalar `Encrypted` payload has no +/// selector at all, and a SteVec document carries selectors only inside its +/// `sv[]` entries, never at the root. +fn is_json_entry(value: &serde_json::Value) -> bool { + value.get("s").is_some() +} + +/// Reshape a single `eql_v3_json_entry` into a one-entry SteVec document. +/// +/// The entry is `{v,i,h,s,c,op}`: document-envelope fields (`v`, `i`, `h`) +/// alongside one SteVec entry's fields (`s`, `c`, the optional array marker +/// `a`, and the optional ordering term `op`). Move the entry fields under +/// `sv:[{…}]` and tag the object as a SteVec (`k:"sv"`), yielding +/// `{v,k:"sv",i,h,sv:[{s,c,a?,op?}]}` — the shape an [`EqlCiphertext`] SteVec +/// document deserialises from and the decrypt path knows how to open. +fn json_entry_into_ste_vec_document(value: &mut serde_json::Value) -> Result<(), Error> { + use serde_json::Value; + + let object = value + .as_object_mut() + .ok_or(EncryptError::ColumnCouldNotBeParsed)?; + + let mut entry = serde_json::Map::new(); + for key in ["s", "c", "a", "op"] { + if let Some(field) = object.remove(key) { + entry.insert(key.to_owned(), field); + } + } + + object.insert("k".to_owned(), Value::String("sv".to_owned())); + object.insert("sv".to_owned(), Value::Array(vec![Value::Object(entry)])); + + Ok(()) +} + +fn log_deserialise_error(err: serde_json::Error) -> Error { + debug!(target: DECRYPT, error = err.to_string()); + err.into() +} + +#[cfg(test)] +mod tests { + use super::{as_ciphertext, rewrite}; + use crate::{ + config::{LogConfig, LogLevel}, + log, + postgresql::Column, + }; + use crate::{error::ProtocolError, Identifier}; + use bytes::{Buf, Bytes}; + use cipherstash_client::schema::{ColumnConfig, ColumnType}; + use pg_proto::DataRow as PgDataRow; + + fn to_message(s: &[u8]) -> PgDataRow { + assert_eq!(s.first(), Some(&b'D')); + let mut body = &s[5..]; + let count = body.get_i16() as usize; + let columns = (0..count) + .map(|_| { + let len = body.get_i32(); + (len >= 0).then(|| { + let len = len as usize; + let value = Bytes::copy_from_slice(&body[..len]); + body.advance(len); + value + }) + }) + .collect(); + PgDataRow { columns } + } + + fn column_config(column: &str) -> Option { + let identifier = Identifier::new("encrypted", column); + let config = ColumnConfig::build("column".to_string()).casts_as(ColumnType::SmallInt); + let column = Column::new(identifier, config, None, eql_mapper::EqlTermVariant::Full); + Some(column) + } + + fn column_config_with_id(column: &str) -> Vec> { + vec![None, column_config(column)] + } + + #[test] + fn rewrite_rejects_more_plaintexts_than_columns() { + let mut row = PgDataRow { + columns: vec![Some(Bytes::from_static(b"one"))], + }; + let plaintexts = vec![None, Some(bytes::BytesMut::from("two"))]; + + let error = rewrite(&mut row, &plaintexts).unwrap_err(); + + assert!(matches!( + error, + crate::error::Error::Protocol(ProtocolError::DataRowColumnCountMismatch { + expected: 2, + received: 1 + }) + )); + } + + // The four `to_ciphertext_*` fixtures below are REAL EQL v3 wire captures + // taken from Postgres -> Proxy `DataRow` messages for the `encrypted` test + // table, via a live encrypt round-trip against ZeroKMS. They exercise + // `DataRow::try_from` + `as_ciphertext` across the binary (jsonb `0x01` + // version header) and text (bare JSON) wire encodings, and NULL columns. + // + // Captured against EQL v3.0.2. The build has since moved to the version + // pinned by `CS_EQL_VERSION` in `mise.toml`, and these still pass — the + // shapes under test (the jsonb version header, the bare-JSON text form, and + // the payload's `i`/`v` fields) have not changed. Regenerate against the + // pinned version, not against 3.0.2, if a future release does change them. + #[test] + pub fn to_ciphertext_with_binary_encoding() { + log::init(LogConfig::with_level(LogLevel::Debug)); + + // `SELECT encrypted_text FROM encrypted WHERE id = $1` (extended/binary): + // the jsonb column arrives as `0x01` + the v3 EqlCiphertextV3 JSON. + let bytes = to_message(b"D\x00\x00\x03\x16\x00\x01\x00\x00\x03\x0c\x01{\"c\": \"mBbL3gJuL?E})+>NeOq5<7N279rs9aRhBwjz3>wOdg{d64myql`6cXIurM_?B|pR<+M8(SeOLoLt~axenSv%=hCOb&m`FC5F;fS-ykq76u4Qgxa(QrcWn^D;Wq5SN5EJ90LtnW_NroxKJj=JLK>\", \"i\": {\"c\": \"encrypted_text\", \"t\": \"encrypted\"}, \"v\": 3, \"bf\": [1512, 1681, 836, 288, 1837, 1131, 415, 1430, 60, 812, 1990, 1211, 1368, 343, 1473, 1980, 598, 1549, 457, 1389, 1557, 941, 494, 1009, 1604, 1033, 2046, 222, 2012, 671, 7, 1525, 265, 901, 743, 543, 1771, 1149, 890, 755, 1974, 1960, 387, 1947, 1298, 130, 1758, 1060, 268, 844, 1375, 746, 1251, 2040], \"hm\": \"96aeaf9852416229d6b33ceb018d9abc90d70cbe7632539d69ef1462c9aa86a0\", \"op\": \"00bf0281ccb68cc6fe496bb1c8277e3484f6392517d5b8425536af7ec00ad7cc40e17e6336568ac4ed98dd659f7581f8a113fe5669b89833d9dd8eadc587a8950b6bd94f872e7f4205a6859e071df47134d3cccf1e53295417\"}"); + let mut data_row = bytes; + + let column_config = vec![column_config("encrypted_text")]; + let encrypted = as_ciphertext(&mut data_row, &column_config); + + assert_eq!(encrypted.len(), 1); + assert!(encrypted[0].is_some()); + assert_eq!( + &column_config[0].as_ref().unwrap().identifier, + encrypted[0].as_ref().unwrap().identifier() + ); + } + + #[test] + pub fn to_ciphertext_with_binary_encoding_and_null() { + log::init(LogConfig::with_level(LogLevel::Debug)); + + // `SELECT encrypted_text, encrypted_bool FROM encrypted WHERE id = $1` + // (binary), encrypted_text set, encrypted_bool NULL. + let bytes = to_message(b"D\x00\x00\x03\x1a\x00\x02\x00\x00\x03\x0c\x01{\"c\": \"mBbL3gJuL?E})+>NeOq5<7N279rs9aRhBwjz3>wOdg{d64myql`6cXIurM_?B|pR<+M8(SeOLoLt~axenSv%=hCOb&m`FC5F;fS-ykq76u4Qgxa(QrcWn^D;Wq5SN5EJ90LtnW_NroxKJj=JLK>\", \"i\": {\"c\": \"encrypted_text\", \"t\": \"encrypted\"}, \"v\": 3, \"bf\": [1512, 1681, 836, 288, 1837, 1131, 415, 1430, 60, 812, 1990, 1211, 1368, 343, 1473, 1980, 598, 1549, 457, 1389, 1557, 941, 494, 1009, 1604, 1033, 2046, 222, 2012, 671, 7, 1525, 265, 901, 743, 543, 1771, 1149, 890, 755, 1974, 1960, 387, 1947, 1298, 130, 1758, 1060, 268, 844, 1375, 746, 1251, 2040], \"hm\": \"96aeaf9852416229d6b33ceb018d9abc90d70cbe7632539d69ef1462c9aa86a0\", \"op\": \"00bf0281ccb68cc6fe496bb1c8277e3484f6392517d5b8425536af7ec00ad7cc40e17e6336568ac4ed98dd659f7581f8a113fe5669b89833d9dd8eadc587a8950b6bd94f872e7f4205a6859e071df47134d3cccf1e53295417\"}\xff\xff\xff\xff"); + let mut data_row = bytes; + + let column_config = vec![ + column_config("encrypted_text"), + column_config("encrypted_bool"), + ]; + let encrypted = as_ciphertext(&mut data_row, &column_config); + + assert_eq!(encrypted.len(), 2); + assert!(encrypted[0].is_some()); + assert!(encrypted[1].is_none()); + } + + #[test] + pub fn to_ciphertext_with_text_encoding() { + log::init(LogConfig::with_level(LogLevel::Debug)); + + // `SELECT encrypted_jsonb FROM encrypted WHERE id = 2` (simple/text): the + // jsonb column arrives as bare JSON text, no version header. + let bytes = to_message(b"D\x00\x00\x027\x00\x01\x00\x00\x02-{\"h\": \"l*AC8+7wO)sD**%APm>F3Bc9FAg#FNCmyISKh%bW{NbL}o`gZpBwFD}ye0IoZJ}<8La$|RV{&NeOq5<7N279rs9aRhBwjz3>wOdg{d64myql`6cXIurM_?B|pR<+M8(SeOLoLt~axenSv%=hCOb&m`FC5F;fS-ykq76u4Qgxa(QrcWn^D;Wq5SN5EJ90LtnW_NroxKJj=JLK>\", \"i\": {\"c\": \"encrypted_text\", \"t\": \"encrypted\"}, \"v\": 3, \"bf\": [1512, 1681, 836, 288, 1837, 1131, 415, 1430, 60, 812, 1990, 1211, 1368, 343, 1473, 1980, 598, 1549, 457, 1389, 1557, 941, 494, 1009, 1604, 1033, 2046, 222, 2012, 671, 7, 1525, 265, 901, 743, 543, 1771, 1149, 890, 755, 1974, 1960, 387, 1947, 1298, 130, 1758, 1060, 268, 844, 1375, 746, 1251, 2040], \"hm\": \"96aeaf9852416229d6b33ceb018d9abc90d70cbe7632539d69ef1462c9aa86a0\", \"op\": \"00bf0281ccb68cc6fe496bb1c8277e3484f6392517d5b8425536af7ec00ad7cc40e17e6336568ac4ed98dd659f7581f8a113fe5669b89833d9dd8eadc587a8950b6bd94f872e7f4205a6859e071df47134d3cccf1e53295417\"}\xff\xff\xff\xff"); + let mut data_row = bytes; + + let column_config = vec![ + column_config("encrypted_text"), + column_config("encrypted_bool"), + ]; + let encrypted = as_ciphertext(&mut data_row, &column_config); + + assert_eq!(encrypted.len(), 2); + assert!(encrypted[0].is_some()); + assert!(encrypted[1].is_none()); + } +} diff --git a/packages/cipherstash-proxy/src/postgresql/rewrite/mod.rs b/packages/cipherstash-proxy/src/postgresql/rewrite/mod.rs new file mode 100644 index 000000000..082b688be --- /dev/null +++ b/packages/cipherstash-proxy/src/postgresql/rewrite/mod.rs @@ -0,0 +1,16 @@ +use bytes::BytesMut; + +pub mod bind; +pub mod data_row; + +pub type Name = bytes::Bytes; +pub const NULL: i32 = -1; +pub const UNSPECIFIED_TYPE_OID: i32 = 0; + +pub fn maybe_json(bytes: &BytesMut) -> bool { + bytes.first() == Some(&b'{') +} + +pub fn maybe_jsonb(bytes: &BytesMut) -> bool { + bytes.len() > 3 && bytes[0] == 1 && bytes[1] == b'{' +} diff --git a/packages/cipherstash-proxy/src/postgresql/startup.rs b/packages/cipherstash-proxy/src/postgresql/startup.rs deleted file mode 100644 index a21e64f95..000000000 --- a/packages/cipherstash-proxy/src/postgresql/startup.rs +++ /dev/null @@ -1,163 +0,0 @@ -use std::time::Duration; - -use bytes::{BufMut, BytesMut}; -use tokio::{ - io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt}, - time::timeout, -}; -use tracing::{debug, error, warn}; - -use crate::{ - connect::AsyncStream, - error::{Error, ProtocolError}, - log::PROTOCOL, - postgresql::{SSL_REQUEST, SSL_RESPONSE_NO, SSL_RESPONSE_YES}, - tls, TandemConfig, SIZE_I32, -}; - -use super::protocol::StartupMessage; - -pub async fn with_tls(stream: AsyncStream, config: &TandemConfig) -> Result { - if config.database_tls_disabled() { - warn!(msg = "Connecting to database without Transport Layer Security (TLS)"); - return Ok(stream); - } - match stream { - AsyncStream::Tcp(mut tcp_stream) => { - let server_supports_ssl = send_ssl_request(&mut tcp_stream).await?; - - match server_supports_ssl { - true => { - let tls_stream = tls::client(tcp_stream, config).await?; - Ok(AsyncStream::Tls(Box::new(tls_stream))) - } - false => { - warn!(msg = "Connecting to database without Transport Layer Security (TLS)"); - Ok(AsyncStream::Tcp(tcp_stream)) - } - } - } - AsyncStream::Tls(_) => { - // Technically unreachable unless the server is misbehaving - warn!(msg = "Database already connected over Transport Layer Security (TLS)"); - Ok(stream) - } - } -} - -/// -/// Reads a Postgres startup message from client with an optional timeout -/// -/// Timeout values are in config -/// -/// -pub async fn read_message( - mut stream: S, - connection_timeout: Option, -) -> Result { - match connection_timeout { - Some(duration) => read_message_with_timeout(stream, duration).await, - None => read(&mut stream).await, - } -} - -/// -/// Reads a Postgres message from client with a timeout -/// -/// Timeout values are in config -/// -/// -async fn read_message_with_timeout( - mut stream: S, - duration: Duration, -) -> Result { - timeout(duration, read(&mut stream)) - .await - .map_err(|_| Error::ConnectionTimeout { duration })? -} - -/// -/// Read the start up message from the client -/// Startup messages are sent by the client to the server to initiate a connection -/// -/// -/// -async fn read(client: &mut C) -> Result -where - C: AsyncRead + Unpin, -{ - let len = client.read_i32().await?; - - let capacity = len as usize; - - let mut bytes = BytesMut::with_capacity(capacity); - bytes.put_i32(len); - bytes.resize(capacity, b'0'); - - let slice_start = SIZE_I32; - client.read_exact(&mut bytes[slice_start..]).await?; - - // code is the first 4 bytes after len - let code_bytes: [u8; 4] = [ - bytes.as_ref()[4], - bytes.as_ref()[5], - bytes.as_ref()[6], - bytes.as_ref()[7], - ]; - - let code = i32::from_be_bytes(code_bytes); - - let message = StartupMessage { - code: code.into(), - bytes, - }; - debug!(target: PROTOCOL, StartupMessage = ?message); - - Ok(message) -} - -/// -/// Send SSLRequest to the stream and return the response -/// Returns true if the server indicates support for TLS -/// -pub async fn send_ssl_request( - stream: &mut T, -) -> Result { - let mut bytes = BytesMut::with_capacity(12); - bytes.put_i32(8); - bytes.put_i32(SSL_REQUEST); - - stream.write_all(&bytes).await?; - - // Server supports TLS - let response = match stream.read_u8().await? { - SSL_RESPONSE_YES => true, - SSL_RESPONSE_NO => false, - code => { - error!(msg = "Unexpected startup message", code = ?(code as char)); - return Err(ProtocolError::UnexpectedStartupMessage.into()); - } - }; - - debug!(target: PROTOCOL, msg = "Database SSLResponse", SSLResponse = ?response); - Ok(response) -} - -/// -/// Send SSLRequest to the stream -/// Returns true if the server indicates support for TLS -/// N for no, S for yeS or tlS -/// The SSLResponse MUST come before the TLS handshake -/// -pub async fn send_ssl_response( - stream: &mut T, - tls: bool, -) -> Result<(), Error> { - let response = if tls { b'S' } else { b'N' }; - - debug!(target: PROTOCOL, msg = "SSLResponse to Client", SSLResponse = ?response); - - stream.write_all(&[response]).await?; - - Ok(()) -} diff --git a/packages/cipherstash-proxy/src/tls/mod.rs b/packages/cipherstash-proxy/src/tls/mod.rs index 82c70cea1..8c1fab46b 100644 --- a/packages/cipherstash-proxy/src/tls/mod.rs +++ b/packages/cipherstash-proxy/src/tls/mod.rs @@ -1,40 +1,10 @@ +use crate::DatabaseConfig; use crate::{config::TlsConfig, error::Error}; -use crate::{DatabaseConfig, TandemConfig}; use rustls::client::danger::ServerCertVerifier; use rustls::ClientConfig; use rustls_pki_types::{pem::PemObject, CertificateDer, PrivateKeyDer, ServerName}; use rustls_platform_verifier::ConfigVerifierExt; use std::sync::Arc; -use tokio::net::TcpStream; -use tokio_rustls::{TlsAcceptor, TlsConnector, TlsStream}; - -/// -/// Create a Server TLS connection -/// The returned type is the higher-level TlsStream that wraps both Client & Server variants -/// -pub async fn client( - stream: TcpStream, - config: &TandemConfig, -) -> Result, Error> { - let tls_config = configure_client(&config.database); - let connector = TlsConnector::from(Arc::new(tls_config)); - let domain = config.database.server_name()?.to_owned(); - let tls_stream = connector.connect(domain, stream).await?; - - Ok(tls_stream.into()) -} - -/// -/// Create a Server TLS connection -/// The returned type is the higher-level TlsStream that wraps both Client & Server variants -/// -pub async fn server(stream: TcpStream, config: &TlsConfig) -> Result, Error> { - let server_config = configure_server(config)?; - let acceptor = TlsAcceptor::from(Arc::new(server_config)); - let tls_stream = acceptor.accept(stream).await?; - - Ok(tls_stream.into()) -} /// /// Configure the server TLS settings @@ -44,6 +14,14 @@ pub async fn server(stream: TcpStream, config: &TlsConfig) -> Result Result { + configure_server_with_leaf(config).map(|(config, _)| config) +} + +/// Builds the server TLS configuration and returns its leaf certificate for +/// pg-proto's RFC 5929 channel-binding transport. +pub fn configure_server_with_leaf( + config: &TlsConfig, +) -> Result<(rustls::ServerConfig, CertificateDer<'static>), Error> { let certs = match config { TlsConfig::Pem { certificate_pem: certificate, @@ -68,11 +46,15 @@ pub fn configure_server(config: &TlsConfig) -> Result PrivateKeyDer::from_pem_file(private_key), }?; + let leaf = certs + .first() + .cloned() + .ok_or(rustls::Error::NoCertificatesPresented)?; let server_config = rustls::ServerConfig::builder() .with_no_client_auth() .with_single_cert(certs, key)?; - Ok(server_config) + Ok((server_config, leaf)) } ///