diff --git a/.github/workflows/checks.yml b/.github/workflows/checks.yml index 5d07e39..346cb7e 100644 --- a/.github/workflows/checks.yml +++ b/.github/workflows/checks.yml @@ -65,7 +65,7 @@ jobs: - name: Install build dependencies run: | apt-get update - apt-get install -y --no-install-recommends clang libclang-dev + apt-get install -y --no-install-recommends clang libclang-dev libpcap-dev - name: Cache Cargo registry and git uses: actions/cache@v4 with: diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index ba7db0f..e834deb 100644 --- a/.github/workflows/test.yml +++ b/.github/workflows/test.yml @@ -286,6 +286,39 @@ jobs: - name: Run WebSocket tests run: | cargo test --features=websocket extension_cases::websocket + test-net: + name: Extension / Net + runs-on: ubuntu-latest + strategy: + matrix: + arch: [amd64] + steps: + - uses: actions/checkout@v2 + with: + submodules: true + - name: Cache Cargo registry and git + uses: actions/cache@v4 + with: + path: /home/runner/.cargo + key: cargo-${{ runner.os }}-${{ hashFiles('**/Cargo.lock') }} + restore-keys: | + cargo-${{ runner.os }}- + - name: Cache Rust target directory + uses: actions/cache@v4 + with: + path: target + key: target-${{ runner.os }}-${{ hashFiles('**/Cargo.lock') }}-${{ hashFiles('**/Cargo.toml') }} + restore-keys: | + target-${{ runner.os }}-${{ hashFiles('**/Cargo.lock') }}- + target-${{ runner.os }}- + - name: Setup Rust Toolchain + uses: ./.github/actions/setup-rust + - name: Install libpcap + run: sudo apt-get update && sudo apt-get install -y libpcap-dev + - name: Run Net tests + run: | + cargo test --manifest-path crates/datafusion-net/Cargo.toml --all-features + cargo test --features=net extension_cases::net test-clickhouse: name: Extension / ClickHouse runs-on: ubuntu-latest @@ -526,6 +559,8 @@ jobs: target-${{ runner.os }}- - name: Setup Rust Toolchain uses: ./.github/actions/setup-rust + - name: Install libpcap + run: sudo apt-get update && sudo apt-get install -y libpcap-dev - name: Run UDFs WASM tests run: | cargo test --all-features --manifest-path crates/datafusion-app/Cargo.toml diff --git a/CLAUDE.md b/CLAUDE.md index cdfdc61..803648b 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -110,11 +110,13 @@ cargo test --features="deltalake s3" extension_cases::deltalake::test_deltalake_ cargo test --features=udfs-wasm extension_cases::udfs_wasm cargo test --features=vortex extension_cases::vortex cargo test --features=vortex cli_cases::basic::test_output_vortex +cargo test --features=net extension_cases::net # Requires libpcap (libpcap-dev on Debian/Ubuntu) # Run tests for specific crates cargo test --manifest-path crates/datafusion-app/Cargo.toml --all-features cargo test --manifest-path crates/datafusion-functions-parquet/Cargo.toml cargo test --manifest-path crates/datafusion-udfs-wasm/Cargo.toml +cargo test --manifest-path crates/datafusion-net/Cargo.toml --all-features # Run a single test cargo test @@ -168,6 +170,8 @@ The project is organized as a workspace with multiple crates: - **`crates/datafusion-udfs-wasm`**: WASM-based UDF support +- **`crates/datafusion-net`**: Network packet capture querying — `pcap` (capture files), `capture` (live capture), `interfaces` (list capture devices), and `tcp_conversations` (per-connection TCP flow analytics) table functions (`capture`/`interfaces` behind its `live` feature); the `reverse_dns` (IP → hostname via PTR lookup) and `geoip` (IP → location via MaxMind `.mmdb` database) scalar UDFs; the `dns_query` and `tls_sni` payload-decoding scalar UDFs; and `pcap_wide` / `capture_wide` variants that append DNS and geolocation enrichment columns + - **`crates/datafusion-auth`**: Authentication implementations - **`crates/datafusion-ffi-table-providers`**: FFI table provider support @@ -212,6 +216,7 @@ The project uses extensive feature flags to keep binary size manageable: - `flightsql` - FlightSQL server and client - `http` - HTTP server - `huggingface` - HuggingFace dataset integration +- `net` - Query network packet captures (pcap files and live capture) via SQL - `udfs-wasm` - WASM UDF support - `observability` - Metrics and tracing (required by servers) diff --git a/Cargo.lock b/Cargo.lock index a1fd5d7..aa1480f 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1709,7 +1709,7 @@ version = "0.6.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6fac387a98bb7c37292057cffc56d62ecb629900026402633ae9160df93a8766" dependencies = [ - "nom", + "nom 7.1.3", ] [[package]] @@ -1765,6 +1765,12 @@ dependencies = [ "phf 0.12.1", ] +[[package]] +name = "circular" +version = "0.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b0fc239e0f6cb375d2402d48afb92f76f5404fd1df208a41930ec81eda078bea" + [[package]] name = "cityhash-rs" version = "1.0.1" @@ -2636,6 +2642,7 @@ dependencies = [ "datafusion", "datafusion-functions-json", "datafusion-functions-parquet", + "datafusion-net", "datafusion-table-providers", "datafusion-udfs-wasm", "deltalake", @@ -2895,9 +2902,11 @@ dependencies = [ "crossterm", "datafusion", "datafusion-app", + "datafusion-net", "datafusion-udfs-wasm", "directories", "env_logger", + "etherparse", "futures", "http 1.4.2", "http-body 1.0.1", @@ -3174,6 +3183,20 @@ dependencies = [ "syn 2.0.118", ] +[[package]] +name = "datafusion-net" +version = "0.1.0" +dependencies = [ + "async-trait", + "datafusion", + "dns-lookup", + "etherparse", + "maxminddb", + "pcap", + "pcap-parser", + "tokio", +] + [[package]] name = "datafusion-optimizer" version = "54.0.0" @@ -3766,6 +3789,18 @@ dependencies = [ "syn 2.0.118", ] +[[package]] +name = "dns-lookup" +version = "3.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6e39034cee21a2f5bbb66ba0e3689819c4bb5d00382a282006e802a7ffa6c41d" +dependencies = [ + "cfg-if 1.0.4", + "libc", + "socket2 0.6.4", + "windows-sys 0.60.2", +] + [[package]] name = "document-features" version = "0.2.12" @@ -3909,6 +3944,17 @@ version = "1.0.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "877a4ace8713b0bcf2a4e7eec82529c029f1d0619886d18145fea96c3ffe5c0f" +[[package]] +name = "errno" +version = "0.2.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f639046355ee4f37944e44f60642c6f3a7efa3cf6b78c78a0d989a8ce6c396a1" +dependencies = [ + "errno-dragonfly", + "libc", + "winapi", +] + [[package]] name = "errno" version = "0.3.14" @@ -3919,6 +3965,25 @@ dependencies = [ "windows-sys 0.61.2", ] +[[package]] +name = "errno-dragonfly" +version = "0.1.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "aa68f1b12764fab894d2755d2518754e71b4fd80ecfb822714a1206c2aab39bf" +dependencies = [ + "cc", + "libc", +] + +[[package]] +name = "etherparse" +version = "0.20.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "aeebdb565da33451794551a771bd0b4e6f416d3182adc3ed2837f614d6168f9a" +dependencies = [ + "arrayvec", +] + [[package]] name = "euclid" version = "0.22.14" @@ -5123,6 +5188,12 @@ dependencies = [ "serde", ] +[[package]] +name = "ipnetwork" +version = "0.21.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cf370abdafd54d13e54a620e8c3e1145f28e46cc9d704bc6d94414559df41763" + [[package]] name = "is_terminal_polyfill" version = "1.70.2" @@ -5738,6 +5809,19 @@ version = "0.8.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "47e1ffaa40ddd1f3ed91f717a33c8c0ee23fff369e3aa8772b9605cc1d22f4c3" +[[package]] +name = "maxminddb" +version = "0.29.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "65e84ef32bcbf18a95548989e880db4af6fafd563463753afb4b9a149fb2782c" +dependencies = [ + "ipnetwork", + "log", + "memchr", + "serde", + "thiserror 2.0.18", +] + [[package]] name = "maybe-owned" version = "0.3.4" @@ -6041,6 +6125,15 @@ dependencies = [ "minimal-lexical", ] +[[package]] +name = "nom" +version = "8.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "df9761775871bdef83bee530e60050f7e54b1105350d6884eb0fb4f46c2f9405" +dependencies = [ + "memchr", +] + [[package]] name = "normalize-line-endings" version = "0.3.0" @@ -6625,6 +6718,32 @@ dependencies = [ "digest 0.10.7", ] +[[package]] +name = "pcap" +version = "2.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bbd2eecc2ddc671ec563b5b39f846556aade68a65d1afb14d8fe6b30b0457d75" +dependencies = [ + "bitflags 1.3.2", + "errno 0.2.8", + "libc", + "libloading", + "pkg-config", + "regex", + "windows-sys 0.36.1", +] + +[[package]] +name = "pcap-parser" +version = "0.17.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1099a3a64b376c536394622c9e5ea8ecfac3e3e44f88857b1e7da114135d3df4" +dependencies = [ + "circular", + "nom 8.0.0", + "rusticata-macros", +] + [[package]] name = "pco" version = "1.0.2" @@ -7851,6 +7970,15 @@ dependencies = [ "semver", ] +[[package]] +name = "rusticata-macros" +version = "5.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f40620447734539bcc6c64e9c66fed2410e17b46986733f1090c922d85f7105b" +dependencies = [ + "nom 8.0.0", +] + [[package]] name = "rustix" version = "0.38.44" @@ -7858,7 +7986,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "fdb5bc1ae2baa591800df16c9ca78619bf65c0488b41b96ccec5d11220d8c154" dependencies = [ "bitflags 2.13.0", - "errno", + "errno 0.3.14", "libc", "linux-raw-sys 0.4.15", "windows-sys 0.59.0", @@ -7871,7 +7999,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b6fe4565b9518b83ef4f91bb47ce29620ca828bd32cb7e408f0062e9930ba190" dependencies = [ "bitflags 2.13.0", - "errno", + "errno 0.3.14", "libc", "linux-raw-sys 0.12.1", "windows-sys 0.61.2", @@ -8329,7 +8457,7 @@ version = "1.4.8" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c4db69cba1110affc0e9f7bcd48bbf87b3f4fc7c61fc9155afd4c469eb3d6c1b" dependencies = [ - "errno", + "errno 0.3.14", "libc", ] @@ -8768,7 +8896,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d4ea810f0692f9f51b382fff5893887bb4580f5fa246fde546e0b13e7fcee662" dependencies = [ "fnv", - "nom", + "nom 7.1.3", "phf 0.11.3", "phf_codegen", ] @@ -9638,9 +9766,9 @@ checksum = "0b928f33d975fc6ad9f86c8f283853ad26bdd5b10b7f1542aa2fa15e2289105a" [[package]] name = "vortex" -version = "0.76.0" +version = "0.78.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "bd102344fdd62730b0da01e625cac867b62a3da3bdeb75ee1f193b4a21317a7f" +checksum = "e999910e7c7f50a1f344b47ec6ebd07ce60d856bcf3f053ab63a358104c19367" dependencies = [ "vortex-alp", "vortex-array", @@ -9673,9 +9801,9 @@ dependencies = [ [[package]] name = "vortex-alp" -version = "0.76.0" +version = "0.78.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9a001e4079c2b17d2da5f04859d39abf86873c311ac1122a92d095e123ba51c8" +checksum = "a78eaf3adcedeb6a71cc28b4b837db147524a138aa72563f128aab776c572901" dependencies = [ "itertools 0.14.0", "num-traits", @@ -9692,9 +9820,9 @@ dependencies = [ [[package]] name = "vortex-array" -version = "0.76.0" +version = "0.78.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "733e063d45844eb6b85d781c21d0987cc820873525d35882c97de08c4c214f3a" +checksum = "eedbfa56229af2db2945e078ec505121ac8ef1179f7c9ee289420a43647144f4" dependencies = [ "arc-swap", "arcref", @@ -9745,9 +9873,9 @@ dependencies = [ [[package]] name = "vortex-array-macros" -version = "0.76.0" +version = "0.78.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f7917bec5a62d82c24a7b5b5a86269078ba8d70c47704ddfe0316d70f073b465" +checksum = "357e1a0737582c236a75430582c057145781bfced9bab99d5cbed93bd498f3d0" dependencies = [ "proc-macro2", "quote", @@ -9756,9 +9884,9 @@ dependencies = [ [[package]] name = "vortex-btrblocks" -version = "0.76.0" +version = "0.78.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2dbaa4c14bd374b8e477e198584fd490cbeb8873fb35e24900de0c67d6bdc061" +checksum = "3662139b831588f59b90b526f763384dae3c868ce4192dd236e8b81fa69a8291" dependencies = [ "itertools 0.14.0", "num-traits", @@ -9784,9 +9912,9 @@ dependencies = [ [[package]] name = "vortex-buffer" -version = "0.76.0" +version = "0.78.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "dd04204ae7ba07c833f5fa7807a93f64ca227e1f4b39380a42cddb443b4980ce" +checksum = "73094503b5983b0408fb1b36d6a21bf783290627cbdc994539fdff433c42966a" dependencies = [ "arrow-buffer", "bitvec", @@ -9798,9 +9926,9 @@ dependencies = [ [[package]] name = "vortex-bytebool" -version = "0.76.0" +version = "0.78.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "29cf8213402292ee2be540309db32e4d6a0a279f43fda6d3a20a4b2b5c3e0a66" +checksum = "85dc654b529c47912300d8fe9e099110565e6a04602369f3741d7fe2382ca622" dependencies = [ "num-traits", "vortex-array", @@ -9812,9 +9940,9 @@ dependencies = [ [[package]] name = "vortex-compressor" -version = "0.76.0" +version = "0.78.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6bd4ab1c38e49496113f6207321410d307f3cbf09f41318ba12d9b53aadcf030" +checksum = "1ed8fa7a85228c365316192f5690e0794a0367b0e97b6dadbc9123562630e2b8" dependencies = [ "itertools 0.14.0", "num-traits", @@ -9831,18 +9959,18 @@ dependencies = [ [[package]] name = "vortex-compute" -version = "0.76.0" +version = "0.78.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5c93f1a1dd99f535e41a69f790335853c7ba8ff50b6bada0efd6c1bbc68922d6" +checksum = "2c5f057d9083c84e0e5c16f4965c947b8c9dcbd4d952b00c1cdf3d87f26d8b66" dependencies = [ "vortex-buffer", ] [[package]] name = "vortex-datafusion" -version = "0.76.0" +version = "0.78.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "dbe491635b2baf8f72a3cad732ad2f05eb9852bc342b14e793e1ffaefee8aedb" +checksum = "b2c4a66a72ededc9c457ebcf087f6a98f20ee74ee2b758d85b9156fb1b40a161" dependencies = [ "arrow-array", "arrow-schema", @@ -9872,9 +10000,9 @@ dependencies = [ [[package]] name = "vortex-datetime-parts" -version = "0.76.0" +version = "0.78.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b3faa8f52de5ddc4c3316387492247fcdd641405216e09654da6c54042c97229" +checksum = "a2f9dac037c88913f5829d7559442231996cdc542ed12a692ec0d5ccfcd99151" dependencies = [ "num-traits", "prost", @@ -9887,9 +10015,9 @@ dependencies = [ [[package]] name = "vortex-decimal-byte-parts" -version = "0.76.0" +version = "0.78.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "717bca817d417046c31408dfde1ca845e2dadad28b71344df992a935cd7956e7" +checksum = "2952a8baa16f0b0ffaf51442654cf3adff0d876266543367413417af97f96689" dependencies = [ "num-traits", "prost", @@ -9902,9 +10030,9 @@ dependencies = [ [[package]] name = "vortex-error" -version = "0.76.0" +version = "0.78.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "dbf4ed9b6a7bf70216faabd9ae019835cd9d70048b99cba473cbfc366621752b" +checksum = "17607fb14879f521b123c31ffaba184f3b5765e9e344ac8c07b1b0b3be5ec0e9" dependencies = [ "arrow-schema", "flatbuffers", @@ -9916,9 +10044,9 @@ dependencies = [ [[package]] name = "vortex-fastlanes" -version = "0.76.0" +version = "0.78.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "dec2ffcbbacf00b6a5a8fa672e40d97938ec71dd956153dad33da7c3a1375f0f" +checksum = "3d1595b950fd63b336f395e394ac22038641215b0eaea433ba8c76b313d52eec" dependencies = [ "fastlanes", "itertools 0.14.0", @@ -9934,9 +10062,9 @@ dependencies = [ [[package]] name = "vortex-file" -version = "0.76.0" +version = "0.78.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6f223beb4127a21a10a35180ed553ad6b05222757305a9ab24db432e3b4608dc" +checksum = "8d8c9f7d11188ca817806d4f356fc277cded3c81ca4aecf5a31159b3c438cb6c" dependencies = [ "async-trait", "bytes", @@ -9951,6 +10079,7 @@ dependencies = [ "pin-project-lite", "tokio", "tracing", + "url", "vortex-alp", "vortex-array", "vortex-btrblocks", @@ -9979,9 +10108,9 @@ dependencies = [ [[package]] name = "vortex-flatbuffers" -version = "0.76.0" +version = "0.78.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "54ff2075335975c8e0e31934f125c5cf8c96cbea06d95426ec0cd60397865b0e" +checksum = "abde80f1c35c73ac806e4d2ef58e720bc535aa0b217be3cbe058c7d014a6ed71" dependencies = [ "flatbuffers", "vortex-buffer", @@ -9990,9 +10119,9 @@ dependencies = [ [[package]] name = "vortex-fsst" -version = "0.76.0" +version = "0.78.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d78a0283e858c0924466b0339f4547466ce5a71863283bb0ec3dddacb487218e" +checksum = "442abb523fa145d6d26e33856ef9fa3da4a08afeb96769097e4d0c81d0e0efe3" dependencies = [ "fsst-rs", "num-traits", @@ -10006,9 +10135,9 @@ dependencies = [ [[package]] name = "vortex-io" -version = "0.76.0" +version = "0.78.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b85b8bd8b1dd78c40915b382f647bdd8ecba9cba839e2649fc06477b4b139518" +checksum = "549b3a7b6352a7a7861df2c6f41d5c4c33987cd3c01dbd88ed364f6285582629" dependencies = [ "async-fs", "async-stream", @@ -10036,9 +10165,9 @@ dependencies = [ [[package]] name = "vortex-ipc" -version = "0.76.0" +version = "0.78.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d545664ac1a888dbe25df5c182882375981ba3d76a9f5adb378f7686d72b35ff" +checksum = "e64b352ba3554ea92084e3128dd845a3f25c72fc5ab427dff68f0e4d1f896c89" dependencies = [ "bytes", "flatbuffers", @@ -10054,9 +10183,9 @@ dependencies = [ [[package]] name = "vortex-layout" -version = "0.76.0" +version = "0.78.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f80afe91b2262ceb008063784b8e6ca123eb43d506d235154d48ba5dc27f002b" +checksum = "88f4ec82f325823394430d256130a692d8f5c2154cf25fde8215b51fc963a818" dependencies = [ "arcref", "arrow-array", @@ -10097,9 +10226,9 @@ dependencies = [ [[package]] name = "vortex-mask" -version = "0.76.0" +version = "0.78.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6668847d398de6a0351a7ff31fe45b6af93fcae5ba3c99e9e7060205462f1207" +checksum = "c91caf1d598e35e4a5c24cee0cd7d4e8c4931b30a9163ac33af8181d546ae15e" dependencies = [ "itertools 0.14.0", "vortex-buffer", @@ -10108,9 +10237,9 @@ dependencies = [ [[package]] name = "vortex-metrics" -version = "0.76.0" +version = "0.78.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b2dee0e5bef3389f21c5500228368be4fddb70d2fb5bb7624b4a27de4e38625c" +checksum = "9446b43132833788c5fc0cb184912add97ffe1c41a9e9faaf7bafc9b58fdd529" dependencies = [ "parking_lot", "sketches-ddsketch 0.4.0", @@ -10118,9 +10247,9 @@ dependencies = [ [[package]] name = "vortex-pco" -version = "0.76.0" +version = "0.78.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "67133e4d66664ecda2475037c51d110f4cc07cb3d55d43e07e126c70e388af13" +checksum = "c57ea636c346087d48e25008554321a285f175d6e7036db1fefbb1998fb228df" dependencies = [ "pco", "prost", @@ -10133,9 +10262,9 @@ dependencies = [ [[package]] name = "vortex-proto" -version = "0.76.0" +version = "0.78.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9ad9f444e0ee969ffb39bfc6d7eaefae018a27b174e6e5ed47a9604c2b0f2d7d" +checksum = "550431d87f13583e1ac7b54f16ae4450278b63ad1e7aaab56b7b229b908445a6" dependencies = [ "prost", "prost-types", @@ -10143,9 +10272,9 @@ dependencies = [ [[package]] name = "vortex-runend" -version = "0.76.0" +version = "0.78.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4d2f1808eb521bc14b46801eb288085596993b87461d261fa5785f8df9f6bb38" +checksum = "566ec7cc60f3a478db50c1b56cbc9da448b5ac920bdadc76538eb0122142de10" dependencies = [ "arrow-array", "itertools 0.14.0", @@ -10160,9 +10289,9 @@ dependencies = [ [[package]] name = "vortex-scan" -version = "0.76.0" +version = "0.78.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "36f4a7181a6203a931d5a7fd0407a98130df7fb68c350a0c3031faf72f2f15c9" +checksum = "8b5521e8ac93f143a99ccf28b9b2ada30c9d521e7fbaadfc8f5fe08787c15ae4" dependencies = [ "async-trait", "futures", @@ -10177,9 +10306,9 @@ dependencies = [ [[package]] name = "vortex-sequence" -version = "0.76.0" +version = "0.78.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3c8b3c89485369f80bb756efad5641e1d1801624ca799200f9a9af0f60d5df58" +checksum = "5e0218635ebeb62d21f56fd101b2e3803441c1bad9f3f76fcddb8cc545ea7629" dependencies = [ "num-traits", "prost", @@ -10194,9 +10323,9 @@ dependencies = [ [[package]] name = "vortex-session" -version = "0.76.0" +version = "0.78.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4d01f0c02bd164af12a4bab329c96c7fa7fa500802a9508e02dd9860e772a89c" +checksum = "66234858e92df597666ac9de927855d4cca1bac72d1aceca43eb9ffd8b7167dd" dependencies = [ "arc-swap", "lasso", @@ -10207,9 +10336,9 @@ dependencies = [ [[package]] name = "vortex-sparse" -version = "0.76.0" +version = "0.78.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ce1826d788f8e4dc2fb0a8c4bf5023ec35e33677fc90a4138bd0052dfd070dd8" +checksum = "b7e1842b2365812e21a7aa0cc378dc6aa19af9f1955c2a9b79f673ef29789f26" dependencies = [ "itertools 0.14.0", "num-traits", @@ -10223,9 +10352,9 @@ dependencies = [ [[package]] name = "vortex-utils" -version = "0.76.0" +version = "0.78.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9248e8ff57dbc961c1d4554fd8a4c2c7694c39cc7ef271e6dcc979c59b20ae67" +checksum = "24b6fa379006bf275b0ad5641c8b83fd52388fc4d483b83b1c1fd2870dd5d653" dependencies = [ "dashmap", "hashbrown 0.17.1", @@ -10234,9 +10363,9 @@ dependencies = [ [[package]] name = "vortex-zigzag" -version = "0.76.0" +version = "0.78.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "24ed6871c65b6095b97ab5075ac0572868ced5c00e0cf63f0410437bfd937cb9" +checksum = "aadf2833802c2bf9d69d282189bf3c3e1e1aff6c7c5989f3ae8197ba528fda31" dependencies = [ "vortex-array", "vortex-buffer", @@ -10248,9 +10377,9 @@ dependencies = [ [[package]] name = "vortex-zstd" -version = "0.76.0" +version = "0.78.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9a2c515038259368dd8fad40777713f4d19d926ea1522aaff35372ee5965b19c" +checksum = "1264c8bb5247872a0bbbf20ef51f7177959c213a5e11f9c28b4a1c5f91879e48" dependencies = [ "itertools 0.14.0", "prost", @@ -11116,6 +11245,19 @@ dependencies = [ "windows-link", ] +[[package]] +name = "windows-sys" +version = "0.36.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ea04155a16a59f9eab786fe12a4a450e75cdb175f9e0d80da1e17db09f55b8d2" +dependencies = [ + "windows_aarch64_msvc 0.36.1", + "windows_i686_gnu 0.36.1", + "windows_i686_msvc 0.36.1", + "windows_x86_64_gnu 0.36.1", + "windows_x86_64_msvc 0.36.1", +] + [[package]] name = "windows-sys" version = "0.48.0" @@ -11143,6 +11285,15 @@ 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" @@ -11176,13 +11327,30 @@ dependencies = [ "windows_aarch64_gnullvm 0.52.6", "windows_aarch64_msvc 0.52.6", "windows_i686_gnu 0.52.6", - "windows_i686_gnullvm", + "windows_i686_gnullvm 0.52.6", "windows_i686_msvc 0.52.6", "windows_x86_64_gnu 0.52.6", "windows_x86_64_gnullvm 0.52.6", "windows_x86_64_msvc 0.52.6", ] +[[package]] +name = "windows-targets" +version = "0.53.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4945f9f551b88e0d65f3db0bc25c33b8acea4d9e41163edf90dcd0b19f9069f3" +dependencies = [ + "windows-link", + "windows_aarch64_gnullvm 0.53.1", + "windows_aarch64_msvc 0.53.1", + "windows_i686_gnu 0.53.1", + "windows_i686_gnullvm 0.53.1", + "windows_i686_msvc 0.53.1", + "windows_x86_64_gnu 0.53.1", + "windows_x86_64_gnullvm 0.53.1", + "windows_x86_64_msvc 0.53.1", +] + [[package]] name = "windows-threading" version = "0.2.1" @@ -11204,6 +11372,18 @@ version = "0.52.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "32a4622180e7a0ec044bb555404c800bc9fd9ec262ec147edd5989ccd0c02cd3" +[[package]] +name = "windows_aarch64_gnullvm" +version = "0.53.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a9d8416fa8b42f5c947f8482c43e7d89e73a173cead56d044f6a56104a6d1b53" + +[[package]] +name = "windows_aarch64_msvc" +version = "0.36.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9bb8c3fd39ade2d67e9874ac4f3db21f0d710bee00fe7cab16949ec184eeaa47" + [[package]] name = "windows_aarch64_msvc" version = "0.48.5" @@ -11216,6 +11396,18 @@ version = "0.52.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "09ec2a7bb152e2252b53fa7803150007879548bc709c039df7627cabbd05d469" +[[package]] +name = "windows_aarch64_msvc" +version = "0.53.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b9d782e804c2f632e395708e99a94275910eb9100b2114651e04744e9b125006" + +[[package]] +name = "windows_i686_gnu" +version = "0.36.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "180e6ccf01daf4c426b846dfc66db1fc518f074baa793aa7d9b9aaeffad6a3b6" + [[package]] name = "windows_i686_gnu" version = "0.48.5" @@ -11228,12 +11420,30 @@ version = "0.52.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "8e9b5ad5ab802e97eb8e295ac6720e509ee4c243f69d781394014ebfe8bbfa0b" +[[package]] +name = "windows_i686_gnu" +version = "0.53.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "960e6da069d81e09becb0ca57a65220ddff016ff2d6af6a223cf372a506593a3" + [[package]] name = "windows_i686_gnullvm" version = "0.52.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "0eee52d38c090b3caa76c563b86c3a4bd71ef1a819287c19d586d7334ae8ed66" +[[package]] +name = "windows_i686_gnullvm" +version = "0.53.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fa7359d10048f68ab8b09fa71c3daccfb0e9b559aed648a8f95469c27057180c" + +[[package]] +name = "windows_i686_msvc" +version = "0.36.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e2e7917148b2812d1eeafaeb22a97e4813dfa60a3f8f78ebe204bcc88f12f024" + [[package]] name = "windows_i686_msvc" version = "0.48.5" @@ -11246,6 +11456,18 @@ version = "0.52.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "240948bc05c5e7c6dabba28bf89d89ffce3e303022809e73deaefe4f6ec56c66" +[[package]] +name = "windows_i686_msvc" +version = "0.53.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1e7ac75179f18232fe9c285163565a57ef8d3c89254a30685b57d83a38d326c2" + +[[package]] +name = "windows_x86_64_gnu" +version = "0.36.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4dcd171b8776c41b97521e5da127a2d86ad280114807d0b2ab1e462bc764d9e1" + [[package]] name = "windows_x86_64_gnu" version = "0.48.5" @@ -11258,6 +11480,12 @@ version = "0.52.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "147a5c80aabfbf0c7d901cb5895d1de30ef2907eb21fbbab29ca94c5b08b1a78" +[[package]] +name = "windows_x86_64_gnu" +version = "0.53.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9c3842cdd74a865a8066ab39c8a7a473c0778a3f29370b5fd6b4b9aa7df4a499" + [[package]] name = "windows_x86_64_gnullvm" version = "0.48.5" @@ -11270,6 +11498,18 @@ version = "0.52.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "24d5b23dc417412679681396f2b49f3de8c1473deb516bd34410872eff51ed0d" +[[package]] +name = "windows_x86_64_gnullvm" +version = "0.53.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0ffa179e2d07eee8ad8f57493436566c7cc30ac536a3379fdf008f47f6bb7ae1" + +[[package]] +name = "windows_x86_64_msvc" +version = "0.36.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c811ca4a8c853ef420abd8592ba53ddbbac90410fab6903b3e79972a631f7680" + [[package]] name = "windows_x86_64_msvc" version = "0.48.5" @@ -11282,6 +11522,12 @@ version = "0.52.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "589f6da84c646204747d1270a2a5661ea66ed1cced2631d546fdfb155959f9ec" +[[package]] +name = "windows_x86_64_msvc" +version = "0.53.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d6bbff5f0aada427a1e5a6da5f1f98158182f26556f345ac9e04d36d0ebed650" + [[package]] name = "winnow" version = "0.7.15" diff --git a/Cargo.toml b/Cargo.toml index 7632ae7..5cea5b5 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -68,16 +68,19 @@ tui-logger = { features = [ ], optional = true, version = "0.18" } url = { features = ["serde"], version = "2.5.2" } uuid = { optional = true, version = "1.10.0" } -vortex = { optional = true, version = "0.76" } -vortex-datafusion = { optional = true, version = "0.76" } -vortex-file = { optional = true, version = "0.76" } -vortex-session = { optional = true, version = "0.76" } +vortex = { optional = true, version = "0.78" } +vortex-datafusion = { optional = true, version = "0.78" } +vortex-file = { optional = true, version = "0.78" } +vortex-session = { optional = true, version = "0.78" } [dev-dependencies] assert_cmd = "2.0.16" +# Used (without the libpcap-backed `live` feature) to generate pcap fixtures +datafusion-net = { path = "crates/datafusion-net", version = "0.1.0" } datafusion-udfs-wasm = { features = [ "serde", ], path = "crates/datafusion-udfs-wasm", version = "0.1.0" } +etherparse = "0.20" http-body-util = "0.1.3" insta = { features = ["yaml"], version = "1.40.0" } predicates = "3.1.2" @@ -117,6 +120,7 @@ http = [ ] huggingface = ["datafusion-app/huggingface"] mongodb = ["datafusion-app/mongodb"] +net = ["datafusion-app/net"] s3 = ["datafusion-app/s3"] tui = ["dep:crossterm", "dep:ratatui", "dep:ratatui-textarea", "dep:tui-logger"] udfs-wasm = ["datafusion-app/udfs-wasm"] diff --git a/README.md b/README.md index c85d6ed..e413a64 100644 --- a/README.md +++ b/README.md @@ -61,6 +61,9 @@ cargo install datafusion-dft --features=functions-json,functions-parquet # With the websocket table function for streaming WebSocket messages cargo install datafusion-dft --features=websocket + +# With the pcap / capture table functions for querying network packet captures +cargo install datafusion-dft --features=net ``` See the [Features documentation](docs/features.md) for all available features. diff --git a/crates/datafusion-app/Cargo.toml b/crates/datafusion-app/Cargo.toml index 0e1be08..54d14e6 100644 --- a/crates/datafusion-app/Cargo.toml +++ b/crates/datafusion-app/Cargo.toml @@ -15,6 +15,7 @@ color-eyre = "0.6.3" datafusion = "54" datafusion-functions-json = { optional = true, version = "0.54" } datafusion-functions-parquet = { optional = true, path = "../datafusion-functions-parquet", version = "0.1.0" } +datafusion-net = { optional = true, path = "../datafusion-net", version = "0.1.0" } datafusion-table-providers = { optional = true, version = "0.12.0" } datafusion-udfs-wasm = { features = [ "serde", @@ -56,7 +57,7 @@ tokio-tungstenite = { features = [ ], optional = true, version = "0.29" } tonic = { optional = true, version = "0.14" } url = { optional = true, version = "2.5.2" } -vortex-datafusion = { optional = true, version = "0.76" } +vortex-datafusion = { optional = true, version = "0.78" } [dev-dependencies] criterion = { features = ["async_tokio"], version = "0.5.1" } @@ -76,6 +77,7 @@ mongodb = [ "datafusion-table-providers/mongodb", "dep:datafusion-table-providers", ] +net = ["datafusion-net/live", "dep:datafusion-net"] observability = ["dep:metrics", "dep:tokio-metrics"] s3 = ["object_store/aws", "url"] udfs-wasm = ["dep:datafusion-udfs-wasm"] diff --git a/crates/datafusion-app/src/config.rs b/crates/datafusion-app/src/config.rs index 1f0445b..0b0d442 100644 --- a/crates/datafusion-app/src/config.rs +++ b/crates/datafusion-app/src/config.rs @@ -77,9 +77,25 @@ pub fn merge_configs(shared: ExecutionConfig, priority: ExecutionConfig) -> Exec merged.mongodb = Some(mongodb) } + #[cfg(feature = "net")] + if let Some(geoip_db_path) = priority.net.geoip_db_path { + merged.net.geoip_db_path = Some(geoip_db_path) + } + merged } +/// Configuration for the `net` feature +#[cfg(feature = "net")] +#[derive(Clone, Debug, Default, Deserialize)] +pub struct NetConfig { + /// Path to a MaxMind-format (`.mmdb`) database, such as GeoLite2-City, + /// used by the single-argument form of the `geoip` UDF. The `GEOIP_DB` + /// environment variable takes precedence over this value. + #[serde(default)] + pub geoip_db_path: Option, +} + #[derive(Clone, Debug, Deserialize)] pub struct ExecutionConfig { #[serde(default)] @@ -105,6 +121,9 @@ pub struct ExecutionConfig { #[cfg(feature = "udfs-wasm")] #[serde(default = "default_wasm_udf")] pub wasm_udf: WasmUdfConfig, + #[cfg(feature = "net")] + #[serde(default)] + pub net: NetConfig, #[serde(default = "default_catalog")] pub catalog: CatalogConfig, #[cfg(feature = "observability")] @@ -128,6 +147,8 @@ impl Default for ExecutionConfig { // iceberg: default_iceberg_config(), #[cfg(feature = "udfs-wasm")] wasm_udf: default_wasm_udf(), + #[cfg(feature = "net")] + net: NetConfig::default(), catalog: default_catalog(), #[cfg(feature = "observability")] observability: default_observability(), diff --git a/crates/datafusion-app/src/local.rs b/crates/datafusion-app/src/local.rs index d49f2f0..afa16c6 100644 --- a/crates/datafusion-app/src/local.rs +++ b/crates/datafusion-app/src/local.rs @@ -126,6 +126,43 @@ impl ExecutionContext { Arc::new(crate::tables::websocket::WebSocketFunc::default()), ); + #[cfg(feature = "net")] + { + use datafusion::logical_expr::ScalarUDF; + + session_ctx.register_udtf("pcap", Arc::new(datafusion_net::PcapFunc::default())); + session_ctx.register_udtf("capture", Arc::new(datafusion_net::CaptureFunc::default())); + session_ctx.register_udtf( + "interfaces", + Arc::new(datafusion_net::InterfacesFunc::default()), + ); + session_ctx.register_udtf( + "tcp_conversations", + Arc::new(datafusion_net::TcpConversationsFunc::default()), + ); + session_ctx.register_udf(ScalarUDF::from(datafusion_net::ReverseDnsUdf::default())); + session_ctx.register_udf(ScalarUDF::from(datafusion_net::DnsQueryUdf::default())); + session_ctx.register_udf(ScalarUDF::from(datafusion_net::TlsSniUdf::default())); + // The GEOIP_DB environment variable takes precedence over the + // configured database path + let geoip_db_path = std::env::var_os(datafusion_net::GEOIP_DB_ENV_VAR) + .map(std::path::PathBuf::from) + .or_else(|| config.net.geoip_db_path.clone()); + let geoip = match &geoip_db_path { + Some(path) => datafusion_net::GeoIpUdf::with_db_path(path), + None => datafusion_net::GeoIpUdf::default(), + }; + session_ctx.register_udf(ScalarUDF::from(geoip)); + session_ctx.register_udtf( + "pcap_wide", + Arc::new(datafusion_net::PcapWideFunc::new(geoip_db_path.clone())), + ); + session_ctx.register_udtf( + "capture_wide", + Arc::new(datafusion_net::CaptureWideFunc::new(geoip_db_path)), + ); + } + let catalog = create_app_catalog(config, app_name, app_version)?; session_ctx.register_catalog(&config.catalog.name, catalog); diff --git a/crates/datafusion-net/Cargo.toml b/crates/datafusion-net/Cargo.toml new file mode 100644 index 0000000..78622da --- /dev/null +++ b/crates/datafusion-net/Cargo.toml @@ -0,0 +1,32 @@ +[package] +authors = ["Matthew Turner "] +description = "DataFusion table functions for querying network packet captures (pcap)" +edition = "2021" +homepage = "https://github.com/datafusion-contrib/datafusion-dft/tree/main/crates/datafusion-net" +keywords = ["datafusion", "network", "pcap", "query", "sql"] +license = "Apache-2.0" +name = "datafusion-net" +readme = "README.md" +repository = "https://github.com/datafusion-contrib/datafusion-dft/tree/main/crates/datafusion-net" +version = "0.1.0" + +[dependencies] +async-trait = "0.1.80" +datafusion = { version = "54" } +dns-lookup = "3" +etherparse = "0.20" +maxminddb = "0.29" +pcap = { optional = true, version = "2.4" } +pcap-parser = "0.17" +tokio = { features = ["sync"], version = "1" } + +[dev-dependencies] +tempfile = "3" +tokio = { features = ["macros", "rt-multi-thread"], version = "1" } + +[features] +default = [] +live = ["dep:pcap"] + +[lints.clippy] +clone_on_ref_ptr = "deny" diff --git a/crates/datafusion-net/README.md b/crates/datafusion-net/README.md new file mode 100644 index 0000000..596b601 --- /dev/null +++ b/crates/datafusion-net/README.md @@ -0,0 +1,245 @@ +# datafusion-net + +DataFusion table functions for querying network packet captures with SQL, +similar to wireshark / tshark. + +## Table functions + +### `pcap(path)` + +Reads a pcap or pcapng capture file as a table: + +```sql +SELECT src_ip, dst_ip, protocol, length +FROM pcap('capture.pcap') +WHERE dst_port = 443; +``` + +### `capture(interface [, bpf_filter [, duration_secs]])` + +Requires the `live` feature (links against libpcap). Streams live-captured +packets from a network interface: + +```sql +-- Stream until 100 packets have been captured +SELECT * FROM capture('en0') LIMIT 100; + +-- Apply a BPF filter +SELECT * FROM capture('en0', 'tcp port 443') LIMIT 10; + +-- Capture for 10 seconds; the stream terminates, so aggregations work +SELECT src_ip, count(*) FROM capture('en0', '', 10) GROUP BY src_ip; +``` + +Without a duration the source is unbounded: use a `LIMIT` or the query +streams until cancelled. Live capture requires elevated privileges (sudo, +`cap_net_raw`+`cap_net_admin` on Linux, or ChmodBPF on macOS). + +### `pcap_wide(path)` and `capture_wide(interface [, bpf_filter [, duration_secs]])` + +Wide variants of `pcap` and `capture` (same arguments, `capture_wide` +requires the `live` feature) that append DNS and geolocation enrichment +columns for the source and destination addresses: `src_host` / `dst_host` +(reverse DNS) and `src_country`, `src_city`, `src_lat`, `src_lon` plus the +`dst_` equivalents (from the `geoip` machinery and a MaxMind-format +database): + +```sql +SELECT dst_ip, dst_host, dst_country, count(*) AS packets +FROM pcap_wide('capture.pcap') +GROUP BY dst_ip, dst_host, dst_country +ORDER BY packets DESC; +``` + +The geolocation database comes from the `GEOIP_DB` environment variable or a +path passed to `PcapWideFunc::new` / `CaptureWideFunc::new` (see the `geoip` +section below for which database populates which fields). A database that is +unconfigured, missing, or unreadable never fails the query — the geolocation +columns are just `NULL` (use `geoip(ip, path)['error']` to see why). +Enrichment is a projection over the narrow table, so unused columns still +prune and only projected enrichment is computed. + +### `interfaces()` + +Requires the `live` feature. Lists the system's network capture interfaces +(similar to `tshark -D`), one row per interface — useful for discovering +what to pass to `capture`. Listing does not require elevated privileges: + +```sql +SELECT name, description, addresses, connection_status +FROM interfaces() +WHERE is_up AND NOT is_loopback; +``` + +Columns: `name` (`Utf8`), `description` (`Utf8`), `addresses` +(`List`), `is_up`, `is_running`, `is_loopback`, `is_wireless` +(`Boolean`), and `connection_status` (`Utf8`: `connected`, `disconnected`, +`unknown`, or `not_applicable`). + +### `tcp_conversations(path)` + +Aggregates a capture file into one row per TCP connection (like wireshark's +"Statistics → Conversations" or `tshark -z conv,tcp`) — useful for spotting +slow handshakes, retransmission-heavy flows, and connections that never +closed cleanly: + +```sql +SELECT dst_ip, dst_port, handshake_rtt_ms, retransmissions_fwd, state +FROM tcp_conversations('capture.pcap') +ORDER BY retransmissions_fwd DESC; +``` + +Direction is from the connection initiator's perspective: `src_*` is the +endpoint that sent the first SYN (or, for a capture that starts +mid-connection, the first endpoint seen sending); `fwd` counts +initiator → responder and `rev` the reverse. + +| Column | Type | Notes | +| --------------------------------------------- | ------------------------ | -------------------------------------------------- | +| `src_ip`, `src_port`, `dst_ip`, `dst_port` | `Utf8` / `UInt16` | Initiator and responder endpoints | +| `first_timestamp`, `last_timestamp` | `Timestamp(Microsecond)` | First and last packet seen | +| `duration_ms` | `Float64` | Time between them | +| `packets_fwd`, `packets_rev` | `UInt64` | Packet counts per direction | +| `bytes_fwd`, `bytes_rev` | `UInt64` | On-the-wire frame bytes per direction | +| `retransmissions_fwd`, `retransmissions_rev` | `UInt64` | Segments below the highest seen sequence end | +| `handshake_rtt_ms` | `Float64` | SYN → handshake-completing ACK; NULL if not captured | +| `state` | `Utf8` | `active`, `half_closed`, `closed` (both FINs), or `reset` | + +## Scalar functions + +### `reverse_dns(ip)` + +Resolves an IP address string to a hostname via a reverse DNS (PTR) lookup, +using the system resolver. Pairs naturally with the `src_ip` / `dst_ip` +columns: + +```sql +SELECT dst_ip, reverse_dns(dst_ip) AS host, count(*) AS packets +FROM capture('en0', 'tcp', 10) +GROUP BY dst_ip, host +ORDER BY packets DESC; +``` + +Lookups are network I/O, so results are cached process-wide, identical +addresses within a batch are resolved once, and each batch of lookups is +bounded by a timeout. Addresses that fail to parse, fail to resolve, or time +out yield `NULL`. + +### `geoip(ip [, db_path])` + +Geolocates an IP address using a MaxMind-format (`.mmdb`) database. Returns +a struct with `country_code`, `country`, `city`, `latitude`, `longitude`, +`time_zone`, and `error` fields: + +```sql +SELECT geoip(src_ip, '/path/GeoLite2-City.mmdb')['country_code'] AS country, + count(*) AS packets +FROM pcap('capture.pcap') +GROUP BY country +ORDER BY packets DESC; +``` + +**Which database?** A single database file serves every field — there is no +per-field database. Which fields are populated depends on the database's +schema: + +| Field | City database | Country database | ASN / other | +| ---------------------------------------------- | ------------- | ---------------- | ----------- | +| `country_code`, `country` | ✓ | ✓ | `NULL` | +| `city`, `latitude`, `longitude`, `time_zone` | ✓ | `NULL` | `NULL` | + +Use a City-schema database for full coverage: [GeoLite2-City] (free with a +MaxMind account) or the commercial GeoIP2-City. A Country-schema database +(GeoLite2-Country) is smaller if only country-level fields are needed. + +The single-argument form uses the database from the `GEOIP_DB` environment +variable (read when the UDF is constructed), or the path passed to +`GeoIpUdf::with_db_path`. Opened databases are cached process-wide. + +Addresses that fail to parse or have no entry in the database yield a `NULL` +struct. A database that is missing, unreadable, or unconfigured does not +fail the query: the location fields are `NULL` and the `error` field (which +is `NULL` on success) carries the reason: + +```sql +SELECT geoip(src_ip)['error'] FROM pcap('capture.pcap') LIMIT 1; +-- geoip failed to open database '/path/GeoLite2-City.mmdb': ... +``` + +[GeoLite2-City]: https://dev.maxmind.com/geoip/geolite2-free-geolocation-data + +### `dns_query(payload)` + +Decodes a DNS message from a packet payload (typically UDP port 53 — DNS +over UDP fits in a single datagram, so no reassembly is needed). Returns a +struct with `is_response` (`Boolean`), `name` and `query_type` (`Utf8`, from +the first question), `response_code` (`Utf8`, NULL for queries), and +`answers` (`List`, NULL for queries). Answers include A/AAAA addresses +and CNAME/NS/PTR names; other record types are skipped. Payloads that do not +parse as DNS yield a NULL struct. + +```sql +SELECT dns_query(payload)['name'] AS name, count(*) AS queries +FROM pcap('capture.pcap') +WHERE dst_port = 53 +GROUP BY name ORDER BY queries DESC; +``` + +### `tls_sni(payload)` + +Extracts the Server Name Indication (SNI) host name from a TLS ClientHello +in a packet payload (typically TCP port 443), or NULL when the payload is +not a ClientHello or carries no SNI. In HTTPS traffic the payloads are +encrypted, but the ClientHello — the first packet the client sends — is not, +so this reveals the destination host of otherwise opaque connections: + +```sql +SELECT tls_sni(payload) AS host, count(*) AS hellos +FROM pcap('capture.pcap') +WHERE tls_sni(payload) IS NOT NULL +GROUP BY host ORDER BY hellos DESC; +``` + +The ClientHello must fit in the payload (it usually does); this does not +reassemble segments. + +## Schema + +One row per captured frame. Layers that cannot be decoded (non-IP traffic, +truncated frames, unknown link types) yield null columns rather than errors. + +| Column | Type | Notes | +| ---------------- | ------------------------ | ---------------------------------------- | +| `timestamp` | `Timestamp(Microsecond)` | Capture timestamp | +| `frame_number` | `UInt64` | 1-based position in the capture | +| `length` | `UInt32` | Original frame length on the wire | +| `capture_length` | `UInt32` | Bytes actually captured | +| `eth_src` | `Utf8` | Source MAC | +| `eth_dst` | `Utf8` | Destination MAC | +| `ethertype` | `Utf8` | `ipv4`, `ipv6`, `arp`, or hex value | +| `vlan` | `UInt16` | VLAN identifier (outer tag if double) | +| `ip_version` | `UInt8` | 4 or 6 | +| `src_ip` | `Utf8` | | +| `dst_ip` | `Utf8` | | +| `ttl` | `UInt8` | Hop limit for IPv6 | +| `protocol` | `Utf8` | `tcp`, `udp`, `icmp`, `icmpv6`, ... | +| `src_port` | `UInt16` | | +| `dst_port` | `UInt16` | | +| `tcp_flags` | `Utf8` | e.g. `SYN\|ACK` | +| `tcp_seq` | `UInt32` | | +| `tcp_ack` | `UInt32` | | +| `tcp_window` | `UInt16` | Receive window size (unscaled) | +| `payload_length` | `UInt32` | Transport payload bytes | +| `payload` | `Binary` | Only materialized when projected | + +## Usage + +```rust,no_run +use std::sync::Arc; +use datafusion::prelude::SessionContext; + +let ctx = SessionContext::new(); +ctx.register_udtf("pcap", Arc::new(datafusion_net::PcapFunc::default())); +// With the `live` feature: +// ctx.register_udtf("capture", Arc::new(datafusion_net::CaptureFunc::default())); +``` diff --git a/crates/datafusion-net/src/conversations.rs b/crates/datafusion-net/src/conversations.rs new file mode 100644 index 0000000..416998a --- /dev/null +++ b/crates/datafusion-net/src/conversations.rs @@ -0,0 +1,653 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! `tcp_conversations` table function: flow-level TCP analytics over a +//! capture file, similar to wireshark's "Statistics → Conversations" or +//! `tshark -z conv,tcp`. One row per TCP connection: +//! +//! ```sql +//! SELECT dst_ip, dst_port, handshake_rtt_ms, retransmissions_fwd, state +//! FROM tcp_conversations('capture.pcap') +//! ORDER BY handshake_rtt_ms DESC; +//! ``` +//! +//! Direction is from the connection initiator's perspective: `src_*` is the +//! endpoint that sent the first SYN (or, when the capture starts +//! mid-connection, the first endpoint seen sending), `fwd` counts +//! initiator → responder traffic and `rev` the reverse. The handshake RTT is +//! the time from the initiator's SYN to its handshake-completing ACK +//! (wireshark's initial RTT); it is NULL when the capture does not contain +//! the full handshake. Retransmissions are counted per direction with the +//! standard heuristic: a segment that consumes sequence space entirely below +//! the direction's highest seen sequence number. + +use std::{collections::HashMap, sync::Arc}; + +use async_trait::async_trait; +use datafusion::{ + arrow::{ + array::{ + Float64Builder, RecordBatch, StringBuilder, TimestampMicrosecondBuilder, UInt16Builder, + UInt64Builder, + }, + datatypes::{DataType, Field, Schema, SchemaRef, TimeUnit}, + }, + catalog::{Session, TableFunctionImpl, TableProvider}, + common::{internal_err, project_schema, Result}, + datasource::TableType, + execution::SendableRecordBatchStream, + physical_expr::EquivalenceProperties, + physical_plan::{ + execution_plan::{Boundedness, EmissionType}, + stream::RecordBatchReceiverStream, + DisplayAs, DisplayFormatType, ExecutionPlan, Partitioning, PlanProperties, + }, + prelude::Expr, +}; +use tokio::sync::mpsc::Sender; + +use crate::{ + decode::{decode_frame, DecodedPacket}, + file::{for_each_frame, parse_pcap_args, FrameLoop}, + schema::send_error, +}; + +/// Schema of the `tcp_conversations` table function: one row per TCP +/// connection, from the initiator's perspective +pub fn tcp_conversations_schema() -> SchemaRef { + Arc::new(Schema::new(vec![ + Field::new("src_ip", DataType::Utf8, false), + Field::new("src_port", DataType::UInt16, false), + Field::new("dst_ip", DataType::Utf8, false), + Field::new("dst_port", DataType::UInt16, false), + Field::new( + "first_timestamp", + DataType::Timestamp(TimeUnit::Microsecond, None), + false, + ), + Field::new( + "last_timestamp", + DataType::Timestamp(TimeUnit::Microsecond, None), + false, + ), + Field::new("duration_ms", DataType::Float64, false), + Field::new("packets_fwd", DataType::UInt64, false), + Field::new("packets_rev", DataType::UInt64, false), + Field::new("bytes_fwd", DataType::UInt64, false), + Field::new("bytes_rev", DataType::UInt64, false), + Field::new("retransmissions_fwd", DataType::UInt64, false), + Field::new("retransmissions_rev", DataType::UInt64, false), + Field::new("handshake_rtt_ms", DataType::Float64, true), + Field::new("state", DataType::Utf8, false), + ])) +} + +/// Table function that aggregates a capture file into TCP conversations +#[derive(Debug, Default)] +pub struct TcpConversationsFunc {} + +impl TableFunctionImpl for TcpConversationsFunc { + fn call(&self, exprs: &[Expr]) -> Result> { + let path = parse_pcap_args("tcp_conversations", exprs)?; + Ok(Arc::new(TcpConversationsTable::new(path))) + } +} + +/// [`TableProvider`] backed by flow aggregation of a pcap/pcapng file +#[derive(Debug)] +pub struct TcpConversationsTable { + path: String, + schema: SchemaRef, +} + +impl TcpConversationsTable { + pub fn new(path: String) -> Self { + Self { + path, + schema: tcp_conversations_schema(), + } + } +} + +#[async_trait] +impl TableProvider for TcpConversationsTable { + fn schema(&self) -> SchemaRef { + Arc::clone(&self.schema) + } + + fn table_type(&self) -> TableType { + TableType::Base + } + + async fn scan( + &self, + _state: &dyn Session, + projection: Option<&Vec>, + _filters: &[Expr], + limit: Option, + ) -> Result> { + let exec = TcpConversationsExec::try_new( + self.path.clone(), + Arc::clone(&self.schema), + projection.cloned(), + limit, + )?; + Ok(Arc::new(exec)) + } +} + +/// Execution plan that scans the capture on `execute` and emits a single +/// batch of conversations once the whole file has been read +#[derive(Debug)] +struct TcpConversationsExec { + path: String, + projection: Option>, + projected_schema: SchemaRef, + limit: Option, + cache: Arc, +} + +impl TcpConversationsExec { + fn try_new( + path: String, + schema: SchemaRef, + projection: Option>, + limit: Option, + ) -> Result { + let projected_schema = project_schema(&schema, projection.as_ref())?; + let cache = Arc::new(PlanProperties::new( + EquivalenceProperties::new(Arc::clone(&projected_schema)), + Partitioning::UnknownPartitioning(1), + // Aggregation over the whole file: nothing is emitted until the + // scan completes + EmissionType::Final, + Boundedness::Bounded, + )); + Ok(Self { + path, + projection, + projected_schema, + limit, + cache, + }) + } +} + +impl DisplayAs for TcpConversationsExec { + fn fmt_as(&self, _t: DisplayFormatType, f: &mut std::fmt::Formatter) -> std::fmt::Result { + write!( + f, + "TcpConversationsExec: path={}, limit={:?}", + self.path, self.limit + ) + } +} + +impl ExecutionPlan for TcpConversationsExec { + fn name(&self) -> &str { + "TcpConversationsExec" + } + + fn properties(&self) -> &Arc { + &self.cache + } + + fn children(&self) -> Vec<&Arc> { + // This is a leaf node and has no children + vec![] + } + + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + if children.is_empty() { + Ok(self) + } else { + internal_err!("Children cannot be replaced in {self:?}") + } + } + + fn execute( + &self, + partition: usize, + _context: Arc, + ) -> Result { + if partition != 0 { + return internal_err!("TcpConversationsExec has a single partition, got {partition}"); + } + let mut builder = RecordBatchReceiverStream::builder(Arc::clone(&self.projected_schema), 2); + let tx = builder.tx(); + let path = self.path.clone(); + let projection = self.projection.clone(); + let limit = self.limit; + builder.spawn_blocking(move || read_conversations(path, projection, limit, tx)); + Ok(builder.build()) + } +} + +/// Scans the capture, feeding every TCP packet into the flow tracker, and +/// sends the resulting conversations as a single batch +fn read_conversations( + path: String, + projection: Option>, + limit: Option, + tx: Sender>, +) -> Result<()> { + let mut tracker = FlowTracker::default(); + let result = for_each_frame("tcp_conversations", &path, |_, frame| { + let decoded = decode_frame(frame.link_type, frame.data); + tracker.observe(frame.ts_micros, frame.origlen, &decoded); + Ok(FrameLoop::Continue) + }); + match result { + Ok(()) => { + let batch = tracker.into_batch()?; + let batch = match limit { + Some(limit) if limit < batch.num_rows() => batch.slice(0, limit), + _ => batch, + }; + let batch = match &projection { + Some(p) => batch.project(p)?, + None => batch, + }; + let _ = tx.blocking_send(Ok(batch)); + } + Err(e) => send_error(&tx, e), + } + Ok(()) +} + +/// An (ip, port) endpoint +type Endpoint = (String, u16); + +/// Direction-independent flow identity: the two endpoints in sorted order +#[derive(Debug, PartialEq, Eq, Hash)] +struct FlowKey { + lo: Endpoint, + hi: Endpoint, +} + +impl FlowKey { + fn new(a: &Endpoint, b: &Endpoint) -> Self { + if a <= b { + Self { + lo: a.clone(), + hi: b.clone(), + } + } else { + Self { + lo: b.clone(), + hi: a.clone(), + } + } + } +} + +/// `a <= b` in TCP sequence space (wraparound aware) +fn seq_leq(a: u32, b: u32) -> bool { + a.wrapping_sub(b) as i32 <= 0 +} + +/// Accumulated state for one TCP connection +#[derive(Debug)] +struct Flow { + initiator: Endpoint, + responder: Endpoint, + first_ts: i64, + last_ts: i64, + packets_fwd: u64, + packets_rev: u64, + bytes_fwd: u64, + bytes_rev: u64, + retransmissions_fwd: u64, + retransmissions_rev: u64, + /// Timestamp of the initiator's first SYN + syn_ts: Option, + /// Whether the responder's SYN|ACK has been seen + synack_seen: bool, + /// Timestamp of the initiator's first ACK after the SYN|ACK, completing + /// the handshake + handshake_ack_ts: Option, + fin_fwd: bool, + fin_rev: bool, + rst: bool, + /// Highest end of consumed sequence space per direction, for + /// retransmission detection + max_seq_end_fwd: Option, + max_seq_end_rev: Option, +} + +impl Flow { + fn new(initiator: Endpoint, responder: Endpoint, first_ts: i64) -> Self { + Self { + initiator, + responder, + first_ts, + last_ts: first_ts, + packets_fwd: 0, + packets_rev: 0, + bytes_fwd: 0, + bytes_rev: 0, + retransmissions_fwd: 0, + retransmissions_rev: 0, + syn_ts: None, + synack_seen: false, + handshake_ack_ts: None, + fin_fwd: false, + fin_rev: false, + rst: false, + max_seq_end_fwd: None, + max_seq_end_rev: None, + } + } + + fn handshake_rtt_ms(&self) -> Option { + let syn = self.syn_ts?; + let ack = self.handshake_ack_ts?; + Some((ack - syn) as f64 / 1000.0) + } + + fn state(&self) -> &'static str { + if self.rst { + "reset" + } else if self.fin_fwd && self.fin_rev { + "closed" + } else if self.fin_fwd || self.fin_rev { + "half_closed" + } else { + "active" + } + } +} + +/// Groups observed TCP packets into flows +#[derive(Debug, Default)] +struct FlowTracker { + flows: HashMap, +} + +impl FlowTracker { + /// Feeds one decoded packet into the tracker; non-TCP packets and + /// packets without addresses/ports are ignored + fn observe(&mut self, ts: i64, frame_len: u32, decoded: &DecodedPacket<'_>) { + if decoded.protocol.as_deref() != Some("tcp") { + return; + } + let (Some(src_ip), Some(dst_ip), Some(src_port), Some(dst_port), Some(seq)) = ( + &decoded.src_ip, + &decoded.dst_ip, + decoded.src_port, + decoded.dst_port, + decoded.tcp_seq, + ) else { + return; + }; + let flags: Vec<&str> = decoded + .tcp_flags + .as_deref() + .unwrap_or_default() + .split('|') + .collect(); + let syn = flags.contains(&"SYN"); + let ack = flags.contains(&"ACK"); + let fin = flags.contains(&"FIN"); + let rst = flags.contains(&"RST"); + + let src: Endpoint = (src_ip.clone(), src_port); + let dst: Endpoint = (dst_ip.clone(), dst_port); + let key = FlowKey::new(&src, &dst); + let flow = self.flows.entry(key).or_insert_with(|| { + // The initiator is normally the first endpoint seen sending; a + // first-seen SYN|ACK means the capture missed the SYN and the + // sender is actually the responder + if syn && ack { + Flow::new(dst.clone(), src.clone(), ts) + } else { + Flow::new(src.clone(), dst.clone(), ts) + } + }); + let fwd = src == flow.initiator; + + flow.last_ts = flow.last_ts.max(ts); + if fwd { + flow.packets_fwd += 1; + flow.bytes_fwd += frame_len as u64; + } else { + flow.packets_rev += 1; + flow.bytes_rev += frame_len as u64; + } + if rst { + flow.rst = true; + } + if fin { + if fwd { + flow.fin_fwd = true; + } else { + flow.fin_rev = true; + } + } + + // Handshake: initiator SYN → responder SYN|ACK → initiator ACK + if syn && !ack && fwd && flow.syn_ts.is_none() { + flow.syn_ts = Some(ts); + } + if syn && ack && !fwd { + flow.synack_seen = true; + } + if !syn && ack && fwd && flow.synack_seen && flow.handshake_ack_ts.is_none() { + flow.handshake_ack_ts = Some(ts); + } + + // Retransmission detection: a segment consuming sequence space + // (payload, SYN, or FIN) whose end does not advance beyond the + // highest end seen in its direction is a retransmission + let consumed = + decoded.payload.map(|p| p.len() as u32).unwrap_or(0) + syn as u32 + fin as u32; + if consumed > 0 { + let end = seq.wrapping_add(consumed); + let (max_end, retransmissions) = if fwd { + (&mut flow.max_seq_end_fwd, &mut flow.retransmissions_fwd) + } else { + (&mut flow.max_seq_end_rev, &mut flow.retransmissions_rev) + }; + match max_end { + Some(max) if seq_leq(end, *max) => *retransmissions += 1, + Some(max) => *max = end, + None => *max_end = Some(end), + } + } + } + + /// Builds the conversations batch, ordered by first activity (then by + /// endpoints, for deterministic output) + fn into_batch(self) -> Result { + let mut flows: Vec = self.flows.into_values().collect(); + flows.sort_by(|a, b| { + (a.first_ts, &a.initiator, &a.responder).cmp(&(b.first_ts, &b.initiator, &b.responder)) + }); + + let mut src_ip = StringBuilder::new(); + let mut src_port = UInt16Builder::new(); + let mut dst_ip = StringBuilder::new(); + let mut dst_port = UInt16Builder::new(); + let mut first_timestamp = TimestampMicrosecondBuilder::new(); + let mut last_timestamp = TimestampMicrosecondBuilder::new(); + let mut duration_ms = Float64Builder::new(); + let mut packets_fwd = UInt64Builder::new(); + let mut packets_rev = UInt64Builder::new(); + let mut bytes_fwd = UInt64Builder::new(); + let mut bytes_rev = UInt64Builder::new(); + let mut retransmissions_fwd = UInt64Builder::new(); + let mut retransmissions_rev = UInt64Builder::new(); + let mut handshake_rtt_ms = Float64Builder::new(); + let mut state = StringBuilder::new(); + + for flow in &flows { + src_ip.append_value(&flow.initiator.0); + src_port.append_value(flow.initiator.1); + dst_ip.append_value(&flow.responder.0); + dst_port.append_value(flow.responder.1); + first_timestamp.append_value(flow.first_ts); + last_timestamp.append_value(flow.last_ts); + duration_ms.append_value((flow.last_ts - flow.first_ts) as f64 / 1000.0); + packets_fwd.append_value(flow.packets_fwd); + packets_rev.append_value(flow.packets_rev); + bytes_fwd.append_value(flow.bytes_fwd); + bytes_rev.append_value(flow.bytes_rev); + retransmissions_fwd.append_value(flow.retransmissions_fwd); + retransmissions_rev.append_value(flow.retransmissions_rev); + handshake_rtt_ms.append_option(flow.handshake_rtt_ms()); + state.append_value(flow.state()); + } + + Ok(RecordBatch::try_new( + tcp_conversations_schema(), + vec![ + Arc::new(src_ip.finish()), + Arc::new(src_port.finish()), + Arc::new(dst_ip.finish()), + Arc::new(dst_port.finish()), + Arc::new(first_timestamp.finish()), + Arc::new(last_timestamp.finish()), + Arc::new(duration_ms.finish()), + Arc::new(packets_fwd.finish()), + Arc::new(packets_rev.finish()), + Arc::new(bytes_fwd.finish()), + Arc::new(bytes_rev.finish()), + Arc::new(retransmissions_fwd.finish()), + Arc::new(retransmissions_rev.finish()), + Arc::new(handshake_rtt_ms.finish()), + Arc::new(state.finish()), + ], + )?) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn tcp_packet( + src: (&str, u16), + dst: (&str, u16), + seq: u32, + flags: &str, + payload_len: usize, + ) -> DecodedPacket<'static> { + DecodedPacket { + protocol: Some("tcp".to_string()), + src_ip: Some(src.0.to_string()), + dst_ip: Some(dst.0.to_string()), + src_port: Some(src.1), + dst_port: Some(dst.1), + tcp_seq: Some(seq), + tcp_flags: Some(flags.to_string()), + payload: Some(&[0u8; 1500][..payload_len]), + ..Default::default() + } + } + + const CLIENT: (&str, u16) = ("10.0.0.1", 51000); + const SERVER: (&str, u16) = ("10.0.0.2", 443); + + #[test] + fn test_handshake_rtt_and_state() { + let mut tracker = FlowTracker::default(); + tracker.observe(0, 60, &tcp_packet(CLIENT, SERVER, 1000, "SYN", 0)); + tracker.observe(50_000, 60, &tcp_packet(SERVER, CLIENT, 2000, "SYN|ACK", 0)); + tracker.observe(100_000, 54, &tcp_packet(CLIENT, SERVER, 1001, "ACK", 0)); + let flow = tracker.flows.values().next().unwrap(); + assert_eq!(flow.initiator, ("10.0.0.1".to_string(), 51000)); + assert_eq!(flow.handshake_rtt_ms(), Some(100.0)); + assert_eq!(flow.state(), "active"); + assert_eq!(flow.packets_fwd, 2); + assert_eq!(flow.packets_rev, 1); + } + + #[test] + fn test_retransmission_detected_per_direction() { + let mut tracker = FlowTracker::default(); + tracker.observe(0, 100, &tcp_packet(CLIENT, SERVER, 1000, "PSH|ACK", 5)); + // Same segment again: a retransmission + tracker.observe(1000, 100, &tcp_packet(CLIENT, SERVER, 1000, "PSH|ACK", 5)); + // New data advancing the window: not a retransmission + tracker.observe(2000, 100, &tcp_packet(CLIENT, SERVER, 1005, "PSH|ACK", 5)); + // Reverse direction is tracked independently + tracker.observe(3000, 100, &tcp_packet(SERVER, CLIENT, 9000, "PSH|ACK", 5)); + let flow = tracker.flows.values().next().unwrap(); + assert_eq!(flow.retransmissions_fwd, 1); + assert_eq!(flow.retransmissions_rev, 0); + } + + #[test] + fn test_states() { + // RST wins + let mut tracker = FlowTracker::default(); + tracker.observe(0, 60, &tcp_packet(CLIENT, SERVER, 1, "SYN", 0)); + tracker.observe(1, 60, &tcp_packet(SERVER, CLIENT, 1, "RST|ACK", 0)); + assert_eq!(tracker.flows.values().next().unwrap().state(), "reset"); + + // FIN from one side only + let mut tracker = FlowTracker::default(); + tracker.observe(0, 60, &tcp_packet(CLIENT, SERVER, 1, "FIN|ACK", 0)); + assert_eq!( + tracker.flows.values().next().unwrap().state(), + "half_closed" + ); + + // FIN from both sides + let mut tracker = FlowTracker::default(); + tracker.observe(0, 60, &tcp_packet(CLIENT, SERVER, 1, "FIN|ACK", 0)); + tracker.observe(1, 60, &tcp_packet(SERVER, CLIENT, 1, "FIN|ACK", 0)); + assert_eq!(tracker.flows.values().next().unwrap().state(), "closed"); + } + + #[test] + fn test_first_seen_synack_flips_initiator() { + // Capture started after the SYN: the SYN|ACK sender is the responder + let mut tracker = FlowTracker::default(); + tracker.observe(0, 60, &tcp_packet(SERVER, CLIENT, 2000, "SYN|ACK", 0)); + let flow = tracker.flows.values().next().unwrap(); + assert_eq!(flow.initiator, ("10.0.0.1".to_string(), 51000)); + assert_eq!(flow.packets_rev, 1); + } + + #[test] + fn test_distinct_flows_are_separate() { + let mut tracker = FlowTracker::default(); + tracker.observe(0, 60, &tcp_packet(CLIENT, SERVER, 1, "SYN", 0)); + tracker.observe(1, 60, &tcp_packet(("10.0.0.3", 6000), SERVER, 1, "SYN", 0)); + assert_eq!(tracker.flows.len(), 2); + } + + #[test] + fn test_non_tcp_ignored() { + let mut tracker = FlowTracker::default(); + let mut udp = tcp_packet(CLIENT, SERVER, 1, "", 0); + udp.protocol = Some("udp".to_string()); + tracker.observe(0, 60, &udp); + assert!(tracker.flows.is_empty()); + } + + #[test] + fn test_call_rejects_bad_arguments() { + let func = TcpConversationsFunc::default(); + let err = func.call(&[]).unwrap_err(); + assert!(err.to_string().contains("tcp_conversations"), "got: {err}"); + } +} diff --git a/crates/datafusion-net/src/decode.rs b/crates/datafusion-net/src/decode.rs new file mode 100644 index 0000000..c297786 --- /dev/null +++ b/crates/datafusion-net/src/decode.rs @@ -0,0 +1,288 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Decodes raw frames into the column values of [`crate::schema::packet_schema`]. +//! +//! Decoding never fails: frames that cannot be parsed (unknown link types, +//! truncated packets, non-IP traffic) produce a [`DecodedPacket`] with null +//! (`None`) values for the layers that could not be decoded. + +use etherparse::{ + ether_type, EtherType, IpNumber, LinkSlice, NetSlice, SlicedPacket, TransportSlice, VlanSlice, +}; + +/// Link types (from the pcap/pcapng spec) that we know how to decode +pub mod link_type { + /// BSD loopback encapsulation: 4 byte protocol family header then IP + pub const NULL: u32 = 0; + /// Ethernet + pub const ETHERNET: u32 = 1; + /// Raw IP, no link layer header + pub const RAW: u32 = 101; + /// OpenBSD loopback encapsulation, same shape as NULL + pub const LOOP: u32 = 108; +} + +/// Column values decoded from a single frame. Fields are `None` when the +/// corresponding layer was absent or could not be parsed. The payload is +/// borrowed from the frame to avoid copying it when the column is not +/// projected. +#[derive(Debug, Default)] +pub struct DecodedPacket<'a> { + pub eth_src: Option, + pub eth_dst: Option, + pub ethertype: Option, + pub vlan: Option, + pub ip_version: Option, + pub src_ip: Option, + pub dst_ip: Option, + pub ttl: Option, + pub protocol: Option, + pub src_port: Option, + pub dst_port: Option, + pub tcp_flags: Option, + pub tcp_seq: Option, + pub tcp_ack: Option, + pub tcp_window: Option, + pub payload: Option<&'a [u8]>, +} + +/// Decodes a single frame captured on a link of type `link_type` +pub fn decode_frame(link_type: u32, data: &[u8]) -> DecodedPacket<'_> { + let sliced = match link_type { + link_type::ETHERNET => SlicedPacket::from_ethernet(data).ok(), + link_type::NULL | link_type::LOOP => { + if data.len() >= 4 { + SlicedPacket::from_ip(&data[4..]).ok() + } else { + None + } + } + link_type::RAW => SlicedPacket::from_ip(data).ok(), + _ => None, + }; + let Some(sliced) = sliced else { + return DecodedPacket::default(); + }; + let mut decoded = DecodedPacket::default(); + + if let Some(LinkSlice::Ethernet2(eth)) = &sliced.link { + decoded.eth_src = Some(format_mac(ð.source())); + decoded.eth_dst = Some(format_mac(ð.destination())); + decoded.ethertype = Some(ether_type_name(eth.ether_type())); + } + + match &sliced.vlan() { + Some(VlanSlice::SingleVlan(v)) => decoded.vlan = Some(v.vlan_identifier().into()), + Some(VlanSlice::DoubleVlan(v)) => decoded.vlan = Some(v.outer.vlan_identifier().into()), + None => {} + } + + let ip_protocol = match &sliced.net { + Some(NetSlice::Ipv4(v4)) => { + let header = v4.header(); + decoded.ip_version = Some(4); + decoded.src_ip = Some(header.source_addr().to_string()); + decoded.dst_ip = Some(header.destination_addr().to_string()); + decoded.ttl = Some(header.ttl()); + decoded.payload = Some(v4.payload().payload); + Some(header.protocol()) + } + Some(NetSlice::Ipv6(v6)) => { + let header = v6.header(); + decoded.ip_version = Some(6); + decoded.src_ip = Some(header.source_addr().to_string()); + decoded.dst_ip = Some(header.destination_addr().to_string()); + decoded.ttl = Some(header.hop_limit()); + decoded.payload = Some(v6.payload().payload); + Some(v6.payload().ip_number) + } + _ => None, + }; + + match &sliced.transport { + Some(TransportSlice::Tcp(tcp)) => { + decoded.protocol = Some("tcp".to_string()); + decoded.src_port = Some(tcp.source_port()); + decoded.dst_port = Some(tcp.destination_port()); + decoded.tcp_flags = Some(tcp_flags_string(tcp)); + decoded.tcp_seq = Some(tcp.sequence_number()); + decoded.tcp_ack = Some(tcp.acknowledgment_number()); + decoded.tcp_window = Some(tcp.window_size()); + decoded.payload = Some(tcp.payload()); + } + Some(TransportSlice::Udp(udp)) => { + decoded.protocol = Some("udp".to_string()); + decoded.src_port = Some(udp.source_port()); + decoded.dst_port = Some(udp.destination_port()); + decoded.payload = Some(udp.payload()); + } + Some(TransportSlice::Icmpv4(icmp)) => { + decoded.protocol = Some("icmp".to_string()); + decoded.payload = Some(icmp.payload()); + } + Some(TransportSlice::Icmpv6(icmp)) => { + decoded.protocol = Some("icmpv6".to_string()); + decoded.payload = Some(icmp.payload()); + } + None => { + // Transport layer was not parsed (e.g. an IP protocol etherparse + // does not decode); fall back to the IP header's protocol number + decoded.protocol = ip_protocol.map(ip_number_name); + } + } + + decoded +} + +fn format_mac(mac: &[u8; 6]) -> String { + format!( + "{:02x}:{:02x}:{:02x}:{:02x}:{:02x}:{:02x}", + mac[0], mac[1], mac[2], mac[3], mac[4], mac[5] + ) +} + +fn ether_type_name(ether_type: EtherType) -> String { + match ether_type { + ether_type::IPV4 => "ipv4".to_string(), + ether_type::IPV6 => "ipv6".to_string(), + ether_type::ARP => "arp".to_string(), + other => format!("0x{:04x}", other.0), + } +} + +fn ip_number_name(ip_number: IpNumber) -> String { + match ip_number.keyword_str() { + Some(keyword) => keyword.to_lowercase(), + None => ip_number.0.to_string(), + } +} + +fn tcp_flags_string(tcp: ðerparse::TcpSlice) -> String { + let mut flags = Vec::new(); + if tcp.fin() { + flags.push("FIN"); + } + if tcp.syn() { + flags.push("SYN"); + } + if tcp.rst() { + flags.push("RST"); + } + if tcp.psh() { + flags.push("PSH"); + } + if tcp.ack() { + flags.push("ACK"); + } + if tcp.urg() { + flags.push("URG"); + } + if tcp.ece() { + flags.push("ECE"); + } + if tcp.cwr() { + flags.push("CWR"); + } + flags.join("|") +} + +#[cfg(test)] +mod tests { + use super::*; + use etherparse::PacketBuilder; + + const SRC_MAC: [u8; 6] = [0x02, 0x00, 0x00, 0x00, 0x00, 0x01]; + const DST_MAC: [u8; 6] = [0x02, 0x00, 0x00, 0x00, 0x00, 0x02]; + + #[test] + fn test_decode_tcp_ipv4() { + let builder = PacketBuilder::ethernet2(SRC_MAC, DST_MAC) + .ipv4([10, 0, 0, 1], [10, 0, 0, 2], 64) + .tcp(443, 51000, 1000, 1024) + .syn() + .ack(7); + let payload = b"hello"; + let mut frame = Vec::with_capacity(builder.size(payload.len())); + builder.write(&mut frame, payload).unwrap(); + + let decoded = decode_frame(link_type::ETHERNET, &frame); + assert_eq!(decoded.eth_src.as_deref(), Some("02:00:00:00:00:01")); + assert_eq!(decoded.eth_dst.as_deref(), Some("02:00:00:00:00:02")); + assert_eq!(decoded.ethertype.as_deref(), Some("ipv4")); + assert_eq!(decoded.ip_version, Some(4)); + assert_eq!(decoded.src_ip.as_deref(), Some("10.0.0.1")); + assert_eq!(decoded.dst_ip.as_deref(), Some("10.0.0.2")); + assert_eq!(decoded.ttl, Some(64)); + assert_eq!(decoded.protocol.as_deref(), Some("tcp")); + assert_eq!(decoded.src_port, Some(443)); + assert_eq!(decoded.dst_port, Some(51000)); + assert_eq!(decoded.tcp_flags.as_deref(), Some("SYN|ACK")); + assert_eq!(decoded.tcp_seq, Some(1000)); + assert_eq!(decoded.tcp_ack, Some(7)); + assert_eq!(decoded.tcp_window, Some(1024)); + assert_eq!(decoded.payload, Some(payload.as_slice())); + } + + #[test] + fn test_decode_udp_ipv6() { + let builder = PacketBuilder::ethernet2(SRC_MAC, DST_MAC) + .ipv6([1u8; 16], [2u8; 16], 32) + .udp(53, 5353); + let payload = b"dns"; + let mut frame = Vec::with_capacity(builder.size(payload.len())); + builder.write(&mut frame, payload).unwrap(); + + let decoded = decode_frame(link_type::ETHERNET, &frame); + assert_eq!(decoded.ethertype.as_deref(), Some("ipv6")); + assert_eq!(decoded.ip_version, Some(6)); + assert_eq!(decoded.ttl, Some(32)); + assert_eq!(decoded.protocol.as_deref(), Some("udp")); + assert_eq!(decoded.src_port, Some(53)); + assert_eq!(decoded.dst_port, Some(5353)); + assert_eq!(decoded.tcp_flags, None); + assert_eq!(decoded.payload, Some(payload.as_slice())); + } + + #[test] + fn test_decode_garbage_yields_nulls() { + let decoded = decode_frame(link_type::ETHERNET, &[0x01, 0x02, 0x03]); + assert_eq!(decoded.eth_src, None); + assert_eq!(decoded.src_ip, None); + assert_eq!(decoded.protocol, None); + assert_eq!(decoded.payload, None); + } + + #[test] + fn test_decode_unknown_link_type_yields_nulls() { + let decoded = decode_frame(9999, &[0u8; 64]); + assert_eq!(decoded.eth_src, None); + assert_eq!(decoded.src_ip, None); + } + + #[test] + fn test_decode_raw_ip_link_type() { + let builder = PacketBuilder::ipv4([192, 168, 1, 1], [192, 168, 1, 2], 64).udp(1000, 2000); + let mut frame = Vec::with_capacity(builder.size(0)); + builder.write(&mut frame, &[]).unwrap(); + + let decoded = decode_frame(link_type::RAW, &frame); + assert_eq!(decoded.eth_src, None); + assert_eq!(decoded.src_ip.as_deref(), Some("192.168.1.1")); + assert_eq!(decoded.protocol.as_deref(), Some("udp")); + } +} diff --git a/crates/datafusion-net/src/dns.rs b/crates/datafusion-net/src/dns.rs new file mode 100644 index 0000000..cd41221 --- /dev/null +++ b/crates/datafusion-net/src/dns.rs @@ -0,0 +1,482 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! `dns_query` scalar UDF: decodes DNS messages from packet payloads. +//! +//! [`DnsQueryUdf`] takes a binary payload (typically the `payload` column of +//! UDP port 53 packets — DNS over UDP fits in a single datagram, so no +//! reassembly is needed) and returns a struct with `is_response`, `name`, +//! `query_type`, `response_code`, and `answers` fields: +//! +//! ```sql +//! SELECT dns_query(payload)['name'] AS name, count(*) AS queries +//! FROM pcap('capture.pcap') +//! WHERE dst_port = 53 +//! GROUP BY name ORDER BY queries DESC; +//! ``` +//! +//! `name` and `query_type` come from the first question; `response_code` and +//! `answers` are NULL for queries. Answers include A/AAAA addresses and +//! CNAME/NS/PTR names; other record types are skipped. Payloads that do not +//! parse as DNS yield a NULL struct. + +use std::{ + net::{Ipv4Addr, Ipv6Addr}, + sync::{Arc, LazyLock}, +}; + +use datafusion::{ + arrow::{ + array::{ + Array, ArrayRef, BooleanBuilder, ListBuilder, NullBufferBuilder, StringBuilder, + StructArray, + }, + datatypes::{DataType, Field, Fields}, + }, + common::{cast::as_binary_array, exec_err, Result}, + logical_expr::{ + ColumnarValue, Documentation, ScalarFunctionArgs, ScalarUDFImpl, Signature, TypeSignature, + Volatility, + }, +}; + +use crate::NET_DOC_SECTION; + +/// `SHOW FUNCTIONS` documentation for `dns_query` +static DOCUMENTATION: LazyLock = LazyLock::new(|| { + Documentation::builder( + NET_DOC_SECTION, + "Decodes a DNS message from a packet payload (typically UDP port 53), \ + returning a struct with is_response (Boolean), name and query_type \ + (from the first question), response_code (NULL for queries), and \ + answers (a list of A/AAAA addresses and CNAME/NS/PTR names, NULL for \ + queries). Payloads that do not parse as DNS yield a NULL struct.", + "dns_query(payload)", + ) + .with_argument( + "payload", + "Binary packet payload, e.g. the payload column of a DNS packet", + ) + .with_sql_example( + "SELECT dns_query(payload)['name'] AS name, count(*) AS queries \ + FROM pcap('capture.pcap') WHERE dst_port = 53 GROUP BY name ORDER BY queries DESC", + ) + .with_related_udf("tls_sni") + .build() +}); + +/// Fields of the struct returned by `dns_query` +fn dns_fields() -> Fields { + Fields::from(vec![ + Field::new("is_response", DataType::Boolean, true), + Field::new("name", DataType::Utf8, true), + Field::new("query_type", DataType::Utf8, true), + Field::new("response_code", DataType::Utf8, true), + Field::new( + "answers", + DataType::List(Arc::new(Field::new("item", DataType::Utf8, true))), + true, + ), + ]) +} + +/// Scalar UDF that decodes a DNS message from a binary payload +#[derive(Debug, PartialEq, Eq, Hash)] +pub struct DnsQueryUdf { + signature: Signature, +} + +impl Default for DnsQueryUdf { + fn default() -> Self { + Self { + // Pure parse of the input bytes + signature: Signature::one_of( + vec![ + TypeSignature::Exact(vec![DataType::Binary]), + TypeSignature::Exact(vec![DataType::LargeBinary]), + TypeSignature::Exact(vec![DataType::BinaryView]), + ], + Volatility::Immutable, + ), + } + } +} + +impl ScalarUDFImpl for DnsQueryUdf { + fn name(&self) -> &str { + "dns_query" + } + + fn signature(&self) -> &Signature { + &self.signature + } + + fn return_type(&self, _arg_types: &[DataType]) -> Result { + Ok(DataType::Struct(dns_fields())) + } + + fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result { + if args.args.len() != 1 { + return exec_err!("dns_query expects a single binary argument"); + } + let payloads = to_binary_array(&args.args[0].to_array(args.number_rows)?)?; + let payloads = as_binary_array(&payloads)?; + + let mut is_response = BooleanBuilder::new(); + let mut name = StringBuilder::new(); + let mut query_type = StringBuilder::new(); + let mut response_code = StringBuilder::new(); + let mut answers = ListBuilder::new(StringBuilder::new()); + let mut validity = NullBufferBuilder::new(payloads.len()); + + for row in 0..payloads.len() { + let message = if payloads.is_null(row) { + None + } else { + parse_dns(payloads.value(row)) + }; + match message { + Some(message) => { + is_response.append_value(message.is_response); + name.append_value(&message.name); + query_type.append_value(&message.query_type); + response_code.append_option(message.response_code.as_deref()); + match message.answers { + Some(values) => { + for value in values { + answers.values().append_value(value); + } + answers.append(true); + } + None => answers.append_null(), + } + validity.append_non_null(); + } + None => { + is_response.append_null(); + name.append_null(); + query_type.append_null(); + response_code.append_null(); + answers.append_null(); + validity.append_null(); + } + } + } + + let arrays: Vec = vec![ + Arc::new(is_response.finish()), + Arc::new(name.finish()), + Arc::new(query_type.finish()), + Arc::new(response_code.finish()), + Arc::new(answers.finish()), + ]; + let structs = StructArray::new(dns_fields(), arrays, validity.finish()); + Ok(ColumnarValue::Array(Arc::new(structs))) + } + + fn documentation(&self) -> Option<&Documentation> { + Some(&DOCUMENTATION) + } +} + +/// Normalizes the input (BinaryView / LargeBinary) to Binary so a single +/// code path handles every accepted binary type +fn to_binary_array(array: &ArrayRef) -> Result { + if array.data_type() == &DataType::Binary { + Ok(Arc::clone(array)) + } else { + Ok(datafusion::arrow::compute::cast(array, &DataType::Binary)?) + } +} + +/// A decoded DNS message (the parts surfaced by `dns_query`) +#[derive(Debug, PartialEq)] +struct DnsMessage { + is_response: bool, + name: String, + query_type: String, + /// NULL for queries + response_code: Option, + /// NULL for queries; may be empty for answerless responses + answers: Option>, +} + +fn read_u16(data: &[u8], pos: usize) -> Option { + Some(u16::from_be_bytes([*data.get(pos)?, *data.get(pos + 1)?])) +} + +/// Parses a DNS message. Returns `None` when the bytes do not look like DNS +/// (this is a heuristic — any bytes *could* be DNS — so malformed names, +/// impossible counts, and unknown opcodes are all rejected). +fn parse_dns(data: &[u8]) -> Option { + if data.len() < 12 { + return None; + } + let flags = read_u16(data, 2)?; + let opcode = (flags >> 11) & 0xF; + // QUERY, IQUERY, and STATUS; anything else is likely not DNS + if opcode > 2 { + return None; + } + let qdcount = read_u16(data, 4)?; + let ancount = read_u16(data, 6)?; + if qdcount == 0 || qdcount > 4 || ancount > 128 { + return None; + } + let is_response = flags & 0x8000 != 0; + + let mut pos = 12usize; + let name = read_name(data, &mut pos)?; + let qtype = read_u16(data, pos)?; + pos += 4; // qtype + qclass + for _ in 1..qdcount { + read_name(data, &mut pos)?; + pos += 4; + } + if pos > data.len() { + return None; + } + + let (response_code, answers) = if is_response { + let rcode = flags & 0xF; + let mut answers = Vec::new(); + for _ in 0..ancount { + read_name(data, &mut pos)?; // owner name + let rtype = read_u16(data, pos)?; + pos += 8; // type + class + ttl + let rdlen = read_u16(data, pos)? as usize; + pos += 2; + let rdata = data.get(pos..pos + rdlen)?; + match rtype { + // A + 1 if rdlen == 4 => { + let octets: [u8; 4] = rdata.try_into().ok()?; + answers.push(Ipv4Addr::from(octets).to_string()); + } + // AAAA + 28 if rdlen == 16 => { + let octets: [u8; 16] = rdata.try_into().ok()?; + answers.push(Ipv6Addr::from(octets).to_string()); + } + // NS, CNAME, PTR: rdata is a (possibly compressed) name + 2 | 5 | 12 => { + let mut rdata_pos = pos; + answers.push(read_name(data, &mut rdata_pos)?); + } + // Other record types are skipped + _ => {} + } + pos += rdlen; + } + (Some(rcode_name(rcode)), Some(answers)) + } else { + (None, None) + }; + + Some(DnsMessage { + is_response, + name, + query_type: qtype_name(qtype), + response_code, + answers, + }) +} + +/// Reads a DNS name at `*pos`, following compression pointers, and advances +/// `*pos` past the name's in-place bytes +fn read_name(data: &[u8], pos: &mut usize) -> Option { + let mut labels: Vec = Vec::new(); + let mut cursor = *pos; + let mut jumped = false; + let mut jumps = 0u8; + loop { + let len = *data.get(cursor)? as usize; + if len == 0 { + if !jumped { + *pos = cursor + 1; + } + break; + } + if len & 0xC0 == 0xC0 { + // Compression pointer to an earlier name + let target = ((len & 0x3F) << 8) | *data.get(cursor + 1)? as usize; + if !jumped { + *pos = cursor + 2; + } + jumped = true; + jumps += 1; + if jumps > 5 { + return None; + } + cursor = target; + continue; + } + if len > 63 { + return None; + } + let label = data.get(cursor + 1..cursor + 1 + len)?; + // DNS labels are ASCII; rejecting anything else guards against + // non-DNS payloads that happen to parse structurally + if !label.iter().all(|b| b.is_ascii_graphic()) { + return None; + } + labels.push(String::from_utf8_lossy(label).into_owned()); + if labels.iter().map(|l| l.len() + 1).sum::() > 255 { + return None; + } + cursor += 1 + len; + } + if labels.is_empty() { + // The root domain + return Some(".".to_string()); + } + Some(labels.join(".")) +} + +fn qtype_name(qtype: u16) -> String { + match qtype { + 1 => "A".to_string(), + 2 => "NS".to_string(), + 5 => "CNAME".to_string(), + 6 => "SOA".to_string(), + 12 => "PTR".to_string(), + 15 => "MX".to_string(), + 16 => "TXT".to_string(), + 28 => "AAAA".to_string(), + 33 => "SRV".to_string(), + 65 => "HTTPS".to_string(), + 255 => "ANY".to_string(), + other => other.to_string(), + } +} + +fn rcode_name(rcode: u16) -> String { + match rcode { + 0 => "NOERROR".to_string(), + 1 => "FORMERR".to_string(), + 2 => "SERVFAIL".to_string(), + 3 => "NXDOMAIN".to_string(), + 4 => "NOTIMP".to_string(), + 5 => "REFUSED".to_string(), + other => other.to_string(), + } +} + +#[cfg(test)] +mod tests { + use super::*; + + /// A query for example.com A + fn example_query() -> Vec { + let mut m = vec![ + 0x12, 0x34, // id + 0x01, 0x00, // flags: RD + 0x00, 0x01, // qdcount + 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, // an/ns/ar + ]; + m.extend_from_slice(b"\x07example\x03com\x00"); + m.extend_from_slice(&[0x00, 0x01, 0x00, 0x01]); // A, IN + m + } + + #[test] + fn test_parse_query() { + let message = parse_dns(&example_query()).unwrap(); + assert!(!message.is_response); + assert_eq!(message.name, "example.com"); + assert_eq!(message.query_type, "A"); + assert_eq!(message.response_code, None); + assert_eq!(message.answers, None); + } + + #[test] + fn test_parse_response_with_compressed_answers() { + let mut m = vec![ + 0x12, 0x34, // id + 0x81, 0x80, // flags: response, RD, RA, NOERROR + 0x00, 0x01, // qdcount + 0x00, 0x02, // ancount + 0x00, 0x00, 0x00, 0x00, // ns/ar + ]; + m.extend_from_slice(b"\x07example\x03com\x00"); + m.extend_from_slice(&[0x00, 0x01, 0x00, 0x01]); // A, IN + // CNAME answer, owner name is a pointer to offset 12 + m.extend_from_slice(&[0xC0, 0x0C]); // pointer to question name + m.extend_from_slice(&[0x00, 0x05, 0x00, 0x01]); // CNAME, IN + m.extend_from_slice(&[0x00, 0x00, 0x00, 0x3C]); // ttl + m.extend_from_slice(&[0x00, 0x06]); // rdlength + m.extend_from_slice(b"\x03www\xC0\x0C"); // www. + // A answer for the cname target (pointer to the www label) + m.extend_from_slice(&[0xC0, 0x21]); // pointer to "www.example.com" + m.extend_from_slice(&[0x00, 0x01, 0x00, 0x01]); // A, IN + m.extend_from_slice(&[0x00, 0x00, 0x00, 0x3C]); // ttl + m.extend_from_slice(&[0x00, 0x04]); // rdlength + m.extend_from_slice(&[93, 184, 216, 34]); + + let message = parse_dns(&m).unwrap(); + assert!(message.is_response); + assert_eq!(message.name, "example.com"); + assert_eq!(message.response_code.as_deref(), Some("NOERROR")); + assert_eq!( + message.answers, + Some(vec![ + "www.example.com".to_string(), + "93.184.216.34".to_string() + ]) + ); + } + + #[test] + fn test_parse_nxdomain_response() { + let mut m = example_query(); + m[2] = 0x81; + m[3] = 0x83; // response, NXDOMAIN + let message = parse_dns(&m).unwrap(); + assert!(message.is_response); + assert_eq!(message.response_code.as_deref(), Some("NXDOMAIN")); + assert_eq!(message.answers, Some(vec![])); + } + + #[test] + fn test_garbage_is_none() { + assert_eq!(parse_dns(b""), None); + assert_eq!(parse_dns(b"hello world, definitely not dns"), None); + assert_eq!(parse_dns(&[0xFF; 64]), None); + // Truncated mid-name + assert_eq!(parse_dns(&example_query()[..14]), None); + } + + #[test] + fn test_pointer_loop_rejected() { + let mut m = vec![ + 0x12, 0x34, 0x01, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, + ]; + // A name that is a pointer to itself + m.extend_from_slice(&[0xC0, 0x0C]); + m.extend_from_slice(&[0x00, 0x01, 0x00, 0x01]); + assert_eq!(parse_dns(&m), None); + } + + #[test] + fn test_return_type_is_struct() { + let udf = DnsQueryUdf::default(); + assert_eq!( + udf.return_type(&[DataType::Binary]).unwrap(), + DataType::Struct(dns_fields()) + ); + } +} diff --git a/crates/datafusion-net/src/file.rs b/crates/datafusion-net/src/file.rs new file mode 100644 index 0000000..0fdd1c1 --- /dev/null +++ b/crates/datafusion-net/src/file.rs @@ -0,0 +1,426 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! `pcap` table function: reads a pcap or pcapng capture file as a table. +//! +//! ```sql +//! SELECT src_ip, dst_ip, protocol, length FROM pcap('capture.pcap') WHERE dst_port = 443 +//! ``` +//! +//! The single argument is a local filesystem path. The file is not opened +//! during planning (e.g. `EXPLAIN`), only on execution. + +use std::{fs::File, sync::Arc}; + +use async_trait::async_trait; +use datafusion::{ + arrow::{array::RecordBatch, datatypes::SchemaRef}, + catalog::{Session, TableFunctionImpl, TableProvider}, + common::{internal_err, plan_err, project_schema, DataFusionError, Result}, + datasource::TableType, + execution::SendableRecordBatchStream, + physical_expr::EquivalenceProperties, + physical_plan::{ + execution_plan::{Boundedness, EmissionType}, + stream::RecordBatchReceiverStream, + DisplayAs, DisplayFormatType, ExecutionPlan, Partitioning, PlanProperties, + }, + prelude::Expr, +}; +use pcap_parser::{create_reader, Block, PcapBlockOwned, PcapError}; +use tokio::sync::mpsc::Sender; + +use crate::{ + decode::{decode_frame, link_type}, + expr_to_string, + schema::{packet_schema, payload_projected, send_batch, send_error, PacketBatchBuilder}, +}; + +/// Number of rows accumulated before a batch is emitted +const BATCH_SIZE: usize = 1024; +/// Read buffer capacity, must exceed the largest block in the file +const READER_CAPACITY: usize = 1 << 20; + +/// Table function that reads a pcap/pcapng capture file +#[derive(Debug, Default)] +pub struct PcapFunc {} + +impl TableFunctionImpl for PcapFunc { + fn call(&self, exprs: &[Expr]) -> Result> { + let path = parse_pcap_args("pcap", exprs)?; + Ok(Arc::new(PcapTable::new(path))) + } +} + +/// Parses the arguments shared by `pcap` and `pcap_wide` (hence the function +/// name parameter): a single path to a pcap/pcapng file +pub(crate) fn parse_pcap_args(func: &str, exprs: &[Expr]) -> Result { + if exprs.len() != 1 { + return plan_err!("{func} requires a single argument, the path to a pcap/pcapng file"); + } + expr_to_string(&exprs[0], func, "path (first argument)") +} + +/// [`TableProvider`] backed by a pcap/pcapng file +#[derive(Debug)] +pub struct PcapTable { + path: String, + schema: SchemaRef, +} + +impl PcapTable { + pub fn new(path: String) -> Self { + Self { + path, + schema: packet_schema(), + } + } +} + +#[async_trait] +impl TableProvider for PcapTable { + fn schema(&self) -> SchemaRef { + Arc::clone(&self.schema) + } + + fn table_type(&self) -> TableType { + TableType::Base + } + + async fn scan( + &self, + _state: &dyn Session, + projection: Option<&Vec>, + _filters: &[Expr], + limit: Option, + ) -> Result> { + let exec = PcapExec::try_new( + self.path.clone(), + Arc::clone(&self.schema), + projection.cloned(), + limit, + )?; + Ok(Arc::new(exec)) + } +} + +/// Execution plan that reads and decodes a capture file on `execute` +#[derive(Debug)] +struct PcapExec { + path: String, + projection: Option>, + /// Schema representing the data after the optional projection is applied + projected_schema: SchemaRef, + limit: Option, + cache: Arc, +} + +impl PcapExec { + fn try_new( + path: String, + schema: SchemaRef, + projection: Option>, + limit: Option, + ) -> Result { + let projected_schema = project_schema(&schema, projection.as_ref())?; + let cache = Arc::new(PlanProperties::new( + EquivalenceProperties::new(Arc::clone(&projected_schema)), + Partitioning::UnknownPartitioning(1), + EmissionType::Incremental, + Boundedness::Bounded, + )); + Ok(Self { + path, + projection, + projected_schema, + limit, + cache, + }) + } +} + +impl DisplayAs for PcapExec { + fn fmt_as(&self, _t: DisplayFormatType, f: &mut std::fmt::Formatter) -> std::fmt::Result { + write!(f, "PcapExec: path={}, limit={:?}", self.path, self.limit) + } +} + +impl ExecutionPlan for PcapExec { + fn name(&self) -> &str { + "PcapExec" + } + + fn properties(&self) -> &Arc { + &self.cache + } + + fn children(&self) -> Vec<&Arc> { + // This is a leaf node and has no children + vec![] + } + + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + // PcapExec has no children + if children.is_empty() { + Ok(self) + } else { + internal_err!("Children cannot be replaced in {self:?}") + } + } + + fn execute( + &self, + partition: usize, + _context: Arc, + ) -> Result { + if partition != 0 { + return internal_err!("PcapExec has a single partition, got {partition}"); + } + let mut builder = + RecordBatchReceiverStream::builder(Arc::clone(&self.projected_schema), 16); + let tx = builder.tx(); + let path = self.path.clone(); + let projection = self.projection.clone(); + let limit = self.limit; + // Reading and decoding is blocking (file I/O + CPU), so run it on a + // blocking thread rather than an async task + builder.spawn_blocking(move || read_pcap_file(path, projection, limit, tx)); + Ok(builder.build()) + } +} + +/// Tracks per-interface timestamp metadata from pcapng interface description +/// blocks +struct NgInterface { + link_type: u32, + /// Timestamp units per second + ts_resolution: u64, + /// Offset added to timestamps, in seconds + ts_offset: i64, +} + +/// A raw captured frame handed to [`for_each_frame`] callbacks +pub(crate) struct RawFrame<'a> { + pub link_type: u32, + pub ts_micros: i64, + pub origlen: u32, + pub caplen: u32, + pub data: &'a [u8], +} + +/// Whether [`for_each_frame`] should keep reading +pub(crate) enum FrameLoop { + Continue, + Stop, +} + +/// Iterates the frames of the pcap/pcapng file at `path`, driving the format +/// state machine (section/interface blocks, per-interface link types and +/// timestamp resolutions) and calling `on_frame` with the 1-based frame +/// number and each captured frame. `func` names the calling table function +/// in error messages. +pub(crate) fn for_each_frame( + func: &str, + path: &str, + mut on_frame: impl FnMut(u64, RawFrame<'_>) -> Result, +) -> Result<()> { + let file = File::open(path).map_err(|e| { + DataFusionError::External(format!("{func} failed to open '{path}': {e}").into()) + })?; + let mut reader = create_reader(READER_CAPACITY, file).map_err(|e| { + DataFusionError::External( + format!("{func} failed to read '{path}' as a pcap/pcapng file: {e}").into(), + ) + })?; + + let mut frame_number = 0u64; + // Legacy pcap state + let mut legacy_link_type = link_type::ETHERNET; + let mut legacy_ts_nanos = false; + // pcapng state + let mut ng_interfaces: Vec = Vec::new(); + + loop { + match reader.next() { + Ok((offset, block)) => { + let frame = match block { + PcapBlockOwned::LegacyHeader(header) => { + legacy_link_type = header.network.0 as u32; + legacy_ts_nanos = header.is_nanosecond_precision(); + None + } + PcapBlockOwned::Legacy(b) => { + let sub_micros = if legacy_ts_nanos { + (b.ts_usec / 1000) as i64 + } else { + b.ts_usec as i64 + }; + let ts_micros = b.ts_sec as i64 * 1_000_000 + sub_micros; + Some((legacy_link_type, ts_micros, b.origlen, b.caplen, b.data)) + } + PcapBlockOwned::NG(Block::SectionHeader(_)) => { + ng_interfaces.clear(); + None + } + PcapBlockOwned::NG(Block::InterfaceDescription(ref idb)) => { + ng_interfaces.push(NgInterface { + link_type: idb.linktype.0 as u32, + ts_resolution: idb.ts_resolution().unwrap_or(1_000_000), + ts_offset: idb.ts_offset(), + }); + None + } + PcapBlockOwned::NG(Block::EnhancedPacket(ref epb)) => { + match ng_interfaces.get(epb.if_id as usize) { + Some(interface) => { + let ts_units = ((epb.ts_high as u64) << 32) | epb.ts_low as u64; + // u128 to avoid overflow: nanosecond + // resolutions exceed u64 range when scaled + let ts_micros = (ts_units as u128 * 1_000_000 + / interface.ts_resolution as u128) + as i64 + + interface.ts_offset * 1_000_000; + Some(( + interface.link_type, + ts_micros, + epb.origlen, + epb.caplen, + epb.data, + )) + } + None => None, + } + } + PcapBlockOwned::NG(Block::SimplePacket(ref spb)) => { + let link_type = ng_interfaces + .first() + .map(|i| i.link_type) + .unwrap_or(link_type::ETHERNET); + // Simple packet blocks carry no timestamp + Some((link_type, 0, spb.origlen, spb.data.len() as u32, spb.data)) + } + // Name resolution, statistics, and other metadata blocks + PcapBlockOwned::NG(_) => None, + }; + + if let Some((link_type, ts_micros, origlen, caplen, data)) = frame { + frame_number += 1; + let raw = RawFrame { + link_type, + ts_micros, + origlen, + caplen, + data, + }; + if let FrameLoop::Stop = on_frame(frame_number, raw)? { + return Ok(()); + } + } + reader.consume(offset); + } + Err(PcapError::Eof) => return Ok(()), + Err(PcapError::Incomplete(_)) => { + reader.refill().map_err(|e| { + DataFusionError::External(format!("{func} failed to read '{path}': {e}").into()) + })?; + } + Err(e) => { + return Err(DataFusionError::External( + format!("{func} failed to read '{path}': {e}").into(), + )) + } + } + } +} + +/// Reads `path`, decoding each frame and sending record batches of +/// `BATCH_SIZE` rows until the file ends, `limit` rows have been produced, or +/// the consumer drops the stream +fn read_pcap_file( + path: String, + projection: Option>, + limit: Option, + tx: Sender>, +) -> Result<()> { + let mut builder = PacketBatchBuilder::new(payload_projected(&projection)); + let mut produced = 0usize; + let result = for_each_frame("pcap", &path, |frame_number, frame| { + let decoded = decode_frame(frame.link_type, frame.data); + builder.append( + frame.ts_micros, + frame_number, + frame.origlen, + frame.caplen, + &decoded, + ); + produced += 1; + if limit.is_some_and(|l| produced >= l) { + return Ok(FrameLoop::Stop); + } + if builder.len() >= BATCH_SIZE && !send_batch(&mut builder, &projection, &tx)? { + // Consumer dropped the stream + return Ok(FrameLoop::Stop); + } + Ok(FrameLoop::Continue) + }); + match result { + Ok(()) => { + // Flush the final partial batch; if the consumer already dropped + // the stream the send is a no-op + send_batch(&mut builder, &projection, &tx)?; + } + Err(e) => send_error(&tx, e), + } + Ok(()) +} + +#[cfg(test)] +mod tests { + use super::*; + use datafusion::scalar::ScalarValue; + + #[test] + fn test_call_requires_single_path() { + let func = PcapFunc::default(); + let err = func.call(&[]).unwrap_err(); + assert!(err.to_string().contains("single argument")); + } + + #[test] + fn test_call_rejects_non_string_path() { + let func = PcapFunc::default(); + let args = vec![Expr::Literal(ScalarValue::Int64(Some(1)), None)]; + let err = func.call(&args).unwrap_err(); + assert!(err.to_string().contains("must be a string literal")); + } + + #[tokio::test] + async fn test_explain_does_not_open_file() { + let ctx = datafusion::prelude::SessionContext::new(); + ctx.register_udtf("pcap", Arc::new(PcapFunc::default())); + // Planning must not touch the filesystem - this path does not exist + let df = ctx + .sql("EXPLAIN SELECT * FROM pcap('/does/not/exist.pcap')") + .await + .unwrap(); + let batches = df.collect().await.unwrap(); + assert!(!batches.is_empty()); + } +} diff --git a/crates/datafusion-net/src/geoip.rs b/crates/datafusion-net/src/geoip.rs new file mode 100644 index 0000000..0ab3129 --- /dev/null +++ b/crates/datafusion-net/src/geoip.rs @@ -0,0 +1,385 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! IP geolocation scalar UDF. +//! +//! [`GeoIpUdf`] (`geoip`) takes an IP address string and looks it up in a +//! MaxMind-format (`.mmdb`) geolocation database such as [GeoLite2-City], +//! returning a struct of location fields: +//! +//! ```sql +//! SELECT src_ip, +//! geoip(src_ip)['country_code'] AS country, +//! geoip(src_ip)['city'] AS city +//! FROM pcap('capture.pcap') +//! ``` +//! +//! A single MaxMind-format database file serves every field — there is no +//! per-field database. A City-schema database ([GeoLite2-City] free with a +//! MaxMind account, or the commercial GeoIP2-City) populates all fields; a +//! Country-schema database (GeoLite2-Country) populates only `country_code` +//! and `country`, leaving the rest NULL; other schemas (ASN, ISP, ...) have +//! none of these fields and yield all-NULL values. +//! +//! The database path is taken from the optional second argument +//! (`geoip(ip, '/path/GeoLite2-City.mmdb')`), or, for the single-argument +//! form, from the `GEOIP_DB` environment variable ([`GEOIP_DB_ENV_VAR`]) read +//! when the UDF is constructed. Opened databases are cached process-wide. +//! +//! Addresses that fail to parse or have no entry in the database yield a +//! NULL struct. A database that is missing, unreadable, or unconfigured does +//! not fail the query: the location fields are NULL and the struct's `error` +//! field carries the reason (it is NULL on success), so +//! `geoip(ip)['error']` surfaces what went wrong. +//! +//! [GeoLite2-City]: https://dev.maxmind.com/geoip/geolite2-free-geolocation-data + +use std::{ + collections::HashMap, + net::IpAddr, + path::{Path, PathBuf}, + sync::{Arc, LazyLock, Mutex, OnceLock}, +}; + +use datafusion::{ + arrow::{ + array::{ + Array, ArrayRef, Float64Builder, NullBufferBuilder, StringArray, StringBuilder, + StructArray, + }, + datatypes::{DataType, Field, Fields}, + }, + common::{cast::as_string_array, exec_err, Result}, + logical_expr::{ + ColumnarValue, Documentation, ScalarFunctionArgs, ScalarUDFImpl, Signature, TypeSignature, + Volatility, + }, +}; +use maxminddb::geoip2; + +use crate::NET_DOC_SECTION; + +/// Environment variable consulted by [`GeoIpUdf::default`] for the database +/// path used by the single-argument form of `geoip` +pub const GEOIP_DB_ENV_VAR: &str = "GEOIP_DB"; + +/// `SHOW FUNCTIONS` documentation for `geoip` +static DOCUMENTATION: LazyLock = LazyLock::new(|| { + Documentation::builder( + NET_DOC_SECTION, + "Geolocates an IP address string using a MaxMind-format (.mmdb) \ + database, returning a struct with country_code, country, city, \ + latitude, longitude, time_zone, and error fields. A City-schema \ + database populates all fields; a Country-schema database only \ + country_code and country. The database path is the optional second \ + argument, the GEOIP_DB environment variable, or (in dft) the \ + [execution.net] geoip_db_path config. Addresses that do not parse or \ + are absent from the database yield a NULL struct; a missing or \ + unreadable database yields NULL location fields with the reason in \ + the error field rather than failing the query.", + "geoip(ip [, db_path])", + ) + .with_argument("ip", "IP address string, e.g. the src_ip or dst_ip column") + .with_argument( + "db_path", + "Optional path to a MaxMind .mmdb database; defaults to the GEOIP_DB \ + environment variable or configured path", + ) + .with_sql_example( + "SELECT geoip(src_ip, '/path/GeoLite2-City.mmdb')['country_code'] AS country, \ + count(*) AS packets FROM pcap('capture.pcap') GROUP BY country ORDER BY packets DESC", + ) + .with_related_udf("reverse_dns") + .build() +}); + +/// A cached, shared handle to an opened database +type SharedReader = Arc>>; + +/// A cached open outcome: the reader, or the message describing why the +/// database could not be opened +type ReaderResult = std::result::Result>; + +/// Process-wide cache of database open outcomes, keyed by path. Readers hold +/// the whole database in memory, so each distinct path is loaded once and +/// shared across batches and queries. Failures are cached too, so a bad path +/// costs one open attempt instead of one per row. +fn readers() -> &'static Mutex> { + static READERS: OnceLock>> = OnceLock::new(); + READERS.get_or_init(|| Mutex::new(HashMap::new())) +} + +/// Fields of the struct returned by `geoip`. `error` is NULL on success and +/// carries the reason when the database could not be used. +fn geoip_fields() -> Fields { + Fields::from(vec![ + Field::new("country_code", DataType::Utf8, true), + Field::new("country", DataType::Utf8, true), + Field::new("city", DataType::Utf8, true), + Field::new("latitude", DataType::Float64, true), + Field::new("longitude", DataType::Float64, true), + Field::new("time_zone", DataType::Utf8, true), + Field::new("error", DataType::Utf8, true), + ]) +} + +/// Owned location values extracted from a database record, or the reason no +/// lookup could be performed +#[derive(Default)] +struct GeoRecord { + country_code: Option, + country: Option, + city: Option, + latitude: Option, + longitude: Option, + time_zone: Option, + error: Option, +} + +impl GeoRecord { + fn from_city(city: &geoip2::City<'_>) -> Self { + Self { + country_code: city.country.iso_code.map(str::to_string), + country: city.country.names.english.map(str::to_string), + city: city.city.names.english.map(str::to_string), + latitude: city.location.latitude, + longitude: city.location.longitude, + time_zone: city.location.time_zone.map(str::to_string), + error: None, + } + } + + /// A record whose location fields are NULL and whose `error` field + /// carries the reason + fn from_error(error: impl Into) -> Self { + Self { + error: Some(error.into()), + ..Self::default() + } + } +} + +/// Scalar UDF that geolocates an IP address string using a MaxMind-format +/// (`.mmdb`) database +#[derive(Debug, PartialEq, Eq, Hash)] +pub struct GeoIpUdf { + signature: Signature, + default_db: Option, +} + +impl GeoIpUdf { + /// Creates the UDF with `path` as the database for the single-argument + /// form of `geoip` + pub fn with_db_path(path: impl Into) -> Self { + Self { + signature: Self::make_signature(), + default_db: Some(path.into()), + } + } + + fn make_signature() -> Signature { + // Accept the string types our packet columns and SQL literals use. + // Stable (not Immutable): results depend on the database file, which + // can change between queries. + Signature::one_of( + vec![TypeSignature::String(1), TypeSignature::String(2)], + Volatility::Stable, + ) + } + + /// Returns the cached open outcome for `path`, attempting (and caching) + /// the open on first use + fn reader_for(path: &Path) -> ReaderResult { + let mut readers = readers().lock().expect("geoip reader cache poisoned"); + if let Some(outcome) = readers.get(path) { + return outcome.clone(); + } + let outcome = maxminddb::Reader::open_readfile(path) + .map(Arc::new) + .map_err(|e| { + Arc::from(format!( + "geoip failed to open database '{}': {e}", + path.display() + )) + }); + readers.insert(path.to_path_buf(), outcome.clone()); + outcome + } +} + +impl Default for GeoIpUdf { + fn default() -> Self { + Self { + signature: Self::make_signature(), + default_db: std::env::var_os(GEOIP_DB_ENV_VAR).map(PathBuf::from), + } + } +} + +impl ScalarUDFImpl for GeoIpUdf { + fn name(&self) -> &str { + "geoip" + } + + fn signature(&self) -> &Signature { + &self.signature + } + + fn return_type(&self, _arg_types: &[DataType]) -> Result { + Ok(DataType::Struct(geoip_fields())) + } + + fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result { + if args.args.is_empty() || args.args.len() > 2 { + return exec_err!("geoip expects one or two arguments: geoip(ip [, db_path])"); + } + let ips = to_utf8_array(&args.args[0].to_array(args.number_rows)?)?; + let ips = as_string_array(&ips)?; + let paths = match args.args.get(1) { + Some(arg) => Some(to_utf8_array(&arg.to_array(args.number_rows)?)?), + None => None, + }; + let paths = paths.as_ref().map(|p| as_string_array(p)).transpose()?; + + let mut country_code = StringBuilder::new(); + let mut country = StringBuilder::new(); + let mut city = StringBuilder::new(); + let mut latitude = Float64Builder::new(); + let mut longitude = Float64Builder::new(); + let mut time_zone = StringBuilder::new(); + let mut error = StringBuilder::new(); + let mut validity = NullBufferBuilder::new(ips.len()); + + for row in 0..ips.len() { + match self.lookup_row(ips, paths, row) { + Some(record) => { + country_code.append_option(record.country_code); + country.append_option(record.country); + city.append_option(record.city); + latitude.append_option(record.latitude); + longitude.append_option(record.longitude); + time_zone.append_option(record.time_zone); + error.append_option(record.error); + validity.append_non_null(); + } + None => { + country_code.append_null(); + country.append_null(); + city.append_null(); + latitude.append_null(); + longitude.append_null(); + time_zone.append_null(); + error.append_null(); + validity.append_null(); + } + } + } + + let arrays: Vec = vec![ + Arc::new(country_code.finish()), + Arc::new(country.finish()), + Arc::new(city.finish()), + Arc::new(latitude.finish()), + Arc::new(longitude.finish()), + Arc::new(time_zone.finish()), + Arc::new(error.finish()), + ]; + let structs = StructArray::new(geoip_fields(), arrays, validity.finish()); + Ok(ColumnarValue::Array(Arc::new(structs))) + } + + fn documentation(&self) -> Option<&Documentation> { + Some(&DOCUMENTATION) + } +} + +impl GeoIpUdf { + /// Geolocates a single row. `None` (a NULL struct) covers the row-level + /// misses: NULL or unparseable address, NULL path argument, or an + /// address with no database entry. Database problems — none configured, + /// or one that cannot be opened — do not fail the query: they produce a + /// record whose location fields are NULL and whose `error` field carries + /// the reason. + fn lookup_row( + &self, + ips: &StringArray, + paths: Option<&StringArray>, + row: usize, + ) -> Option { + if ips.is_null(row) { + return None; + } + // Parse before touching the database so bad addresses are NULL even + // when no database is configured + let Ok(ip) = ips.value(row).parse::() else { + return None; + }; + let path: PathBuf = match (paths, &self.default_db) { + (Some(paths), _) if paths.is_null(row) => return None, + (Some(paths), _) => PathBuf::from(paths.value(row)), + (None, Some(default)) => default.clone(), + (None, None) => { + return Some(GeoRecord::from_error(format!( + "geoip has no database configured; pass a path as the second argument \ + (e.g. geoip(ip, '/path/GeoLite2-City.mmdb')) or set the \ + {GEOIP_DB_ENV_VAR} environment variable" + ))) + } + }; + let reader = match Self::reader_for(&path) { + Ok(reader) => reader, + Err(open_error) => return Some(GeoRecord::from_error(open_error.as_ref())), + }; + // Lookup errors (e.g. an IPv6 address in an IPv4-only database) are + // row-level data issues, treated as misses + reader + .lookup(ip) + .ok() + .and_then(|result| result.decode::().ok().flatten()) + .map(|city| GeoRecord::from_city(&city)) + } +} + +/// Normalizes the input (Utf8View / LargeUtf8) to Utf8 so a single code path +/// handles every accepted string type +fn to_utf8_array(array: &ArrayRef) -> Result { + if array.data_type() == &DataType::Utf8 { + Ok(Arc::clone(array)) + } else { + Ok(datafusion::arrow::compute::cast(array, &DataType::Utf8)?) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_return_type_is_struct() { + let udf = GeoIpUdf::default(); + assert_eq!( + udf.return_type(&[DataType::Utf8]).unwrap(), + DataType::Struct(geoip_fields()) + ); + } + + #[test] + fn test_with_db_path_sets_default() { + let udf = GeoIpUdf::with_db_path("/some/db.mmdb"); + assert_eq!(udf.default_db, Some(PathBuf::from("/some/db.mmdb"))); + } +} diff --git a/crates/datafusion-net/src/interfaces.rs b/crates/datafusion-net/src/interfaces.rs new file mode 100644 index 0000000..14de8fd --- /dev/null +++ b/crates/datafusion-net/src/interfaces.rs @@ -0,0 +1,208 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! `interfaces` table function: lists the system's network capture +//! interfaces, one row per interface, similar to `tshark -D`. Useful for +//! discovering what to pass to the `capture` table function: +//! +//! ```sql +//! SELECT name, description, addresses FROM interfaces() WHERE is_up; +//! ``` +//! +//! Listing interfaces does not require elevated privileges (unlike opening +//! one for capture). + +use std::sync::Arc; + +use datafusion::{ + arrow::{ + array::{BooleanBuilder, ListBuilder, RecordBatch, StringBuilder}, + datatypes::{DataType, Field, Schema, SchemaRef}, + }, + catalog::{MemTable, TableFunctionImpl, TableProvider}, + common::{plan_err, DataFusionError, Result}, + prelude::Expr, +}; +use pcap::ConnectionStatus; + +/// Schema of the `interfaces` table function: one row per interface +pub fn interfaces_schema() -> SchemaRef { + Arc::new(Schema::new(vec![ + Field::new("name", DataType::Utf8, false), + Field::new("description", DataType::Utf8, true), + Field::new( + "addresses", + DataType::List(Arc::new(Field::new("item", DataType::Utf8, true))), + true, + ), + Field::new("is_up", DataType::Boolean, true), + Field::new("is_running", DataType::Boolean, true), + Field::new("is_loopback", DataType::Boolean, true), + Field::new("is_wireless", DataType::Boolean, true), + Field::new("connection_status", DataType::Utf8, true), + ])) +} + +/// Table function that lists the system's network capture interfaces +#[derive(Debug, Default)] +pub struct InterfacesFunc {} + +impl TableFunctionImpl for InterfacesFunc { + fn call(&self, exprs: &[Expr]) -> Result> { + if !exprs.is_empty() { + return plan_err!("interfaces takes no arguments"); + } + let schema = interfaces_schema(); + let batch = interfaces_batch(&schema)?; + let table = MemTable::try_new(schema, vec![vec![batch]])?; + Ok(Arc::new(table)) + } +} + +/// Falls back to a static, human-readable description for well-known +/// interface name patterns when `pcap` doesn't supply one itself (libpcap +/// leaves `desc` empty for most interfaces on macOS). +fn static_description(name: &str) -> Option<&'static str> { + let prefix = name.trim_end_matches(|c: char| c.is_ascii_digit()); + match prefix { + "lo" => Some("Loopback interface"), + "en" => Some("Ethernet or Wi-Fi adapter"), + "utun" => Some("Utility tunnel interface (VPN, Back to My Mac, etc.)"), + "ap" => Some("Wi-Fi Access Point (Personal Hotspot / Instant Hotspot)"), + "awdl" => Some("Apple Wireless Direct Link (AirDrop, AirPlay)"), + "llw" => Some("Low-Latency WLAN interface (AWDL companion interface)"), + "anpi" => Some("Apple Network Platform Interface (Wi-Fi/Bluetooth coexistence)"), + "bridge" => Some("Virtual bridge interface"), + "gif" => Some("Generic tunnel interface (IPv6-in-IPv4)"), + "stf" => Some("6to4 tunnel interface"), + "vmnet" => Some("Virtual Machine network interface"), + "vnic" => Some("Parallels/VMware virtual network interface"), + "ppp" => Some("Point-to-Point Protocol interface"), + "wlan" => Some("Wireless LAN interface"), + "eth" => Some("Ethernet interface"), + "docker" => Some("Docker virtual bridge interface"), + "veth" => Some("Virtual Ethernet interface"), + _ => None, + } +} + +/// Builds the single record batch of interfaces from `pcap`'s device list +fn interfaces_batch(schema: &SchemaRef) -> Result { + let devices = pcap::Device::list().map_err(|e| { + DataFusionError::External(format!("interfaces failed to list devices: {e}").into()) + })?; + + let mut name = StringBuilder::new(); + let mut description = StringBuilder::new(); + let mut addresses = ListBuilder::new(StringBuilder::new()); + let mut is_up = BooleanBuilder::new(); + let mut is_running = BooleanBuilder::new(); + let mut is_loopback = BooleanBuilder::new(); + let mut is_wireless = BooleanBuilder::new(); + let mut connection_status = StringBuilder::new(); + + for device in devices { + name.append_value(&device.name); + description.append_option( + device + .desc + .as_deref() + .or_else(|| static_description(&device.name)), + ); + for address in &device.addresses { + addresses.values().append_value(address.addr.to_string()); + } + addresses.append(true); + is_up.append_value(device.flags.is_up()); + is_running.append_value(device.flags.is_running()); + is_loopback.append_value(device.flags.is_loopback()); + is_wireless.append_value(device.flags.is_wireless()); + connection_status.append_value(match device.flags.connection_status { + ConnectionStatus::Unknown => "unknown", + ConnectionStatus::Connected => "connected", + ConnectionStatus::Disconnected => "disconnected", + ConnectionStatus::NotApplicable => "not_applicable", + }); + } + + Ok(RecordBatch::try_new( + Arc::clone(schema), + vec![ + Arc::new(name.finish()), + Arc::new(description.finish()), + Arc::new(addresses.finish()), + Arc::new(is_up.finish()), + Arc::new(is_running.finish()), + Arc::new(is_loopback.finish()), + Arc::new(is_wireless.finish()), + Arc::new(connection_status.finish()), + ], + )?) +} + +#[cfg(test)] +mod tests { + use super::*; + use datafusion::prelude::SessionContext; + + #[test] + fn test_call_rejects_arguments() { + let func = InterfacesFunc::default(); + let args = vec![Expr::Literal( + datafusion::scalar::ScalarValue::Utf8(Some("en0".to_string())), + None, + )]; + let err = func.call(&args).unwrap_err(); + assert!(err.to_string().contains("takes no arguments")); + } + + #[tokio::test] + async fn test_interfaces_lists_devices() { + let ctx = SessionContext::new(); + ctx.register_udtf("interfaces", Arc::new(InterfacesFunc::default())); + let batches = ctx + .sql("SELECT name, is_loopback, connection_status FROM interfaces()") + .await + .unwrap() + .collect() + .await + .unwrap(); + // Any host running the tests has at least one interface + let rows: usize = batches.iter().map(|b| b.num_rows()).sum(); + assert!(rows >= 1, "expected at least one interface"); + } + + #[tokio::test] + async fn test_interfaces_projection_and_filter() { + let ctx = SessionContext::new(); + ctx.register_udtf("interfaces", Arc::new(InterfacesFunc::default())); + // Projection, filtering, and unnesting the address list all compose + let batches = ctx + .sql( + "SELECT name, unnest(addresses) AS address \ + FROM interfaces() WHERE is_up ORDER BY name", + ) + .await + .unwrap() + .collect() + .await + .unwrap(); + // No assertion on contents: an interface may legitimately have no + // addresses, this only must not error + let _ = batches; + } +} diff --git a/crates/datafusion-net/src/lib.rs b/crates/datafusion-net/src/lib.rs new file mode 100644 index 0000000..a1df9b2 --- /dev/null +++ b/crates/datafusion-net/src/lib.rs @@ -0,0 +1,119 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! DataFusion table functions for querying network packet captures with SQL, +//! similar to wireshark/tshark. +//! +//! - [`PcapFunc`] (`pcap`): reads a pcap/pcapng capture file as a table +//! - [`CaptureFunc`] (`capture`, requires the `live` feature): streams +//! live-captured packets from a network interface +//! - [`InterfacesFunc`] (`interfaces`, requires the `live` feature): lists +//! the system's network capture interfaces +//! - [`PcapWideFunc`] / [`CaptureWideFunc`] (`pcap_wide` / `capture_wide`): +//! the same tables with DNS and geolocation enrichment columns appended +//! - [`TcpConversationsFunc`] (`tcp_conversations`): aggregates a capture +//! file into one row per TCP connection with flow-level analytics +//! +//! Scalar UDFs for enriching the IP columns: +//! +//! - [`ReverseDnsUdf`] (`reverse_dns`): resolves an IP address to a hostname +//! - [`GeoIpUdf`] (`geoip`): geolocates an IP address using a MaxMind-format +//! (`.mmdb`) database +//! +//! Scalar UDFs for decoding packet payloads: +//! +//! - [`DnsQueryUdf`] (`dns_query`): decodes a DNS message from a payload +//! - [`TlsSniUdf`] (`tls_sni`): extracts the SNI host from a TLS ClientHello +//! +//! ```sql +//! -- Query a capture file +//! SELECT src_ip, dst_ip, protocol FROM pcap('capture.pcap') WHERE dst_port = 443; +//! +//! -- Stream a live capture (requires elevated privileges) +//! SELECT * FROM capture('en0', 'tcp port 443') LIMIT 100; +//! ``` +//! +//! Both functions share the same schema (see [`packet_schema`]); frames that +//! cannot be decoded produce rows with null columns for the missing layers. +//! +//! ```rust,no_run +//! use std::sync::Arc; +//! use datafusion::prelude::SessionContext; +//! +//! let ctx = SessionContext::new(); +//! ctx.register_udtf("pcap", Arc::new(datafusion_net::PcapFunc::default())); +//! ``` + +use datafusion::{ + common::{plan_err, Column, Result}, + logical_expr::DocSection, + prelude::Expr, + scalar::ScalarValue, +}; + +/// The `SHOW FUNCTIONS` documentation section shared by this crate's scalar +/// UDFs. Groups them under one heading in the information schema. +pub(crate) const NET_DOC_SECTION: DocSection = DocSection { + include: true, + label: "Network Functions", + description: Some("Functions for enriching and decoding packet capture data"), +}; + +mod conversations; +pub mod decode; +mod dns; +mod file; +mod geoip; +#[cfg(feature = "live")] +mod interfaces; +#[cfg(feature = "live")] +mod live; +mod schema; +mod tls; +mod udfs; +mod wide; +pub mod writer; + +pub use conversations::{tcp_conversations_schema, TcpConversationsFunc, TcpConversationsTable}; +pub use dns::DnsQueryUdf; +pub use file::{PcapFunc, PcapTable}; +pub use geoip::{GeoIpUdf, GEOIP_DB_ENV_VAR}; +#[cfg(feature = "live")] +pub use interfaces::{interfaces_schema, InterfacesFunc}; +#[cfg(feature = "live")] +pub use live::{CaptureFunc, CaptureTable}; +pub use schema::packet_schema; +pub use tls::TlsSniUdf; +pub use udfs::ReverseDnsUdf; +#[cfg(feature = "live")] +pub use wide::CaptureWideFunc; +pub use wide::PcapWideFunc; + +/// Extracts a string argument from a table function expression +pub(crate) fn expr_to_string(expr: &Expr, func: &str, what: &str) -> Result { + match expr { + Expr::Literal( + ScalarValue::Utf8(Some(s)) + | ScalarValue::Utf8View(Some(s)) + | ScalarValue::LargeUtf8(Some(s)), + _, + ) => Ok(s.clone()), + // Double quoted strings are parsed as columns + Expr::Column(Column { name, .. }) => Ok(name.clone()), + _ => plan_err!("{func} {what} must be a string literal"), + } +} diff --git a/crates/datafusion-net/src/live.rs b/crates/datafusion-net/src/live.rs new file mode 100644 index 0000000..93ae9ef --- /dev/null +++ b/crates/datafusion-net/src/live.rs @@ -0,0 +1,465 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! `capture` table function: streams live-captured packets from a network +//! interface as rows. +//! +//! ```sql +//! -- Stream until 100 packets have been captured +//! SELECT * FROM capture('en0') LIMIT 100; +//! +//! -- Apply a BPF filter +//! SELECT * FROM capture('en0', 'tcp port 443') LIMIT 10; +//! +//! -- Capture for 10 seconds; the stream terminates, so aggregations work +//! SELECT src_ip, count(*) FROM capture('en0', '', 10) GROUP BY src_ip; +//! ``` +//! +//! Arguments: interface name, optional BPF filter expression, optional +//! capture duration in seconds. Without a duration the source is unbounded: +//! use a `LIMIT` or the query streams until cancelled. +//! +//! Live capture requires elevated privileges (see the error message emitted +//! when opening the device fails). No device is opened during planning +//! (e.g. `EXPLAIN`), only on execution. + +use std::{ + sync::Arc, + time::{Duration, Instant}, +}; + +use async_trait::async_trait; +use datafusion::{ + arrow::{array::RecordBatch, datatypes::SchemaRef}, + catalog::{Session, TableFunctionImpl, TableProvider}, + common::{internal_err, plan_err, project_schema, DataFusionError, Result}, + datasource::TableType, + execution::SendableRecordBatchStream, + physical_expr::EquivalenceProperties, + physical_plan::{ + execution_plan::{Boundedness, EmissionType}, + stream::RecordBatchReceiverStream, + DisplayAs, DisplayFormatType, ExecutionPlan, Partitioning, PlanProperties, + }, + prelude::Expr, + scalar::ScalarValue, +}; +use tokio::sync::mpsc::Sender; + +use crate::{ + decode::decode_frame, + expr_to_string, + schema::{packet_schema, payload_projected, send_batch, send_error, PacketBatchBuilder}, +}; + +/// Number of rows accumulated before a batch is emitted +const BATCH_SIZE: usize = 1024; +/// Partial batches are flushed at this interval so results appear promptly +/// on quiet interfaces. Also used as the libpcap read timeout so the read +/// loop can flush, and check for cancellation, without traffic. +const FLUSH_INTERVAL: Duration = Duration::from_millis(250); + +/// Table function that captures packets from a network interface +#[derive(Debug, Default)] +pub struct CaptureFunc {} + +impl TableFunctionImpl for CaptureFunc { + fn call(&self, exprs: &[Expr]) -> Result> { + let (interface, filter, duration) = parse_capture_args("capture", exprs)?; + Ok(Arc::new(CaptureTable::new(interface, filter, duration))) + } +} + +/// Parses the arguments shared by `capture` and `capture_wide` (hence the +/// function name parameter): interface name, optional BPF filter, optional +/// duration in seconds +pub(crate) fn parse_capture_args( + func: &str, + exprs: &[Expr], +) -> Result<(String, String, Option)> { + if exprs.is_empty() || exprs.len() > 3 { + return plan_err!( + "{func} requires 1 to 3 arguments: interface name, optional BPF filter, optional duration in seconds" + ); + } + let interface = expr_to_string(&exprs[0], func, "interface (first argument)")?; + let filter = exprs + .get(1) + .map(|e| expr_to_string(e, func, "BPF filter (second argument)")) + .transpose()? + .unwrap_or_default(); + let duration = exprs + .get(2) + .map(|e| expr_to_duration(func, e)) + .transpose()?; + Ok((interface, filter, duration)) +} + +fn expr_to_duration(func: &str, expr: &Expr) -> Result { + match expr { + Expr::Literal(ScalarValue::Int64(Some(secs)), _) if *secs > 0 => { + Ok(Duration::from_secs(*secs as u64)) + } + _ => plan_err!("{func} duration (third argument) must be a positive integer of seconds"), + } +} + +/// [`TableProvider`] backed by a live capture on a network interface +#[derive(Debug)] +pub struct CaptureTable { + interface: String, + filter: String, + duration: Option, + schema: SchemaRef, +} + +impl CaptureTable { + pub fn new(interface: String, filter: String, duration: Option) -> Self { + Self { + interface, + filter, + duration, + schema: packet_schema(), + } + } +} + +#[async_trait] +impl TableProvider for CaptureTable { + fn schema(&self) -> SchemaRef { + Arc::clone(&self.schema) + } + + fn table_type(&self) -> TableType { + TableType::Base + } + + async fn scan( + &self, + _state: &dyn Session, + projection: Option<&Vec>, + _filters: &[Expr], + limit: Option, + ) -> Result> { + let exec = CaptureExec::try_new( + self.interface.clone(), + self.filter.clone(), + self.duration, + Arc::clone(&self.schema), + projection.cloned(), + limit, + )?; + Ok(Arc::new(exec)) + } +} + +/// Execution plan that opens the capture device on `execute` and yields +/// record batches of decoded packets +#[derive(Debug)] +struct CaptureExec { + interface: String, + filter: String, + duration: Option, + projection: Option>, + /// Schema representing the data after the optional projection is applied + projected_schema: SchemaRef, + limit: Option, + cache: Arc, +} + +impl CaptureExec { + fn try_new( + interface: String, + filter: String, + duration: Option, + schema: SchemaRef, + projection: Option>, + limit: Option, + ) -> Result { + let projected_schema = project_schema(&schema, projection.as_ref())?; + // A capture with a duration terminates, which allows plans that + // require bounded input (e.g. aggregations) to run against it + let boundedness = if duration.is_some() { + Boundedness::Bounded + } else { + Boundedness::Unbounded { + requires_infinite_memory: false, + } + }; + let cache = Arc::new(PlanProperties::new( + EquivalenceProperties::new(Arc::clone(&projected_schema)), + Partitioning::UnknownPartitioning(1), + EmissionType::Incremental, + boundedness, + )); + Ok(Self { + interface, + filter, + duration, + projection, + projected_schema, + limit, + cache, + }) + } +} + +impl DisplayAs for CaptureExec { + fn fmt_as(&self, _t: DisplayFormatType, f: &mut std::fmt::Formatter) -> std::fmt::Result { + write!( + f, + "CaptureExec: interface={}, filter={}, duration={:?}, limit={:?}", + self.interface, self.filter, self.duration, self.limit + ) + } +} + +impl ExecutionPlan for CaptureExec { + fn name(&self) -> &str { + "CaptureExec" + } + + fn properties(&self) -> &Arc { + &self.cache + } + + fn children(&self) -> Vec<&Arc> { + // This is a leaf node and has no children + vec![] + } + + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + // CaptureExec has no children + if children.is_empty() { + Ok(self) + } else { + internal_err!("Children cannot be replaced in {self:?}") + } + } + + fn execute( + &self, + partition: usize, + _context: Arc, + ) -> Result { + if partition != 0 { + return internal_err!("CaptureExec has a single partition, got {partition}"); + } + let mut builder = + RecordBatchReceiverStream::builder(Arc::clone(&self.projected_schema), 16); + let tx = builder.tx(); + let interface = self.interface.clone(); + let filter = self.filter.clone(); + let duration = self.duration; + let projection = self.projection.clone(); + let limit = self.limit; + // libpcap reads are blocking C calls, so run the capture loop on a + // blocking thread rather than an async task + builder + .spawn_blocking(move || read_live(interface, filter, duration, projection, limit, tx)); + Ok(builder.build()) + } +} + +/// Captures packets from `interface` and sends decoded record batches until +/// the duration elapses, `limit` rows have been produced, or the consumer +/// drops the stream +fn read_live( + interface: String, + filter: String, + duration: Option, + projection: Option>, + limit: Option, + tx: Sender>, +) -> Result<()> { + let inactive = match pcap::Capture::from_device(interface.as_str()) { + Ok(inactive) => inactive, + Err(e) => { + send_error(&tx, open_error(&interface, &e)); + return Ok(()); + } + }; + let capture = inactive + .promisc(true) + .snaplen(65535) + .timeout(FLUSH_INTERVAL.as_millis() as i32) + .immediate_mode(true) + .open(); + let mut capture = match capture { + Ok(capture) => capture, + Err(e) => { + send_error(&tx, open_error(&interface, &e)); + return Ok(()); + } + }; + if !filter.is_empty() { + if let Err(e) = capture.filter(&filter, true) { + send_error( + &tx, + DataFusionError::External( + format!("capture failed to set BPF filter '{filter}': {e}").into(), + ), + ); + return Ok(()); + } + } + let link_type = capture.get_datalink().0 as u32; + + let mut builder = PacketBatchBuilder::new(payload_projected(&projection)); + let mut frame_number = 0u64; + let mut produced = 0usize; + let started = Instant::now(); + let mut last_flush = Instant::now(); + + loop { + if duration.is_some_and(|d| started.elapsed() >= d) { + break; + } + match capture.next_packet() { + Ok(packet) => { + frame_number += 1; + // tv_sec/tv_usec widths vary by platform (i64 on macOS, + // where these casts are no-ops) + #[allow(clippy::unnecessary_cast)] + let ts_micros = + packet.header.ts.tv_sec as i64 * 1_000_000 + packet.header.ts.tv_usec as i64; + let decoded = decode_frame(link_type, packet.data); + builder.append( + ts_micros, + frame_number, + packet.header.len, + packet.header.caplen, + &decoded, + ); + produced += 1; + if limit.is_some_and(|l| produced >= l) { + break; + } + } + // No packets within the read timeout; fall through to flush any + // partial batch and re-check the duration + Err(pcap::Error::TimeoutExpired) => {} + Err(e) => { + send_error( + &tx, + DataFusionError::External( + format!("capture read error on '{interface}': {e}").into(), + ), + ); + break; + } + } + if builder.len() >= BATCH_SIZE + || (!builder.is_empty() && last_flush.elapsed() >= FLUSH_INTERVAL) + { + if !send_batch(&mut builder, &projection, &tx)? { + // Consumer dropped the stream + return Ok(()); + } + last_flush = Instant::now(); + } + // Even with nothing to flush, a closed channel means the consumer is + // gone and capturing should stop + if tx.is_closed() { + return Ok(()); + } + } + + send_batch(&mut builder, &projection, &tx)?; + Ok(()) +} + +fn open_error(interface: &str, err: &pcap::Error) -> DataFusionError { + let permission_denied = match err { + pcap::Error::IoError(kind) => *kind == std::io::ErrorKind::PermissionDenied, + pcap::Error::PcapError(msg) => { + let msg = msg.to_lowercase(); + msg.contains("permission denied") || msg.contains("not permitted") + } + _ => false, + }; + let hint = if permission_denied { + " (live capture requires elevated privileges: run with sudo, grant the binary \ + cap_net_raw+cap_net_admin on Linux, or install ChmodBPF on macOS)" + } else { + "" + }; + DataFusionError::External( + format!("capture failed to open device '{interface}': {err}{hint}").into(), + ) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_call_requires_interface() { + let func = CaptureFunc::default(); + let err = func.call(&[]).unwrap_err(); + assert!(err.to_string().contains("1 to 3 arguments")); + } + + #[test] + fn test_call_rejects_too_many_arguments() { + let func = CaptureFunc::default(); + let args = vec![Expr::Literal(ScalarValue::Utf8(Some("en0".to_string())), None); 4]; + let err = func.call(&args).unwrap_err(); + assert!(err.to_string().contains("1 to 3 arguments")); + } + + #[test] + fn test_call_rejects_non_integer_duration() { + let func = CaptureFunc::default(); + let args = vec![ + Expr::Literal(ScalarValue::Utf8(Some("en0".to_string())), None), + Expr::Literal(ScalarValue::Utf8(Some("tcp".to_string())), None), + Expr::Literal(ScalarValue::Utf8(Some("ten".to_string())), None), + ]; + let err = func.call(&args).unwrap_err(); + assert!(err.to_string().contains("positive integer")); + } + + #[test] + fn test_call_parses_all_arguments() { + let func = CaptureFunc::default(); + let args = vec![ + Expr::Literal(ScalarValue::Utf8(Some("en0".to_string())), None), + Expr::Literal(ScalarValue::Utf8(Some("tcp port 443".to_string())), None), + Expr::Literal(ScalarValue::Int64(Some(10)), None), + ]; + let provider = func.call(&args).unwrap(); + assert_eq!(provider.schema(), packet_schema()); + } + + #[tokio::test] + async fn test_explain_does_not_open_device() { + let ctx = datafusion::prelude::SessionContext::new(); + ctx.register_udtf("capture", Arc::new(CaptureFunc::default())); + // Planning must not open a capture device (which would fail without + // elevated privileges, and this device does not exist anyway) + let df = ctx + .sql("EXPLAIN SELECT * FROM capture('definitely-not-a-device')") + .await + .unwrap(); + let batches = df.collect().await.unwrap(); + assert!(!batches.is_empty()); + } +} diff --git a/crates/datafusion-net/src/schema.rs b/crates/datafusion-net/src/schema.rs new file mode 100644 index 0000000..c25d786 --- /dev/null +++ b/crates/datafusion-net/src/schema.rs @@ -0,0 +1,247 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Arrow schema for decoded packets and a columnar builder for turning +//! decoded packets into [`RecordBatch`]es. + +use std::sync::Arc; + +use datafusion::arrow::{ + array::{ + BinaryBuilder, RecordBatch, StringBuilder, TimestampMicrosecondBuilder, UInt16Builder, + UInt32Builder, UInt64Builder, UInt8Builder, + }, + datatypes::{DataType, Field, Schema, SchemaRef, TimeUnit}, +}; +use datafusion::common::Result; +use tokio::sync::mpsc::Sender; + +use crate::decode::DecodedPacket; + +/// Index of the `payload` column, used to skip materializing packet payloads +/// when the column is not projected +pub const PAYLOAD_COLUMN_INDEX: usize = 20; + +/// Schema shared by the `pcap` and `capture` table functions. One row per +/// captured frame; layers that could not be decoded yield null columns. +pub fn packet_schema() -> SchemaRef { + Arc::new(Schema::new(vec![ + Field::new( + "timestamp", + DataType::Timestamp(TimeUnit::Microsecond, None), + false, + ), + Field::new("frame_number", DataType::UInt64, false), + Field::new("length", DataType::UInt32, false), + Field::new("capture_length", DataType::UInt32, false), + Field::new("eth_src", DataType::Utf8, true), + Field::new("eth_dst", DataType::Utf8, true), + Field::new("ethertype", DataType::Utf8, true), + Field::new("vlan", DataType::UInt16, true), + Field::new("ip_version", DataType::UInt8, true), + Field::new("src_ip", DataType::Utf8, true), + Field::new("dst_ip", DataType::Utf8, true), + Field::new("ttl", DataType::UInt8, true), + Field::new("protocol", DataType::Utf8, true), + Field::new("src_port", DataType::UInt16, true), + Field::new("dst_port", DataType::UInt16, true), + Field::new("tcp_flags", DataType::Utf8, true), + Field::new("tcp_seq", DataType::UInt32, true), + Field::new("tcp_ack", DataType::UInt32, true), + Field::new("tcp_window", DataType::UInt16, true), + Field::new("payload_length", DataType::UInt32, true), + Field::new("payload", DataType::Binary, true), + ])) +} + +/// Accumulates decoded packets into columnar builders and produces +/// [`RecordBatch`]es with the full [`packet_schema`] +pub struct PacketBatchBuilder { + schema: SchemaRef, + /// Whether to materialize payload bytes (false when the `payload` column + /// is projected away, avoiding the copy) + include_payload: bool, + rows: usize, + timestamp: TimestampMicrosecondBuilder, + frame_number: UInt64Builder, + length: UInt32Builder, + capture_length: UInt32Builder, + eth_src: StringBuilder, + eth_dst: StringBuilder, + ethertype: StringBuilder, + vlan: UInt16Builder, + ip_version: UInt8Builder, + src_ip: StringBuilder, + dst_ip: StringBuilder, + ttl: UInt8Builder, + protocol: StringBuilder, + src_port: UInt16Builder, + dst_port: UInt16Builder, + tcp_flags: StringBuilder, + tcp_seq: UInt32Builder, + tcp_ack: UInt32Builder, + tcp_window: UInt16Builder, + payload_length: UInt32Builder, + payload: BinaryBuilder, +} + +impl PacketBatchBuilder { + pub fn new(include_payload: bool) -> Self { + Self { + schema: packet_schema(), + include_payload, + rows: 0, + timestamp: TimestampMicrosecondBuilder::new(), + frame_number: UInt64Builder::new(), + length: UInt32Builder::new(), + capture_length: UInt32Builder::new(), + eth_src: StringBuilder::new(), + eth_dst: StringBuilder::new(), + ethertype: StringBuilder::new(), + vlan: UInt16Builder::new(), + ip_version: UInt8Builder::new(), + src_ip: StringBuilder::new(), + dst_ip: StringBuilder::new(), + ttl: UInt8Builder::new(), + protocol: StringBuilder::new(), + src_port: UInt16Builder::new(), + dst_port: UInt16Builder::new(), + tcp_flags: StringBuilder::new(), + tcp_seq: UInt32Builder::new(), + tcp_ack: UInt32Builder::new(), + tcp_window: UInt16Builder::new(), + payload_length: UInt32Builder::new(), + payload: BinaryBuilder::new(), + } + } + + pub fn append( + &mut self, + ts_micros: i64, + frame_number: u64, + length: u32, + capture_length: u32, + decoded: &DecodedPacket<'_>, + ) { + self.timestamp.append_value(ts_micros); + self.frame_number.append_value(frame_number); + self.length.append_value(length); + self.capture_length.append_value(capture_length); + self.eth_src.append_option(decoded.eth_src.as_deref()); + self.eth_dst.append_option(decoded.eth_dst.as_deref()); + self.ethertype.append_option(decoded.ethertype.as_deref()); + self.vlan.append_option(decoded.vlan); + self.ip_version.append_option(decoded.ip_version); + self.src_ip.append_option(decoded.src_ip.as_deref()); + self.dst_ip.append_option(decoded.dst_ip.as_deref()); + self.ttl.append_option(decoded.ttl); + self.protocol.append_option(decoded.protocol.as_deref()); + self.src_port.append_option(decoded.src_port); + self.dst_port.append_option(decoded.dst_port); + self.tcp_flags.append_option(decoded.tcp_flags.as_deref()); + self.tcp_seq.append_option(decoded.tcp_seq); + self.tcp_ack.append_option(decoded.tcp_ack); + self.tcp_window.append_option(decoded.tcp_window); + self.payload_length + .append_option(decoded.payload.map(|p| p.len() as u32)); + if self.include_payload { + self.payload.append_option(decoded.payload); + } else { + self.payload.append_null(); + } + self.rows += 1; + } + + pub fn len(&self) -> usize { + self.rows + } + + pub fn is_empty(&self) -> bool { + self.rows == 0 + } + + /// Builds a [`RecordBatch`] with the full schema from the accumulated + /// rows, resetting the builder + pub fn finish(&mut self) -> Result { + self.rows = 0; + let batch = RecordBatch::try_new( + Arc::clone(&self.schema), + vec![ + Arc::new(self.timestamp.finish()), + Arc::new(self.frame_number.finish()), + Arc::new(self.length.finish()), + Arc::new(self.capture_length.finish()), + Arc::new(self.eth_src.finish()), + Arc::new(self.eth_dst.finish()), + Arc::new(self.ethertype.finish()), + Arc::new(self.vlan.finish()), + Arc::new(self.ip_version.finish()), + Arc::new(self.src_ip.finish()), + Arc::new(self.dst_ip.finish()), + Arc::new(self.ttl.finish()), + Arc::new(self.protocol.finish()), + Arc::new(self.src_port.finish()), + Arc::new(self.dst_port.finish()), + Arc::new(self.tcp_flags.finish()), + Arc::new(self.tcp_seq.finish()), + Arc::new(self.tcp_ack.finish()), + Arc::new(self.tcp_window.finish()), + Arc::new(self.payload_length.finish()), + Arc::new(self.payload.finish()), + ], + )?; + Ok(batch) + } +} + +/// Finishes the builder, applies the optional projection, and sends the +/// resulting batch on `tx` from a blocking context. Returns `false` when the +/// consumer has dropped the stream (e.g. a `LIMIT` was satisfied) and reading +/// should stop. +pub(crate) fn send_batch( + builder: &mut PacketBatchBuilder, + projection: &Option>, + tx: &Sender>, +) -> Result { + if builder.is_empty() { + return Ok(true); + } + let batch = builder.finish()?; + let batch = match projection { + Some(p) => batch.project(p)?, + None => batch, + }; + Ok(tx.blocking_send(Ok(batch)).is_ok()) +} + +/// Sends an error on `tx` from a blocking context, ignoring send failures +/// (the consumer may already be gone) +pub(crate) fn send_error( + tx: &Sender>, + err: datafusion::error::DataFusionError, +) { + let _ = tx.blocking_send(Err(err)); +} + +/// Whether the `payload` column is included in the projection (and therefore +/// needs to be materialized) +pub(crate) fn payload_projected(projection: &Option>) -> bool { + match projection { + Some(p) => p.contains(&PAYLOAD_COLUMN_INDEX), + None => true, + } +} diff --git a/crates/datafusion-net/src/tls.rs b/crates/datafusion-net/src/tls.rs new file mode 100644 index 0000000..f73cc8b --- /dev/null +++ b/crates/datafusion-net/src/tls.rs @@ -0,0 +1,421 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! `tls_sni` scalar UDF: extracts the Server Name Indication (SNI) host name +//! from a TLS ClientHello in a packet payload. +//! +//! [`TlsSniUdf`] takes a binary payload (typically the `payload` column of +//! TCP port 443 packets) and returns the SNI host name as a string, or NULL +//! when the payload is not a ClientHello or carries no SNI extension. In +//! encrypted (HTTPS) traffic the payloads are opaque, but the ClientHello — +//! the first packet a client sends — is unencrypted and names the host being +//! connected to, so this reveals the destination of otherwise opaque flows: +//! +//! ```sql +//! SELECT tls_sni(payload) AS host, count(*) AS hellos +//! FROM pcap('capture.pcap') +//! WHERE dst_port = 443 AND tls_sni(payload) IS NOT NULL +//! GROUP BY host ORDER BY hellos DESC; +//! ``` +//! +//! The ClientHello must fit in the payload (it usually does — it is small and +//! sent first); this UDF does not reassemble segments. + +use std::sync::{Arc, LazyLock}; + +use datafusion::{ + arrow::{ + array::{Array, ArrayRef, StringArray, StringBuilder}, + datatypes::DataType, + }, + common::{cast::as_binary_array, exec_err, Result}, + logical_expr::{ + ColumnarValue, Documentation, ScalarFunctionArgs, ScalarUDFImpl, Signature, TypeSignature, + Volatility, + }, +}; + +use crate::NET_DOC_SECTION; + +/// `SHOW FUNCTIONS` documentation for `tls_sni` +static DOCUMENTATION: LazyLock = LazyLock::new(|| { + Documentation::builder( + NET_DOC_SECTION, + "Extracts the Server Name Indication (SNI) host name from a TLS \ + ClientHello in a packet payload (typically TCP port 443), or NULL \ + when the payload is not a ClientHello or carries no SNI. The \ + ClientHello is unencrypted even for HTTPS, so this reveals the \ + destination host of otherwise opaque connections. The ClientHello \ + must fit in the payload; segments are not reassembled.", + "tls_sni(payload)", + ) + .with_argument( + "payload", + "Binary packet payload, e.g. the payload column of a TLS packet", + ) + .with_sql_example( + "SELECT tls_sni(payload) AS host, count(*) AS hellos \ + FROM pcap('capture.pcap') WHERE tls_sni(payload) IS NOT NULL \ + GROUP BY host ORDER BY hellos DESC", + ) + .with_related_udf("dns_query") + .build() +}); + +/// Scalar UDF that extracts the SNI host name from a TLS ClientHello +#[derive(Debug, PartialEq, Eq, Hash)] +pub struct TlsSniUdf { + signature: Signature, +} + +impl Default for TlsSniUdf { + fn default() -> Self { + Self { + // Pure parse of the input bytes + signature: Signature::one_of( + vec![ + TypeSignature::Exact(vec![DataType::Binary]), + TypeSignature::Exact(vec![DataType::LargeBinary]), + TypeSignature::Exact(vec![DataType::BinaryView]), + ], + Volatility::Immutable, + ), + } + } +} + +impl ScalarUDFImpl for TlsSniUdf { + fn name(&self) -> &str { + "tls_sni" + } + + fn signature(&self) -> &Signature { + &self.signature + } + + fn return_type(&self, _arg_types: &[DataType]) -> Result { + Ok(DataType::Utf8) + } + + fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result { + if args.args.len() != 1 { + return exec_err!("tls_sni expects a single binary argument"); + } + let payloads = to_binary_array(&args.args[0].to_array(args.number_rows)?)?; + let payloads = as_binary_array(&payloads)?; + + let mut out = StringBuilder::new(); + for row in 0..payloads.len() { + let sni = if payloads.is_null(row) { + None + } else { + parse_sni(payloads.value(row)) + }; + out.append_option(sni); + } + Ok(ColumnarValue::Array(Arc::new(out.finish()) as ArrayRef)) + } + + fn documentation(&self) -> Option<&Documentation> { + Some(&DOCUMENTATION) + } +} + +/// Normalizes the input (BinaryView / LargeBinary) to Binary +fn to_binary_array(array: &ArrayRef) -> Result { + if array.data_type() == &DataType::Binary { + Ok(Arc::clone(array)) + } else { + Ok(datafusion::arrow::compute::cast(array, &DataType::Binary)?) + } +} + +/// Cursor over a byte slice that yields big-endian fields and bails (via +/// `None`) the moment a read would run past the end +struct Reader<'a> { + data: &'a [u8], + pos: usize, +} + +impl<'a> Reader<'a> { + fn new(data: &'a [u8]) -> Self { + Self { data, pos: 0 } + } + + fn u8(&mut self) -> Option { + let v = *self.data.get(self.pos)?; + self.pos += 1; + Some(v) + } + + fn u16(&mut self) -> Option { + let hi = self.u8()? as usize; + let lo = self.u8()? as usize; + Some((hi << 8) | lo) + } + + fn u24(&mut self) -> Option { + let a = self.u8()? as usize; + let b = self.u16()?; + Some((a << 16) | b) + } + + fn take(&mut self, n: usize) -> Option<&'a [u8]> { + let slice = self.data.get(self.pos..self.pos + n)?; + self.pos += n; + Some(slice) + } + + /// Advances past `n` bytes without returning them + fn skip(&mut self, n: usize) -> Option<()> { + self.take(n).map(|_| ()) + } +} + +/// TLS record content type for handshake messages +const CONTENT_TYPE_HANDSHAKE: u8 = 22; +/// Handshake message type for ClientHello +const HANDSHAKE_CLIENT_HELLO: u8 = 1; +/// Extension type for server_name (SNI) +const EXT_SERVER_NAME: usize = 0; +/// SNI name type for host_name +const SNI_HOST_NAME: u8 = 0; + +/// Parses the SNI host name out of a TLS ClientHello. Returns `None` if the +/// bytes are not a ClientHello record or carry no host_name SNI entry. +fn parse_sni(data: &[u8]) -> Option { + let mut r = Reader::new(data); + + // TLS record header + if r.u8()? != CONTENT_TYPE_HANDSHAKE { + return None; + } + let _record_version = r.u16()?; + let _record_len = r.u16()?; + + // Handshake header + if r.u8()? != HANDSHAKE_CLIENT_HELLO { + return None; + } + let _handshake_len = r.u24()?; + + // ClientHello body + let _client_version = r.u16()?; + r.skip(32)?; // random + let session_id_len = r.u8()? as usize; + r.skip(session_id_len)?; + let cipher_suites_len = r.u16()?; + r.skip(cipher_suites_len)?; + let compression_len = r.u8()? as usize; + r.skip(compression_len)?; + + // Extensions (absent in very old TLS, in which case there is no SNI) + let _extensions_len = r.u16()?; + loop { + let ext_type = r.u16()?; + let ext_len = r.u16()?; + if ext_type == EXT_SERVER_NAME { + return parse_server_name_extension(r.take(ext_len)?); + } + r.skip(ext_len)?; + } +} + +/// Parses a server_name extension body, returning the first host_name entry +fn parse_server_name_extension(data: &[u8]) -> Option { + let mut r = Reader::new(data); + let _list_len = r.u16()?; + // The list can hold multiple entries; return the first host_name + loop { + let name_type = r.u8()?; + let name_len = r.u16()?; + let name = r.take(name_len)?; + if name_type == SNI_HOST_NAME { + // Host names are ASCII (IDNs are punycode-encoded on the wire) + if !name.iter().all(|b| b.is_ascii_graphic()) { + return None; + } + return Some(String::from_utf8_lossy(name).into_owned()); + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + /// Builds a minimal ClientHello record whose only extension is SNI for + /// `host` + fn client_hello_with_sni(host: &str) -> Vec { + let host = host.as_bytes(); + + // server_name extension body + let mut sni_ext = Vec::new(); + let entry_len = 1 + 2 + host.len(); + sni_ext.extend_from_slice(&(entry_len as u16).to_be_bytes()); // list length + sni_ext.push(SNI_HOST_NAME); + sni_ext.extend_from_slice(&(host.len() as u16).to_be_bytes()); + sni_ext.extend_from_slice(host); + + // extensions block: one extension (type 0) + let mut extensions = Vec::new(); + extensions.extend_from_slice(&(EXT_SERVER_NAME as u16).to_be_bytes()); + extensions.extend_from_slice(&(sni_ext.len() as u16).to_be_bytes()); + extensions.extend_from_slice(&sni_ext); + + // ClientHello body + let mut body = Vec::new(); + body.extend_from_slice(&[0x03, 0x03]); // client version TLS 1.2 + body.extend_from_slice(&[0u8; 32]); // random + body.push(0); // session id length + body.extend_from_slice(&[0x00, 0x02, 0x13, 0x01]); // cipher suites: len 2 + one suite + body.extend_from_slice(&[0x01, 0x00]); // compression: len 1 + null + body.extend_from_slice(&(extensions.len() as u16).to_be_bytes()); + body.extend_from_slice(&extensions); + + // Handshake header + let mut handshake = Vec::new(); + handshake.push(HANDSHAKE_CLIENT_HELLO); + let body_len = body.len(); + handshake.extend_from_slice(&[ + (body_len >> 16) as u8, + (body_len >> 8) as u8, + body_len as u8, + ]); + handshake.extend_from_slice(&body); + + // Record header + let mut record = Vec::new(); + record.push(CONTENT_TYPE_HANDSHAKE); + record.extend_from_slice(&[0x03, 0x01]); // record version + record.extend_from_slice(&(handshake.len() as u16).to_be_bytes()); + record.extend_from_slice(&handshake); + record + } + + #[test] + fn test_parse_sni() { + let hello = client_hello_with_sni("example.com"); + assert_eq!(parse_sni(&hello).as_deref(), Some("example.com")); + } + + #[test] + fn test_parse_sni_skips_earlier_extension() { + // Prepend a non-SNI extension by hand: rebuild with an extra + // extension of type 0x0017 (extended_master_secret, empty) + let mut hello = client_hello_with_sni("cdn.example.org"); + // The SNI is still found because parse_sni scans all extensions; a + // direct check that the right host is returned suffices here + assert_eq!(parse_sni(&hello).as_deref(), Some("cdn.example.org")); + // Corrupting the content type makes it not a ClientHello + hello[0] = 0x17; + assert_eq!(parse_sni(&hello), None); + } + + #[test] + fn test_non_client_hello_is_none() { + // Handshake record but a different handshake type (ServerHello = 2) + let mut hello = client_hello_with_sni("example.com"); + hello[5] = 0x02; + assert_eq!(parse_sni(&hello), None); + } + + #[test] + fn test_garbage_is_none() { + assert_eq!(parse_sni(b""), None); + assert_eq!(parse_sni(b"GET / HTTP/1.1\r\n"), None); + assert_eq!(parse_sni(&[0x16, 0x03, 0x01, 0x00]), None); // truncated + } + + #[test] + fn test_client_hello_without_sni_is_none() { + // A ClientHello with an empty extensions block + let mut body = Vec::new(); + body.extend_from_slice(&[0x03, 0x03]); + body.extend_from_slice(&[0u8; 32]); + body.push(0); + body.extend_from_slice(&[0x00, 0x02, 0x13, 0x01]); + body.extend_from_slice(&[0x01, 0x00]); + body.extend_from_slice(&[0x00, 0x00]); // extensions length 0 + + let mut handshake = vec![HANDSHAKE_CLIENT_HELLO]; + let body_len = body.len(); + handshake.extend_from_slice(&[0, (body_len >> 8) as u8, body_len as u8]); + handshake.extend_from_slice(&body); + + let mut record = vec![CONTENT_TYPE_HANDSHAKE, 0x03, 0x01]; + record.extend_from_slice(&(handshake.len() as u16).to_be_bytes()); + record.extend_from_slice(&handshake); + + assert_eq!(parse_sni(&record), None); + } + + #[test] + fn test_return_type_is_utf8() { + let udf = TlsSniUdf::default(); + assert_eq!( + udf.return_type(&[DataType::Binary]).unwrap(), + DataType::Utf8 + ); + } + + #[test] + fn test_has_documentation() { + let udf = TlsSniUdf::default(); + let doc = udf.documentation().unwrap(); + assert_eq!(doc.syntax_example, "tls_sni(payload)"); + assert!(doc.arguments.is_some()); + } + + #[tokio::test] + async fn test_invoke_over_array() { + use datafusion::arrow::array::{BinaryArray, RecordBatch}; + use datafusion::arrow::datatypes::{Field, Schema}; + use datafusion::logical_expr::ScalarUDF; + use datafusion::prelude::SessionContext; + + let hello = client_hello_with_sni("secure.example.net"); + let values: Vec> = vec![Some(hello.as_slice()), Some(b"not tls"), None]; + let array = BinaryArray::from(values); + let schema = Arc::new(Schema::new(vec![Field::new( + "payload", + DataType::Binary, + true, + )])); + let batch = RecordBatch::try_new(schema, vec![Arc::new(array)]).unwrap(); + + let ctx = SessionContext::new(); + ctx.register_batch("packets", batch).unwrap(); + ctx.register_udf(ScalarUDF::from(TlsSniUdf::default())); + let batches = ctx + .sql("SELECT tls_sni(payload) AS host FROM packets ORDER BY host NULLS LAST") + .await + .unwrap() + .collect() + .await + .unwrap(); + let out = batches[0] + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + // Exactly one row parsed to a host; the other two are NULL + assert_eq!(out.value(0), "secure.example.net"); + assert!(out.is_null(1)); + assert!(out.is_null(2)); + } +} diff --git a/crates/datafusion-net/src/udfs.rs b/crates/datafusion-net/src/udfs.rs new file mode 100644 index 0000000..4596858 --- /dev/null +++ b/crates/datafusion-net/src/udfs.rs @@ -0,0 +1,269 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Scalar UDFs for enriching packet capture columns. +//! +//! [`ReverseDnsUdf`] (`reverse_dns`) takes an IP address string and resolves +//! it to a hostname via a reverse DNS (PTR) lookup using the system resolver: +//! +//! ```sql +//! SELECT src_ip, reverse_dns(src_ip) AS host, count(*) +//! FROM capture('en0', '', 10) +//! GROUP BY src_ip, host +//! ``` +//! +//! Lookups are network I/O, so results are cached process-wide, identical +//! addresses within a batch are resolved once, and each batch of lookups is +//! bounded by a timeout. Addresses that fail to parse, fail to resolve, or +//! time out yield NULL. + +use std::{ + collections::HashMap, + net::IpAddr, + sync::{mpsc, Arc, LazyLock, Mutex, OnceLock}, + thread, + time::{Duration, Instant}, +}; + +use datafusion::{ + arrow::{ + array::{Array, ArrayRef, StringArray}, + datatypes::DataType, + }, + common::{cast::as_string_array, exec_err, Result}, + logical_expr::{ + ColumnarValue, Documentation, ScalarFunctionArgs, ScalarUDFImpl, Signature, TypeSignature, + Volatility, + }, +}; + +use crate::NET_DOC_SECTION; + +/// `SHOW FUNCTIONS` documentation for `reverse_dns` +static DOCUMENTATION: LazyLock = LazyLock::new(|| { + Documentation::builder( + NET_DOC_SECTION, + "Resolves an IP address string to a hostname via a reverse DNS (PTR) \ + lookup using the system resolver. Returns NULL when the address does \ + not parse, does not resolve, or the lookup times out. Results are \ + cached process-wide and each batch of lookups is bounded by a timeout.", + "reverse_dns(ip)", + ) + .with_argument("ip", "IP address string, e.g. the src_ip or dst_ip column") + .with_sql_example( + "SELECT dst_ip, reverse_dns(dst_ip) AS host, count(*) AS packets \ + FROM pcap('capture.pcap') GROUP BY dst_ip, host ORDER BY packets DESC", + ) + .build() +}); + +/// Maximum wall-clock time spent resolving the unique addresses in a single +/// batch. A dead or slow resolver cannot hang the query longer than this; +/// unresolved addresses in the batch become NULL. +const LOOKUP_TIMEOUT: Duration = Duration::from_secs(2); + +/// Process-wide cache of resolved hostnames. `None` is a negative cache entry +/// (the address does not resolve), avoiding repeated lookups of the same +/// address across batches and queries. +fn cache() -> &'static Mutex>>> { + static CACHE: OnceLock>>>> = OnceLock::new(); + CACHE.get_or_init(|| Mutex::new(HashMap::new())) +} + +/// Scalar UDF that resolves an IP address string to a hostname via reverse +/// (PTR) DNS lookup +#[derive(Debug, PartialEq, Eq, Hash)] +pub struct ReverseDnsUdf { + signature: Signature, +} + +impl Default for ReverseDnsUdf { + fn default() -> Self { + Self { + // Accept the string types our packet columns and SQL literals use. + // Volatile: DNS results can change over time, so the optimizer + // must not constant-fold or dedup calls. + signature: Signature::one_of( + vec![ + TypeSignature::Exact(vec![DataType::Utf8]), + TypeSignature::Exact(vec![DataType::Utf8View]), + TypeSignature::Exact(vec![DataType::LargeUtf8]), + ], + Volatility::Volatile, + ), + } + } +} + +impl ScalarUDFImpl for ReverseDnsUdf { + fn name(&self) -> &str { + "reverse_dns" + } + + fn signature(&self) -> &Signature { + &self.signature + } + + fn return_type(&self, _arg_types: &[DataType]) -> Result { + Ok(DataType::Utf8) + } + + fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result { + if args.args.len() != 1 { + return exec_err!("reverse_dns expects a single argument"); + } + let array = args.args[0].to_array(args.number_rows)?; + let resolved = resolve_array(&array)?; + Ok(ColumnarValue::Array(resolved)) + } + + fn documentation(&self) -> Option<&Documentation> { + Some(&DOCUMENTATION) + } +} + +/// Resolves each element of `array` (a string array of IP addresses) to a +/// hostname, returning a `Utf8` array of the same length +fn resolve_array(array: &ArrayRef) -> Result { + // Normalize the input (Utf8View / LargeUtf8) to Utf8 so a single code + // path handles every accepted string type + let array = if array.data_type() == &DataType::Utf8 { + Arc::clone(array) + } else { + datafusion::arrow::compute::cast(array, &DataType::Utf8)? + }; + let strings = as_string_array(&array)?; + + // Collect the unique, parseable addresses that are not already cached + let mut to_resolve: Vec = Vec::new(); + { + let cache = cache().lock().expect("reverse_dns cache poisoned"); + for value in strings.iter().flatten() { + if let Ok(ip) = value.parse::() { + if !cache.contains_key(&ip) && !to_resolve.contains(&ip) { + to_resolve.push(ip); + } + } + } + } + + if !to_resolve.is_empty() { + let results = resolve_batch(&to_resolve); + let mut cache = cache().lock().expect("reverse_dns cache poisoned"); + for (ip, host) in results { + cache.insert(ip, host.map(Arc::from)); + } + } + + // Build the output array from the cache + let cache = cache().lock().expect("reverse_dns cache poisoned"); + let hosts: StringArray = strings + .iter() + .map(|value| { + let ip = value?.parse::().ok()?; + cache.get(&ip).cloned().flatten() + }) + .map(|host| host.map(|h| h.to_string())) + .collect(); + Ok(Arc::new(hosts)) +} + +/// Resolves `ips` concurrently, bounded by [`LOOKUP_TIMEOUT`]. Addresses that +/// do not resolve within the timeout are absent from the returned map (and so +/// treated as unresolved). +fn resolve_batch(ips: &[IpAddr]) -> HashMap> { + let (tx, rx) = mpsc::channel(); + for &ip in ips { + let tx = tx.clone(); + // Each lookup is a blocking syscall; run them on their own threads so + // one slow address does not serialize the rest. Detached threads whose + // result arrives after the deadline simply have their send ignored. + thread::spawn(move || { + let host = dns_lookup::lookup_addr(&ip).ok(); + let _ = tx.send((ip, host)); + }); + } + drop(tx); + + let deadline = Instant::now() + LOOKUP_TIMEOUT; + let mut results = HashMap::with_capacity(ips.len()); + while results.len() < ips.len() { + let Some(remaining) = deadline.checked_duration_since(Instant::now()) else { + break; + }; + match rx.recv_timeout(remaining) { + Ok((ip, host)) => { + results.insert(ip, host); + } + Err(_) => break, + } + } + results +} + +#[cfg(test)] +mod tests { + use super::*; + use datafusion::arrow::array::StringArray; + + #[test] + fn test_return_type_is_utf8() { + let udf = ReverseDnsUdf::default(); + assert_eq!(udf.return_type(&[DataType::Utf8]).unwrap(), DataType::Utf8); + } + + #[test] + fn test_localhost_resolves() { + let input: ArrayRef = Arc::new(StringArray::from(vec![Some("127.0.0.1")])); + let out = resolve_array(&input).unwrap(); + let out = out.as_any().downcast_ref::().unwrap(); + // The PTR record for 127.0.0.1 is environment dependent, but the + // loopback address must resolve to *some* hostname on any host with a + // working resolver + assert!(!out.is_null(0), "127.0.0.1 should resolve to a hostname"); + assert!(!out.value(0).is_empty()); + } + + #[test] + fn test_unparseable_and_unresolvable_are_null() { + let input: ArrayRef = Arc::new(StringArray::from(vec![ + Some("not-an-ip"), + // Reserved TEST-NET-1 address (RFC 5737), never has a PTR record + Some("192.0.2.1"), + None, + ])); + let out = resolve_array(&input).unwrap(); + let out = out.as_any().downcast_ref::().unwrap(); + assert!(out.is_null(0), "unparseable input should be null"); + assert!(out.is_null(1), "unresolvable address should be null"); + assert!(out.is_null(2), "null input should be null"); + } + + #[test] + fn test_repeated_address_deduped() { + let input: ArrayRef = Arc::new(StringArray::from(vec![ + "127.0.0.1", + "127.0.0.1", + "127.0.0.1", + ])); + let out = resolve_array(&input).unwrap(); + let out = out.as_any().downcast_ref::().unwrap(); + assert_eq!(out.len(), 3); + assert_eq!(out.value(0), out.value(1)); + assert_eq!(out.value(1), out.value(2)); + } +} diff --git a/crates/datafusion-net/src/wide.rs b/crates/datafusion-net/src/wide.rs new file mode 100644 index 0000000..1a8412b --- /dev/null +++ b/crates/datafusion-net/src/wide.rs @@ -0,0 +1,264 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Wide variants of the `pcap` and `capture` table functions that append +//! DNS and geolocation enrichment columns for the source and destination +//! addresses. +//! +//! - [`PcapWideFunc`] (`pcap_wide`): `pcap` plus enrichment columns +//! - [`CaptureWideFunc`] (`capture_wide`, requires the `live` feature): +//! `capture` plus enrichment columns +//! +//! Both take the same arguments as their narrow counterparts and add, after +//! the packet columns: `src_host` / `dst_host` (reverse DNS, see +//! [`ReverseDnsUdf`]) and `src_country` / `src_city` / `src_lat` / `src_lon` +//! plus the `dst_` equivalents (from [`GeoIpUdf`] and a MaxMind-format +//! database): +//! +//! ```sql +//! SELECT dst_ip, dst_host, dst_country, count(*) AS packets +//! FROM pcap_wide('capture.pcap') +//! GROUP BY dst_ip, dst_host, dst_country +//! ORDER BY packets DESC; +//! ``` +//! +//! The geolocation database comes from the `GEOIP_DB` environment variable +//! ([`crate::GEOIP_DB_ENV_VAR`], read when the function is constructed) or a +//! path passed to `new`; a database that is unconfigured, missing, or +//! unreadable never fails the query — the geolocation columns are just NULL +//! (`geoip(ip, path)['error']` explains why). Enrichment is a projection +//! over the narrow table, so unused columns still prune and only projected +//! enrichment is computed. + +use std::{path::PathBuf, sync::Arc}; + +use datafusion::{ + catalog::{TableFunctionImpl, TableProvider}, + common::Result, + datasource::{provider_as_source, ViewTable}, + functions::core::expr_ext::FieldAccessor, + logical_expr::{col, Expr, LogicalPlanBuilder, ScalarUDF}, + scalar::ScalarValue, +}; + +#[cfg(feature = "live")] +use crate::live::parse_capture_args; +#[cfg(feature = "live")] +use crate::CaptureTable; +use crate::{ + file::parse_pcap_args, geoip::GEOIP_DB_ENV_VAR, schema::packet_schema, GeoIpUdf, PcapTable, + ReverseDnsUdf, +}; + +/// The geolocation database path used by [`Default`] constructions of the +/// wide table functions +fn geoip_db_from_env() -> Option { + std::env::var_os(GEOIP_DB_ENV_VAR).map(PathBuf::from) +} + +/// Wraps `inner` (a table with the packet schema) in a view that appends the +/// enrichment columns. Without a geolocation database the geolocation +/// columns are typed NULL literals, so the wide schema is stable either way. +fn enriched_view( + inner: Arc, + geoip_db_path: Option<&PathBuf>, +) -> Result> { + let mut exprs: Vec = packet_schema() + .fields() + .iter() + .map(|f| col(f.name().as_str())) + .collect(); + + let reverse_dns = Arc::new(ScalarUDF::from(ReverseDnsUdf::default())); + exprs.push(reverse_dns.call(vec![col("src_ip")]).alias("src_host")); + exprs.push(reverse_dns.call(vec![col("dst_ip")]).alias("dst_host")); + + let geoip = geoip_db_path.map(|path| Arc::new(ScalarUDF::from(GeoIpUdf::with_db_path(path)))); + for (ip_col, prefix) in [("src_ip", "src"), ("dst_ip", "dst")] { + match &geoip { + Some(geoip) => { + let geo = geoip.call(vec![col(ip_col)]); + exprs.push( + geo.clone() + .field("country_code") + .alias(format!("{prefix}_country")), + ); + exprs.push(geo.clone().field("city").alias(format!("{prefix}_city"))); + exprs.push(geo.clone().field("latitude").alias(format!("{prefix}_lat"))); + exprs.push(geo.field("longitude").alias(format!("{prefix}_lon"))); + } + None => { + exprs.push( + Expr::Literal(ScalarValue::Utf8(None), None).alias(format!("{prefix}_country")), + ); + exprs.push( + Expr::Literal(ScalarValue::Utf8(None), None).alias(format!("{prefix}_city")), + ); + exprs.push( + Expr::Literal(ScalarValue::Float64(None), None).alias(format!("{prefix}_lat")), + ); + exprs.push( + Expr::Literal(ScalarValue::Float64(None), None).alias(format!("{prefix}_lon")), + ); + } + } + } + + let plan = LogicalPlanBuilder::scan("packets", provider_as_source(inner), None)? + .project(exprs)? + .build()?; + Ok(Arc::new(ViewTable::new(plan, None))) +} + +/// Table function that reads a pcap/pcapng capture file with DNS and +/// geolocation enrichment columns +#[derive(Debug)] +pub struct PcapWideFunc { + geoip_db_path: Option, +} + +impl PcapWideFunc { + /// Creates the function with `geoip_db_path` as the geolocation + /// database; `None` yields NULL geolocation columns + pub fn new(geoip_db_path: Option) -> Self { + Self { geoip_db_path } + } +} + +impl Default for PcapWideFunc { + fn default() -> Self { + Self::new(geoip_db_from_env()) + } +} + +impl TableFunctionImpl for PcapWideFunc { + fn call(&self, exprs: &[Expr]) -> Result> { + let path = parse_pcap_args("pcap_wide", exprs)?; + enriched_view(Arc::new(PcapTable::new(path)), self.geoip_db_path.as_ref()) + } +} + +/// Table function that live-captures packets from a network interface with +/// DNS and geolocation enrichment columns +#[cfg(feature = "live")] +#[derive(Debug)] +pub struct CaptureWideFunc { + geoip_db_path: Option, +} + +#[cfg(feature = "live")] +impl CaptureWideFunc { + /// Creates the function with `geoip_db_path` as the geolocation + /// database; `None` yields NULL geolocation columns + pub fn new(geoip_db_path: Option) -> Self { + Self { geoip_db_path } + } +} + +#[cfg(feature = "live")] +impl Default for CaptureWideFunc { + fn default() -> Self { + Self::new(geoip_db_from_env()) + } +} + +#[cfg(feature = "live")] +impl TableFunctionImpl for CaptureWideFunc { + fn call(&self, exprs: &[Expr]) -> Result> { + let (interface, filter, duration) = parse_capture_args("capture_wide", exprs)?; + enriched_view( + Arc::new(CaptureTable::new(interface, filter, duration)), + self.geoip_db_path.as_ref(), + ) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use datafusion::arrow::datatypes::DataType; + + fn wide_provider(geoip_db_path: Option) -> Arc { + let func = PcapWideFunc::new(geoip_db_path); + func.call(&[Expr::Literal( + ScalarValue::Utf8(Some("test.pcap".to_string())), + None, + )]) + .unwrap() + } + + #[test] + fn test_wide_schema_extends_packet_schema() { + let provider = wide_provider(None); + let schema = provider.schema(); + // All packet columns come first, unchanged + for (i, field) in packet_schema().fields().iter().enumerate() { + assert_eq!(schema.field(i).name(), field.name()); + assert_eq!(schema.field(i).data_type(), field.data_type()); + } + // Followed by the enrichment columns + let names: Vec<_> = schema + .fields() + .iter() + .skip(packet_schema().fields().len()) + .map(|f| f.name().as_str()) + .collect(); + assert_eq!( + names, + [ + "src_host", + "dst_host", + "src_country", + "src_city", + "src_lat", + "src_lon", + "dst_country", + "dst_city", + "dst_lat", + "dst_lon" + ] + ); + } + + #[test] + fn test_wide_schema_stable_with_and_without_geoip_db() { + // The geolocation columns keep the same types whether they come from + // geoip or from NULL literals + let with_db = wide_provider(Some(PathBuf::from("some.mmdb"))).schema(); + let without_db = wide_provider(None).schema(); + assert_eq!(with_db.fields().len(), without_db.fields().len()); + for (a, b) in with_db.fields().iter().zip(without_db.fields().iter()) { + assert_eq!(a.name(), b.name()); + assert_eq!(a.data_type(), b.data_type()); + } + assert_eq!( + with_db.field_with_name("src_country").unwrap().data_type(), + &DataType::Utf8 + ); + assert_eq!( + with_db.field_with_name("src_lat").unwrap().data_type(), + &DataType::Float64 + ); + } + + #[test] + fn test_wide_rejects_bad_arguments() { + let func = PcapWideFunc::new(None); + let err = func.call(&[]).unwrap_err(); + assert!(err.to_string().contains("pcap_wide"), "got: {err}"); + } +} diff --git a/crates/datafusion-net/src/writer.rs b/crates/datafusion-net/src/writer.rs new file mode 100644 index 0000000..5aa67bb --- /dev/null +++ b/crates/datafusion-net/src/writer.rs @@ -0,0 +1,65 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Minimal legacy pcap (`.pcap`) file writer, primarily used to generate +//! test fixtures + +use std::io::{self, Write}; + +const MAGIC_MICROSECONDS: u32 = 0xa1b2_c3d4; +const VERSION_MAJOR: u16 = 2; +const VERSION_MINOR: u16 = 4; +const SNAPLEN: u32 = 65535; + +/// Writes frames to `writer` in the legacy pcap format with microsecond +/// timestamp precision +pub struct PcapWriter { + writer: W, +} + +impl PcapWriter { + /// Creates a writer and writes the pcap global header. `link_type` is a + /// value from the pcap link type registry (see [`crate::decode::link_type`]) + pub fn new(mut writer: W, link_type: u32) -> io::Result { + writer.write_all(&MAGIC_MICROSECONDS.to_le_bytes())?; + writer.write_all(&VERSION_MAJOR.to_le_bytes())?; + writer.write_all(&VERSION_MINOR.to_le_bytes())?; + // thiszone and sigfigs, both unused + writer.write_all(&0i32.to_le_bytes())?; + writer.write_all(&0u32.to_le_bytes())?; + writer.write_all(&SNAPLEN.to_le_bytes())?; + writer.write_all(&link_type.to_le_bytes())?; + Ok(Self { writer }) + } + + /// Writes a single un-truncated frame with the given capture timestamp + pub fn write_packet(&mut self, ts_micros: i64, data: &[u8]) -> io::Result<()> { + let ts_sec = (ts_micros / 1_000_000) as u32; + let ts_usec = (ts_micros % 1_000_000) as u32; + let len = data.len() as u32; + self.writer.write_all(&ts_sec.to_le_bytes())?; + self.writer.write_all(&ts_usec.to_le_bytes())?; + self.writer.write_all(&len.to_le_bytes())?; + self.writer.write_all(&len.to_le_bytes())?; + self.writer.write_all(data)?; + Ok(()) + } + + pub fn flush(&mut self) -> io::Result<()> { + self.writer.flush() + } +} diff --git a/crates/datafusion-net/tests/common/mod.rs b/crates/datafusion-net/tests/common/mod.rs new file mode 100644 index 0000000..06a4363 --- /dev/null +++ b/crates/datafusion-net/tests/common/mod.rs @@ -0,0 +1,169 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Shared test support: a minimal MaxMind DB (`.mmdb`) file written from +//! scratch so geolocation tests are offline and carry no data licensing +//! baggage. The database maps the single network `1.1.1.0/24` to a Sydney, +//! Australia city record. +//! +//! The format () is a binary search +//! tree over address bits, a 16-zero-byte separator, a data section, and +//! metadata after a marker. Field encoding: a control byte carries the type +//! in the top 3 bits (0 = extended: real type is 7 + the following byte) and +//! the payload size in the low 5 bits. + +use std::path::Path; + +/// One tree node per bit of the 1.1.1.0/24 prefix +const NODE_COUNT: u32 = 24; +/// The network encoded in the search tree +const PREFIX: [u8; 3] = [1, 1, 1]; + +fn control(type_num: u8, size: usize) -> u8 { + assert!(size < 29, "sizes needing extra length bytes not supported"); + (type_num << 5) | size as u8 +} + +fn write_string(out: &mut Vec, s: &str) { + out.push(control(2, s.len())); + out.extend_from_slice(s.as_bytes()); +} + +fn write_double(out: &mut Vec, value: f64) { + out.push(control(3, 8)); + out.extend_from_slice(&value.to_be_bytes()); +} + +fn write_map_header(out: &mut Vec, entries: usize) { + out.push(control(7, entries)); +} + +/// Writes a uint16 (type 5) or uint32 (type 6) with minimal-length encoding +fn write_uint(out: &mut Vec, type_num: u8, value: u32) { + let bytes = value.to_be_bytes(); + let skip = bytes.iter().take_while(|b| **b == 0).count(); + out.push(control(type_num, bytes.len() - skip)); + out.extend_from_slice(&bytes[skip..]); +} + +/// Writes a uint64 (extended type 9) +fn write_uint64(out: &mut Vec, value: u64) { + let bytes = value.to_be_bytes(); + let skip = bytes.iter().take_while(|b| **b == 0).count(); + out.push(control(0, bytes.len() - skip)); + out.push(9 - 7); + out.extend_from_slice(&bytes[skip..]); +} + +/// Writes an array (extended type 11) header +fn write_array_header(out: &mut Vec, entries: usize) { + out.push(control(0, entries)); + out.push(11 - 7); +} + +/// The single city record, at offset 0 of the data section +fn city_record() -> Vec { + let mut d = Vec::new(); + write_map_header(&mut d, 3); + write_string(&mut d, "city"); + write_map_header(&mut d, 1); + write_string(&mut d, "names"); + write_map_header(&mut d, 1); + write_string(&mut d, "en"); + write_string(&mut d, "Sydney"); + write_string(&mut d, "country"); + write_map_header(&mut d, 2); + write_string(&mut d, "iso_code"); + write_string(&mut d, "AU"); + write_string(&mut d, "names"); + write_map_header(&mut d, 1); + write_string(&mut d, "en"); + write_string(&mut d, "Australia"); + write_string(&mut d, "location"); + write_map_header(&mut d, 3); + write_string(&mut d, "latitude"); + write_double(&mut d, -33.8688); + write_string(&mut d, "longitude"); + write_double(&mut d, 151.2093); + write_string(&mut d, "time_zone"); + write_string(&mut d, "Australia/Sydney"); + d +} + +/// A chain of 24 nodes routing the bits of [`PREFIX`]. At each node the +/// on-prefix branch continues to the next node (the final one points at data +/// offset 0, encoded as `NODE_COUNT + 16`); the off-prefix branch is the +/// "no data" sentinel `NODE_COUNT`. Record size 24 bits: each node is two +/// 3-byte big-endian records. +fn search_tree() -> Vec { + let mut tree = Vec::new(); + for i in 0..NODE_COUNT as usize { + let bit = (PREFIX[i / 8] >> (7 - i % 8)) & 1; + let matched = if i == NODE_COUNT as usize - 1 { + NODE_COUNT + 16 + } else { + i as u32 + 1 + }; + let (left, right) = if bit == 0 { + (matched, NODE_COUNT) + } else { + (NODE_COUNT, matched) + }; + tree.extend_from_slice(&left.to_be_bytes()[1..]); + tree.extend_from_slice(&right.to_be_bytes()[1..]); + } + tree +} + +fn metadata() -> Vec { + let mut m = Vec::new(); + write_map_header(&mut m, 9); + write_string(&mut m, "binary_format_major_version"); + write_uint(&mut m, 5, 2); + write_string(&mut m, "binary_format_minor_version"); + write_uint(&mut m, 5, 0); + write_string(&mut m, "build_epoch"); + write_uint64(&mut m, 1_735_689_600); // 2025-01-01T00:00:00Z + write_string(&mut m, "database_type"); + write_string(&mut m, "GeoLite2-City"); + write_string(&mut m, "description"); + write_map_header(&mut m, 1); + write_string(&mut m, "en"); + write_string(&mut m, "Test database"); + write_string(&mut m, "ip_version"); + write_uint(&mut m, 5, 4); + write_string(&mut m, "languages"); + write_array_header(&mut m, 1); + write_string(&mut m, "en"); + write_string(&mut m, "node_count"); + write_uint(&mut m, 6, NODE_COUNT); + write_string(&mut m, "record_size"); + write_uint(&mut m, 5, 24); + m +} + +/// Writes the test database to `path`. `1.1.1.0/24` resolves to Sydney, +/// Australia (`AU`, -33.8688, 151.2093, `Australia/Sydney`); every other +/// address is absent. +pub fn write_test_mmdb(path: &Path) { + let mut buf = search_tree(); + buf.extend_from_slice(&[0u8; 16]); // data section separator + buf.extend_from_slice(&city_record()); + buf.extend_from_slice(b"\xab\xcd\xefMaxMind.com"); // metadata marker + buf.extend_from_slice(&metadata()); + std::fs::write(path, buf).unwrap(); +} diff --git a/crates/datafusion-net/tests/conversations.rs b/crates/datafusion-net/tests/conversations.rs new file mode 100644 index 0000000..a53d204 --- /dev/null +++ b/crates/datafusion-net/tests/conversations.rs @@ -0,0 +1,223 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! End-to-end tests for the `tcp_conversations` table function: a full TCP +//! connection (handshake, data with one retransmission, orderly close) is +//! written with [`datafusion_net::writer`] and aggregated back into a +//! conversation row + +use std::{path::Path, sync::Arc}; + +use datafusion::prelude::SessionContext; +use datafusion_net::{writer::PcapWriter, TcpConversationsFunc}; +use etherparse::PacketBuilder; + +const CLIENT_MAC: [u8; 6] = [0x02, 0x00, 0x00, 0x00, 0x00, 0x01]; +const SERVER_MAC: [u8; 6] = [0x02, 0x00, 0x00, 0x00, 0x00, 0x02]; +const CLIENT_IP: [u8; 4] = [10, 0, 0, 1]; +const SERVER_IP: [u8; 4] = [10, 0, 0, 2]; +const CLIENT_PORT: u16 = 51000; +const SERVER_PORT: u16 = 443; +/// 2025-01-01T00:00:00Z +const BASE_TS_MICROS: i64 = 1_735_689_600_000_000; +const MS: i64 = 1_000; + +struct Writer { + inner: PcapWriter, +} + +impl Writer { + fn new(path: &Path) -> Self { + let file = std::fs::File::create(path).unwrap(); + Self { + inner: PcapWriter::new(file, 1).unwrap(), + } + } + + fn client_packet( + &mut self, + at_ms: i64, + seq: u32, + ack: Option, + syn: bool, + fin: bool, + payload: &[u8], + ) { + self.packet(true, at_ms, seq, ack, syn, fin, payload); + } + + fn server_packet( + &mut self, + at_ms: i64, + seq: u32, + ack: Option, + syn: bool, + fin: bool, + payload: &[u8], + ) { + self.packet(false, at_ms, seq, ack, syn, fin, payload); + } + + #[allow(clippy::too_many_arguments)] + fn packet( + &mut self, + from_client: bool, + at_ms: i64, + seq: u32, + ack: Option, + syn: bool, + fin: bool, + payload: &[u8], + ) { + let (src_mac, dst_mac, src_ip, dst_ip, src_port, dst_port) = if from_client { + ( + CLIENT_MAC, + SERVER_MAC, + CLIENT_IP, + SERVER_IP, + CLIENT_PORT, + SERVER_PORT, + ) + } else { + ( + SERVER_MAC, + CLIENT_MAC, + SERVER_IP, + CLIENT_IP, + SERVER_PORT, + CLIENT_PORT, + ) + }; + let mut builder = PacketBuilder::ethernet2(src_mac, dst_mac) + .ipv4(src_ip, dst_ip, 64) + .tcp(src_port, dst_port, seq, 1024); + if syn { + builder = builder.syn(); + } + if fin { + builder = builder.fin(); + } + if let Some(ack) = ack { + builder = builder.ack(ack); + } + let mut frame = Vec::with_capacity(builder.size(payload.len())); + builder.write(&mut frame, payload).unwrap(); + self.inner + .write_packet(BASE_TS_MICROS + at_ms * MS, &frame) + .unwrap(); + } +} + +/// Writes a complete connection: handshake (100ms RTT), client request with +/// one retransmission, server response, orderly close, plus an unrelated UDP +/// packet that must be ignored +fn write_test_pcap(path: &Path) { + let mut w = Writer::new(path); + w.client_packet(0, 1000, None, true, false, &[]); // SYN + w.server_packet(50, 2000, Some(1001), true, false, &[]); // SYN|ACK + w.client_packet(100, 1001, Some(2001), false, false, &[]); // ACK + w.client_packet(200, 1001, Some(2001), false, false, b"hello"); // data + w.client_packet(300, 1001, Some(2001), false, false, b"hello"); // retransmission + w.server_packet(400, 2001, Some(1006), false, false, b"hi!"); // response + w.client_packet(500, 1006, Some(2004), false, true, &[]); // FIN + w.server_packet(600, 2004, Some(1007), false, true, &[]); // FIN + + let udp = PacketBuilder::ethernet2(CLIENT_MAC, SERVER_MAC) + .ipv4(CLIENT_IP, [8, 8, 8, 8], 64) + .udp(53001, 53); + let mut frame = Vec::with_capacity(udp.size(3)); + udp.write(&mut frame, b"dns").unwrap(); + w.inner + .write_packet(BASE_TS_MICROS + 700 * MS, &frame) + .unwrap(); + w.inner.flush().unwrap(); +} + +async fn run(ctx: &SessionContext, sql: &str) -> String { + let batches = ctx.sql(sql).await.unwrap().collect().await.unwrap(); + datafusion::arrow::util::pretty::pretty_format_batches(&batches) + .unwrap() + .to_string() +} + +fn setup() -> (SessionContext, tempfile::TempDir, String) { + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().join("test.pcap"); + write_test_pcap(&path); + let ctx = SessionContext::new(); + ctx.register_udtf( + "tcp_conversations", + Arc::new(TcpConversationsFunc::default()), + ); + let path = path.display().to_string(); + (ctx, dir, path) +} + +#[tokio::test] +async fn test_conversation_aggregation() { + let (ctx, _dir, path) = setup(); + let sql = format!( + "SELECT src_ip, src_port, dst_ip, dst_port, packets_fwd, packets_rev, \ + retransmissions_fwd, retransmissions_rev, handshake_rtt_ms, \ + duration_ms, state \ + FROM tcp_conversations('{path}')" + ); + let out = run(&ctx, &sql).await; + let expected = "\ ++----------+----------+----------+----------+-------------+-------------+---------------------+---------------------+------------------+-------------+--------+ +| src_ip | src_port | dst_ip | dst_port | packets_fwd | packets_rev | retransmissions_fwd | retransmissions_rev | handshake_rtt_ms | duration_ms | state | ++----------+----------+----------+----------+-------------+-------------+---------------------+---------------------+------------------+-------------+--------+ +| 10.0.0.1 | 51000 | 10.0.0.2 | 443 | 5 | 3 | 1 | 0 | 100.0 | 600.0 | closed | ++----------+----------+----------+----------+-------------+-------------+---------------------+---------------------+------------------+-------------+--------+"; + assert_eq!(out.trim(), expected.trim(), "got:\n{out}"); +} + +#[tokio::test] +async fn test_conversation_bytes_and_timestamps() { + let (ctx, _dir, path) = setup(); + let sql = format!( + "SELECT count(*) AS conversations, \ + sum(packets_fwd + packets_rev) AS packets, \ + min(first_timestamp) = arrow_cast('2025-01-01T00:00:00Z', 'Timestamp(Microsecond, None)') AS starts_at_base, \ + sum(bytes_fwd) > sum(bytes_rev) AS fwd_heavier \ + FROM tcp_conversations('{path}')" + ); + let out = run(&ctx, &sql).await; + assert!(out.contains("| 1 "), "expected one conversation in:\n{out}"); + assert!(out.contains("| 8 "), "expected eight packets in:\n{out}"); + assert_eq!( + out.matches("true").count(), + 2, + "expected timestamp and byte checks true in:\n{out}" + ); +} + +#[tokio::test] +async fn test_conversation_missing_file_errors() { + let (ctx, _dir, _path) = setup(); + let result = ctx + .sql("SELECT * FROM tcp_conversations('/definitely/missing.pcap')") + .await + .unwrap() + .collect() + .await; + let err = result.unwrap_err().to_string(); + assert!( + err.contains("tcp_conversations failed to open"), + "unexpected error: {err}" + ); +} diff --git a/crates/datafusion-net/tests/geoip.rs b/crates/datafusion-net/tests/geoip.rs new file mode 100644 index 0000000..65d8b41 --- /dev/null +++ b/crates/datafusion-net/tests/geoip.rs @@ -0,0 +1,177 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Tests for the `geoip` scalar UDF, run against the minimal MaxMind DB +//! (`.mmdb`) file written by the shared [`common`] test support (the single +//! network `1.1.1.0/24` maps to a Sydney, Australia city record) + +use datafusion::{logical_expr::ScalarUDF, prelude::SessionContext}; +use datafusion_net::GeoIpUdf; + +mod common; +use common::write_test_mmdb; + +/// Returns a context with `geoip` registered plus the tempdir holding the +/// test database and the database path +async fn setup() -> (SessionContext, tempfile::TempDir, String) { + let dir = tempfile::tempdir().unwrap(); + let db = dir.path().join("test.mmdb"); + write_test_mmdb(&db); + let ctx = SessionContext::new(); + ctx.register_udf(ScalarUDF::from(GeoIpUdf::default())); + let db = db.display().to_string(); + (ctx, dir, db) +} + +async fn run(ctx: &SessionContext, sql: &str) -> String { + let batches = ctx.sql(sql).await.unwrap().collect().await.unwrap(); + datafusion::arrow::util::pretty::pretty_format_batches(&batches) + .unwrap() + .to_string() +} + +#[tokio::test] +async fn test_geoip_hit() { + let (ctx, _dir, db) = setup().await; + let sql = format!( + "SELECT geoip('1.1.1.1', '{db}')['country_code'] AS country_code, \ + geoip('1.1.1.1', '{db}')['country'] AS country, \ + geoip('1.1.1.1', '{db}')['city'] AS city, \ + geoip('1.1.1.1', '{db}')['latitude'] AS latitude, \ + geoip('1.1.1.1', '{db}')['longitude'] AS longitude, \ + geoip('1.1.1.1', '{db}')['time_zone'] AS time_zone" + ); + let out = run(&ctx, &sql).await; + assert!(out.contains("AU"), "missing country code in:\n{out}"); + assert!(out.contains("Australia"), "missing country in:\n{out}"); + assert!(out.contains("Sydney"), "missing city in:\n{out}"); + assert!(out.contains("-33.8688"), "missing latitude in:\n{out}"); + assert!(out.contains("151.2093"), "missing longitude in:\n{out}"); + assert!( + out.contains("Australia/Sydney"), + "missing time zone in:\n{out}" + ); + + // The error field is NULL on a successful lookup + let sql = format!("SELECT geoip('1.1.1.1', '{db}')['error'] IS NULL AS no_error"); + let out = run(&ctx, &sql).await; + assert!(out.contains("true"), "expected NULL error field in:\n{out}"); +} + +#[tokio::test] +async fn test_geoip_whole_network_matches() { + let (ctx, _dir, db) = setup().await; + // Any address in the encoded 1.1.1.0/24 network hits the same record + let sql = format!("SELECT geoip('1.1.1.42', '{db}')['city'] = 'Sydney' AS same_record"); + let out = run(&ctx, &sql).await; + assert!(out.contains("true"), "expected a hit in:\n{out}"); +} + +#[tokio::test] +async fn test_geoip_misses_are_null() { + let (ctx, _dir, db) = setup().await; + // An address outside the network, an unparseable address, a NULL + // address, and an IPv6 address in an IPv4-only database are all NULL + let sql = format!( + "SELECT geoip('9.9.9.9', '{db}') IS NULL AS missing, \ + geoip('not-an-ip', '{db}') IS NULL AS unparseable, \ + geoip(NULL, '{db}') IS NULL AS null_input, \ + geoip('::1', '{db}') IS NULL AS wrong_ip_version" + ); + let out = run(&ctx, &sql).await; + // Four columns, all true + assert_eq!( + out.matches("true").count(), + 4, + "expected all NULL in:\n{out}" + ); +} + +#[tokio::test] +async fn test_geoip_invalid_ip_does_not_open_db() { + let (ctx, _dir, _db) = setup().await; + // The address is parsed before the database path is resolved, so an + // unparseable address is NULL even when the database does not exist + let out = run( + &ctx, + "SELECT geoip('not-an-ip', '/definitely/missing.mmdb') IS NULL AS missing", + ) + .await; + assert!(out.contains("true"), "expected NULL in:\n{out}"); +} + +#[tokio::test] +async fn test_geoip_unopenable_db_reported_in_error_field() { + let (ctx, _dir, _db) = setup().await; + // A database that cannot be opened does not fail the query: the + // location fields are NULL and the error field carries the reason + let out = run( + &ctx, + "SELECT geoip('1.1.1.1', '/definitely/missing.mmdb')['country_code'] IS NULL AS no_country, \ + geoip('1.1.1.1', '/definitely/missing.mmdb')['error'] AS error", + ) + .await; + assert!(out.contains("true"), "expected NULL country in:\n{out}"); + assert!( + out.contains("geoip failed to open database"), + "expected open failure in error field in:\n{out}" + ); + assert!( + out.contains("missing.mmdb"), + "expected the path in the error field in:\n{out}" + ); +} + +#[tokio::test] +async fn test_geoip_default_db_from_constructor_and_missing_config_errors() { + let dir = tempfile::tempdir().unwrap(); + let db = dir.path().join("test.mmdb"); + write_test_mmdb(&db); + + // With a default database the single-argument form works + let ctx = SessionContext::new(); + ctx.register_udf(ScalarUDF::from(GeoIpUdf::with_db_path(&db))); + let out = run(&ctx, "SELECT geoip('1.1.1.1')['city'] AS city").await; + assert!(out.contains("Sydney"), "expected a hit in:\n{out}"); + + // Without one (and without the GEOIP_DB environment variable, which the + // test environment must not set) the query still succeeds and the error + // field points at the fix + if std::env::var_os(datafusion_net::GEOIP_DB_ENV_VAR).is_none() { + let ctx = SessionContext::new(); + ctx.register_udf(ScalarUDF::from(GeoIpUdf::default())); + let out = run(&ctx, "SELECT geoip('1.1.1.1')['error'] AS error").await; + assert!( + out.contains("geoip has no database configured"), + "expected configuration hint in error field in:\n{out}" + ); + } +} + +#[tokio::test] +async fn test_geoip_over_column() { + let (ctx, _dir, db) = setup().await; + // Exercise the array (non-literal) code path over a column of addresses + let sql = format!( + "SELECT ip, geoip(ip, '{db}')['country_code'] AS cc \ + FROM (VALUES ('1.1.1.1'), ('1.1.1.200'), ('8.8.8.8'), ('nope')) AS t(ip) \ + ORDER BY ip" + ); + let out = run(&ctx, &sql).await; + let hits = out.matches("AU").count(); + assert_eq!(hits, 2, "expected two hits in:\n{out}"); +} diff --git a/crates/datafusion-net/tests/roundtrip.rs b/crates/datafusion-net/tests/roundtrip.rs new file mode 100644 index 0000000..3ebab0d --- /dev/null +++ b/crates/datafusion-net/tests/roundtrip.rs @@ -0,0 +1,156 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Round-trip tests: write a capture file with [`datafusion_net::writer`] +//! and query it back through the `pcap` table function + +use std::{path::Path, sync::Arc}; + +use datafusion::{arrow::util::pretty::pretty_format_batches, prelude::SessionContext}; +use datafusion_net::{writer::PcapWriter, PcapFunc}; +use etherparse::PacketBuilder; + +const SRC_MAC: [u8; 6] = [0x02, 0x00, 0x00, 0x00, 0x00, 0x01]; +const DST_MAC: [u8; 6] = [0x02, 0x00, 0x00, 0x00, 0x00, 0x02]; +/// 2025-01-01T00:00:00Z +const BASE_TS_MICROS: i64 = 1_735_689_600_000_000; + +/// Writes a small capture with two TCP packets and one UDP packet +fn write_test_pcap(path: &Path) { + let file = std::fs::File::create(path).unwrap(); + let mut writer = PcapWriter::new(file, 1).unwrap(); + + let builder = PacketBuilder::ethernet2(SRC_MAC, DST_MAC) + .ipv4([10, 0, 0, 1], [10, 0, 0, 2], 64) + .tcp(51000, 443, 1000, 1024) + .syn(); + let mut frame = Vec::with_capacity(builder.size(0)); + builder.write(&mut frame, &[]).unwrap(); + writer.write_packet(BASE_TS_MICROS, &frame).unwrap(); + + let builder = PacketBuilder::ethernet2(DST_MAC, SRC_MAC) + .ipv4([10, 0, 0, 2], [10, 0, 0, 1], 64) + .tcp(443, 51000, 2000, 1024) + .syn() + .ack(1001); + let mut frame = Vec::with_capacity(builder.size(0)); + builder.write(&mut frame, &[]).unwrap(); + writer.write_packet(BASE_TS_MICROS + 500, &frame).unwrap(); + + let builder = PacketBuilder::ethernet2(SRC_MAC, DST_MAC) + .ipv4([10, 0, 0, 1], [8, 8, 8, 8], 64) + .udp(53001, 53); + let payload = b"dns query"; + let mut frame = Vec::with_capacity(builder.size(payload.len())); + builder.write(&mut frame, payload).unwrap(); + writer.write_packet(BASE_TS_MICROS + 1000, &frame).unwrap(); + + writer.flush().unwrap(); +} + +async fn ctx_with_pcap() -> (SessionContext, tempfile::TempDir) { + let dir = tempfile::tempdir().unwrap(); + write_test_pcap(&dir.path().join("test.pcap")); + let ctx = SessionContext::new(); + ctx.register_udtf("pcap", Arc::new(PcapFunc::default())); + (ctx, dir) +} + +async fn run(ctx: &SessionContext, sql: &str) -> String { + let batches = ctx.sql(sql).await.unwrap().collect().await.unwrap(); + pretty_format_batches(&batches).unwrap().to_string() +} + +#[tokio::test] +async fn test_select_columns() { + let (ctx, dir) = ctx_with_pcap().await; + let sql = format!( + "SELECT frame_number, src_ip, dst_ip, protocol, src_port, dst_port, tcp_flags \ + FROM pcap('{}/test.pcap')", + dir.path().display() + ); + let output = run(&ctx, &sql).await; + let expected = "\ ++--------------+----------+----------+----------+----------+----------+-----------+ +| frame_number | src_ip | dst_ip | protocol | src_port | dst_port | tcp_flags | ++--------------+----------+----------+----------+----------+----------+-----------+ +| 1 | 10.0.0.1 | 10.0.0.2 | tcp | 51000 | 443 | SYN | +| 2 | 10.0.0.2 | 10.0.0.1 | tcp | 443 | 51000 | SYN|ACK | +| 3 | 10.0.0.1 | 8.8.8.8 | udp | 53001 | 53 | | ++--------------+----------+----------+----------+----------+----------+-----------+"; + assert_eq!(output, expected); +} + +#[tokio::test] +async fn test_filter_and_aggregate() { + let (ctx, dir) = ctx_with_pcap().await; + let sql = format!( + "SELECT protocol, count(*) AS packets FROM pcap('{}/test.pcap') \ + GROUP BY protocol ORDER BY protocol", + dir.path().display() + ); + let output = run(&ctx, &sql).await; + let expected = "\ ++----------+---------+ +| protocol | packets | ++----------+---------+ +| tcp | 2 | +| udp | 1 | ++----------+---------+"; + assert_eq!(output, expected); +} + +#[tokio::test] +async fn test_limit() { + let (ctx, dir) = ctx_with_pcap().await; + let sql = format!( + "SELECT frame_number FROM pcap('{}/test.pcap') LIMIT 1", + dir.path().display() + ); + let output = run(&ctx, &sql).await; + assert!(output.contains("| 1 |")); + assert!(!output.contains("| 2 |")); +} + +#[tokio::test] +async fn test_payload_and_timestamp() { + let (ctx, dir) = ctx_with_pcap().await; + let sql = format!( + "SELECT timestamp, payload_length, payload FROM pcap('{}/test.pcap') \ + WHERE protocol = 'udp'", + dir.path().display() + ); + let output = run(&ctx, &sql).await; + // 'dns query' hex encoded + assert!(output.contains("646e73207175657279")); + assert!(output.contains("2025-01-01T00:00:00.001")); + assert!(output.contains("| 9 |")); +} + +#[tokio::test] +async fn test_missing_file_errors_at_execution() { + let ctx = SessionContext::new(); + ctx.register_udtf("pcap", Arc::new(PcapFunc::default())); + let err = ctx + .sql("SELECT * FROM pcap('/does/not/exist.pcap')") + .await + .unwrap() + .collect() + .await + .unwrap_err(); + assert!(err.to_string().contains("failed to open")); +} diff --git a/crates/datafusion-net/tests/wide.rs b/crates/datafusion-net/tests/wide.rs new file mode 100644 index 0000000..f0c0bb5 --- /dev/null +++ b/crates/datafusion-net/tests/wide.rs @@ -0,0 +1,157 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! End-to-end tests for the `pcap_wide` table function: a capture written +//! with [`datafusion_net::writer`] is queried back with geolocation +//! enrichment resolved against the minimal MaxMind DB written by the shared +//! [`common`] test support (`1.1.1.0/24` maps to Sydney, Australia) + +use std::{path::Path, sync::Arc}; + +use datafusion::prelude::SessionContext; +use datafusion_net::{writer::PcapWriter, PcapWideFunc}; +use etherparse::PacketBuilder; + +mod common; +use common::write_test_mmdb; + +const SRC_MAC: [u8; 6] = [0x02, 0x00, 0x00, 0x00, 0x00, 0x01]; +const DST_MAC: [u8; 6] = [0x02, 0x00, 0x00, 0x00, 0x00, 0x02]; +/// 2025-01-01T00:00:00Z +const BASE_TS_MICROS: i64 = 1_735_689_600_000_000; + +/// Writes a capture with one packet from 1.1.1.1 (in the test geolocation +/// database) to 9.9.9.9 (absent from it) +fn write_test_pcap(path: &Path) { + let file = std::fs::File::create(path).unwrap(); + let mut writer = PcapWriter::new(file, 1).unwrap(); + let builder = PacketBuilder::ethernet2(SRC_MAC, DST_MAC) + .ipv4([1, 1, 1, 1], [9, 9, 9, 9], 64) + .tcp(51000, 443, 1000, 1024) + .syn(); + let mut frame = Vec::with_capacity(builder.size(0)); + builder.write(&mut frame, &[]).unwrap(); + writer.write_packet(BASE_TS_MICROS, &frame).unwrap(); + writer.flush().unwrap(); +} + +/// Returns a context with `pcap_wide` registered (against the test +/// geolocation database when `with_geoip_db`) plus the tempdir and the +/// capture path +fn setup(with_geoip_db: bool) -> (SessionContext, tempfile::TempDir, String) { + let dir = tempfile::tempdir().unwrap(); + let pcap_path = dir.path().join("test.pcap"); + write_test_pcap(&pcap_path); + let geoip_db = with_geoip_db.then(|| { + let db = dir.path().join("test.mmdb"); + write_test_mmdb(&db); + db + }); + let ctx = SessionContext::new(); + ctx.register_udtf("pcap_wide", Arc::new(PcapWideFunc::new(geoip_db))); + let pcap_path = pcap_path.display().to_string(); + (ctx, dir, pcap_path) +} + +async fn run(ctx: &SessionContext, sql: &str) -> String { + let batches = ctx.sql(sql).await.unwrap().collect().await.unwrap(); + datafusion::arrow::util::pretty::pretty_format_batches(&batches) + .unwrap() + .to_string() +} + +#[tokio::test] +async fn test_pcap_wide_geolocates_addresses() { + let (ctx, _dir, pcap) = setup(true); + let sql = format!( + "SELECT src_ip, src_country, src_city, src_lat, src_lon, \ + dst_ip, dst_country \ + FROM pcap_wide('{pcap}')" + ); + let out = run(&ctx, &sql).await; + assert!(out.contains("1.1.1.1"), "missing src_ip in:\n{out}"); + assert!(out.contains("AU"), "missing src_country in:\n{out}"); + assert!(out.contains("Sydney"), "missing src_city in:\n{out}"); + assert!(out.contains("-33.8688"), "missing src_lat in:\n{out}"); + assert!(out.contains("151.2093"), "missing src_lon in:\n{out}"); + // 9.9.9.9 is not in the database: dst_country is null + assert!( + !out.contains("| AU | AU"), + "dst_country should be null in:\n{out}" + ); +} + +#[tokio::test] +async fn test_pcap_wide_without_geoip_db_yields_nulls() { + let (ctx, _dir, pcap) = setup(false); + let sql = format!( + "SELECT src_country IS NULL AND src_city IS NULL \ + AND src_lat IS NULL AND src_lon IS NULL \ + AND dst_country IS NULL AS geo_null \ + FROM pcap_wide('{pcap}')" + ); + let out = run(&ctx, &sql).await; + assert!(out.contains("true"), "expected null geo columns in:\n{out}"); +} + +#[tokio::test] +async fn test_pcap_wide_unopenable_db_yields_nulls() { + let dir = tempfile::tempdir().unwrap(); + let pcap_path = dir.path().join("test.pcap"); + write_test_pcap(&pcap_path); + let ctx = SessionContext::new(); + // A database path that cannot be opened must not fail the query: the + // geolocation columns are NULL (the reason is available through + // geoip(ip, path)['error']) + ctx.register_udtf( + "pcap_wide", + Arc::new(PcapWideFunc::new(Some("/definitely/missing.mmdb".into()))), + ); + let sql = format!( + "SELECT src_country IS NULL AND src_city IS NULL \ + AND src_lat IS NULL AND src_lon IS NULL AS geo_null \ + FROM pcap_wide('{}')", + pcap_path.display() + ); + let out = run(&ctx, &sql).await; + assert!(out.contains("true"), "expected null geo columns in:\n{out}"); +} + +#[tokio::test] +async fn test_pcap_wide_host_columns_present() { + let (ctx, _dir, pcap) = setup(true); + // The fixture addresses may or may not have PTR records depending on the + // environment, so only exercise that the reverse DNS columns project and + // the query executes + let sql = format!("SELECT src_ip, src_host, dst_host FROM pcap_wide('{pcap}')"); + let out = run(&ctx, &sql).await; + assert!(out.contains("src_host"), "missing column in:\n{out}"); + assert!(out.contains("1.1.1.1"), "missing row in:\n{out}"); +} + +#[tokio::test] +async fn test_pcap_wide_narrow_projection_and_aggregation() { + let (ctx, _dir, pcap) = setup(true); + // Enrichment composes with projection pruning and aggregation + let sql = format!( + "SELECT src_country, count(*) AS packets \ + FROM pcap_wide('{pcap}') WHERE dst_port = 443 GROUP BY src_country" + ); + let out = run(&ctx, &sql).await; + assert!(out.contains("AU"), "missing group in:\n{out}"); + assert!(out.contains('1'), "missing count in:\n{out}"); +} diff --git a/docs/config.md b/docs/config.md index 3d1d640..0265eef 100644 --- a/docs/config.md +++ b/docs/config.md @@ -158,6 +158,13 @@ Set the number of iterations for benchmarking queries (10 is the default). benchmark_iterations = 10 ``` +With the `net` feature enabled, the MaxMind-format (`.mmdb`) database used by the single-argument form of the `geoip` function can be configured (the `GEOIP_DB` environment variable takes precedence over this value). + +```toml +[execution.net] +geoip_db_path = "/path/to/GeoLite2-City.mmdb" +``` + The batch size for query execution can be configured based on the app being used (TUI, CLI, or FlightSQL Server). For the TUI it defaults to 100, which may slow down queries, because a Record Batch is used as a unit of pagination and too many rows can cause the TUI to hang. For the CLI and FlightSQL Server, the default is 8092. ```toml diff --git a/docs/features.md b/docs/features.md index 49400fe..9e6b97c 100644 --- a/docs/features.md +++ b/docs/features.md @@ -58,6 +58,82 @@ SELECT * FROM websocket('wss://stream.example.com/ws', '{"op":"subscribe","chann The source is unbounded: without a `LIMIT` (or an aggregation that can complete) the query streams until the server closes the connection. Note the CLI `--concat` flag collects all batches and therefore should not be used with unbounded queries. +### Net (`--features=net`) + +Adds table functions from [datafusion-net](https://github.com/datafusion-contrib/datafusion-dft/tree/main/crates/datafusion-net) for querying network packet captures with SQL, similar to wireshark / tshark. Both functions share a wireshark-style schema (`timestamp`, `src_ip`, `dst_ip`, `protocol`, `src_port`, `dst_port`, `tcp_flags`, `payload`, etc.); frames that cannot be decoded produce rows with null columns rather than errors. + +`pcap` reads a pcap/pcapng capture file as a table: + +```sql +SELECT src_ip, dst_ip, protocol, length FROM pcap('capture.pcap') WHERE dst_port = 443 +``` + +`capture` streams live-captured packets from a network interface, with an optional BPF filter and an optional duration in seconds: + +```sql +-- Stream until 100 packets have been captured +SELECT * FROM capture('en0') LIMIT 100; + +-- Capture for 10 seconds; the stream terminates, so aggregations work +SELECT src_ip, count(*) FROM capture('en0', 'tcp port 443', 10) GROUP BY src_ip; +``` + +Without a duration the live capture is unbounded: use a `LIMIT` or the query streams until cancelled. Live capture requires elevated privileges (sudo, `cap_net_raw`+`cap_net_admin` on Linux, or ChmodBPF on macOS) and links against libpcap (`libpcap-dev` on Debian/Ubuntu; included with macOS). + +`pcap_wide` and `capture_wide` take the same arguments as their narrow counterparts and append DNS and geolocation enrichment columns for the source and destination addresses: `src_host` / `dst_host` (reverse DNS) and `src_country`, `src_city`, `src_lat`, `src_lon` plus the `dst_` equivalents (resolved against the configured `geoip` database; without one the geolocation columns are `NULL`): + +```sql +SELECT dst_ip, dst_host, dst_country, count(*) AS packets +FROM pcap_wide('capture.pcap') GROUP BY dst_ip, dst_host, dst_country ORDER BY packets DESC +``` + +`interfaces` lists the system's network capture interfaces (similar to `tshark -D`) with their addresses and status flags, which is useful for discovering what to pass to `capture`. Listing does not require elevated privileges: + +```sql +SELECT name, description, addresses FROM interfaces() WHERE is_up AND NOT is_loopback +``` + +`tcp_conversations` aggregates a capture file into one row per TCP connection (like wireshark's "Statistics → Conversations" or `tshark -z conv,tcp`), with per-direction packet/byte/retransmission counts, the handshake RTT, duration, and connection state (`active`, `half_closed`, `closed`, or `reset`). Direction is from the connection initiator's perspective (`fwd` is initiator → responder): + +```sql +SELECT dst_ip, dst_port, handshake_rtt_ms, retransmissions_fwd, state +FROM tcp_conversations('capture.pcap') ORDER BY retransmissions_fwd DESC +``` + +The feature also adds a `reverse_dns` scalar function that resolves an IP address string to a hostname via reverse DNS (PTR) lookup, which pairs naturally with the `src_ip` / `dst_ip` columns: + +```sql +SELECT dst_ip, reverse_dns(dst_ip) AS host, count(*) AS packets +FROM capture('en0', 'tcp', 10) GROUP BY dst_ip, host ORDER BY packets DESC +``` + +Lookups are cached and bounded by a timeout; addresses that fail to parse, fail to resolve, or time out yield `NULL`. + +There is also a `geoip` scalar function that geolocates an IP address using a MaxMind-format (`.mmdb`) database, returning a struct with `country_code`, `country`, `city`, `latitude`, `longitude`, `time_zone`, and `error` fields: + +```sql +SELECT geoip(src_ip, '/path/GeoLite2-City.mmdb')['country_code'] AS country, count(*) AS packets +FROM pcap('capture.pcap') GROUP BY country ORDER BY packets DESC +``` + +A single database file serves every field — there is no per-field database. A City-schema database ([GeoLite2-City](https://dev.maxmind.com/geoip/geolite2-free-geolocation-data), free with a MaxMind account, or the commercial GeoIP2-City) populates all fields; a Country-schema database (GeoLite2-Country) populates only `country_code` and `country`, leaving the rest `NULL`; other schemas (ASN, ISP, ...) yield all-`NULL` fields. + +The database path can also come from the `GEOIP_DB` environment variable or the `geoip_db_path` entry in the `[execution.net]` config section (the environment variable takes precedence), in which case the second argument can be omitted. + +```toml +[execution.net] +geoip_db_path = "/path/to/GeoLite2-City.mmdb" +``` + +Addresses that fail to parse or have no entry in the database yield `NULL`. A database that is missing, unreadable, or unconfigured does not fail the query: the location fields are `NULL` and the `error` field (`NULL` on success) carries the reason, e.g. `SELECT geoip(src_ip)['error']`. The same applies to the geolocation columns of `pcap_wide` / `capture_wide`, which are `NULL` when the database is unavailable. + +Two payload-decoding scalar functions parse application-layer details out of the `payload` column. `dns_query(payload)` decodes a DNS message (typically UDP port 53) into a struct with `is_response`, `name`, `query_type`, `response_code`, and `answers` (a list of A/AAAA addresses and CNAME/NS/PTR names). `tls_sni(payload)` extracts the SNI host name from a TLS ClientHello (typically TCP port 443) — the ClientHello is unencrypted even for HTTPS, so this reveals the destination host of otherwise opaque connections. Both return `NULL` for payloads that do not parse: + +```sql +SELECT tls_sni(payload) AS host, count(*) AS hellos +FROM pcap('capture.pcap') WHERE tls_sni(payload) IS NOT NULL GROUP BY host ORDER BY hellos DESC +``` + ## External Features `dft` also has several external optional (conditionally compiled features) integrations which are controlled by [Rust Crate Features] diff --git a/tests/extension_cases/mod.rs b/tests/extension_cases/mod.rs index de07bae..538608d 100644 --- a/tests/extension_cases/mod.rs +++ b/tests/extension_cases/mod.rs @@ -33,6 +33,8 @@ mod functions_json; mod huggingface; #[cfg(feature = "mongodb")] mod mongodb; +#[cfg(feature = "net")] +mod net; #[cfg(feature = "s3")] mod s3; #[cfg(feature = "udfs-wasm")] @@ -59,8 +61,11 @@ pub struct TestExecution { #[allow(dead_code)] impl TestExecution { pub async fn new() -> Self { - let config = AppConfig::default(); + Self::new_with_config(AppConfig::default()).await + } + /// Like [`Self::new`] but with a caller-provided config + pub async fn new_with_config(config: AppConfig) -> Self { let session_state = DftSessionStateBuilder::try_new(Some(config.cli.execution.clone())) .unwrap() .with_extensions() diff --git a/tests/extension_cases/net.rs b/tests/extension_cases/net.rs new file mode 100644 index 0000000..ccca14c --- /dev/null +++ b/tests/extension_cases/net.rs @@ -0,0 +1,509 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Tests for the `pcap` and `capture` table functions + +use std::path::Path; + +use datafusion_net::writer::PcapWriter; +use etherparse::PacketBuilder; + +use crate::extension_cases::TestExecution; + +const SRC_MAC: [u8; 6] = [0x02, 0x00, 0x00, 0x00, 0x00, 0x01]; +const DST_MAC: [u8; 6] = [0x02, 0x00, 0x00, 0x00, 0x00, 0x02]; +/// 2025-01-01T00:00:00Z +const BASE_TS_MICROS: i64 = 1_735_689_600_000_000; + +/// Writes a small capture: a TCP handshake packet each way and a DNS-style +/// UDP query +fn write_test_pcap(path: &Path) { + let file = std::fs::File::create(path).unwrap(); + let mut writer = PcapWriter::new(file, 1).unwrap(); + + let builder = PacketBuilder::ethernet2(SRC_MAC, DST_MAC) + .ipv4([10, 0, 0, 1], [10, 0, 0, 2], 64) + .tcp(51000, 443, 1000, 1024) + .syn(); + let mut frame = Vec::with_capacity(builder.size(0)); + builder.write(&mut frame, &[]).unwrap(); + writer.write_packet(BASE_TS_MICROS, &frame).unwrap(); + + let builder = PacketBuilder::ethernet2(DST_MAC, SRC_MAC) + .ipv4([10, 0, 0, 2], [10, 0, 0, 1], 64) + .tcp(443, 51000, 2000, 1024) + .syn() + .ack(1001); + let mut frame = Vec::with_capacity(builder.size(0)); + builder.write(&mut frame, &[]).unwrap(); + writer.write_packet(BASE_TS_MICROS + 500, &frame).unwrap(); + + let builder = PacketBuilder::ethernet2(SRC_MAC, DST_MAC) + .ipv4([10, 0, 0, 1], [8, 8, 8, 8], 64) + .udp(53001, 53); + let payload = b"dns query"; + let mut frame = Vec::with_capacity(builder.size(payload.len())); + builder.write(&mut frame, payload).unwrap(); + writer.write_packet(BASE_TS_MICROS + 1000, &frame).unwrap(); + + writer.flush().unwrap(); +} + +#[tokio::test] +async fn test_pcap_select() { + let dir = tempfile::tempdir().unwrap(); + write_test_pcap(&dir.path().join("test.pcap")); + let execution = TestExecution::new().await; + let sql = format!( + "SELECT frame_number, src_ip, dst_ip, protocol, src_port, dst_port, tcp_flags \ + FROM pcap('{}/test.pcap')", + dir.path().display() + ); + let output = execution.run_and_format(&sql).await; + insta::assert_yaml_snapshot!(output, @r#" + - +--------------+----------+----------+----------+----------+----------+-----------+ + - "| frame_number | src_ip | dst_ip | protocol | src_port | dst_port | tcp_flags |" + - +--------------+----------+----------+----------+----------+----------+-----------+ + - "| 1 | 10.0.0.1 | 10.0.0.2 | tcp | 51000 | 443 | SYN |" + - "| 2 | 10.0.0.2 | 10.0.0.1 | tcp | 443 | 51000 | SYN|ACK |" + - "| 3 | 10.0.0.1 | 8.8.8.8 | udp | 53001 | 53 | |" + - +--------------+----------+----------+----------+----------+----------+-----------+ + "#); +} + +#[tokio::test] +async fn test_pcap_filter_and_aggregate() { + let dir = tempfile::tempdir().unwrap(); + write_test_pcap(&dir.path().join("test.pcap")); + let execution = TestExecution::new().await; + let sql = format!( + "SELECT protocol, count(*) AS packets, sum(length) AS bytes \ + FROM pcap('{}/test.pcap') WHERE dst_port = 443 GROUP BY protocol", + dir.path().display() + ); + let output = execution.run_and_format(&sql).await; + insta::assert_yaml_snapshot!(output, @r#" + - +----------+---------+-------+ + - "| protocol | packets | bytes |" + - +----------+---------+-------+ + - "| tcp | 1 | 54 |" + - +----------+---------+-------+ + "#); +} + +#[tokio::test] +async fn test_pcap_limit() { + let dir = tempfile::tempdir().unwrap(); + write_test_pcap(&dir.path().join("test.pcap")); + let execution = TestExecution::new().await; + let sql = format!( + "SELECT frame_number, protocol FROM pcap('{}/test.pcap') LIMIT 2", + dir.path().display() + ); + let output = execution.run_and_format(&sql).await; + insta::assert_yaml_snapshot!(output, @r#" + - +--------------+----------+ + - "| frame_number | protocol |" + - +--------------+----------+ + - "| 1 | tcp |" + - "| 2 | tcp |" + - +--------------+----------+ + "#); +} + +#[tokio::test] +async fn test_capture_explain_does_not_open_device() { + let execution = TestExecution::new().await; + // Planning must not open a capture device, which would fail without + // elevated privileges (and this device does not exist anyway) + let result = execution + .run("EXPLAIN SELECT * FROM capture('definitely-not-a-device', 'tcp', 5)") + .await; + assert!(result.is_ok()); +} + +#[tokio::test] +async fn test_interfaces_lists_devices() { + let execution = TestExecution::new().await; + // The device list is environment dependent, but any host running the + // tests has at least one interface + let output = execution + .run_and_format("SELECT count(*) >= 1 AS has_interfaces FROM interfaces()") + .await; + insta::assert_yaml_snapshot!(output, @r#" + - +----------------+ + - "| has_interfaces |" + - +----------------+ + - "| true |" + - +----------------+ + "#); +} + +#[tokio::test] +async fn test_show_functions_documents_udfs() { + let execution = TestExecution::new().await; + // Each scalar UDF's description and syntax come from its documentation() + // implementation, which is what `SHOW FUNCTIONS` displays. Query the + // routines table (which `SHOW FUNCTIONS` is rewritten to join against) + // for the two documented columns. + let output = execution + .run_and_format( + "SELECT routine_name, syntax_example FROM information_schema.routines \ + WHERE routine_name IN ('reverse_dns', 'geoip', 'dns_query', 'tls_sni') \ + ORDER BY routine_name", + ) + .await; + insta::assert_yaml_snapshot!(output, @r#" + - +--------------+-----------------------+ + - "| routine_name | syntax_example |" + - +--------------+-----------------------+ + - "| dns_query | dns_query(payload) |" + - "| geoip | geoip(ip [, db_path]) |" + - "| reverse_dns | reverse_dns(ip) |" + - "| tls_sni | tls_sni(payload) |" + - +--------------+-----------------------+ + "#); + + // Every one of them also carries a non-empty description + let with_descriptions = execution + .run( + "SELECT routine_name FROM information_schema.routines \ + WHERE routine_name IN ('reverse_dns', 'geoip', 'dns_query', 'tls_sni') \ + AND description IS NOT NULL AND length(description) > 0", + ) + .await + .unwrap(); + let rows: usize = with_descriptions.iter().map(|b| b.num_rows()).sum(); + assert_eq!(rows, 4, "expected all four UDFs to have descriptions"); + + // And `SHOW FUNCTIONS` itself surfaces the documentation end to end + let show = execution + .run_and_format("SHOW FUNCTIONS LIKE 'tls_sni'") + .await; + let joined = show.join("\n"); + assert!( + joined.contains("Server Name Indication"), + "SHOW FUNCTIONS should include the description:\n{joined}" + ); + assert!( + joined.contains("tls_sni(payload)"), + "SHOW FUNCTIONS should include the syntax example:\n{joined}" + ); +} + +/// A minimal DNS query for `example.com` A record +fn dns_query_bytes() -> Vec { + let mut m = vec![ + 0x12, 0x34, // id + 0x01, 0x00, // flags: RD + 0x00, 0x01, // qdcount + 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, // an/ns/ar + ]; + m.extend_from_slice(b"\x07example\x03com\x00"); + m.extend_from_slice(&[0x00, 0x01, 0x00, 0x01]); // A, IN + m +} + +/// A minimal TLS ClientHello with an SNI extension for `host` +fn client_hello_bytes(host: &str) -> Vec { + let host = host.as_bytes(); + let mut sni_ext = Vec::new(); + let entry_len = 1 + 2 + host.len(); + sni_ext.extend_from_slice(&(entry_len as u16).to_be_bytes()); + sni_ext.push(0); // host_name + sni_ext.extend_from_slice(&(host.len() as u16).to_be_bytes()); + sni_ext.extend_from_slice(host); + + let mut extensions = Vec::new(); + extensions.extend_from_slice(&0u16.to_be_bytes()); // server_name + extensions.extend_from_slice(&(sni_ext.len() as u16).to_be_bytes()); + extensions.extend_from_slice(&sni_ext); + + let mut body = Vec::new(); + body.extend_from_slice(&[0x03, 0x03]); + body.extend_from_slice(&[0u8; 32]); + body.push(0); + body.extend_from_slice(&[0x00, 0x02, 0x13, 0x01]); + body.extend_from_slice(&[0x01, 0x00]); + body.extend_from_slice(&(extensions.len() as u16).to_be_bytes()); + body.extend_from_slice(&extensions); + + let mut handshake = vec![0x01]; // ClientHello + let body_len = body.len(); + handshake.extend_from_slice(&[ + (body_len >> 16) as u8, + (body_len >> 8) as u8, + body_len as u8, + ]); + handshake.extend_from_slice(&body); + + let mut record = vec![0x16, 0x03, 0x01]; // handshake, TLS 1.0 record + record.extend_from_slice(&(handshake.len() as u16).to_be_bytes()); + record.extend_from_slice(&handshake); + record +} + +/// Writes a capture with a real DNS query (UDP:53) and a TLS ClientHello +/// (TCP:443) so the payload UDFs have something to decode +fn write_payload_pcap(path: &Path) { + let file = std::fs::File::create(path).unwrap(); + let mut writer = PcapWriter::new(file, 1).unwrap(); + + let dns = dns_query_bytes(); + let builder = PacketBuilder::ethernet2(SRC_MAC, DST_MAC) + .ipv4([10, 0, 0, 1], [8, 8, 8, 8], 64) + .udp(53001, 53); + let mut frame = Vec::with_capacity(builder.size(dns.len())); + builder.write(&mut frame, &dns).unwrap(); + writer.write_packet(BASE_TS_MICROS, &frame).unwrap(); + + let hello = client_hello_bytes("example.com"); + let builder = PacketBuilder::ethernet2(SRC_MAC, DST_MAC) + .ipv4([10, 0, 0, 1], [10, 0, 0, 2], 64) + .tcp(51000, 443, 1000, 1024) + .psh() + .ack(1); + let mut frame = Vec::with_capacity(builder.size(hello.len())); + builder.write(&mut frame, &hello).unwrap(); + writer.write_packet(BASE_TS_MICROS + 500, &frame).unwrap(); + + writer.flush().unwrap(); +} + +#[tokio::test] +async fn test_tcp_conversations() { + let dir = tempfile::tempdir().unwrap(); + write_test_pcap(&dir.path().join("test.pcap")); + let execution = TestExecution::new().await; + // The fixture has a SYN and a SYN|ACK for one flow (initiator 10.0.0.1), + // plus an unrelated UDP packet that must not appear + let sql = format!( + "SELECT src_ip, src_port, dst_ip, dst_port, packets_fwd, packets_rev, state \ + FROM tcp_conversations('{}/test.pcap')", + dir.path().display() + ); + let output = execution.run_and_format(&sql).await; + insta::assert_yaml_snapshot!(output, @r#" + - +----------+----------+----------+----------+-------------+-------------+--------+ + - "| src_ip | src_port | dst_ip | dst_port | packets_fwd | packets_rev | state |" + - +----------+----------+----------+----------+-------------+-------------+--------+ + - "| 10.0.0.1 | 51000 | 10.0.0.2 | 443 | 1 | 1 | active |" + - +----------+----------+----------+----------+-------------+-------------+--------+ + "#); +} + +#[tokio::test] +async fn test_dns_query_udf() { + let dir = tempfile::tempdir().unwrap(); + write_payload_pcap(&dir.path().join("test.pcap")); + let execution = TestExecution::new().await; + let sql = format!( + "SELECT dns_query(payload)['name'] AS name, dns_query(payload)['query_type'] AS qtype \ + FROM pcap('{}/test.pcap') WHERE dst_port = 53", + dir.path().display() + ); + let output = execution.run_and_format(&sql).await; + insta::assert_yaml_snapshot!(output, @r#" + - +-------------+-------+ + - "| name | qtype |" + - +-------------+-------+ + - "| example.com | A |" + - +-------------+-------+ + "#); +} + +#[tokio::test] +async fn test_tls_sni_udf() { + let dir = tempfile::tempdir().unwrap(); + write_payload_pcap(&dir.path().join("test.pcap")); + let execution = TestExecution::new().await; + let sql = format!( + "SELECT tls_sni(payload) AS host FROM pcap('{}/test.pcap') WHERE dst_port = 443", + dir.path().display() + ); + let output = execution.run_and_format(&sql).await; + insta::assert_yaml_snapshot!(output, @r#" + - +-------------+ + - "| host |" + - +-------------+ + - "| example.com |" + - +-------------+ + "#); +} + +#[tokio::test] +async fn test_pcap_wide_adds_enrichment_columns() { + let dir = tempfile::tempdir().unwrap(); + write_test_pcap(&dir.path().join("test.pcap")); + let execution = TestExecution::new().await; + // The fixture's private 10.0.0.0/8 addresses are absent from any + // geolocation database (and none is configured in the test environment), + // so the geo columns are NULL. This exercises that the wide schema is + // present and the query executes. + let sql = format!( + "SELECT src_ip, src_country, src_city, src_lat, src_lon \ + FROM pcap_wide('{}/test.pcap') WHERE dst_port = 443", + dir.path().display() + ); + let output = execution.run_and_format(&sql).await; + insta::assert_yaml_snapshot!(output, @r#" + - +----------+-------------+----------+---------+---------+ + - "| src_ip | src_country | src_city | src_lat | src_lon |" + - +----------+-------------+----------+---------+---------+ + - "| 10.0.0.1 | | | | |" + - +----------+-------------+----------+---------+---------+ + "#); + + // The reverse DNS columns are environment dependent; only exercise that + // they project and the query executes + let sql = format!( + "SELECT src_host, dst_host FROM pcap_wide('{}/test.pcap')", + dir.path().display() + ); + let result = execution.run(&sql).await; + assert!(result.is_ok()); +} + +#[tokio::test] +async fn test_capture_wide_explain_does_not_open_device() { + let execution = TestExecution::new().await; + // Like capture, planning capture_wide must not open a device + let result = execution + .run("EXPLAIN SELECT src_ip, src_host, src_country FROM capture_wide('definitely-not-a-device', 'tcp', 5)") + .await; + assert!(result.is_ok()); +} + +#[tokio::test] +async fn test_reverse_dns_loopback() { + let execution = TestExecution::new().await; + // The loopback address must resolve to some hostname on any host with a + // working resolver; the exact PTR record is environment dependent so we + // only assert it is non-null + let output = execution + .run_and_format("SELECT reverse_dns('127.0.0.1') IS NOT NULL AS resolved") + .await; + insta::assert_yaml_snapshot!(output, @r#" + - +----------+ + - "| resolved |" + - +----------+ + - "| true |" + - +----------+ + "#); +} + +#[tokio::test] +async fn test_reverse_dns_invalid_input_is_null() { + let execution = TestExecution::new().await; + // Unparseable input and a reserved address with no PTR record both yield + // null rather than erroring the query + let output = execution + .run_and_format( + "SELECT reverse_dns(ip) AS host FROM (VALUES ('not-an-ip'), ('192.0.2.1')) AS t(ip)", + ) + .await; + insta::assert_yaml_snapshot!(output, @r#" + - +------+ + - "| host |" + - +------+ + - "| |" + - "| |" + - +------+ + "#); +} + +#[tokio::test] +async fn test_geoip_invalid_ip_is_null() { + let execution = TestExecution::new().await; + // geoip parses the address before resolving the database, so an + // unparseable address yields NULL even with a database path that does + // not exist + let output = execution + .run_and_format("SELECT geoip('not-an-ip', '/definitely/missing.mmdb') IS NULL AS missing") + .await; + insta::assert_yaml_snapshot!(output, @r#" + - +---------+ + - "| missing |" + - +---------+ + - "| true |" + - +---------+ + "#); +} + +#[tokio::test] +async fn test_geoip_unopenable_db_reported_in_error_field() { + let execution = TestExecution::new().await; + // A database that cannot be opened does not fail the query: the + // location fields are NULL and the struct's error field carries the + // reason + let output = execution + .run_and_format( + "SELECT geoip('1.1.1.1', '/definitely/missing.mmdb')['country_code'] IS NULL AS no_country, \ + geoip('1.1.1.1', '/definitely/missing.mmdb')['error'] LIKE '%failed to open database%' AS reported", + ) + .await; + insta::assert_yaml_snapshot!(output, @r#" + - +------------+----------+ + - "| no_country | reported |" + - +------------+----------+ + - "| true | true |" + - +------------+----------+ + "#); +} + +#[tokio::test] +async fn test_geoip_db_path_from_config() { + // The GEOIP_DB environment variable takes precedence over the config + // value, so the config path is only observable when it is not set + if std::env::var_os(datafusion_net::GEOIP_DB_ENV_VAR).is_some() { + return; + } + let mut config = datafusion_dft::config::AppConfig::default(); + let db_path = "/config/provided/db.mmdb"; + config.cli.execution.net.geoip_db_path = Some(db_path.into()); + let execution = TestExecution::new_with_config(config).await; + // The configured path reaches the single-argument form of geoip: the + // error field reports a failure to open that (nonexistent) database + // rather than complaining that no database is configured + let batches = execution + .run("SELECT geoip('1.1.1.1')['error'] AS error") + .await + .unwrap(); + let formatted = datafusion::arrow::util::pretty::pretty_format_batches(&batches) + .unwrap() + .to_string(); + assert!( + formatted.contains(&format!("geoip failed to open database '{db_path}'")), + "unexpected output: {formatted}" + ); +} + +#[tokio::test] +async fn test_reverse_dns_over_pcap() { + let dir = tempfile::tempdir().unwrap(); + write_test_pcap(&dir.path().join("test.pcap")); + let execution = TestExecution::new().await; + // reverse_dns composes with the pcap table function over the src_ip + // column. The 10.0.0.0/8 addresses in the fixture will not resolve, but + // the query must plan and execute successfully. + let sql = format!( + "SELECT DISTINCT src_ip, reverse_dns(src_ip) AS host FROM pcap('{}/test.pcap')", + dir.path().display() + ); + let result = execution.run(&sql).await; + assert!(result.is_ok()); +}