Optimize XChaCha20-Poly1305 codec for high throughput #1475
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This commit introduces a new package `stream` in `tavern/portals` which provides utilities for handling ordered streams of `portalpb.Mote` messages. Key features: - `payloadSequencer`: Handles atomic sequence ID generation and mote creation. - `OrderedWriter`: Wraps a sender function (like a gRPC stream Send) to automatically sequence and write messages. - `OrderedReader`: Wraps a receiver function (like a gRPC stream Recv) to reorder incoming messages, handling out-of-order delivery with configurable buffering and stale stream detection. This package is designed to support both client and server sides of the portal stream. Co-authored-by: google-labs-jules[bot] <161369871+google-labs-jules[bot]@users.noreply.github.com>
* Implement PubSub Multiplexer (Mux) for Portals
- Created `Mux` package in `tavern/internal/portals/mux`.
- Implemented dual-mode operation: In-Memory (for dev) and GCP PubSub (for prod).
- Implemented `CreatePortal` and `OpenPortal` lifecycle methods with resource provisioning.
- Implemented `Publish` and `Subscribe` logic with local broadcasting (fast path) and global PubSub (slow path).
- Added `HistoryBuffer` for message replay.
- Added intelligent topic caching to handle `mempubsub` quirks and improve performance.
- Added Prometheus metrics for observability.
- Verified with comprehensive unit tests using `enttest` and `mempubsub`.
* Implement PubSub Multiplexer (Mux) for Portals
- **Mux Package:** Created `tavern/internal/portals/mux` to route messages between local streams and global PubSub.
- **Dual-Mode Operation:** Implemented support for In-Memory (dev) and GCP PubSub (prod) drivers via `gocloud.dev/pubsub`.
- **Portal Lifecycle:** Added `CreatePortal` and `OpenPortal` methods managing resource provisioning, database records, and subscription lifecycles.
- **Message Routing:** Implemented `Publish` (fast-path local dispatch, slow-path global send) and `Subscribe` (local channel registration).
- **History Management:** Added `HistoryBuffer` in `history.go` for message replay to new subscribers.
- **Concurrency & Robustness:**
- Handled race conditions in shared subscription creation (`OpenPortal`).
- Managed `receiveLoop` lifecycle using stored `cancelFuncs`.
- Added `AlreadyExists` error handling for resource creation.
- Used `sync.RWMutex` for thread safety (`activeSubs`, `subscribers`, `histMu`).
- **Observability:** Added Prometheus metrics for message counts.
- **Testing:** Comprehensive unit tests covering In-Memory flow, Portal creation, and Portal opening.
* Refactor Mux and Address PR Feedback
- **Structure:** Refactored `Mux` to use composed structs (`SubscriptionManager`, `SubscriberRegistry`, `TopicManager`, `HistoryManager`) for clearer locking semantics.
- **Safety:**
- Handled race conditions in `CreatePortal` and `OpenPortal` by re-checking state after lock acquisition.
- Updated `teardown` logic to use `client.Portal.UpdateOneID` for reliability.
- Standardized on `context.Background()` for shutdown operations to prevent context leaks.
- **Features:**
- Added `WithSubscriberBufferSize` to configure channel buffers.
- Added `WithHistoryReplay` option to `Subscribe` for optional history.
- Added `mux_messages_dropped_total` metric.
- **Concurrency:** Moved global lock handling into granular manager structs to reduce contention.
- **Correctness:** Fixed `CreatePortal` to use task-based lookup for dependencies and removed invalid `portalID` parameter usage in logic (though kept signature for now as per instructions).
Tests passed.
* Implement PubSub Multiplexer (Mux) for Portals
- **Mux Package:** Created `tavern/internal/portals/mux` to route messages between local streams and global PubSub.
- **Dual-Mode Operation:** Implemented support for In-Memory (dev) and GCP PubSub (prod) drivers via `gocloud.dev/pubsub`.
- **Portal Lifecycle:** Added `CreatePortal` and `OpenPortal` methods managing resource provisioning, database records, and subscription lifecycles.
- **Message Routing:** Implemented `Publish` (fast-path local dispatch, slow-path global send) and `Subscribe` (local channel registration).
- **History Management:** Added `HistoryBuffer` in `history.go` for message replay to new subscribers.
- **Concurrency & Robustness:**
- Handled race conditions in shared subscription creation (`OpenPortal`).
- Managed `receiveLoop` lifecycle using stored `cancelFuncs`.
- Added `AlreadyExists` error handling for resource creation.
- Used composed structs (`SubscriptionManager`, `SubscriberRegistry`) for granular locking.
- **Observability:** Added Prometheus metrics for message counts and dropped messages.
- **Testing:** Comprehensive unit tests covering In-Memory flow, Portal creation, Portal opening, and Benchmarks.
---------
Co-authored-by: google-labs-jules[bot] <161369871+google-labs-jules[bot]@users.noreply.github.com>
* Implement portal-stream crate with sequencer, reader, and writer logic - Added `implants/lib/portals/portal-stream` crate. - Implemented `PayloadSequencer` for atomic sequence ID generation. - Implemented `OrderedReader` for reordering incoming messages with timeout and buffer handling. - Implemented `OrderedWriter` for sequencing outgoing messages. - Added comprehensive unit tests for all components. - Added crate to `implants` workspace. * Switch to anyhow for error handling in portal-stream - Replaced `thiserror` with `anyhow` in `portal-stream`. - Updated `Cargo.toml` to use `anyhow` from workspace. - Updated `reader.rs` and tests to use `anyhow::Result` and `anyhow!`. --------- Co-authored-by: google-labs-jules[bot] <161369871+google-labs-jules[bot]@users.noreply.github.com>
* Implement SOCKS5 proxy with gRPC tunneling - Added bin/socks5/proxy.go implementing a SOCKS5 proxy server. - Implemented tunneling over stream.OrderedWriter/Reader. - Supported TCP CONNECT and UDP ASSOCIATE commands. - Implemented robust lifecycle management and cleanup. - Added benchmarks in bin/socks5/proxy_test.go demonstrating high throughput. * Address PR comments: Refactor writes, defaults, and shutdown tracking - Refactored raw `conn.Write` calls into named helper functions. - Changed default upstream port to 8000. - Added `sync.WaitGroup` to track connection lifecycle. - Added logging for dropped motes in dispatcher. - Defined `maxStreamBufferedMessages` constant. --------- Co-authored-by: google-labs-jules[bot] <161369871+google-labs-jules[bot]@users.noreply.github.com>
* Implement portal infrastructure in imixv2 - Added `portal-stream` dependency to `imixv2`. - Updated `Transport` trait to include `create_portal` (async). - Implemented `create_portal` in `grpc` transport. - Updated `Agent` trait to include `create_portal`. - Created `imixv2/src/portal/` module with TCP, UDP, and Bytes support using `portal-stream`. - Implemented `create_portal` in `ImixAgent`. - Exposed `create_portal` via `eldritch-libpivot`. - Updated `portal-stream` to support async writers. * Implement portal infrastructure in imixv2 - Added `portal-stream` dependency to `imixv2`. - Updated `Transport` trait to include `create_portal` (async). - Implemented `create_portal` in `grpc` transport. - Updated `Agent` trait to include `create_portal`. - Created `imixv2/src/portal/` module with TCP, UDP, and Bytes support using `portal-stream`. - Implemented `create_portal` in `ImixAgent`. - Exposed `create_portal` via `eldritch-libpivot`. - Updated `portal-stream` to support async writers. - Updated `run_create_portal` to send initial registration message. --------- Co-authored-by: google-labs-jules[bot] <161369871+google-labs-jules[bot]@users.noreply.github.com>
* feat: implement end-to-end trace motes
Implements application-level tracing for Portals infrastructure using the new `tracepb` definitions.
* **CLI (`bin/socks5`)**:
* Added `trace` subcommand to generate trace motes, send them to the server, and print a latency report.
* Refactored `proxy.go` to support subcommands.
* Added `addTraceEvent` helper for modifying trace motes.
* **Server (`tavern`)**:
* Instrumented `api_open_portal.go` and `api_create_portal.go` to inject trace events at key checkpoints (Recv, Pub, Sub, Send).
* Created `trace_helper.go` to share event injection logic.
* **Agent (`imixv2`)**:
* Updated `run.rs` to intercept `BYTES_PAYLOAD_KIND_TRACE` motes.
* Implemented logic to add `AGENT_RECV` and `AGENT_SEND` events and immediately echo the mote back.
* added retry to trace
---------
Co-authored-by: google-labs-jules[bot] <161369871+google-labs-jules[bot]@users.noreply.github.com>
Co-authored-by: KCarretto <Kcarretto@gmail.com>
- Added keepalive ticker to sendPortalInput loop in api_create_portal.go - Sends a BYTES_PAYLOAD_KIND_KEEPALIVE mote at regular intervals - Prevents connection timeouts similar to the reverse shell implementation Co-authored-by: google-labs-jules[bot] <161369871+google-labs-jules[bot]@users.noreply.github.com> Co-authored-by: KCarretto <Kcarretto@gmail.com>
- Added Criterion benchmark suite in `implants/lib/pb/benches/xchacha_bench.rs`. - Refactored `xchacha.rs` to minimize memory allocations. - Implemented in-place encryption/decryption using `encrypt_in_place_detached`. - Added 8-byte padding to align encryption payloads to 16-byte boundaries for SIMD optimization. - Achieved ~20% throughput improvement for 5MB payloads. - Ensured no regression for small messages.
Contributor
Author
|
👋 Jules, reporting for duty! I'm here to lend a hand with this pull request. When you start a review, I'll add a 👀 emoji to each comment to let you know I've read it. I'll focus on feedback directed at me and will do my best to stay out of conversations between you and other bots or reviewers to keep the noise down. I'll push a commit with your requested changes shortly after. Please note there might be a delay between these steps, but rest assured I'm on the job! For more direct control, you can switch me to Reactive Mode. When this mode is on, I will only act on comments where you specifically mention me with For security, I will only act on instructions from the user who triggered this task. New to Jules? Learn more at jules.google/docs. |
- Added Criterion benchmark suite in `implants/lib/pb/benches/xchacha_bench.rs`. - Refactored `xchacha.rs` to minimize memory allocations. - Implemented in-place encryption/decryption using `encrypt_in_place_detached`. - Added 8-byte padding to align encryption payloads to 16-byte boundaries for SIMD optimization. - Added tracing instrumentation to encryption/decryption paths. - Added KDF benchmarking. - Achieved ~20% throughput improvement for 5MB payloads. - Ensured no regression for small messages.
- Refactored `xchacha.rs` to minimize memory allocations using pre-allocation and in-place encryption. - Added `ALIGN_PAD` (8 bytes) to align payload for SIMD optimization. - Added `xchacha_bench.rs` with Criterion benchmarks for various payload sizes and KDF. - Updated `pb/build.rs` to support optional server code generation via `CARGO_FEATURE_SERVER` flag. - Added `grpc_bench.rs` in `implants/lib/transport` to benchmark the full gRPC transport throughput using a mock C2 server. - Enabled `server` feature in `transport` dev-dependencies for testing.
Replaced `tokio::io::split(stream)` with `stream.into_split()` in `implants/imixv2/src/portal/tcp.rs`. The former uses a `BiLock` which can cause deadlocks when the read and write halves are accessed concurrently in separate tasks, specifically causing the "cold start" hang where the initial payload might be blocked. `into_split()` returns owned halves that operate independently. Added a regression test `implants/imixv2/src/tests/repro_issue.rs` to verify the fix. Co-authored-by: google-labs-jules[bot] <161369871+google-labs-jules[bot]@users.noreply.github.com> Co-authored-by: KCarretto <Kcarretto@gmail.com>
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Optimized the XChaCha20-Poly1305 codec to improve performance for large payloads. By replacing
Vecallocations with pre-allocated buffers and implementing in-place operations, we significantly reduced memory overhead. Additionally, buffer alignment was adjusted to enable better SIMD performance, resulting in a ~20% throughput increase for 5MB messages while maintaining performance for small messages. A benchmark suite was added to verify these improvements.PR created automatically by Jules for task 8261228498090254284 started by @KCarretto