Gateway side things

This commit is contained in:
durch
2025-10-23 00:00:45 +02:00
parent 96e3ff2af9
commit 542e56044a
58 changed files with 11689 additions and 14 deletions
+1
View File
@@ -63,3 +63,4 @@ nym-api/redocly/formatted-openapi.json
**/settings.sql
**/enter_db.sh
.beads
+15 -1
View File
@@ -12,6 +12,16 @@ Nym is a privacy platform that uses mixnet technology to protect against metadat
- Validators for network consensus
- Various service providers and integrations
## Navigation Aids
This repository includes comprehensive navigation documents for efficient code exploration:
- **[CODEMAP.md](./CODEMAP.md)**: Structural overview of the entire repository with directory hierarchy, package descriptions, and navigation hints. Use this to quickly understand the codebase layout and find specific components.
- **[FUNCTION_LEXICON.md](./FUNCTION_LEXICON.md)**: Comprehensive catalog of key functions, signatures, and API patterns across all major modules. Use this to quickly find available functions and understand their usage patterns.
When working with this codebase, start by consulting these documents to understand the structure and available APIs before diving into specific files.
## Build Commands
### Rust Components
@@ -150,7 +160,7 @@ dotenv -f envs/sandbox.env -- cargo run -p nym-api
## Architecture
The Nym platform consists of various components organized as a monorepo:
The Nym platform consists of various components organized as a monorepo. For a detailed structural overview with directory hierarchy and navigation hints, see [CODEMAP.md](./CODEMAP.md).
1. **Core Mixnet Infrastructure**:
- `nym-node`: Core binary supporting mixnode and gateway modes
@@ -422,6 +432,8 @@ The system uses SQLite databases with tables like:
## Development Workflows
**Note**: Before diving into specific workflows, consult [CODEMAP.md](./CODEMAP.md) to understand the repository structure and [FUNCTION_LEXICON.md](./FUNCTION_LEXICON.md) to discover available APIs and functions.
### Running a Node
To run the mixnode or gateway:
@@ -450,6 +462,8 @@ To monitor the health of your node:
## Common Libraries
For a comprehensive catalog of functions and APIs available in these libraries, see [FUNCTION_LEXICON.md](./FUNCTION_LEXICON.md).
- `common/types`: Shared data types across all components
- `common/crypto`: Cryptographic primitives and wrappers
- `common/client-core`: Core client functionality
+456
View File
@@ -0,0 +1,456 @@
# Nym Repository Codemap
<!-- AIDEV-NOTE: This codemap provides structural navigation for the Nym privacy platform monorepo -->
<!-- Last updated: 2024-10-22 (branch: drazen/lp-reg) -->
## Quick Navigation Index
| Component | Location | Purpose |
|-----------|----------|---------|
| [Main Executables](#main-executables) | Root directories | Core binaries and services |
| [Client Implementations](#client-implementations) | `/clients/` | Various client types |
| [Common Libraries](#common-libraries) | `/common/` | 70+ shared modules |
| [Smart Contracts](#smart-contracts) | `/contracts/` | CosmWasm contracts |
| [SDKs](#sdks) | `/sdk/` | Multi-language SDKs |
| [WASM Modules](#wasm-modules) | `/wasm/` | Browser implementations |
| [Service Providers](#service-providers) | `/service-providers/` | Exit nodes & routers |
| [Tools](#tools-and-utilities) | `/tools/` | CLI tools & utilities |
| [Configuration](#configuration-and-environments) | `/envs/` | Environment configs |
## Repository Structure Overview
```
nym/
├── Cargo.toml # Workspace manifest (170+ members)
├── Cargo.lock # Locked dependencies
├── Makefile # Build automation
├── CLAUDE.md # Development guidelines
├── envs/ # Environment configurations
│ ├── local.env # Local development
│ ├── sandbox.env # Test network
│ ├── mainnet.env # Production
│ └── canary.env # Pre-release
├── assets/ # Images, logos, fonts
├── docker/ # Docker configurations
└── scripts/ # Deployment & setup scripts
```
<!-- AIDEV-NOTE: Navigation hint - Use envs/ for network-specific configurations -->
## Main Executables
### Core Network Nodes
#### **nym-node** (v1.19.0) - Universal Node Binary
- **Path**: `/nym-node/`
- **Entry**: `src/main.rs`
- **Modes**: `mixnode`, `gateway`
- **Key Modules**:
- `cli/` - Command-line interface
- `config/` - Configuration management
- `node/` - Core node logic
- `wireguard/` - WireGuard VPN integration
- `throughput_tester/` - Performance testing
<!-- AIDEV-NOTE: Complex area - nym-node replaces legacy gateway and mixnode binaries -->
#### **nym-api** - Network API Server
- **Path**: `/nym-api/`
- **Entry**: `src/main.rs`
- **Database**: PostgreSQL with SQLx
- **Migrations**: `/migrations/` (25+ migration files)
- **Key Subsystems**:
- `circulating_supply_api/` - Token supply tracking
- `ecash/` - E-cash credential management
- `epoch_operations/` - Epoch advancement
- `network_monitor/` - Health monitoring
- `node_performance/` - Performance metrics
- `nym_nodes/` - Node registry
#### **gateway** (Legacy, v1.1.36)
- **Path**: `/gateway/`
- **Status**: Being phased out for nym-node
- **New**: `src/node/lp_listener/` (branch: drazen/lp-reg)
### Supporting Services
| Service | Path | Purpose |
|---------|------|---------|
| `nym-network-monitor` | `/nym-network-monitor/` | Network reliability testing |
| `nym-validator-rewarder` | `/nym-validator-rewarder/` | Reward calculation |
| `nyx-chain-watcher` | `/nyx-chain-watcher/` | Blockchain monitoring |
| `nym-credential-proxy` | `/nym-credential-proxy/` | Credential services |
| `nym-statistics-api` | `/nym-statistics-api/` | Statistics aggregation |
| `nym-node-status-api` | `/nym-node-status-api/` | Node status tracking |
## Client Implementations
### Directory: `/clients/`
```
clients/
├── native/ # Native Rust client
│ └── websocket-requests/ # WebSocket protocol
├── socks5/ # SOCKS5 proxy client
├── validator/ # Blockchain validator client
└── webassembly/ # Browser-based client
```
<!-- AIDEV-NOTE: Pattern reference - All clients use common/client-core for shared functionality -->
## Common Libraries
### Directory: `/common/` (70+ modules)
### Core Infrastructure
| Module | Purpose | Key Types |
|--------|---------|-----------|
| `nym-common` | Shared utilities | Constants, helpers |
| `types` | Common data types | NodeId, MixId |
| `config` | Configuration system | Config traits |
| `commands` | CLI structures | Command builders |
| `bin-common` | Binary utilities | Logging, banners |
### Cryptography & Security
| Module | Purpose | Dependencies |
|--------|---------|-------------|
| `crypto` | Crypto primitives | Ed25519, X25519 |
| `credentials` | Credential system | BLS12-381 |
| `credentials-interface` | Interface definitions | - |
| `credential-verification` | Validation logic | - |
| `pemstore` | PEM storage | - |
### Network Protocol (Sphinx)
<!-- AIDEV-NOTE: Complex area - Sphinx is the core privacy protocol -->
```
nymsphinx/
├── types/ # Core types
├── chunking/ # Message fragmentation
├── forwarding/ # Packet forwarding
├── routing/ # Route selection
├── addressing/ # Address handling
├── anonymous-replies/ # SURB system
├── acknowledgements/ # ACK handling
├── cover/ # Cover traffic
├── params/ # Protocol parameters
└── framing/ # Wire format
```
### New Components (Branch: drazen/lp-reg)
<!-- AIDEV-NOTE: Current branch changes - These are new additions -->
| Module | Path | Status |
|--------|------|--------|
| `nym-lp` | `/common/nym-lp/` | New LP protocol |
| `nym-lp-common` | `/common/nym-lp-common/` | LP utilities |
| `nym-kcp` | `/common/nym-kcp/` | KCP protocol |
### Client Systems
```
client-core/
├── config-types/ # Configuration types
├── gateways-storage/ # Gateway persistence
└── surb-storage/ # SURB storage
client-libs/
├── gateway-client/ # Gateway connection
├── mixnet-client/ # Mixnet interaction
└── validator-client/ # Blockchain queries
```
### Additional Common Modules
**Storage & Data**:
- `statistics/` - Statistical collection
- `topology/` - Network topology
- `node-tester-utils/` - Testing utilities
- `ticketbooks-merkle/` - Merkle trees
**Advanced Features**:
- `dkg/` - Distributed Key Generation
- `ecash-signer-check/` - E-cash validation
- `nym_offline_compact_ecash/` - Offline e-cash
**Blockchain**:
- `ledger/` - Ledger operations
- `nyxd-scraper/` - Chain scraping
- `cosmwasm-smart-contracts/` - Contract interfaces
**Utilities**:
- `task/` - Async task management
- `async-file-watcher/` - File watching
- `nym-cache/` - Caching layer
- `nym-metrics/` - Metrics (Prometheus)
- `bandwidth-controller/` - Bandwidth accounting
## Smart Contracts
### Directory: `/contracts/`
<!-- AIDEV-NOTE: Navigation hint - All contracts use CosmWasm 2.2.2 -->
```
contracts/
├── Cargo.toml # Workspace config
├── .cargo/config.toml # WASM build config
├── coconut-dkg/ # DKG contract
├── ecash/ # E-cash contract
├── mixnet/ # Node registry
├── vesting/ # Token vesting
├── nym-pool/ # Liquidity pool
├── multisig/ # Multi-sig wallet
├── performance/ # Performance tracking
└── mixnet-vesting-integration-tests/
```
### Contract Build Process
```bash
make contracts # Build all
make contract-schema # Generate schemas
make wasm-opt-contracts # Optimize
```
## SDKs
### Directory: `/sdk/`
```
sdk/
├── rust/
│ └── nym-sdk/ # Primary Rust SDK
├── typescript/
│ ├── packages/ # NPM packages
│ ├── codegen/ # Code generation
│ └── examples/ # Usage examples
└── ffi/
├── cpp/ # C++ bindings
├── go/ # Go bindings
└── shared/ # Shared FFI code
```
## WASM Modules
### Directory: `/wasm/`
| Module | Purpose | Build Command |
|--------|---------|---------------|
| `client` | Browser client | `make` in directory |
| `mix-fetch` | Privacy fetch API | `make` in directory |
| `node-tester` | Network testing | `make` in directory |
| `zknym-lib` | Zero-knowledge lib | `make` in directory |
<!-- AIDEV-NOTE: Pattern reference - WASM modules compile from Rust using wasm-pack -->
## Service Providers
### Directory: `/service-providers/`
```
service-providers/
├── network-requester/ # Exit node for external requests
├── ip-packet-router/ # IP packet routing (VPN-like)
└── common/ # Shared utilities
```
## Tools and Utilities
### Directory: `/tools/`
### Public Tools
| Tool | Path | Purpose |
|------|------|---------|
| `nym-cli` | `/tools/nym-cli/` | Node management CLI |
| `nym-id-cli` | `/tools/nym-id-cli/` | Identity management |
| `nymvisor` | `/tools/nymvisor/` | Process supervisor |
| `nym-nr-query` | `/tools/nym-nr-query/` | Network queries |
| `echo-server` | `/tools/echo-server/` | Testing server |
### Internal Tools
```
internal/
├── mixnet-connectivity-check/ # Network diagnostics
├── contract-state-importer/ # Migration tools
├── validator-status-check/ # Validator health
├── ssl-inject/ # SSL injection
├── testnet-manager/ # Testnet management
└── sdk-version-bump/ # Version management
```
## Configuration and Environments
### Environment Files: `/envs/`
<!-- AIDEV-NOTE: Navigation hint - Always use dotenv -f envs/[env].env for proper configuration -->
| Environment | File | API Endpoint | Use Case |
|------------|------|--------------|----------|
| Local | `local.env` | localhost | Development |
| Sandbox | `sandbox.env` | sandbox-nym-api1.nymtech.net | Testing |
| Mainnet | `mainnet.env` | validator.nymtech.net | Production |
| Canary | `canary.env` | - | Pre-release |
### Key Environment Variables
```bash
NETWORK_NAME # Network identifier
NYM_API # API endpoint
NYXD # Blockchain RPC
MIXNET_CONTRACT_ADDRESS # Contract addresses
MNEMONIC # Test mnemonic (NEVER in production)
RUST_LOG # Logging level
DATABASE_URL # PostgreSQL connection
```
## Build System
### Primary Build Commands
```bash
make build # Debug build
make build-release # Release build
make test # Run tests
make clippy # Lint code
make fmt # Format code
make contracts # Build contracts
make sdk-wasm-build # Build WASM
```
### Workspace Configuration
<!-- AIDEV-NOTE: Complex area - Root Cargo.toml manages 170+ workspace members -->
**Root Cargo.toml Structure**:
- `[workspace]` - Lists all 170+ members
- `[workspace.dependencies]` - Shared dependency versions
- `[workspace.lints]` - Shared lint rules
- `[profile.*]` - Build profiles
## Database Structure
### SQLx Usage Pattern
- **Compile-time verified**: All queries checked at build
- **Migration files**: In package `/migrations/` directories
- **Query cache**: `.sqlx/` directory
### Key Tables (nym-api)
```sql
-- Network monitoring
mixnode_status
gateway_status
routes
monitor_run
-- Node registry
nym_nodes
node_descriptions
-- Performance
node_uptime
node_performance
```
## Current Branch Context (drazen/lp-reg)
### New Additions
- `/common/nym-lp/` - Low-level protocol implementation
- `/common/nym-lp-common/` - LP common utilities
- `/common/nym-kcp/` - KCP protocol
- `/gateway/src/node/lp_listener/` - LP listener
### Modified Files
```
M Cargo.lock
M Cargo.toml
M common/registration/
M common/wireguard/
M gateway/
M nym-node/
M nym-node/nym-node-metrics/
```
## Navigation Patterns
<!-- AIDEV-NOTE: Navigation hint - Use these patterns to quickly find code -->
### Finding Code by Type
| Code Type | Look In |
|-----------|---------|
| Main executables | Root directories with `src/main.rs` |
| Libraries | `/common/` with descriptive names |
| Contracts | `/contracts/[name]/src/contract.rs` |
| Tests | Colocated with source, `#[cfg(test)]` |
| Configurations | `/envs/` and `config/` subdirs |
| Database queries | Files with `.sql` or SQLx macros |
| API endpoints | `/nym-api/src/` subdirectories |
| CLI commands | `/cli/commands/` in executables |
### Common Import Locations
```rust
// Crypto
use nym_crypto::asymmetric::{ed25519, x25519};
// Network
use nym_sphinx::forwarding::packet::MixPacket;
use nym_topology::NymTopology;
// Client
use nym_client_core::client::Client;
// Configuration
use nym_network_defaults::NymNetworkDetails;
// Contracts
use nym_mixnet_contract_common::*;
```
## Module Relationships
<!-- AIDEV-NOTE: Complex area - Understanding dependencies helps navigation -->
### Dependency Graph (Simplified)
```
nym-node
├── common/nym-common
├── common/crypto
├── common/nymsphinx
├── common/topology
├── common/client-libs/validator-client
└── common/wireguard
nym-api
├── common/nym-common
├── nym-api-requests
├── common/client-libs/validator-client
├── common/credentials
└── sqlx (database)
clients/native
├── common/client-core
├── common/client-libs/gateway-client
├── common/nymsphinx
└── common/credentials
```
## Development Workflows
### Adding New Feature
1. Check `/envs/` for configuration
2. Find similar code in `/common/`
3. Implement in appropriate module
4. Add tests colocated with code
5. Update `/nym-api/` if needed
6. Run `make test` and `make clippy`
### Debugging Network Issues
1. Start with `/nym-network-monitor/`
2. Check `/common/topology/` for routing
3. Review `/common/nymsphinx/` for protocol
4. Examine logs with `RUST_LOG=debug`
### Contract Development
1. Create in `/contracts/[name]/`
2. Use existing contracts as templates
3. Build with `make contracts`
4. Test with `cw-multi-test`
---
<!-- AIDEV-NOTE: This codemap is optimized for LLM navigation. Use Ctrl+F to quickly find components -->
Generated
+84 -11
View File
@@ -165,6 +165,15 @@ version = "0.1.6"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "4b46cbb362ab8752921c97e041f5e366ee6297bd428a31275b9fcf1e380f7299"
[[package]]
name = "ansi_term"
version = "0.12.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "d52a9bb7ec0cf484c551830a7ce27bd20d67eac647e1befb56b0be4ee39a55d2"
dependencies = [
"winapi",
]
[[package]]
name = "anstream"
version = "0.6.19"
@@ -991,6 +1000,12 @@ version = "0.3.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "e3b5ca7a04898ad4bcd41c90c5285445ff5b791899bb1b0abdd2a2aa791211d7"
[[package]]
name = "byte_string"
version = "1.0.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "11aade7a05aa8c3a351cedc44c3fc45806430543382fcc4743a9b757a2a0b4ed"
[[package]]
name = "bytecodec"
version = "0.4.15"
@@ -2053,6 +2068,20 @@ dependencies = [
"serde",
]
[[package]]
name = "dashmap"
version = "6.1.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "5041cc499144891f3790297212f32a74fb938e5136a14943f338ef9e0ae276cf"
dependencies = [
"cfg-if",
"crossbeam-utils",
"hashbrown 0.14.5",
"lock_api",
"once_cell",
"parking_lot_core",
]
[[package]]
name = "data-encoding"
version = "2.9.0"
@@ -2353,7 +2382,7 @@ dependencies = [
"bytecodec",
"bytes",
"clap",
"dashmap",
"dashmap 5.5.3",
"dirs",
"futures",
"nym-bin-common",
@@ -4842,7 +4871,7 @@ dependencies = [
"cw2",
"cw3",
"cw4",
"dashmap",
"dashmap 5.5.3",
"dotenvy",
"futures",
"humantime-serde",
@@ -5267,7 +5296,7 @@ version = "0.1.0"
dependencies = [
"anyhow",
"async-trait",
"dashmap",
"dashmap 5.5.3",
"nym-crypto",
"nym-sphinx",
"nym-task",
@@ -5804,7 +5833,8 @@ dependencies = [
"bincode",
"bip39",
"bs58",
"dashmap",
"bytes",
"dashmap 5.5.3",
"defguard_wireguard_rs",
"fastrand 2.3.0",
"futures",
@@ -5821,10 +5851,13 @@ dependencies = [
"nym-gateway-storage",
"nym-id",
"nym-ip-packet-router",
"nym-kcp",
"nym-lp",
"nym-mixnet-client",
"nym-network-defaults",
"nym-network-requester",
"nym-node-metrics",
"nym-registration-common",
"nym-sdk",
"nym-service-provider-requests-common",
"nym-sphinx",
@@ -6204,6 +6237,19 @@ dependencies = [
"url",
]
[[package]]
name = "nym-kcp"
version = "0.1.0"
dependencies = [
"ansi_term",
"byte_string",
"bytes",
"env_logger",
"log",
"thiserror 2.0.12",
"tokio-util",
]
[[package]]
name = "nym-ledger"
version = "0.1.0"
@@ -6215,11 +6261,37 @@ dependencies = [
"thiserror 2.0.12",
]
[[package]]
name = "nym-lp"
version = "0.1.0"
dependencies = [
"ansi_term",
"bincode",
"bs58",
"bytes",
"criterion",
"dashmap 6.1.0",
"nym-lp-common",
"nym-sphinx",
"parking_lot",
"rand 0.8.5",
"rand_chacha 0.3.1",
"serde",
"sha2 0.10.9",
"snow",
"thiserror 2.0.12",
"utoipa",
]
[[package]]
name = "nym-lp-common"
version = "0.1.0"
[[package]]
name = "nym-metrics"
version = "0.1.0"
dependencies = [
"dashmap",
"dashmap 5.5.3",
"lazy_static",
"prometheus",
"tracing",
@@ -6229,7 +6301,7 @@ dependencies = [
name = "nym-mixnet-client"
version = "0.1.0"
dependencies = [
"dashmap",
"dashmap 5.5.3",
"futures",
"nym-crypto",
"nym-noise",
@@ -6328,7 +6400,7 @@ dependencies = [
"anyhow",
"axum",
"clap",
"dashmap",
"dashmap 5.5.3",
"futures",
"log",
"nym-bin-common",
@@ -6496,7 +6568,7 @@ dependencies = [
name = "nym-node-metrics"
version = "0.1.0"
dependencies = [
"dashmap",
"dashmap 5.5.3",
"futures",
"nym-metrics",
"nym-statistics-common",
@@ -6816,6 +6888,7 @@ dependencies = [
"nym-crypto",
"nym-ip-packet-requests",
"nym-sphinx",
"serde",
"tokio-util",
]
@@ -6830,7 +6903,7 @@ dependencies = [
"bytecodec",
"bytes",
"clap",
"dashmap",
"dashmap 5.5.3",
"dirs",
"dotenvy",
"futures",
@@ -7109,7 +7182,7 @@ dependencies = [
name = "nym-sphinx-chunking"
version = "0.1.0"
dependencies = [
"dashmap",
"dashmap 5.5.3",
"log",
"nym-crypto",
"nym-metrics",
@@ -7911,7 +7984,7 @@ checksum = "8b3a2a91fdbfdd4d212c0dcc2ab540de2c2bcbbd90be17de7a7daf8822d010c1"
dependencies = [
"async-trait",
"crossbeam-channel",
"dashmap",
"dashmap 5.5.3",
"fnv",
"futures-channel",
"futures-executor",
+3
View File
@@ -72,6 +72,9 @@ members = [
"common/nym-cache",
"common/nym-connection-monitor",
"common/nym-id",
"common/nym-kcp",
"common/nym-lp",
"common/nym-lp-common",
"common/nym-metrics",
"common/nym_offline_compact_ecash",
"common/nymnoise",
+909
View File
@@ -0,0 +1,909 @@
# Nym Function Lexicon
<!-- AIDEV-NOTE: This lexicon catalogs key functions, signatures, and API patterns across the Nym codebase -->
<!-- Last updated: 2024-10-22 (branch: drazen/lp-reg) -->
## Quick Reference Index
| Category | Section | Key Operations |
|----------|---------|----------------|
| [Node Operations](#1-node-operations) | Mixnode & Gateway | Initialization, key management, tasks |
| [Sphinx Protocol](#2-sphinx-packet-protocol) | Packet Processing | Message creation, chunking, routing |
| [Client APIs](#3-client-apis) | Client Operations | Connection, sending, receiving |
| [Network Topology](#4-network-topology) | Routing | Topology queries, route selection |
| [Blockchain](#5-blockchain-operations) | Validator Client | Queries, transactions, contracts |
| [REST APIs](#6-rest-api-endpoints) | HTTP Handlers | API routes and responses |
| [Credentials](#7-credential--ecash) | E-cash | Credential creation, verification |
| [Smart Contracts](#8-smart-contracts) | CosmWasm | Entry points, messages |
| [Common Patterns](#9-common-patterns) | Conventions | Naming, errors, async |
---
## 1. Node Operations
### nym-node Core Functions
<!-- AIDEV-NOTE: Complex area - nym-node unifies mixnode and gateway functionality -->
**Module**: `nym-node/src/node/mod.rs`
```rust
// Node initialization
pub async fn initialise_node(
config: &Config,
rng: &mut impl CryptoRng + RngCore,
) -> Result<NodeData, NymNodeError>
// Key management
pub fn load_x25519_wireguard_keypair(
paths: &KeysPaths,
) -> Result<x25519::KeyPair, NymNodeError>
pub fn load_ed25519_identity_keypair(
paths: &KeysPaths,
) -> Result<ed25519::KeyPair, NymNodeError>
// Gateway-specific initialization
impl GatewayTasksData {
pub async fn new(
config: &GatewayTasksConfig,
client_storage: ClientStorage,
) -> Result<GatewayTasksData, GatewayError>
pub fn initialise(
config: &GatewayTasksConfig,
force_init: bool,
) -> Result<(), GatewayError>
}
// Service provider initialization
impl ServiceProvidersData {
pub fn initialise_client_keys<R: RngCore + CryptoRng>(
rng: &mut R,
gateway_paths: &GatewayPaths,
) -> Result<ed25519::KeyPair, GatewayError>
pub async fn initialise_network_requester<R>(
rng: &mut R,
config: &Config,
) -> Result<Option<LocalNetworkRequester>, GatewayError>
}
```
### Gateway Task Builder Pattern
**Module**: `gateway/src/node/mod.rs`
```rust
pub struct GatewayTasksBuilder {
// Builder methods
pub fn new(
identity_keypair: Arc<ed25519::KeyPair>,
config: Config,
client_storage: ClientStorage,
) -> GatewayTasksBuilder
pub fn set_network_requester_opts(
&mut self,
opts: Option<LocalNetworkRequesterOpts>
) -> &mut Self
pub fn set_ip_packet_router_opts(
&mut self,
opts: Option<LocalIpPacketRouterOpts>
) -> &mut Self
pub async fn build_and_run(
self,
shutdown: TaskManager,
) -> Result<(), GatewayError>
}
```
<!-- AIDEV-NOTE: Pattern reference - Builder pattern is common for complex initialization -->
---
## 2. Sphinx Packet Protocol
### Message Construction & Processing
**Module**: `common/nymsphinx/src/message.rs`
```rust
// Core message types
pub enum NymMessage {
Plain(Vec<u8>),
Repliable(RepliableMessage),
Reply(ReplyMessage),
}
impl NymMessage {
// Constructors
pub fn new_plain(msg: Vec<u8>) -> NymMessage
pub fn new_repliable(msg: RepliableMessage) -> NymMessage
pub fn new_reply(msg: ReplyMessage) -> NymMessage
pub fn new_additional_surbs_request(
recipient: Recipient,
amount: u32
) -> NymMessage
// Processing
pub fn pad_to_full_packet_lengths(
self,
plaintext_per_packet: usize
) -> PaddedMessage
pub fn split_into_fragments<R: Rng>(
self,
rng: &mut R,
packet_size: PacketSize,
) -> Vec<Fragment>
pub fn remove_padding(self) -> Result<NymMessage, NymMessageError>
// Queries
pub fn is_reply_surb_request(&self) -> bool
pub fn available_sphinx_plaintext_per_packet(
&self,
packet_size: PacketSize
) -> usize
pub fn required_packets(&self, packet_size: PacketSize) -> usize
}
```
### Payload Building & Preparation
**Module**: `common/nymsphinx/src/preparer.rs`
```rust
pub struct NymPayloadBuilder {
// Main preparation methods
pub async fn prepare_chunk_for_sending(
&mut self,
message: NymMessage,
topology: &NymTopology,
) -> Result<Vec<MixPacket>, NymPayloadBuilderError>
pub async fn prepare_reply_chunk_for_sending(
&mut self,
reply: NymMessage,
reply_surb: ReplySurb,
) -> Result<Vec<MixPacket>, NymPayloadBuilderError>
// SURB generation
pub fn generate_reply_surbs(
&mut self,
amount: u32,
topology: &NymTopology,
) -> Result<Vec<SurbAck>, NymPayloadBuilderError>
// Fragment splitting
pub fn pad_and_split_message(
&mut self,
message: NymMessage,
) -> Result<Vec<Fragment>, NymPayloadBuilderError>
}
// Builder constructors
pub fn build_regular<R: CryptoRng + Rng>(
rng: R,
sender_address: Option<Recipient>,
) -> NymPayloadBuilder
pub fn build_reply(
sender_address: Recipient,
sender_tag: AnonymousSenderTag,
) -> NymPayloadBuilder
```
### Chunking & Fragmentation
**Module**: `common/nymsphinx/chunking/src/lib.rs`
<!-- AIDEV-NOTE: Complex area - Chunking splits messages into Sphinx-sized packets -->
```rust
// Main chunking function
pub fn split_into_sets(
message: &[u8],
max_plaintext_size: usize,
max_fragments_per_set: usize,
) -> Result<Vec<Vec<Fragment>>, ChunkingError>
// Fragment monitoring (optional feature)
pub mod monitoring {
pub fn enable()
pub fn enabled() -> bool
pub fn fragment_received(fragment: &Fragment)
pub fn fragment_sent(
fragment: &Fragment,
client_nonce: i32,
destination: PublicKey
)
}
```
---
## 3. Client APIs
### Gateway Client
**Module**: `common/client-libs/gateway-client/src/lib.rs`
```rust
pub struct GatewayClient {
// Connection management
pub async fn connect(
config: GatewayClientConfig,
) -> Result<GatewayClient, GatewayClientError>
pub async fn authenticate(
&mut self,
credentials: Credentials,
) -> Result<(), GatewayClientError>
// Message operations
pub async fn send_mix_packet(
&self,
packet: MixPacket,
) -> Result<(), GatewayClientError>
pub async fn receive_messages(
&mut self,
) -> Result<Vec<ReconstructedMessage>, GatewayClientError>
}
// Packet routing
pub struct PacketRouter {
pub fn new(
mix_tx: MixnetMessageSender,
ack_tx: AcknowledgementSender,
) -> PacketRouter
pub async fn route_packet(
&self,
packet: MixPacket,
) -> Result<(), PacketRouterError>
}
```
### Mixnet Client
**Module**: `common/client-libs/mixnet-client/src/lib.rs`
```rust
pub struct Client {
// Core client operations
pub async fn new(config: Config) -> Result<Client, ClientError>
pub async fn send_message(
&mut self,
recipient: Recipient,
message: Vec<u8>,
) -> Result<(), ClientError>
pub async fn receive_message(
&mut self,
) -> Result<ReconstructedMessage, ClientError>
// Connection management
pub fn is_connected(&self) -> bool
pub async fn reconnect(&mut self) -> Result<(), ClientError>
}
// Send without response trait
pub trait SendWithoutResponse {
fn send_without_response(
&self,
packet: MixPacket,
) -> io::Result<()>
}
```
<!-- AIDEV-NOTE: Pattern reference - Async/await is standard for network operations -->
### Client Core Initialization
**Module**: `common/client-core/src/init.rs`
```rust
// Key generation
pub fn generate_new_client_keys<R: CryptoRng + Rng>(
rng: &mut R,
) -> (ed25519::KeyPair, x25519::KeyPair)
// Storage initialization
pub async fn init_storage(
paths: &ClientPaths,
) -> Result<ClientStorage, ClientCoreError>
// Configuration setup
pub fn setup_client_config(
id: &str,
network: Network,
) -> Result<Config, ClientCoreError>
```
---
## 4. Network Topology
### Topology Management
**Module**: `common/topology/src/lib.rs`
```rust
pub struct NymTopology {
// Query methods
pub fn mixnodes(&self) -> &[RoutingNode]
pub fn gateways(&self) -> &[RoutingNode]
pub fn layer_nodes(&self, layer: MixLayer) -> Vec<&RoutingNode>
// Route selection
pub fn random_route<R: Rng>(
&self,
rng: &mut R,
) -> Option<Vec<RoutingNode>>
pub fn get_node_by_id(&self, node_id: NodeId) -> Option<&RoutingNode>
}
// Route provider
pub struct NymRouteProvider {
pub fn new(topology: NymTopology) -> NymRouteProvider
pub fn random_route<R: Rng>(
&self,
rng: &mut R,
) -> Option<Vec<RoutingNode>>
}
// Topology provider trait
pub trait TopologyProvider {
async fn get_topology(&self) -> Result<NymTopology, NymTopologyError>
async fn refresh_topology(&mut self) -> Result<(), NymTopologyError>
}
```
### Routing Node
**Module**: `common/topology/src/node.rs`
```rust
pub struct RoutingNode {
pub fn node_id(&self) -> NodeId
pub fn identity_key(&self) -> &ed25519::PublicKey
pub fn sphinx_key(&self) -> &x25519::PublicKey
pub fn mix_host(&self) -> &SocketAddr
pub fn clients_ws_address(&self) -> Option<&Url>
}
```
---
## 5. Blockchain Operations
### Validator Client
**Module**: `common/client-libs/validator-client/src/client.rs`
<!-- AIDEV-NOTE: Complex area - Handles all blockchain interactions -->
```rust
pub struct Client<C, S = NoSigner> {
// Contract queries
pub async fn query_contract_state<T>(
&self,
contract: &str,
query: T,
) -> Result<ContractStateResponse, ValidatorClientError>
where T: Into<Binary>
// Transaction execution (requires signer)
pub async fn execute_contract_message<M>(
&self,
contract: &str,
msg: M,
funds: Vec<Coin>,
) -> Result<TxResponse, ValidatorClientError>
where M: Into<Binary>
// Specific contract operations
pub async fn bond_mixnode(
&self,
mixnode: MixNode,
cost_params: MixNodeCostParams,
pledge: Coin,
) -> Result<TxResponse, ValidatorClientError>
pub async fn unbond_mixnode(&self) -> Result<TxResponse, ValidatorClientError>
pub async fn delegate_to_mixnode(
&self,
mix_id: MixId,
amount: Coin,
) -> Result<TxResponse, ValidatorClientError>
}
// Nyxd-specific client
pub type DirectSigningHttpRpcNyxdClient =
nyxd::NyxdClient<HttpRpcClient, DirectSecp256k1HdWallet>;
```
### Contract Queries
**Module**: `common/client-libs/validator-client/src/nyxd/contract_traits/`
```rust
// Mixnet contract queries
pub trait MixnetQueryClient {
async fn get_mixnodes(&self) -> Result<Vec<MixNodeDetails>, NyxdError>
async fn get_gateways(&self) -> Result<Vec<Gateway>, NyxdError>
async fn get_current_epoch(&self) -> Result<Epoch, NyxdError>
async fn get_rewarded_set(&self) -> Result<EpochRewardedSet, NyxdError>
}
// Vesting contract queries
pub trait VestingQueryClient {
async fn get_vesting_details(&self, address: &str)
-> Result<VestingDetails, NyxdError>
}
// E-cash contract queries
pub trait EcashQueryClient {
async fn get_deposit(&self, id: DepositId)
-> Result<Deposit, NyxdError>
}
```
---
## 6. REST API Endpoints
### nym-api Main Routes
**Module**: `nym-api/src/main.rs` and submodules
<!-- AIDEV-NOTE: Navigation hint - Each module contains router setup and handlers -->
```rust
// Main API setup
#[tokio::main]
async fn main() -> Result<(), anyhow::Error> {
// Router configuration
let app = Router::new()
.merge(api_routes())
.merge(swagger_ui())
.layer(cors_layer())
.layer(trace_layer());
}
// Core API routes (various modules)
pub fn api_routes() -> Router {
Router::new()
.nest("/v1/status", status_routes())
.nest("/v1/mixnodes", mixnode_routes())
.nest("/v1/gateways", gateway_routes())
.nest("/v1/network", network_routes())
.nest("/v1/ecash", ecash_routes())
}
```
### Status Routes
**Module**: `nym-api/src/status/mod.rs`
```rust
pub async fn status_handler() -> impl IntoResponse {
Json(ApiStatusResponse {
status: "ok",
uptime: get_uptime(),
})
}
pub async fn health_check() -> impl IntoResponse {
StatusCode::OK
}
```
### Network Monitor Routes
**Module**: `nym-api/src/network_monitor/mod.rs`
```rust
pub async fn get_monitor_report(
State(state): State<AppState>,
) -> Result<Json<MonitorReport>, ApiError> {
// Returns network reliability report
}
pub async fn get_node_reliability(
Path(node_id): Path<NodeId>,
State(state): State<AppState>,
) -> Result<Json<NodeReliability>, ApiError> {
// Returns specific node reliability
}
```
### E-cash API
**Module**: `nym-api/src/ecash/mod.rs`
```rust
pub async fn verify_credential(
Json(credential): Json<Credential>,
State(state): State<AppState>,
) -> Result<Json<VerificationResponse>, ApiError> {
// Verifies e-cash credentials
}
pub async fn issue_credential(
Json(request): Json<IssuanceRequest>,
State(state): State<AppState>,
) -> Result<Json<IssuedCredential>, ApiError> {
// Issues new e-cash credentials
}
```
---
## 7. Credential & E-cash
### Credential Operations
**Module**: `common/credentials/src/ecash/mod.rs`
```rust
// Credential spending
pub struct CredentialSpendingData {
pub fn new(
ticketbook: IssuedTicketBook,
gateway_identity: ed25519::PublicKey,
) -> CredentialSpendingData
pub fn prepare_for_spending(
&self,
request_id: i64,
) -> PreparedCredential
}
// Credential signing
pub struct CredentialSigningData {
pub fn sign_credential(
&self,
blinded_credential: BlindedCredential,
) -> Result<BlindedSignature, CredentialError>
}
// Aggregation utilities
pub fn aggregate_verification_keys(
keys: Vec<VerificationKey>,
) -> AggregatedVerificationKey
pub fn obtain_aggregate_wallet(
verification_keys: Vec<VerificationKey>,
commitments: Vec<Commitment>,
) -> Result<AggregateWallet, CredentialError>
```
### Ticketbook Operations
**Module**: `common/credentials/src/ecash/bandwidth/mod.rs`
<!-- AIDEV-NOTE: Complex area - Ticketbooks contain bandwidth credentials -->
```rust
pub struct IssuedTicketBook {
pub fn new(
tickets: Vec<IssuedTicket>,
expiration: OffsetDateTime,
) -> IssuedTicketBook
pub fn total_bandwidth(&self) -> Bandwidth
pub fn is_expired(&self) -> bool
pub fn consume_ticket(&mut self) -> Option<IssuedTicket>
}
pub struct ImportableTicketBook {
pub fn try_from_base58(s: &str) -> Result<Self, CredentialError>
pub fn into_issued(self) -> Result<IssuedTicketBook, CredentialError>
}
```
---
## 8. Smart Contracts
### Mixnet Contract Entry Points
**Module**: `contracts/mixnet/src/contract.rs`
```rust
#[entry_point]
pub fn instantiate(
deps: DepsMut,
env: Env,
info: MessageInfo,
msg: InstantiateMsg,
) -> Result<Response, ContractError>
#[entry_point]
pub fn execute(
deps: DepsMut,
env: Env,
info: MessageInfo,
msg: ExecuteMsg,
) -> Result<Response, ContractError>
#[entry_point]
pub fn query(
deps: Deps,
env: Env,
msg: QueryMsg,
) -> StdResult<Binary>
```
### Execute Message Handlers
**Module**: `contracts/mixnet/src/contract.rs`
```rust
// Node operations
fn try_bond_mixnode(
deps: DepsMut,
env: Env,
info: MessageInfo,
mixnode: MixNode,
) -> Result<Response, ContractError>
fn try_unbond_mixnode(
deps: DepsMut,
env: Env,
info: MessageInfo,
) -> Result<Response, ContractError>
// Delegation operations
fn try_delegate(
deps: DepsMut,
env: Env,
info: MessageInfo,
mix_id: MixId,
) -> Result<Response, ContractError>
fn try_undelegate(
deps: DepsMut,
env: Env,
info: MessageInfo,
mix_id: MixId,
) -> Result<Response, ContractError>
// Reward operations
fn try_reward_mixnode(
deps: DepsMut,
env: Env,
mix_id: MixId,
performance: Performance,
) -> Result<Response, ContractError>
```
### Query Message Handlers
```rust
fn query_mixnode(deps: Deps, mix_id: MixId) -> StdResult<MixnodeDetails>
fn query_gateways(deps: Deps) -> StdResult<Vec<Gateway>>
fn query_rewarded_set(deps: Deps, epoch: Epoch) -> StdResult<EpochRewardedSet>
fn query_current_epoch(deps: Deps) -> StdResult<Epoch>
```
---
## 9. Common Patterns
### Function Naming Conventions
<!-- AIDEV-NOTE: Pattern reference - Consistent naming helps code discovery -->
```rust
// Constructors
pub fn new(...) -> Self // Standard constructor
pub fn with_defaults() -> Self // Constructor with defaults
pub fn from_config(config: Config) -> Self // From configuration
// Async initialization
pub async fn init(...) -> Result<T> // Async initialization
pub async fn initialise(...) -> Result<T> // British spelling variant
pub async fn setup(...) -> Result<T> // Setup function
// Builder pattern
pub fn builder() -> TBuilder // Create builder
pub fn set_field(mut self, val: T) -> Self // Builder setter
pub fn build(self) -> Result<T> // Build final object
// Getters
pub fn field(&self) -> &T // Immutable reference
pub fn field_mut(&mut self) -> &mut T // Mutable reference
pub fn into_inner(self) -> T // Consume and return inner
// Queries
pub fn is_valid(&self) -> bool // Boolean check
pub fn has_field(&self) -> bool // Existence check
pub fn contains(&self, item: &T) -> bool // Contains check
// Transformations
pub fn to_type(&self) -> Type // Convert to type
pub fn into_type(self) -> Type // Consume and convert
pub fn try_into_type(self) -> Result<Type> // Fallible conversion
```
### Error Handling Patterns
```rust
// Custom error types with thiserror
#[derive(Error, Debug)]
pub enum ModuleError {
#[error("Network error: {0}")]
Network(#[from] NetworkError),
#[error("Invalid configuration: {reason}")]
InvalidConfig { reason: String },
#[error(transparent)]
Other(#[from] anyhow::Error),
}
// Result type alias
pub type Result<T> = std::result::Result<T, ModuleError>;
// Error conversion
impl From<io::Error> for ModuleError {
fn from(err: io::Error) -> Self {
ModuleError::Io(err)
}
}
```
### Async Patterns
```rust
// Async trait (with async-trait crate)
#[async_trait]
pub trait AsyncOperation {
async fn perform(&self) -> Result<()>;
}
// Spawning tasks
tokio::spawn(async move {
// Task code
});
// Channels for communication
let (tx, mut rx) = mpsc::channel(100);
// Select on multiple futures
tokio::select! {
result = future1 => { /* handle */ },
result = future2 => { /* handle */ },
_ = shutdown.recv() => { /* shutdown */ },
}
```
### Storage Patterns
```rust
// SQLx queries
sqlx::query!(
"SELECT * FROM nodes WHERE id = ?",
node_id
)
.fetch_optional(&pool)
.await?;
// In-memory caching
use dashmap::DashMap;
let cache: DashMap<Key, Value> = DashMap::new();
// File storage
use std::fs;
fs::write(path, data)?;
let content = fs::read_to_string(path)?;
```
---
## 10. Import Reference
### Standard Imports by Category
```rust
// Nym crypto
use nym_crypto::asymmetric::{ed25519, x25519};
use nym_crypto::symmetric::stream_cipher;
// Sphinx protocol
use nym_sphinx::forwarding::packet::MixPacket;
use nym_sphinx::framing::codec::NymCodec;
use nym_sphinx::addressing::nodes::NymNodeRoutingAddress;
use nym_sphinx::params::{PacketSize, DEFAULT_PACKET_SIZE};
// Client libraries
use nym_client_core::client::Client;
use nym_gateway_client::GatewayClient;
use nym_validator_client::ValidatorClient;
// Topology
use nym_topology::{NymTopology, RoutingNode};
use nym_mixnet_contract_common::NodeId;
// Configuration
use nym_network_defaults::NymNetworkDetails;
use nym_config::defaults::NymNetwork;
// Async runtime
use tokio::sync::{mpsc, RwLock, Mutex};
use tokio::time::{sleep, Duration};
use futures::{StreamExt, SinkExt};
// Error handling
use thiserror::Error;
use anyhow::{anyhow, Result, Context};
// Logging
use tracing::{debug, info, warn, error, instrument};
// Serialization
use serde::{Deserialize, Serialize};
use serde_json::{json, Value};
// Web framework (API)
use axum::{Router, extract::{Path, Query, State}, response::IntoResponse};
use axum::Json;
```
---
## 11. Feature Flags
### Common Feature Gates
```rust
// Client-specific features
#[cfg(feature = "client")]
#[cfg(feature = "cli")]
// Platform-specific
#[cfg(not(target_arch = "wasm32"))]
#[cfg(target_arch = "wasm32")]
// Testing
#[cfg(test)]
#[cfg(feature = "testing")]
#[cfg(feature = "contract-testing")]
// Storage backends
#[cfg(feature = "fs-surb-storage")]
#[cfg(feature = "fs-credentials-storage")]
// Network features
#[cfg(feature = "http-client")]
#[cfg(feature = "websocket")]
```
---
## Quick Lookup Tables
### Async vs Sync Functions
| Operation Type | Typically Async | Typically Sync |
|---------------|-----------------|----------------|
| Network I/O | ✓ | |
| Database queries | ✓ | |
| Contract execution | ✓ | |
| Cryptographic ops | | ✓ |
| Message construction | | ✓ |
| Configuration parsing | | ✓ |
| Topology queries | Both | Both |
### Return Type Patterns
| Pattern | Usage | Example |
|---------|-------|---------|
| `Result<T, E>` | Fallible operations | `connect() -> Result<Client>` |
| `Option<T>` | May not exist | `get_node() -> Option<Node>` |
| `impl Trait` | Return trait impl | `handler() -> impl IntoResponse` |
| `Box<dyn Trait>` | Dynamic dispatch | `create() -> Box<dyn Storage>` |
| Direct type | Infallible ops | `new() -> Self` |
### Module Organization
| Module Type | Location Pattern | Naming Convention |
|------------|------------------|-------------------|
| Binary entry | `/src/main.rs` | - |
| Library root | `/src/lib.rs` | - |
| Submodules | `/src/module/mod.rs` | snake_case |
| Tests | `/src/module/tests.rs` | #[cfg(test)] |
| Errors | `/src/error.rs` | ModuleError |
| Config | `/src/config.rs` | Config struct |
---
<!-- AIDEV-NOTE: This lexicon provides rapid function lookup. Use Ctrl+F to search for specific operations -->
+81
View File
@@ -0,0 +1,81 @@
# CLAUDE.md - nym-kcp
KCP (Fast and Reliable ARQ Protocol) implementation providing reliability over UDP for the Nym network. This crate ensures ordered, reliable delivery of packets.
## Architecture Overview
### Core Components
**KcpDriver** (src/driver.rs)
- High-level interface for KCP operations
- Manages single KCP session and I/O buffer
- Handles packet encoding/decoding
**KcpSession** (src/session.rs)
- Core KCP state machine
- Manages send/receive windows, RTT, congestion control
- Implements ARQ (Automatic Repeat Request) logic
**KcpPacket** (src/packet.rs)
- Wire format: conv(4B) | cmd(1B) | frg(1B) | wnd(2B) | ts(4B) | sn(4B) | una(4B) | len(4B) | data
- Commands: PSH (data), ACK, WND (window probe), ERR
## Key Concepts
### Conversation ID (conv)
- Unique identifier for each KCP connection
- Generated from hash of destination in nym-lp-node
- Must match on both ends for successful communication
### Packet Flow
1. **Send Path**: `send()` → Queue in send buffer → `fetch_outgoing()` → Wire
2. **Receive Path**: Wire → `input()` → Process ACKs/data → Application buffer
3. **Update Loop**: Call `update()` regularly to handle timeouts/retransmissions
### Reliability Mechanisms
- **Sequence Numbers (sn)**: Track packet ordering
- **Fragment Numbers (frg)**: Handle message fragmentation
- **UNA (Unacknowledged)**: Cumulative ACK up to this sequence
- **Selective ACK**: Via individual ACK packets
- **Fast Retransmit**: Triggered by duplicate ACKs
- **RTO Calculation**: Smoothed RTT with variance
## Configuration Parameters
```rust
// In KcpSession
MSS: 1400 // Maximum segment size
WINDOW_SIZE: 128 // Send/receive window
RTO_MIN: 100ms // Minimum retransmission timeout
RTO_MAX: 60000ms // Maximum retransmission timeout
FAST_RESEND: 2 // Fast retransmit threshold
```
## Common Operations
### Processing Incoming Data
```rust
driver.input(data)?; // Decode and process packets
let packets = driver.fetch_outgoing(); // Get any response packets
```
### Sending Data
```rust
driver.send(&data); // Queue for sending
driver.update(current_time); // Trigger flush
let packets = driver.fetch_outgoing(); // Get packets to send
```
## Debugging Tips
- Enable `trace!` logs to see packet-level details
- Monitor `ts_flush` vs `ts_current` for timing issues
- Check `snd_wnd` and `rcv_wnd` for flow control problems
- Watch for "fast retransmit" messages indicating packet loss
## Integration Notes
- AIDEV-NOTE: MSS must account for Sphinx packet overhead
- AIDEV-NOTE: Window size affects memory usage and throughput
- Update frequency impacts latency vs CPU usage tradeoff
- Conv ID must be consistent across session lifecycle
+27
View File
@@ -0,0 +1,27 @@
[package]
name = "nym-kcp"
version = "0.1.0"
edition = "2021"
[lib]
name = "nym_kcp"
path = "src/lib.rs"
[[bin]]
name = "wire_format"
path = "bin/wire_format/main.rs"
[[bin]]
name = "session"
path = "bin/session/main.rs"
[dependencies]
tokio-util = { workspace = true, features = ["codec"] }
byte_string = "1.0"
bytes = { workspace = true }
thiserror = { workspace = true }
log = { workspace = true }
ansi_term = "0.12"
[dev-dependencies]
env_logger = "0.11"
+80
View File
@@ -0,0 +1,80 @@
use bytes::BytesMut;
use log::info;
use nym_kcp::{packet::KcpPacket, session::KcpSession};
fn main() -> Result<(), Box<dyn std::error::Error>> {
// Create two KcpSessions, simulating two endpoints
let mut local_sess = KcpSession::new(42);
let mut remote_sess = KcpSession::new(42);
// Set an MSS (max segment size) smaller than our data to force fragmentation
local_sess.set_mtu(40);
remote_sess.set_mtu(40);
// Some data larger than 30 bytes to demonstrate multi-fragment
let big_data = b"The quick brown fox jumps over the lazy dog. This is a test.";
// --- LOCAL sends data ---
info!(
"Local: sending data: {:?}",
String::from_utf8_lossy(big_data)
);
local_sess.send(big_data);
// Update local session's logic at time=0
local_sess.update(100);
// LOCAL fetches outgoing (to be sent across the network)
let outgoing_pkts = local_sess.fetch_outgoing();
info!("Local: outgoing pkts: {:?}", outgoing_pkts);
// Here you'd normally encrypt and send them. Well just encode them into a buffer.
// Then that buffer is "transferred" to the remote side.
let mut wire_buf = BytesMut::new();
for pkt in &outgoing_pkts {
pkt.encode(&mut wire_buf);
}
// --- REMOTE receives data ---
// The remote side "decrypts" (here we just clone) and decodes
let mut remote_in = wire_buf.clone();
// Decode zero or more KcpPackets from remote_in
while let Some(decoded_pkt) = KcpPacket::decode(&mut remote_in)? {
info!(
"Decoded packet, sn: {}, frg: {}",
decoded_pkt.sn(),
decoded_pkt.frg()
);
remote_sess.input(&decoded_pkt);
}
// Update remote session to process newly received data
remote_sess.update(100);
// The remote session likely generated ACK packets
let ack_pkts = remote_sess.fetch_outgoing();
// --- LOCAL receives ACKs ---
// The local side decodes them
let mut ack_buf = BytesMut::new();
for pkt in &ack_pkts {
pkt.encode(&mut ack_buf);
}
while let Some(decoded_pkt) = KcpPacket::decode(&mut ack_buf)? {
local_sess.input(&decoded_pkt);
}
// Update local again with some arbitrary time, e.g. 50 ms later
local_sess.update(100);
// Just for completeness, local might produce more packets, though typically it's just empty now
let _ = local_sess.fetch_outgoing();
// --- REMOTE reads reassembled data ---
let incoming = remote_sess.fetch_incoming();
info!("Remote: incoming pkts: {:?}", incoming);
Ok(())
}
+83
View File
@@ -0,0 +1,83 @@
use std::{
fs::File,
io::{BufRead as _, BufReader},
};
use bytes::BytesMut;
use log::info;
use nym_kcp::{
codec::KcpCodec,
packet::{KcpCommand, KcpPacket},
};
use tokio_util::codec::{Decoder as _, Encoder as _};
fn main() -> Result<(), Box<dyn std::error::Error>> {
// 1) Open a file and read lines
let file = File::open("bin/wire_format/packets.txt")?;
let reader = BufReader::new(file);
// 2) Create our KcpCodec
let mut codec = KcpCodec {};
// We'll use out_buf for encoded data from *all* lines
let mut out_buf = BytesMut::new();
let mut input_lines = vec![];
// Read lines & encode them all
for (i, line) in reader.lines().enumerate() {
let line = line?;
info!("Original line #{}: {}", i + 1, line);
// Construct a KcpPacket
let pkt = KcpPacket::new(
42,
KcpCommand::Push,
0,
128,
0,
i as u32,
0,
line.as_bytes().to_vec(),
);
input_lines.push(pkt.clone_data());
// Encode (serialize) the packet into out_buf
codec.encode(pkt, &mut out_buf)?;
}
// === Simulate encryption & transmission ===
// In reality, you might do `encrypt(&out_buf)` and then
// send it over the network. We'll just clone here:
let mut received_buf = out_buf.clone();
// 3) Now decode (deserialize) all packets at once
// For demonstration, read them back out
let mut count = 0;
let mut decoded_lines = vec![];
#[allow(clippy::while_let_loop)]
loop {
match codec.decode(&mut received_buf)? {
Some(decoded_pkt) => {
count += 1;
// Convert packet data back to a string
let decoded_str = String::from_utf8_lossy(decoded_pkt.data());
info!("Decoded line #{}: {}", decoded_pkt.sn() + 1, decoded_str);
decoded_lines.push(decoded_pkt.clone_data());
}
None => break,
}
}
for (i, j) in input_lines.iter().zip(decoded_lines.iter()) {
assert_eq!(i, j);
}
info!("Decoded {} lines total.", count);
Ok(())
}
@@ -0,0 +1,10 @@
packet 1
packet 2
packet 3
packet 4
packet 5
packet 6
packet 7
packet 8
packet 9
packet 10
@@ -0,0 +1,85 @@
# Nym-KCP Code Review Summary (2024-07-31)
Based on an initial code review, the following potential issues and areas for improvement were identified in the `nym-kcp` crate:
## Potential Bugs / Protocol Deviations
1. **Simplified Windowing (`session.rs: move_queue_to_buf`):**
* **Issue:** ~~Currently only considers the local send window (`snd_wnd`), ignoring the remote receive window (`rmt_wnd`).~~
* **Status:** Confirmed OK. The implementation correctly uses `cwnd = min(snd_wnd, rmt_wnd)`.
* **Impact:** ~~Violates KCP congestion control principles (`cwnd = min(snd_wnd, rmt_wnd)`). Can potentially overwhelm the receiver.~~ **(Initial concern resolved)**
2. **Naive RTO Backoff (`session.rs: flush_outgoing`):**
* **Issue:** ~~Uses a simple linear increase (`rto += max(rto, rx_rto)`) instead of standard exponential backoff.~~
* **Status:** Resolved. Changed to exponential backoff (`rto *= 2`) clamped to 60s.
* **Impact:** ~~Slower recovery from packet loss/congestion compared to standard KCP.~~
3. **Less Robust UNA Update (`session.rs: parse_una`):**
* **Issue:** ~~Uses `self.snd_una = una` instead of `max(self.snd_una, una)`. ~~
* **Status:** Resolved. Changed to use `cmp::max(self.snd_una, una)`.
* **Impact:** ~~Less resilient to out-of-order packets carrying older UNA values.~~
## Areas for Improvement / Robustness
4. **Limited Testing (`session.rs: tests`):**
* **Issue:** Only one test case focusing on out-of-order fragment reassembly.
* **Impact:** Insufficient coverage for loss, retransmissions, windowing, edge cases. Low confidence in overall robustness.
5. **Unimplemented Wask/Wins (`session.rs: input`):**
* **Issue:** `KcpCommand::Wask` and `KcpCommand::Wins` are not handled.
* **Impact:** Session cannot probe or react to dynamic changes in the peer's receive window.
6. **Concurrency Locking (`driver.rs`):**
* **Issue:** `Arc<Mutex<>>` with `try_lock` and exponential backoff loop.
* **Impact:** Potential performance bottleneck under high contention; hardcoded retry limit.
7. **Fragment Reassembly Complexity (`session.rs: move_buf_to_queue`):**
* **Issue:** Logic for reassembling fragments, while plausible, is complex and needs thorough testing.
* **Impact:** Potential for subtle bugs related to sequence numbers, buffer state.
## Next Steps
* ~~Address the windowing logic deviation (Priority 1).~~ (Confirmed OK)
* Enhance test suite significantly.
* Implement Wask/Wins handling.
* ~~Refine RTO backoff mechanism.~~ (Resolved)
* (Optional) Test robustness of UNA update logic against out-of-order packets.
## Code Fixes
* **RTO Backoff:** Updated `flush_outgoing` to use exponential backoff (`rto *= 2`) for segment retransmissions, clamped to a maximum (60s), instead of the previous linear increase. Addresses Review Item #2.
* **UNA Update:** Updated `parse_una` to use `cmp::max(self.snd_una, una)` for more robust handling of out-of-order packets. Addresses Review Item #3.
* **Windowing Logic:** Confirmed that `move_queue_to_buf` correctly calculates `cwnd = min(snd_wnd, rmt_wnd)`. Initial concern in Review Item #1 was based on misunderstanding or outdated code.
## Proposed Testing Enhancements
1. **Windowing Behavior Tests:**
* Verify `cwnd = min(snd_wnd, rmt_wnd)` limit on outgoing segments.
* Verify `Write` trait returns `ErrorKind::WouldBlock` when `cwnd` is full.
2. **Retransmission & RTO Tests:**
* Simulate packet loss and verify retransmission occurs after RTO.
* Verify RTO backoff mechanism (current naive, future standard).
* Verify ACK prevents scheduled retransmission.
3. **ACK & UNA Processing Tests:**
* Verify UNA correctly clears acknowledged segments from `snd_buf`.
* Verify specific ACK removes the correct segment and updates RTT.
* Test robustness against out-of-order ACKs/UNA (requires `parse_una` fix).
4. **More Fragmentation/Reassembly Tests:**
* Test diverse out-of-order delivery patterns.
* Test handling of duplicate fragments.
* Test loss of fragments and subsequent retransmission/reassembly.
## Testing Progress (2024-08-01)
The following tests have been implemented in `session.rs` based on the proposed enhancements:
* `test_congestion_window_limits_send_buffer`: Verifies that the number of segments moved to `snd_buf` respects `cwnd = min(snd_wnd, rmt_wnd)`. (Addresses Windowing Behavior Test 1)
* `test_segment_retransmission_after_rto`: Verifies that a segment is retransmitted if its RTO expires without an ACK. (Addresses Retransmission Test 1)
* `test_ack_removes_segment_from_send_buffer`: Verifies that receiving a specific ACK removes the corresponding segment from `snd_buf`. (Addresses ACK Processing Test 2, first part)
* `test_ack_updates_rtt`: Verifies that receiving a specific ACK updates the session's RTT estimate and RTO. (Addresses ACK Processing Test 2, second part)
* `test_una_clears_send_buffer`: Verifies that receiving a packet with a UNA value clears all segments with `sn < una` from `snd_buf`. (Addresses ACK Processing Test 1)
## Testing Progress (2024-08-02)
* `test_write_fills_send_queue_when_window_full`: Verifies that `Write` limits accepted data based on `snd_wnd` and `update` respects `cwnd` when moving segments. (Partially addresses Windowing Behavior Test 2)
* `test_ack_prevents_retransmission`: Verifies that a segment is not retransmitted if it is ACKed before its RTO expires. (Addresses Retransmission Test 3)
* `test_duplicate_fragment_handling`: Verifies that the receiver correctly ignores duplicate fragments during reassembly. (Addresses Fragmentation Test 2)
* `test_fragment_loss_and_reassembly`: Verifies that a lost fragment is retransmitted after RTO and the receiver can reassemble the message upon receiving it. (Addresses Fragmentation Test 3)
+30
View File
@@ -0,0 +1,30 @@
use std::io;
use bytes::BytesMut;
use tokio_util::codec::{Decoder, Encoder};
use super::packet::KcpPacket;
/// Our codec for encoding/decoding KCP packets
#[derive(Debug, Default)]
pub struct KcpCodec;
impl Decoder for KcpCodec {
type Item = KcpPacket;
type Error = io::Error;
fn decode(&mut self, src: &mut BytesMut) -> Result<Option<Self::Item>, Self::Error> {
// We simply delegate to `KcpPacket::decode`
KcpPacket::decode(src).map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))
}
}
impl Encoder<KcpPacket> for KcpCodec {
type Error = io::Error;
fn encode(&mut self, item: KcpPacket, dst: &mut BytesMut) -> Result<(), Self::Error> {
// We just call `item.encode` to append the bytes
item.encode(dst);
Ok(())
}
}
+60
View File
@@ -0,0 +1,60 @@
use bytes::BytesMut;
use log::{debug, trace};
use crate::{error::KcpError, packet::KcpPacket, session::KcpSession};
pub struct KcpDriver {
session: KcpSession,
buffer: BytesMut,
}
impl KcpDriver {
pub fn conv_id(&self) -> Result<u32, KcpError> {
Ok(self.session.conv)
}
pub fn send(&mut self, data: &[u8]) {
self.session.send(data);
}
pub fn input(&mut self, data: &[u8]) -> Result<Vec<KcpPacket>, KcpError> {
self.buffer.extend_from_slice(data);
let mut pkts = Vec::new();
while let Ok(Some(pkt)) = KcpPacket::decode(&mut self.buffer) {
debug!(
"Decoded packet, cmd: {}, sn: {}, frg: {}",
pkt.command(),
pkt.sn(),
pkt.frg()
);
self._input(&pkt)?;
pkts.push(pkt);
}
Ok(pkts)
}
fn _input(&mut self, pkt: &KcpPacket) -> Result<(), KcpError> {
self.session.input(pkt);
Ok(())
}
pub fn fetch_outgoing(&mut self) -> Vec<KcpPacket> {
trace!(
"ts_flush: {}, ts_current: {}",
self.session.ts_flush(),
self.session.ts_current()
);
self.session.fetch_outgoing()
}
pub fn update(&mut self, tick: u64) {
self.session.update(tick as u32);
}
pub fn new(session: KcpSession) -> Self {
KcpDriver {
session,
buffer: BytesMut::new(),
}
}
}
+10
View File
@@ -0,0 +1,10 @@
use thiserror::Error;
#[derive(Error, Debug)]
pub enum KcpError {
#[error("Invalid KCP command value: {0}")]
InvalidCommand(u8),
#[error("IO error: {0}")]
Io(#[from] std::io::Error),
}
+5
View File
@@ -0,0 +1,5 @@
pub mod codec;
pub mod driver;
pub mod error;
pub mod packet;
pub mod session;
+219
View File
@@ -0,0 +1,219 @@
use bytes::{Buf, BufMut, BytesMut};
use log::{debug, trace};
use super::error::KcpError;
pub const KCP_HEADER: usize = 24;
/// Typed enumeration for KCP commands.
#[repr(u8)]
#[derive(Debug, Copy, Clone, PartialEq, Eq)]
pub enum KcpCommand {
Push = 81, // cmd: push data
Ack = 82, // cmd: ack
Wask = 83, // cmd: window probe (ask)
Wins = 84, // cmd: window size (tell)
}
impl std::fmt::Display for KcpCommand {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
KcpCommand::Push => write!(f, "Push"),
KcpCommand::Ack => write!(f, "Ack"),
KcpCommand::Wask => write!(f, "Window Probe (ask)"),
KcpCommand::Wins => write!(f, "Window Size (tell)"),
}
}
}
impl TryFrom<u8> for KcpCommand {
type Error = KcpError;
fn try_from(value: u8) -> Result<Self, Self::Error> {
match value {
81 => Ok(KcpCommand::Push),
82 => Ok(KcpCommand::Ack),
83 => Ok(KcpCommand::Wask),
84 => Ok(KcpCommand::Wins),
_ => Err(KcpError::InvalidCommand(value)),
}
}
}
#[allow(clippy::from_over_into)]
impl Into<u8> for KcpCommand {
fn into(self) -> u8 {
self as u8
}
}
/// A single KCP packet (on-wire format).
#[derive(Debug, Clone)]
pub struct KcpPacket {
conv: u32,
cmd: KcpCommand,
frg: u8,
wnd: u16,
ts: u32,
sn: u32,
una: u32,
data: Vec<u8>,
}
#[allow(clippy::too_many_arguments)]
impl KcpPacket {
pub fn new(
conv: u32,
cmd: KcpCommand,
frg: u8,
wnd: u16,
ts: u32,
sn: u32,
una: u32,
data: Vec<u8>,
) -> Self {
Self {
conv,
cmd,
frg,
wnd,
ts,
sn,
una,
data,
}
}
pub fn command(&self) -> KcpCommand {
self.cmd
}
pub fn data(&self) -> &[u8] {
&self.data
}
pub fn clone_data(&self) -> Vec<u8> {
self.data.clone()
}
pub fn conv(&self) -> u32 {
self.conv
}
pub fn cmd(&self) -> KcpCommand {
self.cmd
}
pub fn frg(&self) -> u8 {
self.frg
}
pub fn wnd(&self) -> u16 {
self.wnd
}
pub fn ts(&self) -> u32 {
self.ts
}
pub fn sn(&self) -> u32 {
self.sn
}
pub fn una(&self) -> u32 {
self.una
}
}
impl Default for KcpPacket {
fn default() -> Self {
// We must pick some default command, e.g. `Push`.
// Or omit `Default` if you don't need it.
KcpPacket {
conv: 0,
cmd: KcpCommand::Push,
frg: 0,
wnd: 0,
ts: 0,
sn: 0,
una: 0,
data: Vec::new(),
}
}
}
impl KcpPacket {
/// Attempt to decode a `KcpPacket` from `src`.
/// Returns Ok(Some(pkt)) if fully available, Ok(None) if not enough data,
/// or Err(...) if there's an invalid command or other error.
pub fn decode(src: &mut BytesMut) -> Result<Option<Self>, KcpError> {
trace!("Decoding buffer with len: {}", src.len());
if src.len() < KCP_HEADER {
// Not enough for even the header, this is usually fine, more data will arrive
debug!("Not enough data for header");
return Ok(None);
}
// Peek into the first 28 bytes
let mut header = &src[..KCP_HEADER];
let conv = header.get_u32_le();
let cmd_byte = header.get_u8();
let frg = header.get_u8();
let wnd = header.get_u16_le();
let ts = header.get_u32_le();
let sn = header.get_u32_le();
let una = header.get_u32_le();
let len = header.get_u32_le() as usize;
let total_needed = KCP_HEADER + len;
if src.len() < total_needed {
// We don't have the full packet yet
debug!(
"Not enough data for packet, want {}, have {}",
total_needed,
src.len()
);
return Ok(None);
}
// Convert the raw u8 into our KcpCommand enum
let cmd = KcpCommand::try_from(cmd_byte)?;
// Now we can read out the data portion
let data = src[KCP_HEADER..KCP_HEADER + len].to_vec();
// Advance the buffer so it no longer contains this packet
src.advance(total_needed);
Ok(Some(Self {
conv,
cmd,
frg,
wnd,
ts,
sn,
una,
data,
}))
}
/// Encode this packet into `dst`.
pub fn encode(&self, dst: &mut BytesMut) {
let total_len = KCP_HEADER + self.data.len();
trace!("Encoding packet: {:?}, len: {}", self, total_len);
dst.reserve(total_len);
dst.put_u32_le(self.conv);
dst.put_u8(self.cmd.into()); // Convert enum -> u8
dst.put_u8(self.frg);
dst.put_u16_le(self.wnd);
dst.put_u32_le(self.ts);
dst.put_u32_le(self.sn);
dst.put_u32_le(self.una);
dst.put_u32_le(self.data.len() as u32);
dst.extend_from_slice(&self.data);
trace!("Encoded packet: {:?}, len: {}", dst, dst.len());
}
}
File diff suppressed because it is too large Load Diff
+6
View File
@@ -0,0 +1,6 @@
[package]
name = "nym-lp-common"
version = "0.1.0"
edition = "2021"
[dependencies]
+28
View File
@@ -0,0 +1,28 @@
use std::fmt;
use std::fmt::Write;
pub fn format_debug_bytes(bytes: &[u8]) -> Result<String, fmt::Error> {
let mut out = String::new();
const LINE_LEN: usize = 16;
for (i, chunk) in bytes.chunks(LINE_LEN).enumerate() {
let line_prefix = format!("[{}:{}]", 1 + i * LINE_LEN, i * LINE_LEN + chunk.len());
write!(out, "{line_prefix:12}")?;
let mut line = String::new();
for b in chunk {
line.push_str(format!("{:02x} ", b).as_str());
}
write!(
out,
"{line:48} {}",
chunk
.iter()
.map(|&b| b as char)
.map(|c| if c.is_alphanumeric() { c } else { '.' })
.collect::<String>()
)?;
writeln!(out)?;
}
Ok(out)
}
+30
View File
@@ -0,0 +1,30 @@
[package]
name = "nym-lp"
version = "0.1.0"
edition = "2021"
[dependencies]
bincode = { workspace = true }
thiserror = { workspace = true }
parking_lot = "0.12"
snow = "0.9.6"
bs58 = "0.5.1"
serde = { workspace = true }
bytes = { workspace = true }
dashmap = "6.1.0"
sha2 = "0.10"
ansi_term = "0.12"
utoipa = { workspace = true, features = ["macros", "non_strict_integers"] }
nym-lp-common = { path = "../nym-lp-common" }
nym-sphinx = { path = "../nymsphinx" }
[dev-dependencies]
criterion = { version = "0.5", features = ["html_reports"] }
rand = "0.8"
rand_chacha = "0.3"
[[bench]]
name = "replay_protection"
harness = false
+71
View File
@@ -0,0 +1,71 @@
# Nym Lewes Protocol
The Lewes Protocol (LP) is a secure network communication protocol implemented in Rust. This README provides an overview of the protocol's session management and replay protection mechanisms.
## Architecture Overview
```
+-----------------+ +----------------+ +---------------+
| Transport Layer |<--->| LP Session |<--->| LP Codec |
| (UDP/TCP) | | - Replay prot. | | - Enc/dec only|
+-----------------+ | - Crypto state | +---------------+
+----------------+
```
## Packet Structure
The protocol uses a structured packet format:
```
+------------------+-------------------+------------------+
| Header (16B) | Message | Trailer (16B) |
| - Version (1B) | - Type (2B) | - Authentication |
| - Reserved (3B) | - Content | - tag/MAC |
| - SenderIdx (4B) | | |
| - Counter (8B) | | |
+------------------+-------------------+------------------+
```
- Header contains protocol version, sender identification, and counter for replay protection
- Message carries the actual payload with a type identifier
- Trailer provides authentication and integrity verification (16 bytes)
- Total packet size is constrained by MTU (1500 bytes), accounting for network overhead
## Sessions
- Sessions are managed by `LPSession` and `SessionManager` classes
- Each session has unique receiving and sending indices to identify connections
- Sessions maintain:
- Cryptographic state (currently mocked implementations)
- Counter for outgoing packets
- Replay protection mechanism for incoming packets
## Session Management
- `SessionManager` handles session lifecycle (creation, retrieval, removal)
- Sessions are stored in a thread-safe HashMap indexed by receiving index
- The manager generates unique indices for new sessions
- Sessions are Arc-wrapped for safe concurrent access
## Replay Protection
- Implemented in the `ReceivingKeyCounterValidator` class
- Uses a bitmap-based approach to track received packet counters
- The bitmap allows reordering of up to 1024 packets (configurable)
- SIMD optimizations are used when available for performance
## Replay Protection Process
1. Quick validation (`will_accept_branchless`):
- Checks if counter is valid before expensive operations
- Detects duplicates, out-of-window packets
2. Marking packets (`mark_did_receive_branchless`):
- Updates the bitmap to mark counter as received
- Updates statistics and sliding window as needed
3. Window Sliding:
- Automatically slides the acceptance window when new higher counters arrive
- Clears bits for old counters that fall outside the window
This architecture effectively prevents replay attacks while allowing some packet reordering, an essential feature for secure network protocols.
+235
View File
@@ -0,0 +1,235 @@
use criterion::{black_box, criterion_group, criterion_main, BenchmarkId, Criterion, Throughput};
use nym_lp::replay::ReceivingKeyCounterValidator;
use parking_lot::Mutex;
use rand::{Rng, SeedableRng};
use rand_chacha::ChaCha8Rng;
use std::sync::Arc;
fn bench_sequential_counters(c: &mut Criterion) {
let mut group = c.benchmark_group("replay_sequential");
group.sample_size(1000);
for size in [100, 1000, 10000] {
group.throughput(Throughput::Elements(size as u64));
group.bench_with_input(
BenchmarkId::new("sequential_counters", size),
&size,
|b, &size| {
let validator = ReceivingKeyCounterValidator::default();
let counters: Vec<u64> = (0..size).collect();
b.iter(|| {
let mut validator = validator.clone();
for &counter in &counters {
let _ = black_box(validator.will_accept_branchless(counter));
let _ = black_box(validator.mark_did_receive_branchless(counter));
}
});
},
);
}
group.finish();
}
fn bench_out_of_order_counters(c: &mut Criterion) {
let mut group = c.benchmark_group("replay_out_of_order");
group.sample_size(1000);
for size in [100, 1000, 10000] {
group.throughput(Throughput::Elements(size as u64));
group.bench_with_input(
BenchmarkId::new("out_of_order_counters", size),
&size,
|b, &size| {
let validator = ReceivingKeyCounterValidator::default();
// Create random counters within a valid window
let mut rng = ChaCha8Rng::seed_from_u64(42);
let counters: Vec<u64> = (0..size).map(|_| rng.gen_range(0..1024)).collect();
b.iter(|| {
let mut validator = validator.clone();
for &counter in &counters {
let _ = black_box(validator.will_accept_branchless(counter));
let _ = black_box(validator.mark_did_receive_branchless(counter));
}
});
},
);
}
group.finish();
}
fn bench_thread_safety(c: &mut Criterion) {
let mut group = c.benchmark_group("replay_thread_safety");
group.sample_size(1000);
for size in [100, 1000, 10000] {
group.throughput(Throughput::Elements(size as u64));
group.bench_with_input(
BenchmarkId::new("thread_safe_validator", size),
&size,
|b, &size| {
let validator = Arc::new(Mutex::new(ReceivingKeyCounterValidator::default()));
let counters: Vec<u64> = (0..size).collect();
b.iter(|| {
for &counter in &counters {
let result = {
let guard = validator.lock();
black_box(guard.will_accept_branchless(counter))
};
if result.is_ok() {
let mut guard = validator.lock();
let _ = black_box(guard.mark_did_receive_branchless(counter));
}
}
});
},
);
}
group.finish();
}
fn bench_window_sliding(c: &mut Criterion) {
let mut group = c.benchmark_group("replay_window_sliding");
group.sample_size(100);
for window_size in [128, 512, 1024] {
group.throughput(Throughput::Elements(window_size as u64));
group.bench_with_input(
BenchmarkId::new("window_sliding", window_size),
&window_size,
|b, &window_size| {
b.iter(|| {
let mut validator = ReceivingKeyCounterValidator::default();
// First fill the window with sequential packets
for i in 0..window_size {
let _ = black_box(validator.mark_did_receive_branchless(i));
}
// Then jump ahead to force window sliding
let _ = black_box(validator.mark_did_receive_branchless(window_size * 3));
// Try some packets in the new window
for i in (window_size * 2 + 1)..(window_size * 3) {
let _ = black_box(validator.will_accept_branchless(i));
}
});
},
);
}
group.finish();
}
/// Benchmark operations that would benefit from SIMD optimization
fn bench_core_operations(c: &mut Criterion) {
let mut group = c.benchmark_group("replay_core_operations");
group.sample_size(1000);
// Create validators with different states
let mut empty_validator = ReceivingKeyCounterValidator::default();
let mut half_full_validator = ReceivingKeyCounterValidator::default();
let mut full_validator = ReceivingKeyCounterValidator::default();
// Fill validators with different patterns
for i in 0..512 {
half_full_validator.mark_did_receive_branchless(i).unwrap();
}
for i in 0..1024 {
full_validator.mark_did_receive_branchless(i).unwrap();
}
// Benchmark clearing operations
group.bench_function("clear_empty_window", |b| {
b.iter(|| {
let mut validator = empty_validator.clone();
// Force window sliding that will clear bitmap
black_box(validator.mark_did_receive_branchless(2000).unwrap());
})
});
group.bench_function("clear_half_full_window", |b| {
b.iter(|| {
let mut validator = half_full_validator.clone();
// Force window sliding that will clear bitmap
black_box(validator.mark_did_receive_branchless(2000).unwrap());
})
});
group.bench_function("clear_full_window", |b| {
b.iter(|| {
let mut validator = full_validator.clone();
// Force window sliding that will clear bitmap
black_box(validator.mark_did_receive_branchless(2000).unwrap());
})
});
group.finish();
}
/// Benchmark thread safety with different thread counts
fn bench_concurrency_scaling(c: &mut Criterion) {
let mut group = c.benchmark_group("replay_concurrency_scaling");
group.sample_size(50);
for thread_count in [1, 2, 4, 8] {
group.bench_with_input(
BenchmarkId::new("mutex_threads", thread_count),
&thread_count,
|b, &thread_count| {
b.iter(|| {
let validator = Arc::new(Mutex::new(ReceivingKeyCounterValidator::default()));
let mut handles = Vec::new();
for t in 0..thread_count {
let validator_clone = Arc::clone(&validator);
let handle = std::thread::spawn(move || {
let mut success_count = 0;
for i in 0..100 {
let counter = t * 1000 + i;
let mut guard = validator_clone.lock();
if guard.mark_did_receive_branchless(counter as u64).is_ok() {
success_count += 1;
}
}
success_count
});
handles.push(handle);
}
let mut total = 0;
for handle in handles {
total += handle.join().unwrap();
}
black_box(total)
})
},
);
}
group.finish();
}
criterion_group!(
replay_benches,
bench_sequential_counters,
bench_out_of_order_counters,
bench_thread_safety,
bench_window_sliding,
bench_core_operations,
bench_concurrency_scaling
);
criterion_main!(replay_benches);
+395
View File
@@ -0,0 +1,395 @@
// Copyright 2025 - Nym Technologies SA <contact@nymtech.net>
// SPDX-License-Identifier: Apache-2.0
use crate::message::{ClientHelloData, LpMessage, MessageType};
use crate::packet::{LpHeader, LpPacket, TRAILER_LEN};
use crate::LpError;
use bytes::BytesMut;
/// Parses a complete Lewes Protocol packet from a byte slice (e.g., a UDP datagram payload).
///
/// Assumes the input `src` contains exactly one complete packet. It does not handle
/// stream fragmentation or provide replay protection checks (these belong at the session level).
pub fn parse_lp_packet(src: &[u8]) -> Result<LpPacket, LpError> {
// Minimum size check: LpHeader + Type + Trailer (for 0-payload message)
let min_size = LpHeader::SIZE + 2 + TRAILER_LEN;
if src.len() < min_size {
return Err(LpError::InsufficientBufferSize);
}
// Parse LpHeader
let header = LpHeader::parse(&src[..LpHeader::SIZE])?; // Uses the new LpHeader::parse
// Parse Message Type
let type_start = LpHeader::SIZE;
let type_end = type_start + 2;
let mut message_type_bytes = [0u8; 2];
message_type_bytes.copy_from_slice(&src[type_start..type_end]);
let message_type_raw = u16::from_le_bytes(message_type_bytes);
let message_type = MessageType::from_u16(message_type_raw)
.ok_or_else(|| LpError::invalid_message_type(message_type_raw))?;
// Calculate payload size based on total length
let total_size = src.len();
let message_size = total_size - min_size; // Size of the payload part
// Extract payload based on message type
let message_start = type_end;
let message_end = message_start + message_size;
let payload_slice = &src[message_start..message_end]; // Bounds already checked by min_size and total_size calculation
let message = match message_type {
MessageType::Busy => {
if message_size != 0 {
return Err(LpError::InvalidPayloadSize {
expected: 0,
actual: message_size,
});
}
LpMessage::Busy
}
MessageType::Handshake => {
// No size validation needed here for Handshake, it's variable
LpMessage::Handshake(payload_slice.to_vec())
}
MessageType::EncryptedData => {
// No size validation needed here for EncryptedData, it's variable
LpMessage::EncryptedData(payload_slice.to_vec())
}
MessageType::ClientHello => {
// ClientHello has structured data
// Deserialize ClientHelloData from payload
let data: ClientHelloData = bincode::deserialize(payload_slice)
.map_err(|e| LpError::DeserializationError(e.to_string()))?;
LpMessage::ClientHello(data)
}
};
// Extract trailer
let trailer_start = message_end;
let trailer_end = trailer_start + TRAILER_LEN;
// Check if trailer_end exceeds src length (shouldn't happen if min_size check passed and calculation is correct, but good for safety)
if trailer_end > total_size {
// This indicates an internal logic error or buffer manipulation issue
return Err(LpError::InsufficientBufferSize); // Or a more specific internal error
}
let trailer_slice = &src[trailer_start..trailer_end];
let mut trailer = [0u8; TRAILER_LEN];
trailer.copy_from_slice(trailer_slice);
// Create and return the packet
Ok(LpPacket {
header,
message,
trailer,
})
}
/// Serializes an LpPacket into the provided BytesMut buffer.
pub fn serialize_lp_packet(item: &LpPacket, dst: &mut BytesMut) -> Result<(), LpError> {
// Reserve approximate size - consider making this more accurate if needed
dst.reserve(LpHeader::SIZE + 2 + item.message.len() + TRAILER_LEN);
item.encode(dst); // Use the existing encode method on LpPacket
Ok(())
}
// Add a new error variant for invalid message types (Moved from previous impl LpError block)
impl LpError {
pub fn invalid_message_type(message_type: u16) -> Self {
LpError::InvalidMessageType(message_type)
}
}
#[cfg(test)]
mod tests {
// Import standalone functions
use super::{parse_lp_packet, serialize_lp_packet};
// Keep necessary imports
use crate::message::{LpMessage, MessageType};
use crate::packet::{LpHeader, LpPacket, TRAILER_LEN};
use crate::LpError;
use bytes::BytesMut;
// Helper function to create a test packet's BytesMut representation directly
fn create_test_packet_bytes(counter: u64, message: LpMessage, trailer_fill: u8) -> BytesMut {
let packet = LpPacket {
header: LpHeader {
protocol_version: 1,
session_id: 42,
counter,
},
message,
trailer: [trailer_fill; TRAILER_LEN],
};
let mut buf = BytesMut::new();
serialize_lp_packet(&packet, &mut buf).unwrap();
buf
}
// === Updated Encode/Decode Tests ===
#[test]
fn test_serialize_parse_busy() {
let mut dst = BytesMut::new();
// Create a Busy packet
let packet = LpPacket {
header: LpHeader {
protocol_version: 1,
session_id: 42,
counter: 123,
},
message: LpMessage::Busy,
trailer: [0; TRAILER_LEN],
};
// Serialize the packet
serialize_lp_packet(&packet, &mut dst).unwrap();
// Parse the packet
let decoded = parse_lp_packet(&dst).unwrap();
// Verify the packet fields
assert_eq!(decoded.header.protocol_version, 1);
assert_eq!(decoded.header.session_id, 42);
assert_eq!(decoded.header.counter, 123);
assert!(matches!(decoded.message, LpMessage::Busy));
assert_eq!(decoded.trailer, [0; TRAILER_LEN]);
}
#[test]
fn test_serialize_parse_handshake() {
let mut dst = BytesMut::new();
// Create a Handshake message packet
let payload = vec![42u8; 80]; // Example payload size
let packet = LpPacket {
header: LpHeader {
protocol_version: 1,
session_id: 42,
counter: 123,
},
message: LpMessage::Handshake(payload.clone()),
trailer: [0; TRAILER_LEN],
};
// Serialize the packet
serialize_lp_packet(&packet, &mut dst).unwrap();
// Parse the packet
let decoded = parse_lp_packet(&dst).unwrap();
// Verify the packet fields
assert_eq!(decoded.header.protocol_version, 1);
assert_eq!(decoded.header.session_id, 42);
assert_eq!(decoded.header.counter, 123);
// Verify message type and data
match decoded.message {
LpMessage::Handshake(decoded_payload) => {
assert_eq!(decoded_payload, payload);
}
_ => panic!("Expected Handshake message"),
}
assert_eq!(decoded.trailer, [0; TRAILER_LEN]);
}
#[test]
fn test_serialize_parse_encrypted_data() {
let mut dst = BytesMut::new();
// Create an EncryptedData message packet
let payload = vec![43u8; 124]; // Example payload size
let packet = LpPacket {
header: LpHeader {
protocol_version: 1,
session_id: 42,
counter: 123,
},
message: LpMessage::EncryptedData(payload.clone()),
trailer: [0; TRAILER_LEN],
};
// Serialize the packet
serialize_lp_packet(&packet, &mut dst).unwrap();
// Parse the packet
let decoded = parse_lp_packet(&dst).unwrap();
// Verify the packet fields
assert_eq!(decoded.header.protocol_version, 1);
assert_eq!(decoded.header.session_id, 42);
assert_eq!(decoded.header.counter, 123);
// Verify message type and data
match decoded.message {
LpMessage::EncryptedData(decoded_payload) => {
assert_eq!(decoded_payload, payload);
}
_ => panic!("Expected EncryptedData message"),
}
assert_eq!(decoded.trailer, [0; TRAILER_LEN]);
}
// === Updated Incomplete Data Tests ===
#[test]
fn test_parse_incomplete_header() {
// Create a buffer with incomplete header
let mut buf = BytesMut::new();
buf.extend_from_slice(&[1, 0, 0, 0]); // Only 4 bytes, not enough for LpHeader::SIZE
// Attempt to parse - expect error
let result = parse_lp_packet(&buf);
assert!(result.is_err());
assert!(matches!(
result.unwrap_err(),
LpError::InsufficientBufferSize
));
}
#[test]
fn test_parse_incomplete_message_type() {
// Create a buffer with complete header but incomplete message type
let mut buf = BytesMut::new();
buf.extend_from_slice(&[1, 0, 0, 0]); // Version + reserved
buf.extend_from_slice(&42u32.to_le_bytes()); // Sender index
buf.extend_from_slice(&123u64.to_le_bytes()); // Counter
buf.extend_from_slice(&[0]); // Only 1 byte of message type (need 2)
// Buffer length = 16 + 1 = 17. Min size = 16 + 2 + 16 = 34.
let result = parse_lp_packet(&buf);
assert!(result.is_err());
assert!(matches!(
result.unwrap_err(),
LpError::InsufficientBufferSize
));
}
#[test]
fn test_parse_incomplete_message_data() {
// Create a buffer simulating Handshake but missing trailer and maybe partial payload
let mut buf = BytesMut::new();
buf.extend_from_slice(&[1, 0, 0, 0]); // Version + reserved
buf.extend_from_slice(&42u32.to_le_bytes()); // Sender index
buf.extend_from_slice(&123u64.to_le_bytes()); // Counter
buf.extend_from_slice(&MessageType::Handshake.to_u16().to_le_bytes()); // Handshake type
buf.extend_from_slice(&[42; 40]); // 40 bytes of payload data
// Buffer length = 16 + 2 + 40 = 58. Min size = 16 + 2 + 16 = 34.
// Payload size calculated as 58 - 34 = 24.
// Trailer expected at index 16 + 2 + 24 = 42.
// Trailer read attempts src[42..58].
// This *should* parse successfully based on the logic, but the trailer is garbage.
// Let's rethink: parse_lp_packet assumes the *entire slice* is the packet.
// If the slice doesn't end exactly where the trailer should, it's an error.
// In this case, total length is 58. LpHeader(16) + Type(2) + Trailer(16) = 34. Payload = 58-34=24.
// Trailer starts at 16+2+24 = 42. Ends at 42+16=58. It fits exactly.
// This test *still* doesn't test incompleteness correctly for the datagram parser.
// Let's test a buffer that's *too short* even for header+type+trailer+min_payload
let mut buf_too_short = BytesMut::new();
buf_too_short.extend_from_slice(&[1, 0, 0, 0]); // Version + reserved
buf_too_short.extend_from_slice(&42u32.to_le_bytes()); // Sender index
buf_too_short.extend_from_slice(&123u64.to_le_bytes()); // Counter
buf_too_short.extend_from_slice(&MessageType::Handshake.to_u16().to_le_bytes()); // Handshake type
// No payload, no trailer. Length = 16+2=18. Min size = 34.
let result_too_short = parse_lp_packet(&buf_too_short);
assert!(result_too_short.is_err());
assert!(matches!(
result_too_short.unwrap_err(),
LpError::InsufficientBufferSize
));
// Test a buffer missing PART of the trailer
let mut buf_partial_trailer = BytesMut::new();
buf_partial_trailer.extend_from_slice(&[1, 0, 0, 0]); // Version + reserved
buf_partial_trailer.extend_from_slice(&42u32.to_le_bytes()); // Sender index
buf_partial_trailer.extend_from_slice(&123u64.to_le_bytes()); // Counter
buf_partial_trailer.extend_from_slice(&MessageType::Handshake.to_u16().to_le_bytes()); // Handshake type
let payload = vec![42u8; 20]; // Assume 20 byte payload
buf_partial_trailer.extend_from_slice(&payload);
buf_partial_trailer.extend_from_slice(&[0; TRAILER_LEN - 1]); // Missing last byte of trailer
// Total length = 16 + 2 + 20 + 15 = 53. Min size = 34. This passes.
// Payload size = 53 - 34 = 19. <--- THIS IS WRONG. The parser assumes the length dictates payload.
// Let's fix the parser logic slightly.
// The point is, parse_lp_packet expects a COMPLETE datagram. Providing less bytes
// than LpHeader + Type + Trailer should fail. Providing *more* is also an issue unless
// the length calculation works out perfectly. The most direct test is just < min_size.
// Renaming test to reflect this.
}
#[test]
fn test_parse_buffer_smaller_than_minimum() {
// Test a buffer that's smaller than the smallest possible packet (LpHeader+Type+Trailer)
let mut buf_too_short = BytesMut::new();
buf_too_short.extend_from_slice(&[1, 0, 0, 0]); // Version + reserved
buf_too_short.extend_from_slice(&42u32.to_le_bytes()); // Sender index
buf_too_short.extend_from_slice(&123u64.to_le_bytes()); // Counter
buf_too_short.extend_from_slice(&MessageType::Busy.to_u16().to_le_bytes()); // Type
buf_too_short.extend_from_slice(&[0; TRAILER_LEN - 1]); // Missing last byte of trailer
// Length = 16 + 2 + 15 = 33. Min Size = 34.
let result_too_short = parse_lp_packet(&buf_too_short);
assert!(
result_too_short.is_err(),
"Expected error for buffer size 33, min 34"
);
assert!(matches!(
result_too_short.unwrap_err(),
LpError::InsufficientBufferSize
));
}
#[test]
fn test_parse_invalid_message_type() {
// Create a buffer with invalid message type
let mut buf = BytesMut::new();
buf.extend_from_slice(&[1, 0, 0, 0]); // Version + reserved
buf.extend_from_slice(&42u32.to_le_bytes()); // Sender index
buf.extend_from_slice(&123u64.to_le_bytes()); // Counter
buf.extend_from_slice(&255u16.to_le_bytes()); // Invalid message type
// Need payload and trailer to meet min_size requirement
let payload_size = 10; // Arbitrary
buf.extend_from_slice(&vec![0u8; payload_size]); // Some data
buf.extend_from_slice(&[0; TRAILER_LEN]); // Trailer
// Attempt to parse
let result = parse_lp_packet(&buf);
assert!(result.is_err());
match result {
Err(LpError::InvalidMessageType(255)) => {} // Expected error
Err(e) => panic!("Expected InvalidMessageType error, got {:?}", e),
Ok(_) => panic!("Expected error, but got Ok"),
}
}
#[test]
fn test_parse_incorrect_payload_size_for_busy() {
// Create a Busy packet but *with* a payload (which is invalid)
let mut buf = BytesMut::new();
buf.extend_from_slice(&[1, 0, 0, 0]); // Version + reserved
buf.extend_from_slice(&42u32.to_le_bytes()); // Sender index
buf.extend_from_slice(&123u64.to_le_bytes()); // Counter
buf.extend_from_slice(&MessageType::Busy.to_u16().to_le_bytes()); // Busy type
buf.extend_from_slice(&[42; 1]); // <<< Invalid 1-byte payload for Busy
buf.extend_from_slice(&[0; TRAILER_LEN]); // Trailer
// Total size = 16 + 2 + 1 + 16 = 35. Min size = 34.
// Calculated payload size = 35 - 34 = 1.
let result = parse_lp_packet(&buf);
assert!(result.is_err());
assert!(matches!(
result.unwrap_err(),
LpError::InvalidPayloadSize {
expected: 0,
actual: 1
}
));
}
// Test multiple packets simulation isn't relevant for datagram parsing
// #[test]
// fn test_multiple_packets_in_buffer() { ... }
}
+73
View File
@@ -0,0 +1,73 @@
// Copyright 2025 - Nym Technologies SA <contact@nymtech.net>
// SPDX-License-Identifier: Apache-2.0
use crate::{noise_protocol::NoiseError, replay::ReplayError};
use thiserror::Error;
#[derive(Error, Debug)]
pub enum LpError {
#[error("IO Error: {0}")]
IoError(#[from] std::io::Error),
#[error("Snow Error: {0}")]
SnowKeyError(#[from] snow::Error),
#[error("Snow Pattern Error: {0}")]
SnowPatternError(String),
#[error("Noise Protocol Error: {0}")]
NoiseError(#[from] NoiseError),
#[error("Replay detected: {0}")]
Replay(#[from] ReplayError),
#[error("Invalid packet format: {0}")]
InvalidPacketFormat(String),
#[error("Invalid message type: {0}")]
InvalidMessageType(u16),
#[error("Payload too large: {0}")]
PayloadTooLarge(usize),
#[error("Insufficient buffer size provided")]
InsufficientBufferSize,
#[error("Attempted operation on closed session")]
SessionClosed,
#[error("Internal error: {0}")]
Internal(String),
#[error("Invalid state transition: tried input {input:?} in state {state:?}")]
InvalidStateTransition { state: String, input: String },
#[error("Invalid payload size: expected {expected}, got {actual}")]
InvalidPayloadSize { expected: usize, actual: usize },
#[error("Deserialization error: {0}")]
DeserializationError(String),
#[error(transparent)]
InvalidBase58String(#[from] bs58::decode::Error),
/// Session ID from incoming packet does not match any known session.
#[error("Received packet with unknown session ID: {0}")]
UnknownSessionId(u32),
/// Invalid state transition attempt in the state machine.
#[error("Invalid input '{input}' for current state '{state}'")]
InvalidStateTransitionAttempt { state: String, input: String },
/// Session is closed.
#[error("Session is closed")]
LpSessionClosed,
/// Session is processing an input event.
#[error("Session is processing an input event")]
LpSessionProcessing,
/// State machine not found.
#[error("State machine not found for lp_id: {lp_id}")]
StateMachineNotFound { lp_id: u32 },
}
+165
View File
@@ -0,0 +1,165 @@
use std::fmt::{self, Display, Formatter};
use std::ops::Deref;
use std::str::FromStr;
use nym_sphinx::{PrivateKey as SphinxPrivateKey, PublicKey as SphinxPublicKey};
use serde::Serialize;
use utoipa::ToSchema;
use crate::LpError;
pub struct PrivateKey(SphinxPrivateKey);
impl Deref for PrivateKey {
type Target = SphinxPrivateKey;
fn deref(&self) -> &Self::Target {
&self.0
}
}
impl Default for PrivateKey {
fn default() -> Self {
Self::new()
}
}
impl PrivateKey {
pub fn new() -> Self {
let private_key = SphinxPrivateKey::random();
Self(private_key)
}
pub fn to_base58_string(&self) -> String {
bs58::encode(self.0.to_bytes()).into_string()
}
pub fn from_base58_string(s: &str) -> Result<Self, LpError> {
let bytes: [u8; 32] = bs58::decode(s).into_vec()?.try_into().unwrap();
Ok(PrivateKey(SphinxPrivateKey::from(bytes)))
}
pub fn public_key(&self) -> PublicKey {
let public_key = SphinxPublicKey::from(&self.0);
PublicKey(public_key)
}
}
pub struct PublicKey(SphinxPublicKey);
impl Deref for PublicKey {
type Target = SphinxPublicKey;
fn deref(&self) -> &Self::Target {
&self.0
}
}
impl PublicKey {
pub fn to_base58_string(&self) -> String {
bs58::encode(self.0.as_bytes()).into_string()
}
pub fn from_base58_string(s: &str) -> Result<Self, LpError> {
let bytes: [u8; 32] = bs58::decode(s).into_vec()?.try_into().unwrap();
Ok(PublicKey(SphinxPublicKey::from(bytes)))
}
pub fn from_bytes(bytes: &[u8; 32]) -> Result<Self, LpError> {
Ok(PublicKey(SphinxPublicKey::from(*bytes)))
}
pub fn as_bytes(&self) -> &[u8; 32] {
self.0.as_bytes()
}
}
impl Default for PublicKey {
fn default() -> Self {
let private_key = PrivateKey::default();
private_key.public_key()
}
}
pub struct Keypair {
private_key: PrivateKey,
public_key: PublicKey,
}
impl Default for Keypair {
fn default() -> Self {
Self::new()
}
}
impl Keypair {
pub fn new() -> Self {
let private_key = PrivateKey::default();
let public_key = private_key.public_key();
Self {
private_key,
public_key,
}
}
pub fn private_key(&self) -> &PrivateKey {
&self.private_key
}
pub fn public_key(&self) -> &PublicKey {
&self.public_key
}
}
impl From<KeypairReadable> for Keypair {
fn from(keypair: KeypairReadable) -> Self {
Self {
private_key: PrivateKey::from_base58_string(&keypair.private).unwrap(),
public_key: PublicKey::from_base58_string(&keypair.public).unwrap(),
}
}
}
impl From<&Keypair> for KeypairReadable {
fn from(keypair: &Keypair) -> Self {
Self {
private: keypair.private_key.to_base58_string(),
public: keypair.public_key.to_base58_string(),
}
}
}
impl FromStr for PrivateKey {
type Err = LpError;
fn from_str(s: &str) -> Result<Self, Self::Err> {
PrivateKey::from_base58_string(s)
}
}
impl Display for PrivateKey {
fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
write!(f, "{}", self.to_base58_string())
}
}
impl Display for PublicKey {
fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
write!(f, "{}", self.to_base58_string())
}
}
#[derive(Serialize, serde::Deserialize, Clone, ToSchema, Debug)]
pub struct KeypairReadable {
private: String,
public: String,
}
impl KeypairReadable {
pub fn private_key(&self) -> Result<PrivateKey, LpError> {
PrivateKey::from_base58_string(&self.private)
}
pub fn public_key(&self) -> Result<PublicKey, LpError> {
PublicKey::from_base58_string(&self.public)
}
}
+318
View File
@@ -0,0 +1,318 @@
// Copyright 2025 - Nym Technologies SA <contact@nymtech.net>
// SPDX-License-Identifier: Apache-2.0
pub mod codec;
pub mod error;
pub mod keypair;
pub mod message;
pub mod noise_protocol;
pub mod packet;
pub mod replay;
pub mod session;
mod session_integration;
pub mod session_manager;
use std::hash::{DefaultHasher, Hasher as _};
pub use error::LpError;
use keypair::PublicKey;
pub use message::{ClientHelloData, LpMessage};
pub use packet::LpPacket;
pub use replay::{ReceivingKeyCounterValidator, ReplayError};
pub use session::LpSession;
pub use session_manager::SessionManager;
// Add the new state machine module
pub mod state_machine;
pub use state_machine::LpStateMachine;
pub const NOISE_PATTERN: &str = "Noise_XKpsk3_25519_ChaChaPoly_SHA256";
pub const NOISE_PSK_INDEX: u8 = 3;
#[cfg(test)]
pub fn sessions_for_tests() -> (LpSession, LpSession) {
use crate::{keypair::Keypair, make_lp_id};
let keypair_1 = Keypair::default();
let keypair_2 = Keypair::default();
let id = make_lp_id(&keypair_1.public_key(), &keypair_2.public_key());
let initiator_session = LpSession::new(
id,
true,
&keypair_1.private_key().to_bytes(),
&keypair_2.public_key().to_bytes(),
&[0u8; 32],
)
.expect("Test session creation failed");
let responder_session = LpSession::new(
id,
false,
&keypair_2.private_key().to_bytes(),
&keypair_1.public_key().to_bytes(),
&[0u8; 32],
)
.expect("Test session creation failed");
(initiator_session, responder_session)
}
#[cfg(test)]
mod tests {
use crate::keypair::Keypair;
use crate::message::LpMessage;
use crate::packet::{LpHeader, LpPacket, TRAILER_LEN};
use crate::session_manager::SessionManager;
use crate::{make_lp_id, sessions_for_tests, LpError};
use bytes::BytesMut;
// Import the new standalone functions
use crate::codec::{parse_lp_packet, serialize_lp_packet};
#[test]
fn test_replay_protection_integration() {
// Create session
let session = sessions_for_tests().0;
// === Packet 1 (Counter 0 - Should succeed) ===
let packet1 = LpPacket {
header: LpHeader {
protocol_version: 1,
session_id: 42, // Matches session's sending_index assumption for this test
counter: 0,
},
message: LpMessage::Busy,
trailer: [0u8; TRAILER_LEN],
};
// Serialize packet
let mut buf1 = BytesMut::new();
serialize_lp_packet(&packet1, &mut buf1).unwrap();
// Parse packet
let parsed_packet1 = parse_lp_packet(&buf1).unwrap();
// Perform replay check (should pass)
session
.receiving_counter_quick_check(parsed_packet1.header.counter)
.expect("Initial packet failed replay check");
// Mark received (simulating successful processing)
session
.receiving_counter_mark(parsed_packet1.header.counter)
.expect("Failed to mark initial packet received");
// === Packet 2 (Counter 0 - Replay, should fail check) ===
let packet2 = LpPacket {
header: LpHeader {
protocol_version: 1,
session_id: 42,
counter: 0, // Same counter as before (replay)
},
message: LpMessage::Busy,
trailer: [0u8; TRAILER_LEN],
};
// Serialize packet
let mut buf2 = BytesMut::new();
serialize_lp_packet(&packet2, &mut buf2).unwrap();
// Parse packet
let parsed_packet2 = parse_lp_packet(&buf2).unwrap();
// Perform replay check (should fail)
let replay_result = session.receiving_counter_quick_check(parsed_packet2.header.counter);
assert!(replay_result.is_err());
match replay_result.unwrap_err() {
LpError::Replay(e) => {
assert!(matches!(e, crate::replay::ReplayError::DuplicateCounter));
}
e => panic!("Expected replay error, got {:?}", e),
}
// Do not mark received as it failed validation
// === Packet 3 (Counter 1 - Should succeed) ===
let packet3 = LpPacket {
header: LpHeader {
protocol_version: 1,
session_id: 42,
counter: 1, // Incremented counter
},
message: LpMessage::Busy,
trailer: [0u8; TRAILER_LEN],
};
// Serialize packet
let mut buf3 = BytesMut::new();
serialize_lp_packet(&packet3, &mut buf3).unwrap();
// Parse packet
let parsed_packet3 = parse_lp_packet(&buf3).unwrap();
// Perform replay check (should pass)
session
.receiving_counter_quick_check(parsed_packet3.header.counter)
.expect("Packet 3 failed replay check");
// Mark received
session
.receiving_counter_mark(parsed_packet3.header.counter)
.expect("Failed to mark packet 3 received");
// Verify validator state directly on the session
let state = session.current_packet_cnt();
assert_eq!(state.0, 2); // Next expected counter (correct - was 1, now expects 2)
assert_eq!(state.1, 2); // Total marked received (correct - packets 1 and 3)
}
#[test]
fn test_session_manager_integration() {
// Create session manager
let local_manager = SessionManager::new();
let remote_manager = SessionManager::new();
let local_keypair = Keypair::default();
let remote_keypair = Keypair::default();
let lp_id = make_lp_id(&local_keypair.public_key(), &remote_keypair.public_key());
// Create a session via manager
let _ = local_manager
.create_session_state_machine(
&local_keypair,
&remote_keypair.public_key(),
true,
&[2u8; 32],
)
.unwrap();
let _ = remote_manager
.create_session_state_machine(
&remote_keypair,
&local_keypair.public_key(),
false,
&[2u8; 32],
)
.unwrap();
// === Packet 1 (Counter 0 - Should succeed) ===
let packet1 = LpPacket {
header: LpHeader {
protocol_version: 1,
session_id: lp_id,
counter: 0,
},
message: LpMessage::Busy,
trailer: [0u8; TRAILER_LEN],
};
// Serialize
let mut buf1 = BytesMut::new();
serialize_lp_packet(&packet1, &mut buf1).unwrap();
// Parse
let parsed_packet1 = parse_lp_packet(&buf1).unwrap();
// Process via SessionManager method (which should handle checks + marking)
// NOTE: We might need a method on SessionManager/LpSession like `process_incoming_packet`
// that encapsulates parse -> check -> process_noise -> mark.
// For now, we simulate the steps using the retrieved session.
// Perform replay check
local_manager
.receiving_counter_quick_check(lp_id, parsed_packet1.header.counter)
.expect("Packet 1 check failed");
// Mark received
local_manager
.receiving_counter_mark(lp_id, parsed_packet1.header.counter)
.expect("Packet 1 mark failed");
// === Packet 2 (Counter 1 - Should succeed on same session) ===
let packet2 = LpPacket {
header: LpHeader {
protocol_version: 1,
session_id: lp_id,
counter: 1,
},
message: LpMessage::Busy,
trailer: [0u8; TRAILER_LEN],
};
// Serialize
let mut buf2 = BytesMut::new();
serialize_lp_packet(&packet2, &mut buf2).unwrap();
// Parse
let parsed_packet2 = parse_lp_packet(&buf2).unwrap();
// Perform replay check
local_manager
.receiving_counter_quick_check(lp_id, parsed_packet2.header.counter)
.expect("Packet 2 check failed");
// Mark received
local_manager
.receiving_counter_mark(lp_id, parsed_packet2.header.counter)
.expect("Packet 2 mark failed");
// === Packet 3 (Counter 0 - Replay, should fail check) ===
let packet3 = LpPacket {
header: LpHeader {
protocol_version: 1,
session_id: lp_id,
counter: 0, // Replay of first packet
},
message: LpMessage::Busy,
trailer: [0u8; TRAILER_LEN],
};
// Serialize
let mut buf3 = BytesMut::new();
serialize_lp_packet(&packet3, &mut buf3).unwrap();
// Parse
let parsed_packet3 = parse_lp_packet(&buf3).unwrap();
// Perform replay check (should fail)
let replay_result =
local_manager.receiving_counter_quick_check(lp_id, parsed_packet3.header.counter);
assert!(replay_result.is_err());
match replay_result.unwrap_err() {
LpError::Replay(e) => {
assert!(matches!(e, crate::replay::ReplayError::DuplicateCounter));
}
e => panic!("Expected replay error for packet 3, got {:?}", e),
}
// Do not mark received
}
}
/// Generates a deterministic u32 session ID for the Lewes Protocol
/// based on two public keys. The order of the keys does not matter.
///
/// Uses a different internal delimiter than `make_conv_id` to avoid
/// potential collisions if the same key pairs were used in both contexts.
fn make_id(key1_bytes: &[u8], key2_bytes: &[u8], sep: u8) -> u32 {
let mut hasher = DefaultHasher::new();
// Ensure consistent order for hashing to make the ID order-independent.
// This guarantees make_lp_id(a, b) == make_lp_id(b, a).
if key1_bytes < key2_bytes {
hasher.write(key1_bytes);
// Use a delimiter specific to Lewes Protocol ID generation
// (0xCC chosen arbitrarily, could be any value different from 0xFF)
hasher.write_u8(sep);
hasher.write(key2_bytes);
} else {
hasher.write(key2_bytes);
hasher.write_u8(sep);
hasher.write(key1_bytes);
}
// Truncate the u64 hash result to u32
(hasher.finish() & 0xFFFF_FFFF) as u32
}
pub fn make_lp_id(key1_bytes: &PublicKey, key2_bytes: &PublicKey) -> u32 {
make_id(key1_bytes.as_bytes(), key2_bytes.as_bytes(), 0xCC)
}
pub fn make_conv_id(src: &[u8], dst: &[u8]) -> u32 {
make_id(src, dst, 0xFF)
}
+158
View File
@@ -0,0 +1,158 @@
use std::fmt::{self, Display};
// Copyright 2025 - Nym Technologies SA <contact@nymtech.net>
// SPDX-License-Identifier: Apache-2.0
use bytes::{BufMut, BytesMut};
use serde::{Deserialize, Serialize};
/// Data structure for the ClientHello message
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ClientHelloData {
/// Client's LP x25519 public key (32 bytes)
pub client_lp_public_key: [u8; 32],
/// Protocol version for future compatibility
pub protocol_version: u8,
}
#[derive(Debug, Copy, Clone, PartialEq, Eq)]
#[repr(u16)]
pub enum MessageType {
Busy = 0x0000,
Handshake = 0x0001,
EncryptedData = 0x0002,
ClientHello = 0x0003,
}
impl MessageType {
pub(crate) fn from_u16(value: u16) -> Option<Self> {
match value {
0x0000 => Some(MessageType::Busy),
0x0001 => Some(MessageType::Handshake),
0x0002 => Some(MessageType::EncryptedData),
0x0003 => Some(MessageType::ClientHello),
_ => None,
}
}
pub fn to_u16(&self) -> u16 {
match self {
MessageType::Busy => 0x0000,
MessageType::Handshake => 0x0001,
MessageType::EncryptedData => 0x0002,
MessageType::ClientHello => 0x0003,
}
}
}
#[derive(Debug, Clone)]
pub enum LpMessage {
Busy,
Handshake(Vec<u8>),
EncryptedData(Vec<u8>),
ClientHello(ClientHelloData),
}
impl Display for LpMessage {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
match self {
LpMessage::Busy => write!(f, "Busy"),
LpMessage::Handshake(_) => write!(f, "Handshake"),
LpMessage::EncryptedData(_) => write!(f, "EncryptedData"),
LpMessage::ClientHello(_) => write!(f, "ClientHello"),
}
}
}
impl LpMessage {
pub fn payload(&self) -> &[u8] {
match self {
LpMessage::Busy => &[],
LpMessage::Handshake(payload) => payload,
LpMessage::EncryptedData(payload) => payload,
LpMessage::ClientHello(_) => &[], // Structured data, serialized in encode_content
}
}
pub fn is_empty(&self) -> bool {
match self {
LpMessage::Busy => true,
LpMessage::Handshake(payload) => payload.is_empty(),
LpMessage::EncryptedData(payload) => payload.is_empty(),
LpMessage::ClientHello(_) => false, // Always has data
}
}
pub fn len(&self) -> usize {
match self {
LpMessage::Busy => 0,
LpMessage::Handshake(payload) => payload.len(),
LpMessage::EncryptedData(payload) => payload.len(),
LpMessage::ClientHello(_) => 33, // 32 bytes key + 1 byte version
}
}
pub fn typ(&self) -> MessageType {
match self {
LpMessage::Busy => MessageType::Busy,
LpMessage::Handshake(_) => MessageType::Handshake,
LpMessage::EncryptedData(_) => MessageType::EncryptedData,
LpMessage::ClientHello(_) => MessageType::ClientHello,
}
}
pub fn encode_content(&self, dst: &mut BytesMut) {
match self {
LpMessage::Busy => { /* No content */ }
LpMessage::Handshake(payload) => {
dst.put_slice(payload);
}
LpMessage::EncryptedData(payload) => {
dst.put_slice(payload);
}
LpMessage::ClientHello(data) => {
// Serialize ClientHelloData using bincode
let serialized = bincode::serialize(data)
.expect("Failed to serialize ClientHelloData");
dst.put_slice(&serialized);
}
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::packet::{LpHeader, TRAILER_LEN};
use crate::LpPacket;
#[test]
fn encoding() {
let message = LpMessage::EncryptedData(vec![11u8; 124]);
let resp_header = LpHeader {
protocol_version: 1,
session_id: 0,
counter: 0,
};
let packet = LpPacket {
header: resp_header,
message,
trailer: [80; TRAILER_LEN],
};
// Just print packet for debug, will be captured in test output
println!("{packet:?}");
// Verify message type
assert!(matches!(packet.message.typ(), MessageType::EncryptedData));
// Verify correct data in message
match &packet.message {
LpMessage::EncryptedData(data) => {
assert_eq!(*data, vec![11u8; 124]);
}
_ => panic!("Wrong message type"),
}
}
}
+298
View File
@@ -0,0 +1,298 @@
// Copyright 2025 - Nym Technologies SA <contact@nymtech.net>
// SPDX-License-Identifier: Apache-2.0
//! Sans-IO Noise protocol state machine, adapted from noise-psq.
use snow::{params::NoiseParams, TransportState};
use thiserror::Error;
// --- Error Definition ---
/// Errors related to the Noise protocol state machine.
#[derive(Error, Debug)]
pub enum NoiseError {
#[error("encountered a Noise decryption error")]
DecryptionError,
#[error("encountered a Noise Protocol error - {0}")]
ProtocolError(snow::Error),
#[error("operation is invalid in the current protocol state")]
IncorrectStateError,
#[error("Other Noise-related error: {0}")]
Other(String),
}
impl From<snow::Error> for NoiseError {
fn from(err: snow::Error) -> Self {
match err {
snow::Error::Decrypt => NoiseError::DecryptionError,
err => NoiseError::ProtocolError(err),
}
}
}
// --- Protocol State and Structs ---
/// Represents the possible states of the Noise protocol machine.
#[derive(Debug)]
pub enum NoiseProtocolState {
/// The protocol is currently performing the handshake.
/// Contains the Snow handshake state.
Handshaking(Box<snow::HandshakeState>),
/// The handshake is complete, and the protocol is in transport mode.
/// Contains the Snow transport state.
Transport(TransportState),
/// The protocol has encountered an unrecoverable error.
/// Stores the error description.
Failed(String),
}
/// The core sans-io Noise protocol state machine.
#[derive(Debug)]
pub struct NoiseProtocol {
state: NoiseProtocolState,
// We might need buffers for incoming/outgoing data later if we add internal buffering
// read_buffer: Vec<u8>,
// write_buffer: Vec<u8>,
}
/// Represents the outcome of processing received bytes via `read_message`.
#[derive(Debug, PartialEq)]
pub enum ReadResult {
/// A handshake or transport message was successfully processed, but yielded no application data
/// and did not complete the handshake.
NoOp,
/// A complete application data message was decrypted.
DecryptedData(Vec<u8>),
/// The handshake successfully completed during this read operation.
HandshakeComplete,
// NOTE: NeedMoreBytes variant removed as read_message expects full frames.
}
// --- Implementation ---
impl NoiseProtocol {
/// Creates a new `NoiseProtocol` instance in the Handshaking state.
///
/// Takes an initialized `snow::HandshakeState` (e.g., from `snow::Builder`).
pub fn new(initial_state: snow::HandshakeState) -> Self {
NoiseProtocol {
state: NoiseProtocolState::Handshaking(Box::new(initial_state)),
}
}
/// Processes a single, complete incoming Noise message frame.
///
/// Assumes the caller handles buffering and framing to provide one full message.
/// Returns the result of processing the message.
pub fn read_message(&mut self, input: &[u8]) -> Result<ReadResult, NoiseError> {
// Allocate a buffer large enough for the maximum possible Noise message size.
// TODO: Consider reusing a buffer for efficiency.
let mut buffer = vec![0u8; 65535]; // Max Noise message size
match &mut self.state {
NoiseProtocolState::Handshaking(handshake_state) => {
match handshake_state.read_message(input, &mut buffer) {
Ok(_) => {
if handshake_state.is_handshake_finished() {
// Transition to Transport state.
let current_state = std::mem::replace(
&mut self.state,
// Temporary placeholder needed for mem::replace
NoiseProtocolState::Failed(
NoiseError::IncorrectStateError.to_string(),
),
);
if let NoiseProtocolState::Handshaking(state_to_convert) = current_state
{
match state_to_convert.into_transport_mode() {
Ok(transport_state) => {
self.state = NoiseProtocolState::Transport(transport_state);
Ok(ReadResult::HandshakeComplete)
}
Err(e) => {
let err = NoiseError::from(e);
self.state = NoiseProtocolState::Failed(err.to_string());
Err(err)
}
}
} else {
// Should be unreachable
let err = NoiseError::IncorrectStateError;
self.state = NoiseProtocolState::Failed(err.to_string());
Err(err)
}
} else {
// Handshake continues
Ok(ReadResult::NoOp)
}
}
Err(e) => {
let err = NoiseError::from(e);
self.state = NoiseProtocolState::Failed(err.to_string());
Err(err)
}
}
}
NoiseProtocolState::Transport(transport_state) => {
match transport_state.read_message(input, &mut buffer) {
Ok(len) => Ok(ReadResult::DecryptedData(buffer[..len].to_vec())),
Err(e) => {
let err = NoiseError::from(e);
self.state = NoiseProtocolState::Failed(err.to_string());
Err(err)
}
}
}
NoiseProtocolState::Failed(_) => Err(NoiseError::IncorrectStateError),
}
}
/// Checks if there are pending handshake messages to send.
///
/// If in Handshaking state and it's our turn, generates the message.
/// Transitions state to Transport if the handshake completes after this message.
/// Returns `None` if not in Handshaking state or not our turn.
pub fn get_bytes_to_send(&mut self) -> Option<Result<Vec<u8>, NoiseError>> {
match &mut self.state {
NoiseProtocolState::Handshaking(handshake_state) => {
if handshake_state.is_my_turn() {
let mut buffer = vec![0u8; 65535];
match handshake_state.write_message(&[], &mut buffer) {
// Empty payload for handshake msg
Ok(len) => {
if handshake_state.is_handshake_finished() {
// Transition to Transport state.
let current_state = std::mem::replace(
&mut self.state,
NoiseProtocolState::Failed(
NoiseError::IncorrectStateError.to_string(),
),
);
if let NoiseProtocolState::Handshaking(state_to_convert) =
current_state
{
match state_to_convert.into_transport_mode() {
Ok(transport_state) => {
self.state =
NoiseProtocolState::Transport(transport_state);
Some(Ok(buffer[..len].to_vec())) // Return final handshake msg
}
Err(e) => {
let err = NoiseError::from(e);
self.state =
NoiseProtocolState::Failed(err.to_string());
Some(Err(err))
}
}
} else {
// Should be unreachable
let err = NoiseError::IncorrectStateError;
self.state = NoiseProtocolState::Failed(err.to_string());
Some(Err(err))
}
} else {
// Handshake continues
Some(Ok(buffer[..len].to_vec()))
}
}
Err(e) => {
let err = NoiseError::from(e);
self.state = NoiseProtocolState::Failed(err.to_string());
Some(Err(err))
}
}
} else {
// Not our turn
None
}
}
NoiseProtocolState::Transport(_) | NoiseProtocolState::Failed(_) => {
// No handshake messages to send in these states
None
}
}
}
/// Encrypts an application data payload for sending during the Transport phase.
///
/// Returns the ciphertext (payload + 16-byte tag).
/// Errors if not in Transport state or encryption fails.
pub fn write_message(&mut self, payload: &[u8]) -> Result<Vec<u8>, NoiseError> {
match &mut self.state {
NoiseProtocolState::Transport(transport_state) => {
let mut buffer = vec![0u8; payload.len() + 16]; // Payload + tag
match transport_state.write_message(payload, &mut buffer) {
Ok(len) => Ok(buffer[..len].to_vec()),
Err(e) => {
let err = NoiseError::from(e);
self.state = NoiseProtocolState::Failed(err.to_string());
Err(err)
}
}
}
NoiseProtocolState::Handshaking(_) | NoiseProtocolState::Failed(_) => {
Err(NoiseError::IncorrectStateError)
}
}
}
/// Returns true if the protocol is in the transport phase (handshake complete).
pub fn is_transport(&self) -> bool {
matches!(self.state, NoiseProtocolState::Transport(_))
}
/// Returns true if the protocol has failed.
pub fn is_failed(&self) -> bool {
matches!(self.state, NoiseProtocolState::Failed(_))
}
/// Check if the handshake has finished and the protocol is in transport mode.
pub fn is_handshake_finished(&self) -> bool {
matches!(self.state, NoiseProtocolState::Transport(_))
}
}
pub fn create_noise_state(
local_private_key: &[u8],
remote_public_key: &[u8],
psk: &[u8],
) -> Result<NoiseProtocol, NoiseError> {
let pattern_name = "Noise_XKpsk3_25519_ChaChaPoly_SHA256";
let psk_index = 3;
let noise_params: NoiseParams = pattern_name.parse().unwrap();
let builder = snow::Builder::new(noise_params.clone());
// Using dummy remote key as it's not needed for state creation itself
// In a real scenario, the key would depend on initiator/responder role
let handshake_state = builder
.local_private_key(local_private_key)
.remote_public_key(remote_public_key) // Use own public as dummy remote
.psk(psk_index, psk)
.build_initiator()?;
Ok(NoiseProtocol::new(handshake_state))
}
pub fn create_noise_state_responder(
local_private_key: &[u8],
remote_public_key: &[u8],
psk: &[u8],
) -> Result<NoiseProtocol, NoiseError> {
let pattern_name = "Noise_XKpsk3_25519_ChaChaPoly_SHA256";
let psk_index = 3;
let noise_params: NoiseParams = pattern_name.parse().unwrap();
let builder = snow::Builder::new(noise_params.clone());
// Using dummy remote key as it's not needed for state creation itself
// In a real scenario, the key would depend on initiator/responder role
let handshake_state = builder
.local_private_key(local_private_key)
.remote_public_key(remote_public_key) // Use own public as dummy remote
.psk(psk_index, psk)
.build_responder()?;
Ok(NoiseProtocol::new(handshake_state))
}
+195
View File
@@ -0,0 +1,195 @@
// Copyright 2025 - Nym Technologies SA <contact@nymtech.net>
// SPDX-License-Identifier: Apache-2.0
use crate::message::LpMessage;
use crate::replay::ReceivingKeyCounterValidator;
use crate::LpError;
use bytes::{BufMut, BytesMut};
use nym_lp_common::format_debug_bytes;
use parking_lot::Mutex;
use std::fmt::Write;
use std::fmt::{Debug, Formatter};
use std::sync::Arc;
#[allow(dead_code)]
pub(crate) const UDP_HEADER_LEN: usize = 8;
#[allow(dead_code)]
pub(crate) const IP_HEADER_LEN: usize = 40; // v4 - 20, v6 - 40
#[allow(dead_code)]
pub(crate) const MTU: usize = 1500;
#[allow(dead_code)]
pub(crate) const UDP_OVERHEAD: usize = UDP_HEADER_LEN + IP_HEADER_LEN;
#[allow(dead_code)]
pub const TRAILER_LEN: usize = 16;
#[allow(dead_code)]
pub(crate) const UDP_PAYLOAD_SIZE: usize = MTU - UDP_OVERHEAD - TRAILER_LEN;
#[derive(Clone)]
pub struct LpPacket {
pub(crate) header: LpHeader,
pub(crate) message: LpMessage,
pub(crate) trailer: [u8; TRAILER_LEN],
}
impl Debug for LpPacket {
fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
write!(f, "{}", format_debug_bytes(&self.debug_bytes())?)
}
}
impl LpPacket {
pub fn new(header: LpHeader, message: LpMessage) -> Self {
Self {
header,
message,
trailer: [0; TRAILER_LEN],
}
}
/// Compute a hash of the message payload
///
/// This can be used for message integrity verification or deduplication
pub fn hash_payload(&self) -> [u8; 32] {
use sha2::{Digest, Sha256};
let mut hasher = Sha256::new();
let mut buffer = BytesMut::new();
// Include message type and content in the hash
buffer.put_slice(&(self.message.typ() as u16).to_le_bytes());
self.message.encode_content(&mut buffer);
hasher.update(&buffer);
hasher.finalize().into()
}
pub fn hash_payload_hex(&self) -> String {
let hash = self.hash_payload();
hash.iter()
.fold(String::with_capacity(hash.len() * 2), |mut acc, byte| {
let _ = write!(acc, "{:02x}", byte);
acc
})
}
pub fn message(&self) -> &LpMessage {
&self.message
}
pub fn header(&self) -> &LpHeader {
&self.header
}
pub(crate) fn debug_bytes(&self) -> Vec<u8> {
let mut bytes = BytesMut::new();
self.encode(&mut bytes);
bytes.freeze().to_vec()
}
pub(crate) fn encode(&self, dst: &mut BytesMut) {
self.header.encode(dst);
dst.put_slice(&(self.message.typ() as u16).to_le_bytes());
self.message.encode_content(dst);
dst.put_slice(&self.trailer)
}
/// Validate packet counter against a replay protection validator
///
/// This performs a quick check to see if the packet counter is valid before
/// any expensive processing is done.
pub fn validate_counter(
&self,
validator: &Arc<Mutex<ReceivingKeyCounterValidator>>,
) -> Result<(), LpError> {
let guard = validator.lock();
guard.will_accept_branchless(self.header.counter)?;
Ok(())
}
/// Mark packet as received in the replay protection validator
///
/// This should be called after a packet has been successfully processed.
pub fn mark_received(
&self,
validator: &Arc<Mutex<ReceivingKeyCounterValidator>>,
) -> Result<(), LpError> {
let mut guard = validator.lock();
guard.mark_did_receive_branchless(self.header.counter)?;
Ok(())
}
}
// VERSION [1B] || RESERVED [3B] || SENDER_INDEX [4B] || COUNTER [8B]
#[derive(Debug, Clone)]
pub struct LpHeader {
pub protocol_version: u8,
pub session_id: u32,
pub counter: u64,
}
impl LpHeader {
pub const SIZE: usize = 16;
}
impl LpHeader {
pub fn new(session_id: u32, counter: u64) -> Self {
Self {
protocol_version: 1,
session_id,
counter,
}
}
pub fn encode(&self, dst: &mut BytesMut) {
// protocol version
dst.put_u8(self.protocol_version);
// reserved
dst.put_slice(&[0, 0, 0]);
// sender index
dst.put_slice(&self.session_id.to_le_bytes());
// counter
dst.put_slice(&self.counter.to_le_bytes());
}
pub fn parse(src: &[u8]) -> Result<Self, LpError> {
if src.len() < Self::SIZE {
return Err(LpError::InsufficientBufferSize);
}
let protocol_version = src[0];
// Skip reserved bytes [1..4]
let mut session_id_bytes = [0u8; 4];
session_id_bytes.copy_from_slice(&src[4..8]);
let session_id = u32::from_le_bytes(session_id_bytes);
let mut counter_bytes = [0u8; 8];
counter_bytes.copy_from_slice(&src[8..16]);
let counter = u64::from_le_bytes(counter_bytes);
Ok(LpHeader {
protocol_version,
session_id,
counter,
})
}
/// Get the counter value from the header
pub fn counter(&self) -> u64 {
self.counter
}
/// Get the sender index from the header
pub fn session_id(&self) -> u32 {
self.session_id
}
}
// subsequent data: MessageType || Data
+68
View File
@@ -0,0 +1,68 @@
// Copyright 2025 - Nym Technologies SA <contact@nymtech.net>
// SPDX-License-Identifier: Apache-2.0
//! Error types for replay protection.
use thiserror::Error;
/// Errors that can occur during replay protection validation.
#[derive(Debug, Error)]
pub enum ReplayError {
/// The counter value is invalid (e.g., too far in the future)
#[error("Invalid counter value")]
InvalidCounter,
/// The packet has already been received (replay attack)
#[error("Duplicate counter value")]
DuplicateCounter,
/// The packet is outside the replay window
#[error("Packet outside replay window")]
OutOfWindow,
}
/// Result type for replay protection operations
pub type ReplayResult<T> = Result<T, ReplayError>;
#[cfg(test)]
mod tests {
use super::*;
use crate::error::LpError;
#[test]
fn test_replay_error_variants() {
let invalid = ReplayError::InvalidCounter;
let duplicate = ReplayError::DuplicateCounter;
let out_of_window = ReplayError::OutOfWindow;
assert_eq!(invalid.to_string(), "Invalid counter value");
assert_eq!(duplicate.to_string(), "Duplicate counter value");
assert_eq!(out_of_window.to_string(), "Packet outside replay window");
}
#[test]
fn test_replay_error_conversion() {
let replay_error = ReplayError::InvalidCounter;
let lp_error: LpError = replay_error.into();
match lp_error {
LpError::Replay(e) => {
assert!(matches!(e, ReplayError::InvalidCounter));
}
_ => panic!("Expected Replay variant"),
}
}
#[test]
fn test_replay_result() {
let ok_result: ReplayResult<()> = Ok(());
let err_result: ReplayResult<()> = Err(ReplayError::InvalidCounter);
assert!(ok_result.is_ok());
assert!(err_result.is_err());
assert!(matches!(
err_result.unwrap_err(),
ReplayError::InvalidCounter
));
}
}
+15
View File
@@ -0,0 +1,15 @@
// Copyright 2025 - Nym Technologies SA <contact@nymtech.net>
// SPDX-License-Identifier: Apache-2.0
//! Replay protection module for the Lewes Protocol.
//!
//! This module implements BoringTun-style replay protection to prevent
//! replay attacks and ensure packet ordering. It uses a bitmap-based
//! approach to track received packets and validate their sequence.
pub mod error;
pub mod simd;
pub mod validator;
pub use error::ReplayError;
pub use validator::ReceivingKeyCounterValidator;
+278
View File
@@ -0,0 +1,278 @@
// Copyright 2025 - Nym Technologies SA <contact@nymtech.net>
// SPDX-License-Identifier: Apache-2.0
//! ARM NEON implementation of bitmap operations.
use super::BitmapOps;
#[cfg(target_feature = "neon")]
use std::arch::aarch64::{vceqq_u64, vdupq_n_u64, vgetq_lane_u64, vld1q_u64, vst1q_u64};
/// ARM NEON bitmap operations implementation
pub struct ArmBitmapOps;
impl BitmapOps for ArmBitmapOps {
#[inline(always)]
fn clear_words(bitmap: &mut [u64], start_idx: usize, num_words: usize) {
debug_assert!(start_idx + num_words <= bitmap.len());
#[cfg(target_feature = "neon")]
unsafe {
// Process 2 words at a time with NEON
// Safety:
// - vdupq_n_u64 is safe to call with any u64 value
let zero_vec = vdupq_n_u64(0);
let mut idx = start_idx;
let end_idx = start_idx + num_words;
// Process aligned blocks of 2 words
while idx + 2 <= end_idx {
// Safety:
// - bitmap[idx..] is valid for reads/writes of at least 2 u64 words (16 bytes)
// - We've validated with the debug_assert that start_idx + num_words <= bitmap.len()
// - We check that idx + 2 <= end_idx to ensure we have 2 complete words
vst1q_u64(bitmap[idx..].as_mut_ptr(), zero_vec);
idx += 2;
}
// Handle remaining words (0 or 1)
while idx < end_idx {
bitmap[idx] = 0;
idx += 1;
}
}
#[cfg(not(target_feature = "neon"))]
{
// Fallback to scalar implementation
for i in start_idx..(start_idx + num_words) {
bitmap[i] = 0;
}
}
}
#[inline(always)]
fn is_range_zero(bitmap: &[u64], start_idx: usize, num_words: usize) -> bool {
debug_assert!(start_idx + num_words <= bitmap.len());
#[cfg(target_feature = "neon")]
unsafe {
// Process 2 words at a time with NEON
// Safety:
// - vdupq_n_u64 is safe to call with any u64 value
let zero_vec = vdupq_n_u64(0);
let mut idx = start_idx;
let end_idx = start_idx + num_words;
// Process aligned blocks of 2 words
while idx + 2 <= end_idx {
// Safety:
// - bitmap[idx..] is valid for reads of at least 2 u64 words (16 bytes)
// - We've validated with the debug_assert that start_idx + num_words <= bitmap.len()
// - We check that idx + 2 <= end_idx to ensure we have 2 complete words
let data_vec = vld1q_u64(bitmap[idx..].as_ptr());
// Safety:
// - vceqq_u64 is safe when given valid vector values from vld1q_u64 and vdupq_n_u64
// - vgetq_lane_u64 is safe with valid indices (0 and 1) for a 2-lane vector
let cmp_result = vceqq_u64(data_vec, zero_vec);
let mask1 = vgetq_lane_u64(cmp_result, 0);
let mask2 = vgetq_lane_u64(cmp_result, 1);
if (mask1 & mask2) != u64::MAX {
return false;
}
idx += 2;
}
// Handle remaining words (0 or 1)
while idx < end_idx {
if bitmap[idx] != 0 {
return false;
}
idx += 1;
}
true
}
#[cfg(not(target_feature = "neon"))]
{
// Fallback to scalar implementation
bitmap[start_idx..(start_idx + num_words)]
.iter()
.all(|&w| w == 0)
}
}
#[inline(always)]
fn set_bit(bitmap: &mut [u64], bit_idx: u64) {
let word_idx = (bit_idx / 64) as usize;
let bit_pos = bit_idx % 64;
bitmap[word_idx] |= 1u64 << bit_pos;
}
#[inline(always)]
fn clear_bit(bitmap: &mut [u64], bit_idx: u64) {
let word_idx = (bit_idx / 64) as usize;
let bit_pos = bit_idx % 64;
bitmap[word_idx] &= !(1u64 << bit_pos);
}
#[inline(always)]
fn check_bit(bitmap: &[u64], bit_idx: u64) -> bool {
let word_idx = (bit_idx / 64) as usize;
let bit_pos = bit_idx % 64;
(bitmap[word_idx] & (1u64 << bit_pos)) != 0
}
}
/// We also implement optimized versions for specific operations that could
/// benefit from NEON but don't fit the general trait pattern
///
/// Atomic operations for the bitmap
pub mod atomic {
#[cfg(target_feature = "neon")]
use std::arch::aarch64::{vdupq_n_u64, vld1q_u64, vorrq_u64, vst1q_u64};
/// Check and set bit, returning the previous state
/// This function is not actually atomic! It's just a non-atomic optimization
/// For actual atomic operations, the caller must provide proper synchronization
#[inline(always)]
pub fn check_and_set_bit(bitmap: &mut [u64], bit_idx: u64) -> bool {
let word_idx = (bit_idx / 64) as usize;
let bit_pos = bit_idx % 64;
let mask = 1u64 << bit_pos;
// Get old value
let old_word = bitmap[word_idx];
// Set bit regardless of current state
bitmap[word_idx] |= mask;
// Return true if bit was already set (duplicate)
(old_word & mask) != 0
}
/// Set a range of bits efficiently using NEON
///
/// # Safety
///
/// This function is unsafe because it:
/// - Uses SIMD intrinsics that require the NEON CPU feature to be available
/// - Accesses bitmap memory through raw pointers
/// - Does not perform bounds checking beyond what's required for SIMD operations
///
/// Caller must ensure:
/// - The NEON feature is available on the current CPU
/// - `bitmap` has sufficient size to hold indices up to `end_bit/64`
/// - `start_bit` and `end_bit` are valid bit indices within the bitmap
/// - No other thread is concurrently modifying the same memory
#[inline(always)]
#[cfg(target_feature = "neon")]
pub unsafe fn set_bits_range(bitmap: &mut [u64], start_bit: u64, end_bit: u64) {
// Process whole words where possible
let start_word = (start_bit / 64) as usize;
let end_word = (end_bit / 64) as usize;
if start_word == end_word {
// Special case: all bits in the same word
let start_mask = u64::MAX << (start_bit % 64);
let end_mask = u64::MAX >> (63 - (end_bit % 64));
bitmap[start_word] |= start_mask & end_mask;
return;
}
// Handle partial words at the beginning and end
if start_bit % 64 != 0 {
let start_mask = u64::MAX << (start_bit % 64);
bitmap[start_word] |= start_mask;
}
if (end_bit + 1) % 64 != 0 {
let end_mask = u64::MAX >> (63 - (end_bit % 64));
bitmap[end_word] |= end_mask;
}
// Handle complete words in the middle using NEON
let first_full_word = if start_bit % 64 == 0 {
start_word
} else {
start_word + 1
};
let last_full_word = if (end_bit + 1) % 64 == 0 {
end_word
} else {
end_word - 1
};
if first_full_word <= last_full_word {
// Use NEON to set words faster
// Safety: vdupq_n_u64 is safe to call with any u64 value
let ones_vec = vdupq_n_u64(u64::MAX);
let mut idx = first_full_word;
while idx + 2 <= last_full_word + 1 {
// Safety:
// - bitmap[idx..] is valid for reads/writes of at least 2 u64 words (16 bytes)
// - We check that idx + 2 <= last_full_word + 1 to ensure we have 2 complete words
let current_vec = vld1q_u64(bitmap[idx..].as_ptr());
// Safety: vorrq_u64 is safe when given valid vector values
let result_vec = vorrq_u64(current_vec, ones_vec);
vst1q_u64(bitmap[idx..].as_mut_ptr(), result_vec);
idx += 2;
}
// Handle remaining words
while idx <= last_full_word {
bitmap[idx] = u64::MAX;
idx += 1;
}
}
}
/// Set a range of bits efficiently (scalar fallback)
#[inline(always)]
#[cfg(not(target_feature = "neon"))]
pub fn set_bits_range(bitmap: &mut [u64], start_bit: u64, end_bit: u64) {
// Process whole words where possible
let start_word = (start_bit / 64) as usize;
let end_word = (end_bit / 64) as usize;
if start_word == end_word {
// Special case: all bits in the same word
let start_mask = u64::MAX << (start_bit % 64);
let end_mask = u64::MAX >> (63 - (end_bit % 64));
bitmap[start_word] |= start_mask & end_mask;
return;
}
// Handle partial words at the beginning and end
if start_bit % 64 != 0 {
let start_mask = u64::MAX << (start_bit % 64);
bitmap[start_word] |= start_mask;
}
if (end_bit + 1) % 64 != 0 {
let end_mask = u64::MAX >> (63 - (end_bit % 64));
bitmap[end_word] |= end_mask;
}
// Handle complete words in the middle
let first_full_word = if start_bit % 64 == 0 {
start_word
} else {
start_word + 1
};
let last_full_word = if (end_bit + 1) % 64 == 0 {
end_word
} else {
end_word - 1
};
for word_idx in first_full_word..=last_full_word {
bitmap[word_idx] = u64::MAX;
}
}
}
+71
View File
@@ -0,0 +1,71 @@
// Copyright 2025 - Nym Technologies SA <contact@nymtech.net>
// SPDX-License-Identifier: Apache-2.0
//! SIMD optimizations for the replay protection bitmap operations.
//!
//! This module provides architecture-specific SIMD implementations with a common interface.
// Re-export the appropriate implementation
#[cfg(target_arch = "x86_64")]
mod x86;
#[cfg(target_arch = "x86_64")]
pub use self::x86::*;
#[cfg(target_arch = "aarch64")]
mod arm;
#[cfg(target_arch = "aarch64")]
pub use self::arm::*;
// Fallback scalar implementation for all other architectures
#[cfg(not(any(target_arch = "x86_64", target_arch = "aarch64")))]
mod scalar;
#[cfg(not(any(target_arch = "x86_64", target_arch = "aarch64")))]
pub use self::scalar::*;
/// Trait defining SIMD operations for bitmap manipulation
pub trait BitmapOps {
/// Clear a range of words in the bitmap
fn clear_words(bitmap: &mut [u64], start_idx: usize, num_words: usize);
/// Check if a range of words in the bitmap is all zeros
fn is_range_zero(bitmap: &[u64], start_idx: usize, num_words: usize) -> bool;
/// Set a specific bit in the bitmap
fn set_bit(bitmap: &mut [u64], bit_idx: u64);
/// Clear a specific bit in the bitmap
fn clear_bit(bitmap: &mut [u64], bit_idx: u64);
/// Check if a specific bit is set in the bitmap
fn check_bit(bitmap: &[u64], bit_idx: u64) -> bool;
}
/// Get the optimal number of words to process in a SIMD operation
/// for the current architecture
#[inline(always)]
pub fn optimal_simd_width() -> usize {
// This value is specialized for each architecture in their respective modules
OPTIMAL_SIMD_WIDTH
}
/// Constant indicating the optimal SIMD processing width in number of u64 words
/// for the current architecture
#[cfg(target_arch = "x86_64")]
#[cfg(target_feature = "avx2")]
pub const OPTIMAL_SIMD_WIDTH: usize = 4; // 256 bits = 4 u64 words
#[cfg(target_arch = "x86_64")]
#[cfg(all(not(target_feature = "avx2"), target_feature = "sse2"))]
pub const OPTIMAL_SIMD_WIDTH: usize = 2; // 128 bits = 2 u64 words
#[cfg(target_arch = "aarch64")]
#[cfg(target_feature = "neon")]
pub const OPTIMAL_SIMD_WIDTH: usize = 2; // 128 bits = 2 u64 words
// Fallback for non-SIMD platforms or when features aren't available
#[cfg(not(any(
all(target_arch = "x86_64", target_feature = "avx2"),
all(target_arch = "x86_64", target_feature = "sse2"),
all(target_arch = "aarch64", target_feature = "neon")
)))]
pub const OPTIMAL_SIMD_WIDTH: usize = 1; // Scalar fallback
+114
View File
@@ -0,0 +1,114 @@
// Copyright 2025 - Nym Technologies SA <contact@nymtech.net>
// SPDX-License-Identifier: Apache-2.0
//! Scalar (non-SIMD) implementation of bitmap operations.
//! Used as a fallback when SIMD instructions are unavailable.
use super::BitmapOps;
/// Scalar (non-SIMD) bitmap operations implementation
pub struct ScalarBitmapOps;
impl BitmapOps for ScalarBitmapOps {
#[inline(always)]
fn clear_words(bitmap: &mut [u64], start_idx: usize, num_words: usize) {
for i in start_idx..(start_idx + num_words) {
bitmap[i] = 0;
}
}
#[inline(always)]
fn is_range_zero(bitmap: &[u64], start_idx: usize, num_words: usize) -> bool {
for i in start_idx..(start_idx + num_words) {
if bitmap[i] != 0 {
return false;
}
}
true
}
#[inline(always)]
fn set_bit(bitmap: &mut [u64], bit_idx: u64) {
let word_idx = (bit_idx / 64) as usize;
let bit_pos = (bit_idx % 64) as u64;
bitmap[word_idx] |= 1u64 << bit_pos;
}
#[inline(always)]
fn clear_bit(bitmap: &mut [u64], bit_idx: u64) {
let word_idx = (bit_idx / 64) as usize;
let bit_pos = (bit_idx % 64) as u64;
bitmap[word_idx] &= !(1u64 << bit_pos);
}
#[inline(always)]
fn check_bit(bitmap: &[u64], bit_idx: u64) -> bool {
let word_idx = (bit_idx / 64) as usize;
let bit_pos = (bit_idx % 64) as u64;
(bitmap[word_idx] & (1u64 << bit_pos)) != 0
}
}
/// Scalar implementations of other bitmap utilities
pub mod atomic {
/// Check and set bit, returning the previous state
/// This function is not actually atomic! It's just a normal operation
#[inline(always)]
pub fn check_and_set_bit(bitmap: &mut [u64], bit_idx: u64) -> bool {
let word_idx = (bit_idx / 64) as usize;
let bit_pos = (bit_idx % 64) as u64;
let mask = 1u64 << bit_pos;
// Get old value
let old_word = bitmap[word_idx];
// Set bit regardless of current state
bitmap[word_idx] |= mask;
// Return true if bit was already set (duplicate)
(old_word & mask) != 0
}
/// Set a range of bits efficiently
#[inline(always)]
pub fn set_bits_range(bitmap: &mut [u64], start_bit: u64, end_bit: u64) {
// Process whole words where possible
let start_word = (start_bit / 64) as usize;
let end_word = (end_bit / 64) as usize;
if start_word == end_word {
// Special case: all bits in the same word
let start_mask = u64::MAX << (start_bit % 64);
let end_mask = u64::MAX >> (63 - (end_bit % 64));
bitmap[start_word] |= start_mask & end_mask;
return;
}
// Handle partial words at the beginning and end
if start_bit % 64 != 0 {
let start_mask = u64::MAX << (start_bit % 64);
bitmap[start_word] |= start_mask;
}
if (end_bit + 1) % 64 != 0 {
let end_mask = u64::MAX >> (63 - (end_bit % 64));
bitmap[end_word] |= end_mask;
}
// Handle complete words in the middle
let first_full_word = if start_bit % 64 == 0 {
start_word
} else {
start_word + 1
};
let last_full_word = if (end_bit + 1) % 64 == 0 {
end_word
} else {
end_word - 1
};
for word_idx in first_full_word..=last_full_word {
bitmap[word_idx] = u64::MAX;
}
}
}
+489
View File
@@ -0,0 +1,489 @@
// Copyright 2025 - Nym Technologies SA <contact@nymtech.net>
// SPDX-License-Identifier: Apache-2.0
//! x86/x86_64 SIMD implementation of bitmap operations.
//! Provides optimized implementations using SSE2 and AVX2 intrinsics.
use super::BitmapOps;
// Track execution counts for debugging
static mut AVX2_CLEAR_COUNT: usize = 0;
static mut SSE2_CLEAR_COUNT: usize = 0;
static mut SCALAR_CLEAR_COUNT: usize = 0;
// Import the appropriate SIMD intrinsics
#[cfg(target_feature = "avx2")]
use std::arch::x86_64::{
__m256i, _mm256_cmpeq_epi64, _mm256_load_si256, _mm256_loadu_si256, _mm256_movemask_epi8,
_mm256_or_si256, _mm256_set1_epi64x, _mm256_setzero_si256, _mm256_store_si256,
_mm256_storeu_si256, _mm256_testz_si256,
};
#[cfg(target_feature = "sse2")]
use std::arch::x86_64::{
__m128i, _mm_cmpeq_epi64, _mm_load_si128, _mm_loadu_si128, _mm_movemask_epi8, _mm_or_si128,
_mm_set1_epi64x, _mm_setzero_si128, _mm_store_si128, _mm_storeu_si128, _mm_testz_si128,
};
/// x86/x86_64 SIMD bitmap operations implementation
pub struct X86BitmapOps;
impl BitmapOps for X86BitmapOps {
#[inline(always)]
fn clear_words(bitmap: &mut [u64], start_idx: usize, num_words: usize) {
debug_assert!(start_idx + num_words <= bitmap.len());
// First try AVX2 (256-bit, 4 words at a time)
#[cfg(target_feature = "avx2")]
unsafe {
// Track execution count
AVX2_CLEAR_COUNT += 1;
// Process 4 words at a time with AVX2
let zero_vec = _mm256_setzero_si256();
let mut idx = start_idx;
let end_idx = start_idx + num_words;
// Process aligned blocks of 4 words
while idx + 4 <= end_idx {
// Safety:
// - bitmap[idx..] is valid for reads/writes of at least 4 u64 words (32 bytes)
// - We've validated with the debug_assert that start_idx + num_words <= bitmap.len()
// - We check that idx + 4 <= end_idx to ensure we have 4 complete words
// - The unaligned _storeu_ variant is used to handle any alignment
_mm256_storeu_si256(bitmap[idx..].as_mut_ptr() as *mut __m256i, zero_vec);
idx += 4;
}
// Handle remaining words with SSE2 or scalar ops
if idx < end_idx {
if idx + 2 <= end_idx {
// Use SSE2 for 2 words
// Safety: Same as above, but for 2 words (16 bytes) instead of 4
let sse_zero = _mm_setzero_si128();
_mm_storeu_si128(bitmap[idx..].as_mut_ptr() as *mut __m128i, sse_zero);
idx += 2;
}
// Handle any remaining words
while idx < end_idx {
bitmap[idx] = 0;
idx += 1;
}
}
return;
}
// If AVX2 is unavailable, try SSE2 (128-bit, 2 words at a time)
#[cfg(all(target_feature = "sse2", not(target_feature = "avx2")))]
unsafe {
// Track execution count
SSE2_CLEAR_COUNT += 1;
// Process 2 words at a time with SSE2
let zero_vec = _mm_setzero_si128();
let mut idx = start_idx;
let end_idx = start_idx + num_words;
// Process aligned blocks of 2 words
while idx + 2 <= end_idx {
// Safety:
// - bitmap[idx..] is valid for reads/writes of at least 2 u64 words (16 bytes)
// - We've validated with the debug_assert that start_idx + num_words <= bitmap.len()
// - We check that idx + 2 <= end_idx to ensure we have 2 complete words
// - The unaligned _storeu_ variant is used to handle any alignment
_mm_storeu_si128(bitmap[idx..].as_mut_ptr() as *mut __m128i, zero_vec);
idx += 2;
}
// Handle remaining word (if any)
if idx < end_idx {
bitmap[idx] = 0;
}
return;
}
// Fallback to scalar implementation if no SIMD features available
unsafe {
// Safety: Just increments a static counter, with no possibility of data races
// as long as this function isn't called concurrently
SCALAR_CLEAR_COUNT += 1;
}
// Scalar fallback
for i in start_idx..(start_idx + num_words) {
bitmap[i] = 0;
}
}
#[inline(always)]
fn is_range_zero(bitmap: &[u64], start_idx: usize, num_words: usize) -> bool {
debug_assert!(start_idx + num_words <= bitmap.len());
// First try AVX2 (256-bit, 4 words at a time)
#[cfg(target_feature = "avx2")]
unsafe {
let mut idx = start_idx;
let end_idx = start_idx + num_words;
// Process aligned blocks of 4 words
while idx + 4 <= end_idx {
// Safety:
// - bitmap[idx..] is valid for reads of at least 4 u64 words (32 bytes)
// - We've validated with the debug_assert that start_idx + num_words <= bitmap.len()
// - We check that idx + 4 <= end_idx to ensure we have 4 complete words
// - The unaligned _loadu_ variant is used to handle any alignment
let data_vec = _mm256_loadu_si256(bitmap[idx..].as_ptr() as *const __m256i);
// Check if any bits are non-zero
// Safety: _mm256_testz_si256 is safe when given valid __m256i values,
// which data_vec is guaranteed to be
if !_mm256_testz_si256(data_vec, data_vec) {
return false;
}
idx += 4;
}
// Handle remaining words with SSE2 or scalar ops
if idx < end_idx {
if idx + 2 <= end_idx {
// Use SSE2 for 2 words
// Safety:
// - bitmap[idx..] is valid for reads of at least 2 u64 words (16 bytes)
// - We check that idx + 2 <= end_idx to ensure we have 2 complete words
let data_vec = _mm_loadu_si128(bitmap[idx..].as_ptr() as *const __m128i);
// Safety: _mm_testz_si128 is safe when given valid __m128i values
if !_mm_testz_si128(data_vec, data_vec) {
return false;
}
idx += 2;
}
// Handle any remaining words
while idx < end_idx {
if bitmap[idx] != 0 {
return false;
}
idx += 1;
}
}
return true;
}
// If AVX2 is unavailable, try SSE2 (128-bit, 2 words at a time)
#[cfg(all(target_feature = "sse2", not(target_feature = "avx2")))]
unsafe {
let mut idx = start_idx;
let end_idx = start_idx + num_words;
// Process aligned blocks of 2 words
while idx + 2 <= end_idx {
// Safety:
// - bitmap[idx..] is valid for reads of at least 2 u64 words (16 bytes)
// - We've validated with the debug_assert that start_idx + num_words <= bitmap.len()
// - We check that idx + 2 <= end_idx to ensure we have 2 complete words
// - The unaligned _loadu_ variant is used to handle any alignment
let data_vec = _mm_loadu_si128(bitmap[idx..].as_ptr() as *const __m128i);
// Check if any bits are non-zero (SSE4.1 would have _mm_testz_si128,
// but for SSE2 compatibility we need to use a different approach)
#[cfg(target_feature = "sse4.1")]
{
// Safety: _mm_testz_si128 is safe when given valid __m128i values
if !_mm_testz_si128(data_vec, data_vec) {
return false;
}
}
#[cfg(not(target_feature = "sse4.1"))]
{
// Compare with zero vector using SSE2 only
// Safety: All operations are valid with the data_vec value
let zero_vec = _mm_setzero_si128();
let cmp = _mm_cmpeq_epi64(data_vec, zero_vec);
// The movemask gives us a bit for each byte, set if the high bit of the byte is set
// For all-zero comparison, all 16 bits should be set (0xFFFF)
let mask = _mm_movemask_epi8(cmp);
if mask != 0xFFFF {
return false;
}
}
idx += 2;
}
// Handle remaining word (if any)
if idx < end_idx && bitmap[idx] != 0 {
return false;
}
return true;
}
// Scalar fallback
bitmap[start_idx..(start_idx + num_words)]
.iter()
.all(|&word| word == 0)
}
#[inline(always)]
fn set_bit(bitmap: &mut [u64], bit_idx: u64) {
let word_idx = (bit_idx / 64) as usize;
let bit_pos = (bit_idx % 64) as u64;
bitmap[word_idx] |= 1u64 << bit_pos;
}
#[inline(always)]
fn clear_bit(bitmap: &mut [u64], bit_idx: u64) {
let word_idx = (bit_idx / 64) as usize;
let bit_pos = (bit_idx % 64) as u64;
bitmap[word_idx] &= !(1u64 << bit_pos);
}
#[inline(always)]
fn check_bit(bitmap: &[u64], bit_idx: u64) -> bool {
let word_idx = (bit_idx / 64) as usize;
let bit_pos = (bit_idx % 64) as u64;
(bitmap[word_idx] & (1u64 << bit_pos)) != 0
}
}
/// Additional x86 optimized operations not covered by the trait
pub mod atomic {
use super::*;
/// Check and set bit, returning the previous state
/// This function is not actually atomic! It's just a non-atomic optimization
#[inline(always)]
pub fn check_and_set_bit(bitmap: &mut [u64], bit_idx: u64) -> bool {
let word_idx = (bit_idx / 64) as usize;
let bit_pos = (bit_idx % 64) as u64;
let mask = 1u64 << bit_pos;
// Get old value
let old_word = bitmap[word_idx];
// Set bit regardless of current state
bitmap[word_idx] |= mask;
// Return true if bit was already set (duplicate)
(old_word & mask) != 0
}
/// Set multiple bits at once using SIMD when possible
///
/// # Safety
///
/// This function is unsafe because it:
/// - Uses SIMD intrinsics that require the AVX2 CPU feature to be available
/// - Accesses bitmap memory through raw pointers
/// - Does not perform bounds checking beyond what's required for SIMD operations
///
/// Caller must ensure:
/// - The AVX2 feature is available on the current CPU
/// - `bitmap` has sufficient size to hold indices up to `end_bit/64`
/// - `start_bit` and `end_bit` are valid bit indices within the bitmap
/// - No other thread is concurrently modifying the same memory
#[inline(always)]
#[cfg(target_feature = "avx2")]
pub unsafe fn set_bits_range(bitmap: &mut [u64], start_bit: u64, end_bit: u64) {
// Process whole words where possible
let start_word = (start_bit / 64) as usize;
let end_word = (end_bit / 64) as usize;
// Special case: all bits in the same word
if start_word == end_word {
let start_mask = u64::MAX << (start_bit % 64);
let end_mask = u64::MAX >> (63 - (end_bit % 64));
bitmap[start_word] |= start_mask & end_mask;
return;
}
// Handle partial words at the beginning and end
if start_bit % 64 != 0 {
let start_mask = u64::MAX << (start_bit % 64);
bitmap[start_word] |= start_mask;
}
if (end_bit + 1) % 64 != 0 {
let end_mask = u64::MAX >> (63 - (end_bit % 64));
bitmap[end_word] |= end_mask;
}
// Handle complete words in the middle using AVX2
let first_full_word = if start_bit % 64 == 0 {
start_word
} else {
start_word + 1
};
let last_full_word = if (end_bit + 1) % 64 == 0 {
end_word
} else {
end_word - 1
};
if first_full_word <= last_full_word {
// Use AVX2 to set multiple words at once
// Safety: _mm256_set1_epi64x is safe to call with any i64 value
let ones = _mm256_set1_epi64x(-1); // All bits set to 1
let mut i = first_full_word;
while i + 4 <= last_full_word + 1 {
// Safety:
// - bitmap[i..] is valid for reads/writes of at least 4 u64 words (32 bytes)
// - We check that i + 4 <= last_full_word + 1 to ensure we have 4 complete words
// - The unaligned _loadu/_storeu variants are used to handle any alignment
let current = _mm256_loadu_si256(bitmap[i..].as_ptr() as *const __m256i);
let result = _mm256_or_si256(current, ones);
_mm256_storeu_si256(bitmap[i..].as_mut_ptr() as *mut __m256i, result);
i += 4;
}
// Use SSE2 for remaining pairs of words
if i + 2 <= last_full_word + 1 {
// Safety:
// - bitmap[i..] is valid for reads/writes of at least 2 u64 words (16 bytes)
// - We check that i + 2 <= last_full_word + 1 to ensure we have 2 complete words
// - The unaligned _loadu/_storeu variants are used to handle any alignment
let sse_ones = _mm_set1_epi64x(-1);
let current = _mm_loadu_si128(bitmap[i..].as_ptr() as *const __m128i);
let result = _mm_or_si128(current, sse_ones);
_mm_storeu_si128(bitmap[i..].as_mut_ptr() as *mut __m128i, result);
i += 2;
}
// Handle any remaining words
while i <= last_full_word {
bitmap[i] = u64::MAX;
i += 1;
}
}
}
/// Set multiple bits at once using SSE2 (when AVX2 not available)
///
/// # Safety
///
/// This function is unsafe because it:
/// - Uses SIMD intrinsics that require the SSE2 CPU feature to be available
/// - Accesses bitmap memory through raw pointers
/// - Does not perform bounds checking beyond what's required for SIMD operations
///
/// Caller must ensure:
/// - The SSE2 feature is available on the current CPU
/// - `bitmap` has sufficient size to hold indices up to `end_bit/64`
/// - `start_bit` and `end_bit` are valid bit indices within the bitmap
/// - No other thread is concurrently modifying the same memory
#[inline(always)]
#[cfg(all(target_feature = "sse2", not(target_feature = "avx2")))]
pub unsafe fn set_bits_range(bitmap: &mut [u64], start_bit: u64, end_bit: u64) {
// Process whole words where possible
let start_word = (start_bit / 64) as usize;
let end_word = (end_bit / 64) as usize;
// Special case: all bits in the same word
if start_word == end_word {
let start_mask = u64::MAX << (start_bit % 64);
let end_mask = u64::MAX >> (63 - (end_bit % 64));
bitmap[start_word] |= start_mask & end_mask;
return;
}
// Handle partial words at the beginning and end
if start_bit % 64 != 0 {
let start_mask = u64::MAX << (start_bit % 64);
bitmap[start_word] |= start_mask;
}
if (end_bit + 1) % 64 != 0 {
let end_mask = u64::MAX >> (63 - (end_bit % 64));
bitmap[end_word] |= end_mask;
}
// Handle complete words in the middle using SSE2
let first_full_word = if start_bit % 64 == 0 {
start_word
} else {
start_word + 1
};
let last_full_word = if (end_bit + 1) % 64 == 0 {
end_word
} else {
end_word - 1
};
if first_full_word <= last_full_word {
// Use SSE2 to set multiple words at once
// Safety: _mm_set1_epi64x is safe to call with any i64 value
let ones = _mm_set1_epi64x(-1); // All bits set to 1
let mut i = first_full_word;
while i + 2 <= last_full_word + 1 {
// Safety:
// - bitmap[i..] is valid for reads/writes of at least 2 u64 words (16 bytes)
// - We check that i + 2 <= last_full_word + 1 to ensure we have 2 complete words
// - The unaligned _loadu/_storeu variants are used to handle any alignment
let current = _mm_loadu_si128(bitmap[i..].as_ptr() as *const __m128i);
let result = _mm_or_si128(current, ones);
_mm_storeu_si128(bitmap[i..].as_mut_ptr() as *mut __m128i, result);
i += 2;
}
// Handle any remaining words
while i <= last_full_word {
bitmap[i] = u64::MAX;
i += 1;
}
}
}
/// Set multiple bits at once using scalar operations (fallback)
#[inline(always)]
#[cfg(not(any(target_feature = "avx2", target_feature = "sse2")))]
pub fn set_bits_range(bitmap: &mut [u64], start_bit: u64, end_bit: u64) {
// Process whole words where possible
let start_word = (start_bit / 64) as usize;
let end_word = (end_bit / 64) as usize;
// Special case: all bits in the same word
if start_word == end_word {
let start_mask = u64::MAX << (start_bit % 64);
let end_mask = u64::MAX >> (63 - (end_bit % 64));
bitmap[start_word] |= start_mask & end_mask;
return;
}
// Handle partial words at the beginning and end
if start_bit % 64 != 0 {
let start_mask = u64::MAX << (start_bit % 64);
bitmap[start_word] |= start_mask;
}
if (end_bit + 1) % 64 != 0 {
let end_mask = u64::MAX >> (63 - (end_bit % 64));
bitmap[end_word] |= end_mask;
}
// Handle complete words in the middle
let first_full_word = if start_bit % 64 == 0 {
start_word
} else {
start_word + 1
};
let last_full_word = if (end_bit + 1) % 64 == 0 {
end_word
} else {
end_word - 1
};
for i in first_full_word..=last_full_word {
bitmap[i] = u64::MAX;
}
}
}
+876
View File
@@ -0,0 +1,876 @@
// Copyright 2025 - Nym Technologies SA <contact@nymtech.net>
// SPDX-License-Identifier: Apache-2.0
//! Replay protection validator implementation.
//!
//! This module implements the core replay protection logic using a bitmap-based
//! approach to track received packets and validate their sequence.
use crate::replay::error::{ReplayError, ReplayResult};
use crate::replay::simd::{self, BitmapOps};
// Determine the appropriate SIMD implementation at compile time
#[cfg(target_arch = "aarch64")]
#[cfg(target_feature = "neon")]
use crate::replay::simd::ArmBitmapOps as SimdImpl;
#[cfg(target_arch = "x86_64")]
#[cfg(target_feature = "avx2")]
use crate::replay::simd::X86BitmapOps as SimdImpl;
#[cfg(target_arch = "x86_64")]
#[cfg(all(not(target_feature = "avx2"), target_feature = "sse2"))]
use crate::replay::simd::X86BitmapOps as SimdImpl;
#[cfg(not(any(
all(target_arch = "x86_64", target_feature = "avx2"),
all(target_arch = "x86_64", target_feature = "sse2"),
all(target_arch = "aarch64", target_feature = "neon")
)))]
use crate::replay::simd::ScalarBitmapOps as SimdImpl;
/// Size of a word in the bitmap (64 bits)
const WORD_SIZE: usize = 64;
/// Number of words in the bitmap (allows reordering of 64*16 = 1024 packets)
const N_WORDS: usize = 16;
/// Total number of bits in the bitmap
const N_BITS: usize = WORD_SIZE * N_WORDS;
/// Validator for receiving key counters to prevent replay attacks.
///
/// This structure maintains a bitmap of received packets and validates
/// incoming packet counters to ensure they are not replayed.
#[derive(Debug, Clone, Default)]
pub struct ReceivingKeyCounterValidator {
/// Next expected counter value
next: u64,
/// Total number of received packets
receive_cnt: u64,
/// Bitmap for tracking received packets
bitmap: [u64; N_WORDS],
}
impl ReceivingKeyCounterValidator {
/// Creates a new validator with the given initial counter value.
pub fn new(initial_counter: u64) -> Self {
Self {
next: initial_counter,
receive_cnt: 0,
bitmap: [0; N_WORDS],
}
}
/// Sets a bit in the bitmap to mark a counter as received.
#[inline(always)]
fn set_bit(&mut self, idx: u64) {
SimdImpl::set_bit(&mut self.bitmap, idx % (N_BITS as u64));
}
/// Clears a bit in the bitmap.
#[inline(always)]
fn clear_bit(&mut self, idx: u64) {
SimdImpl::clear_bit(&mut self.bitmap, idx % (N_BITS as u64));
}
/// Clears the word that contains the given index.
#[inline(always)]
#[allow(dead_code)]
fn clear_word(&mut self, idx: u64) {
let bit_idx = idx % (N_BITS as u64);
let word = (bit_idx / (WORD_SIZE as u64)) as usize;
SimdImpl::clear_words(&mut self.bitmap, word, 1);
}
/// Returns true if the bit is set, false otherwise.
#[inline(always)]
fn check_bit_branchless(&self, idx: u64) -> bool {
SimdImpl::check_bit(&self.bitmap, idx % (N_BITS as u64))
}
/// Performs a quick check to determine if a counter will be accepted.
///
/// This is a fast check that can be done before more expensive operations.
///
/// Returns:
/// - `Ok(())` if the counter is acceptable
/// - `Err(ReplayError::InvalidCounter)` if the counter is invalid (too far back)
/// - `Err(ReplayError::DuplicateCounter)` if the counter has already been received
#[inline(always)]
pub fn will_accept_branchless(&self, counter: u64) -> ReplayResult<()> {
// Calculate conditions
let is_growing = counter >= self.next;
// Handle potential overflow when adding N_BITS to counter
let too_far_back = if counter > u64::MAX - (N_BITS as u64) {
// If adding N_BITS would overflow, it can't be too far back
false
} else {
counter + (N_BITS as u64) < self.next
};
let duplicate = self.check_bit_branchless(counter);
// Using Option to avoid early returns
let result = if is_growing {
Some(Ok(()))
} else if too_far_back {
Some(Err(ReplayError::OutOfWindow))
} else if duplicate {
Some(Err(ReplayError::DuplicateCounter))
} else {
Some(Ok(()))
};
// Unwrap the option (always Some)
result.unwrap()
}
/// Special case function for clearing the entire bitmap
/// Used for the fast path when we know the bitmap must be entirely cleared
#[inline(always)]
fn clear_window_fast(&mut self) {
SimdImpl::clear_words(&mut self.bitmap, 0, N_WORDS);
}
/// Checks if the bitmap is completely empty (all zeros)
/// This is used for fast path optimization
#[inline(always)]
fn is_bitmap_empty(&self) -> bool {
SimdImpl::is_range_zero(&self.bitmap, 0, N_WORDS)
}
/// Marks a counter as received and updates internal state.
///
/// This method should be called after a packet has been validated
/// and processed successfully.
///
/// Returns:
/// - `Ok(())` if the counter was successfully marked
/// - `Err(ReplayError::InvalidCounter)` if the counter is invalid (too far back)
/// - `Err(ReplayError::DuplicateCounter)` if the counter has already been received
#[inline(always)]
pub fn mark_did_receive_branchless(&mut self, counter: u64) -> ReplayResult<()> {
// Calculate conditions once - using saturating operations to prevent overflow
// For the too_far_back check, we need to avoid overflowing when adding N_BITS to counter
let too_far_back = if counter > u64::MAX - (N_BITS as u64) {
// If adding N_BITS would overflow, it can't be too far back
false
} else {
counter + (N_BITS as u64) < self.next
};
let is_sequential = counter == self.next;
let is_out_of_order = counter < self.next;
// Early return for out-of-window condition
if too_far_back {
return Err(ReplayError::OutOfWindow);
}
// Check for duplicate (only matters for out-of-order packets)
let duplicate = is_out_of_order && self.check_bit_branchless(counter);
if duplicate {
return Err(ReplayError::DuplicateCounter);
}
// Fast path for far ahead counters with empty bitmap
let far_ahead = counter.saturating_sub(self.next) >= (N_BITS as u64);
if far_ahead && self.is_bitmap_empty() {
// No need to clear anything, just set the new bit
self.set_bit(counter);
self.next = counter.saturating_add(1);
self.receive_cnt += 1;
return Ok(());
}
// Handle bitmap clearing for ahead counters that aren't sequential
if !is_sequential && !is_out_of_order {
self.clear_window(counter);
}
// Set the bit and update counters
self.set_bit(counter);
// Update next counter safely - avoid overflow
self.next = if is_sequential {
counter.saturating_add(1)
} else {
self.next.max(counter.saturating_add(1))
};
self.receive_cnt += 1;
Ok(())
}
/// Returns the current packet count statistics.
///
/// Returns a tuple of `(next, receive_cnt)` where:
/// - `next` is the next expected counter value
/// - `receive_cnt` is the total number of received packets
pub fn current_packet_cnt(&self) -> (u64, u64) {
(self.next, self.receive_cnt)
}
#[inline(always)]
#[allow(dead_code)]
fn check_and_set_bit_branchless(&mut self, idx: u64) -> bool {
let bit_idx = idx % (N_BITS as u64);
simd::atomic::check_and_set_bit(&mut self.bitmap, bit_idx)
}
#[inline(always)]
#[allow(dead_code)]
fn increment_counter_branchless(&mut self, condition: bool) {
// Add either 1 or 0 based on condition
self.receive_cnt += condition as u64;
}
#[inline(always)]
pub fn mark_sequential_branchless(&mut self, counter: u64) -> ReplayResult<()> {
// Check if sequential
let is_sequential = counter == self.next;
// Set the bit
self.set_bit(counter);
// Conditionally update next counter using saturating add to prevent overflow
self.next = self.next.saturating_add(is_sequential as u64);
// Always increment receive count if we got here
self.receive_cnt += 1;
Ok(())
}
// Helper function for window clearing with SIMD optimization
#[inline(always)]
fn clear_window(&mut self, counter: u64) {
// Handle potential overflow safely
// If counter is very large (close to u64::MAX), we need special handling
let counter_distance = counter.saturating_sub(self.next);
let far_ahead = counter_distance >= (N_BITS as u64);
// Fast path: Complete window clearing for far ahead counters
if far_ahead {
// Check if window is already clear for fast path optimization
if !self.is_bitmap_empty() {
// Use SIMD to clear the entire bitmap at once
self.clear_window_fast();
}
return;
}
// Prepare for partial window clearing
let mut i = self.next;
// Get SIMD processing width (platform optimized)
let simd_width = simd::optimal_simd_width();
// Pre-alignment clearing
if i % (WORD_SIZE as u64) != 0 {
let current_word = (i % (N_BITS as u64) / (WORD_SIZE as u64)) as usize;
// Check if we need to clear this word
if self.bitmap[current_word] != 0 {
// Safely handle potential overflow by checking before each increment
while i % (WORD_SIZE as u64) != 0 && i < counter {
self.clear_bit(i);
// Prevent overflow on increment
if i == u64::MAX {
break;
}
i += 1;
}
} else {
// Fast forward to the next word boundary
let words_to_skip = (WORD_SIZE as u64) - (i % (WORD_SIZE as u64));
if words_to_skip > u64::MAX - i {
// Would overflow, just set to MAX
i = u64::MAX;
} else {
i += words_to_skip;
}
}
}
// Word-aligned clearing with SIMD where possible
while i <= counter.saturating_sub(WORD_SIZE as u64) {
let current_word = (i % (N_BITS as u64) / (WORD_SIZE as u64)) as usize;
// Check if we have enough consecutive words to use SIMD
if current_word + simd_width <= N_WORDS
&& i % (simd_width as u64 * WORD_SIZE as u64) == 0
{
// Use SIMD to clear multiple words at once if any need clearing
let needs_clearing =
!SimdImpl::is_range_zero(&self.bitmap, current_word, simd_width);
if needs_clearing {
SimdImpl::clear_words(&mut self.bitmap, current_word, simd_width);
}
// Skip the words we just processed
let words_to_skip = simd_width as u64 * WORD_SIZE as u64;
if words_to_skip > u64::MAX - i {
i = u64::MAX;
break;
}
i += words_to_skip;
} else {
// Process single word
if self.bitmap[current_word] != 0 {
self.bitmap[current_word] = 0;
}
// Check for potential overflow before incrementing
if i > u64::MAX - (WORD_SIZE as u64) {
i = u64::MAX;
break;
}
i += WORD_SIZE as u64;
}
}
// Post-alignment clearing (bit by bit for remaining bits)
if i < counter {
let final_word = (i % (N_BITS as u64) / (WORD_SIZE as u64)) as usize;
let is_final_word_empty = self.bitmap[final_word] == 0;
// Skip clearing if word is already empty
if !is_final_word_empty {
while i < counter {
self.clear_bit(i);
// Prevent overflow on increment
if i == u64::MAX {
break;
}
i += 1;
}
}
}
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_replay_counter_basic() {
let mut validator = ReceivingKeyCounterValidator::default();
// Check initial state
assert_eq!(validator.next, 0);
assert_eq!(validator.receive_cnt, 0);
// Test sequential counters
assert!(validator.mark_did_receive_branchless(0).is_ok());
assert!(validator.mark_did_receive_branchless(0).is_err());
assert!(validator.mark_did_receive_branchless(1).is_ok());
assert!(validator.mark_did_receive_branchless(1).is_err());
}
#[test]
fn test_replay_counter_out_of_order() {
let mut validator = ReceivingKeyCounterValidator::default();
// Process some sequential packets
assert!(validator.mark_did_receive_branchless(0).is_ok());
assert!(validator.mark_did_receive_branchless(1).is_ok());
assert!(validator.mark_did_receive_branchless(2).is_ok());
// Out-of-order packet that hasn't been seen yet
assert!(validator.mark_did_receive_branchless(1).is_err()); // Already seen
assert!(validator.mark_did_receive_branchless(10).is_ok()); // New packet, ahead of next
// Next should now be 11
assert_eq!(validator.next, 11);
// Can still accept packets in the valid window
assert!(validator.will_accept_branchless(9).is_ok());
assert!(validator.will_accept_branchless(8).is_ok());
// But duplicates are rejected
assert!(validator.will_accept_branchless(10).is_err());
}
#[test]
fn test_replay_counter_full() {
let mut validator = ReceivingKeyCounterValidator::default();
// Process a bunch of sequential packets
for i in 0..64 {
assert!(validator.mark_did_receive_branchless(i).is_ok());
assert!(validator.mark_did_receive_branchless(i).is_err());
}
// Test out of order within window
assert!(validator.mark_did_receive_branchless(15).is_err()); // Already seen
assert!(validator.mark_did_receive_branchless(63).is_err()); // Already seen
// Test for packets within bitmap range
for i in 64..(N_BITS as u64) + 128 {
assert!(validator.mark_did_receive_branchless(i).is_ok());
assert!(validator.mark_did_receive_branchless(i).is_err());
}
}
#[test]
fn test_replay_counter_window_sliding() {
let mut validator = ReceivingKeyCounterValidator::default();
// Jump far ahead to force window sliding
let far_ahead = (N_BITS as u64) * 3;
assert!(validator.mark_did_receive_branchless(far_ahead).is_ok());
// Everything too far back should be rejected
for i in 0..=(N_BITS as u64) * 2 {
assert!(matches!(
validator.will_accept_branchless(i),
Err(ReplayError::OutOfWindow)
));
assert!(validator.mark_did_receive_branchless(i).is_err());
}
// Values in window but less than far_ahead should be accepted
for i in (N_BITS as u64) * 2 + 1..far_ahead {
assert!(validator.will_accept_branchless(i).is_ok());
}
// The far_ahead value itself should be rejected now (duplicate)
assert!(matches!(
validator.will_accept_branchless(far_ahead),
Err(ReplayError::DuplicateCounter)
));
// Test receiving packets in reverse order within window
for i in ((N_BITS as u64) * 2 + 1..far_ahead).rev() {
assert!(validator.mark_did_receive_branchless(i).is_ok());
assert!(validator.mark_did_receive_branchless(i).is_err());
}
}
#[test]
fn test_out_of_order_tracking() {
let mut validator = ReceivingKeyCounterValidator::default();
// Jump ahead
assert!(validator.mark_did_receive_branchless(1000).is_ok());
// Test some more additions
assert!(validator.mark_did_receive_branchless(1000 + 70).is_ok());
assert!(validator.mark_did_receive_branchless(1000 + 71).is_ok());
assert!(validator.mark_did_receive_branchless(1000 + 72).is_ok());
assert!(validator
.mark_did_receive_branchless(1000 + 72 + 125)
.is_ok());
assert!(validator.mark_did_receive_branchless(1000 + 63).is_ok());
// Check duplicates
assert!(validator.mark_did_receive_branchless(1000 + 70).is_err());
assert!(validator.mark_did_receive_branchless(1000 + 71).is_err());
assert!(validator.mark_did_receive_branchless(1000 + 72).is_err());
}
#[test]
fn test_counter_stats() {
let mut validator = ReceivingKeyCounterValidator::default();
// Initial state
let (next, count) = validator.current_packet_cnt();
assert_eq!(next, 0);
assert_eq!(count, 0);
// After receiving some packets
assert!(validator.mark_did_receive_branchless(0).is_ok());
assert!(validator.mark_did_receive_branchless(1).is_ok());
assert!(validator.mark_did_receive_branchless(2).is_ok());
let (next, count) = validator.current_packet_cnt();
assert_eq!(next, 3);
assert_eq!(count, 3);
// After an out of order packet
assert!(validator.mark_did_receive_branchless(10).is_ok());
let (next, count) = validator.current_packet_cnt();
assert_eq!(next, 11);
assert_eq!(count, 4);
// After a packet from the past (within window)
assert!(validator.mark_did_receive_branchless(5).is_ok());
let (next, count) = validator.current_packet_cnt();
assert_eq!(next, 11); // Next doesn't change
assert_eq!(count, 5); // Count increases
}
#[test]
fn test_window_boundary_edge_cases() {
let mut validator = ReceivingKeyCounterValidator::default();
// First process a sequence of packets
for i in 0..100 {
assert!(validator.mark_did_receive_branchless(i).is_ok());
}
// The window should now span from 100 to 100+N_BITS
// Test packet near the upper edge of the window
let upper_edge = 100 + (N_BITS as u64) - 1;
assert!(validator.will_accept_branchless(upper_edge).is_ok());
assert!(validator.mark_did_receive_branchless(upper_edge).is_ok());
// Test packet just outside the upper edge (should be accepted)
let just_outside_upper = 100 + (N_BITS as u64);
assert!(validator.will_accept_branchless(just_outside_upper).is_ok());
// Test packet near the lower edge of the window
let lower_edge = 100 + 1; // +1 because we've already processed 100
assert!(validator.will_accept_branchless(lower_edge).is_ok());
// Test packet just outside the lower edge (should be rejected)
if upper_edge >= (N_BITS as u64) * 2 {
// Only test this if we're far enough along to have a lower bound
let just_outside_lower = 100 - (N_BITS as u64);
assert!(matches!(
validator.will_accept_branchless(just_outside_lower),
Err(ReplayError::OutOfWindow)
));
}
}
#[test]
fn test_multiple_window_shifts() {
let mut validator = ReceivingKeyCounterValidator::default();
// First jump - process packet far ahead
let first_jump = 2000;
assert!(validator.mark_did_receive_branchless(first_jump).is_ok());
// Verify next counter is updated
let (next, _) = validator.current_packet_cnt();
assert_eq!(next, first_jump + 1);
// Second large jump, even further ahead
let second_jump = first_jump + 5000;
assert!(validator.mark_did_receive_branchless(second_jump).is_ok());
// Verify next counter is updated again
let (next, _) = validator.current_packet_cnt();
assert_eq!(next, second_jump + 1);
// Test packets within the new window
let mid_window = second_jump - 500;
assert!(validator.will_accept_branchless(mid_window).is_ok());
// Test packets outside the new window
let outside_window = first_jump + 100;
assert!(matches!(
validator.will_accept_branchless(outside_window),
Err(ReplayError::OutOfWindow)
));
}
#[test]
fn test_interleaved_packets_at_boundaries() {
let mut validator = ReceivingKeyCounterValidator::default();
// Jump ahead to establish a large window
let jump = 2000;
assert!(validator.mark_did_receive_branchless(jump).is_ok());
// Process a sequence at the upper boundary
for i in 0..10 {
let upper_packet = jump + 100 + i;
assert!(validator.mark_did_receive_branchless(upper_packet).is_ok());
}
// Process a sequence at the lower boundary
for i in 0..10 {
let lower_packet = jump - (N_BITS as u64) + 100 + i;
// These might fail if they're outside the window, that's ok
let _ = validator.mark_did_receive_branchless(lower_packet);
}
// Process alternating packets at both ends
for i in 0..5 {
let upper = jump + 200 + i;
let lower = jump - (N_BITS as u64) + 200 + i;
assert!(validator.will_accept_branchless(upper).is_ok());
let lower_result = validator.will_accept_branchless(lower);
// Lower might be accepted or rejected, depending on exactly where the window is
if lower_result.is_ok() {
assert!(validator.mark_did_receive_branchless(lower).is_ok());
}
assert!(validator.mark_did_receive_branchless(upper).is_ok());
}
}
#[test]
fn test_exact_window_size_with_full_bitmap() {
let mut validator = ReceivingKeyCounterValidator::default();
// Fill the entire bitmap with non-sequential packets
// This tests both window size and bitmap capacity
// Generate a random but reproducible pattern
let mut positions = Vec::new();
for i in 0..N_BITS {
positions.push((i * 7) % N_BITS);
}
// Mark packets in this pattern
for pos in &positions {
assert!(validator.mark_did_receive_branchless(*pos as u64).is_ok());
}
// Try to mark them again (should all fail as duplicates)
for pos in &positions {
assert!(matches!(
validator.mark_did_receive_branchless(*pos as u64),
Err(ReplayError::DuplicateCounter)
));
}
// Force window to slide
let far_ahead = (N_BITS as u64) * 2;
assert!(validator.mark_did_receive_branchless(far_ahead).is_ok());
// Old packets should now be outside the window
for pos in &positions {
if *pos as u64 + (N_BITS as u64) < far_ahead {
assert!(matches!(
validator.will_accept_branchless(*pos as u64),
Err(ReplayError::OutOfWindow)
));
}
}
}
use std::sync::{Arc, Barrier};
use std::thread;
#[test]
fn test_concurrent_access() {
let validator = Arc::new(std::sync::Mutex::new(
ReceivingKeyCounterValidator::default(),
));
let num_threads = 8;
let operations_per_thread = 1000;
let barrier = Arc::new(Barrier::new(num_threads));
// Create thread handles
let mut handles = vec![];
for thread_id in 0..num_threads {
let validator_clone = Arc::clone(&validator);
let barrier_clone = Arc::clone(&barrier);
let handle = thread::spawn(move || {
// Wait for all threads to be ready
barrier_clone.wait();
let mut successes = 0;
let mut duplicates = 0;
let mut out_of_window = 0;
for i in 0..operations_per_thread {
// Generate a somewhat random but reproducible counter value
// Different threads will sometimes try to insert the same value
let counter = (i * 7 + thread_id * 13) as u64;
let mut guard = validator_clone.lock().unwrap();
match guard.mark_did_receive_branchless(counter) {
Ok(()) => successes += 1,
Err(ReplayError::DuplicateCounter) => duplicates += 1,
Err(ReplayError::OutOfWindow) => out_of_window += 1,
_ => {}
}
}
(successes, duplicates, out_of_window)
});
handles.push(handle);
}
// Collect results
let mut total_successes = 0;
let mut total_duplicates = 0;
let mut total_out_of_window = 0;
for handle in handles {
let (successes, duplicates, out_of_window) = handle.join().unwrap();
total_successes += successes;
total_duplicates += duplicates;
total_out_of_window += out_of_window;
}
// Verify that all operations were accounted for
assert_eq!(
total_successes + total_duplicates + total_out_of_window,
num_threads * operations_per_thread
);
// Verify that some operations were successful and some were duplicates
assert!(total_successes > 0);
assert!(total_duplicates > 0);
// Check final state of the validator
let final_state = validator.lock().unwrap();
let (next, receive_cnt) = final_state.current_packet_cnt();
// Verify that the received count matches our successful operations
assert_eq!(receive_cnt, total_successes as u64);
}
#[test]
fn test_memory_usage() {
use std::mem::{size_of, size_of_val};
// Test small validator
let validator_default = ReceivingKeyCounterValidator::default();
let size_default = size_of_val(&validator_default);
// Expected size calculation
let expected_size = size_of::<u64>() * 2 + // next + receive_cnt
size_of::<u64>() * N_WORDS; // bitmap
assert_eq!(size_default, expected_size);
println!("Default validator size: {} bytes", size_default);
// Memory efficiency calculation (bits tracked per byte of memory)
let bits_per_byte = N_BITS as f64 / size_default as f64;
println!(
"Memory efficiency: {:.2} bits tracked per byte of memory",
bits_per_byte
);
// Verify minimum memory needed for different window sizes
for window_size in [64, 128, 256, 512, 1024, 2048] {
let words_needed = (window_size + WORD_SIZE - 1) / WORD_SIZE; // Ceiling division
let memory_needed = size_of::<u64>() * 2 + size_of::<u64>() * words_needed;
println!(
"Window size {}: {} bytes minimum",
window_size, memory_needed
);
}
}
#[test]
#[cfg(any(
target_feature = "sse2",
target_feature = "avx2",
target_feature = "neon"
))]
fn test_simd_operations() {
// This test verifies that SIMD-optimized operations would produce
// the same results as the scalar implementation
// Create a validator with a known state
let mut validator = ReceivingKeyCounterValidator::default();
// Fill bitmap with a pattern
for i in 0..64 {
validator.set_bit(i);
}
// Create a copy for comparison
let original_bitmap = validator.bitmap;
// Simulate SIMD clear (4 words at a time)
#[cfg(target_feature = "avx2")]
{
use std::arch::x86_64::{_mm256_setzero_si256, _mm256_storeu_si256};
// Clear words 0-3 using AVX2
unsafe {
let zero_vec = _mm256_setzero_si256();
_mm256_storeu_si256(validator.bitmap.as_mut_ptr() as *mut _, zero_vec);
}
// Verify first 4 words are cleared
assert_eq!(validator.bitmap[0], 0);
assert_eq!(validator.bitmap[1], 0);
assert_eq!(validator.bitmap[2], 0);
assert_eq!(validator.bitmap[3], 0);
// Verify other words are unchanged
for i in 4..N_WORDS {
assert_eq!(validator.bitmap[i], original_bitmap[i]);
}
}
#[cfg(target_feature = "sse2")]
{
use std::arch::x86_64::{_mm_setzero_si128, _mm_storeu_si128};
// Reset validator
validator.bitmap = original_bitmap;
// Clear words 0-1 using SSE2
unsafe {
let zero_vec = _mm_setzero_si128();
_mm_storeu_si128(validator.bitmap.as_mut_ptr() as *mut _, zero_vec);
}
// Verify first 2 words are cleared
assert_eq!(validator.bitmap[0], 0);
assert_eq!(validator.bitmap[1], 0);
// Verify other words are unchanged
for i in 2..N_WORDS {
assert_eq!(validator.bitmap[i], original_bitmap[i]);
}
}
// No SIMD available, make this test a no-op
#[cfg(not(any(
target_feature = "sse2",
target_feature = "avx2",
target_feature = "neon"
)))]
{
println!("No SIMD features available, skipping SIMD test");
}
}
#[test]
fn test_clear_window_overflow() {
let mut validator = ReceivingKeyCounterValidator::default();
// Set a very large next value, close to u64::MAX
validator.next = u64::MAX - 1000;
// Try to clear window with an even higher counter
// This should exercise the potentially problematic code
let counter = u64::MAX - 500;
// Call clear_window directly (this is what we suspect has issues)
validator.clear_window(counter);
// If we got here without a panic, at least it's not crashing
// Let's verify the bitmap state is reasonable
let any_non_zero = validator.bitmap.iter().any(|&word| word != 0);
assert!(!any_non_zero, "Bitmap should be cleared");
// Try the full function which uses clear_window internally
assert!(validator.mark_did_receive_branchless(counter).is_ok());
// Verify it was marked
assert!(matches!(
validator.will_accept_branchless(counter),
Err(ReplayError::DuplicateCounter)
));
}
}
+658
View File
@@ -0,0 +1,658 @@
// Copyright 2025 - Nym Technologies SA <contact@nymtech.net>
// SPDX-License-Identifier: Apache-2.0
//! Session management for the Lewes Protocol.
//!
//! This module implements session management functionality, including replay protection
//! and Noise protocol state handling.
use crate::noise_protocol::{NoiseError, NoiseProtocol, ReadResult};
use crate::packet::LpHeader;
use crate::replay::ReceivingKeyCounterValidator;
use crate::{LpError, LpMessage, LpPacket};
use parking_lot::Mutex;
use snow::Builder;
use std::sync::atomic::{AtomicU64, Ordering};
/// A session in the Lewes Protocol, handling connection state with Noise.
///
/// Sessions manage connection state, including LP replay protection and Noise cryptography.
/// Each session has a unique receiving index and sending index for connection identification.
#[derive(Debug)]
pub struct LpSession {
id: u32,
/// Flag indicating if this session acts as the Noise protocol initiator.
is_initiator: bool,
/// Noise protocol state machine
noise_state: Mutex<NoiseProtocol>,
/// Counter for outgoing packets
sending_counter: AtomicU64,
/// Validator for incoming packet counters to prevent replay attacks
receiving_counter: Mutex<ReceivingKeyCounterValidator>,
}
impl LpSession {
pub fn id(&self) -> u32 {
self.id
}
pub fn noise_state(&self) -> &Mutex<NoiseProtocol> {
&self.noise_state
}
/// Returns true if this session was created as the initiator.
pub fn is_initiator(&self) -> bool {
self.is_initiator
}
/// Creates a new session and initializes the Noise protocol state.
///
/// # Arguments
///
/// * `receiving_index` - Index used for receiving packets (becomes session ID).
/// * `sending_index` - Index used for sending packets to the peer.
/// * `is_initiator` - True if this side initiates the Noise handshake.
/// * `local_static_key` - This side's static private key (e.g., X25519).
/// * `remote_static_key` - The peer's static public key (required for initiator in some patterns like XK).
/// * `psk` - The pre-shared key established out-of-band.
/// * `pattern_name` - The Noise protocol pattern string (e.g., "Noise_XKpsk3_25519_ChaChaPoly_SHA256").
/// * `psk_index` - The index/position where the PSK is mixed in according to the pattern.
pub fn new(
id: u32,
is_initiator: bool,
local_private_key: &[u8],
remote_public_key: &[u8],
psk: &[u8],
) -> Result<Self, LpError> {
// AIDEV-NOTE: XKpsk3 pattern requires remote static key known upfront (XK)
// and PSK mixed at position 3. This provides forward secrecy with PSK authentication.
let pattern_name = "Noise_XKpsk3_25519_ChaChaPoly_SHA256";
let psk_index = 3;
let params = pattern_name.parse()?;
let builder = Builder::new(params);
let builder = builder.local_private_key(local_private_key);
let builder = builder.remote_public_key(remote_public_key);
let builder = builder.psk(psk_index, psk);
let initial_state = if is_initiator {
builder.build_initiator().map_err(LpError::SnowKeyError)?
} else {
builder.build_responder().map_err(LpError::SnowKeyError)?
};
let noise_protocol = NoiseProtocol::new(initial_state);
Ok(Self {
id,
is_initiator,
noise_state: Mutex::new(noise_protocol),
sending_counter: AtomicU64::new(0),
receiving_counter: Mutex::new(ReceivingKeyCounterValidator::default()),
})
}
pub fn next_packet(&self, message: LpMessage) -> Result<LpPacket, LpError> {
let counter = self.next_counter();
let header = LpHeader::new(self.id(), counter);
let packet = LpPacket::new(header, message);
Ok(packet)
}
/// Generates the next counter value for outgoing packets.
pub fn next_counter(&self) -> u64 {
self.sending_counter.fetch_add(1, Ordering::Relaxed)
}
/// Performs a quick validation check for an incoming packet counter.
///
/// This should be called before performing any expensive operations like
/// decryption/Noise processing to efficiently filter out potential replay attacks.
///
/// # Arguments
///
/// * `counter` - The counter value to check
///
/// # Returns
///
/// * `Ok(())` if the counter is likely valid
/// * `Err(LpError::Replay)` if the counter is invalid or a potential replay
pub fn receiving_counter_quick_check(&self, counter: u64) -> Result<(), LpError> {
// AIDEV-NOTE: Branchless implementation uses SIMD when available for constant-time
// operations, preventing timing attacks. Check before crypto to save CPU cycles.
let counter_validator = self.receiving_counter.lock();
counter_validator
.will_accept_branchless(counter)
.map_err(LpError::Replay)
}
/// Marks a counter as received after successful packet processing.
///
/// This should be called after a packet has been successfully decoded and processed
/// (including Noise decryption/handshake step) to update the replay protection state.
///
/// # Arguments
///
/// * `counter` - The counter value to mark as received
///
/// # Returns
///
/// * `Ok(())` if the counter was successfully marked
/// * `Err(LpError::Replay)` if the counter cannot be marked (duplicate, too old, etc.)
pub fn receiving_counter_mark(&self, counter: u64) -> Result<(), LpError> {
let mut counter_validator = self.receiving_counter.lock();
counter_validator
.mark_did_receive_branchless(counter)
.map_err(LpError::Replay)
}
/// Returns current packet statistics for monitoring.
///
/// # Returns
///
/// A tuple containing:
/// * The next expected counter value for incoming packets
/// * The total number of received packets
pub fn current_packet_cnt(&self) -> (u64, u64) {
let counter_validator = self.receiving_counter.lock();
counter_validator.current_packet_cnt()
}
/// Prepares the next handshake message to be sent, if any.
///
/// This should be called by the driver/IO layer to check if the Noise protocol
/// state machine requires a message to be sent to the peer.
///
/// # Returns
///
/// * `Ok(None)` if no message needs to be sent currently (e.g., waiting for peer, or handshake complete).
/// * `Err(NoiseError)` if there's an error within the Noise protocol state.
pub fn prepare_handshake_message(&self) -> Option<Result<LpMessage, LpError>> {
let mut noise_state = self.noise_state.lock();
if let Some(message) = noise_state.get_bytes_to_send() {
match message {
Ok(message) => Some(Ok(LpMessage::Handshake(message))),
Err(e) => Some(Err(LpError::NoiseError(e))),
}
} else {
None
}
}
/// Processes a received handshake message from the peer.
///
/// This should be called by the driver/IO layer after receiving a potential
/// handshake message payload from an LP packet.
///
/// # Arguments
///
/// * `noise_payload` - The raw bytes received from the peer, purported to be a Noise handshake message.
///
/// # Returns
///
/// * `Ok(ReadResult)` detailing the outcome (e.g., handshake complete, no-op).
/// * `Err(NoiseError)` if the message is invalid or causes a Noise protocol error.
pub fn process_handshake_message(&self, message: &LpMessage) -> Result<ReadResult, NoiseError> {
let mut noise_state = self.noise_state.lock();
match message {
LpMessage::Handshake(payload) => {
// The sans-io NoiseProtocol::read_message expects only the payload.
noise_state.read_message(payload)
}
_ => Err(NoiseError::IncorrectStateError),
}
}
/// Checks if the Noise handshake phase is complete.
pub fn is_handshake_complete(&self) -> bool {
self.noise_state.lock().is_handshake_finished()
}
/// Encrypts application data payload using the established Noise transport session.
///
/// This should only be called after the handshake is complete (`is_handshake_complete` returns true).
///
/// # Arguments
///
/// * `payload` - The application data to encrypt.
///
/// # Returns
///
/// * `Ok(Vec<u8>)` containing the encrypted Noise message ciphertext.
/// * `Err(NoiseError)` if the session is not in transport mode or encryption fails.
pub fn encrypt_data(&self, payload: &[u8]) -> Result<LpMessage, NoiseError> {
let mut noise_state = self.noise_state.lock();
// Explicitly check if handshake is finished before trying to write
if !noise_state.is_handshake_finished() {
return Err(NoiseError::IncorrectStateError);
}
let payload = noise_state.write_message(payload)?;
Ok(LpMessage::EncryptedData(payload))
}
/// Decrypts an incoming Noise message containing application data.
///
/// This should only be called after the handshake is complete (`is_handshake_complete` returns true)
/// and when an `LPMessage::EncryptedData` is received.
///
/// # Arguments
///
/// * `noise_ciphertext` - The encrypted Noise message received from the peer.
///
/// # Returns
///
/// * `Ok(Vec<u8>)` containing the decrypted application data payload.
/// * `Err(NoiseError)` if the session is not in transport mode, decryption fails, or the message is not data.
pub fn decrypt_data(&self, noise_ciphertext: &LpMessage) -> Result<Vec<u8>, NoiseError> {
let mut noise_state = self.noise_state.lock();
// Explicitly check if handshake is finished before trying to read
if !noise_state.is_handshake_finished() {
return Err(NoiseError::IncorrectStateError);
}
let payload = noise_ciphertext.payload();
match noise_state.read_message(payload)? {
ReadResult::DecryptedData(data) => Ok(data),
_ => Err(NoiseError::IncorrectStateError),
}
}
}
#[cfg(test)]
mod tests {
use snow::{params::NoiseParams, Keypair};
use super::*;
use crate::{replay::ReplayError, sessions_for_tests, NOISE_PATTERN};
// Helper function to generate keypairs for tests
fn generate_keypair() -> Keypair {
let params: NoiseParams = NOISE_PATTERN.parse().unwrap();
snow::Builder::new(params).generate_keypair().unwrap()
}
// Helper function to create a session with real keys for handshake tests
fn create_handshake_test_session(
is_initiator: bool,
local_keys: &Keypair,
remote_pub_key: &[u8],
psk: &[u8],
) -> LpSession {
// Use a dummy ID for testing, the important part is is_initiator
let test_id = if is_initiator { 1 } else { 2 };
LpSession::new(
test_id,
is_initiator,
&local_keys.private,
remote_pub_key,
psk,
)
.expect("Test session creation failed")
}
#[test]
fn test_session_creation() {
let session = sessions_for_tests().0;
// Initial counter should be zero
let counter = session.next_counter();
assert_eq!(counter, 0);
// Counter should increment
let counter = session.next_counter();
assert_eq!(counter, 1);
}
#[test]
fn test_replay_protection_sequential() {
let session = sessions_for_tests().1;
// Sequential counters should be accepted
assert!(session.receiving_counter_quick_check(0).is_ok());
assert!(session.receiving_counter_mark(0).is_ok());
assert!(session.receiving_counter_quick_check(1).is_ok());
assert!(session.receiving_counter_mark(1).is_ok());
// Duplicates should be rejected
assert!(session.receiving_counter_quick_check(0).is_err());
let err = session.receiving_counter_mark(0).unwrap_err();
match err {
LpError::Replay(replay_error) => {
assert!(matches!(replay_error, ReplayError::DuplicateCounter));
}
_ => panic!("Expected replay error"),
}
}
#[test]
fn test_replay_protection_out_of_order() {
let session = sessions_for_tests().1;
// Receive packets in order
assert!(session.receiving_counter_mark(0).is_ok());
assert!(session.receiving_counter_mark(1).is_ok());
assert!(session.receiving_counter_mark(2).is_ok());
// Skip ahead
assert!(session.receiving_counter_mark(10).is_ok());
// Can still receive out-of-order packets within window
assert!(session.receiving_counter_quick_check(5).is_ok());
assert!(session.receiving_counter_mark(5).is_ok());
// But duplicates are still rejected
assert!(session.receiving_counter_quick_check(5).is_err());
assert!(session.receiving_counter_mark(5).is_err());
}
#[test]
fn test_packet_stats() {
let session = sessions_for_tests().1;
// Initial stats
let (next, received) = session.current_packet_cnt();
assert_eq!(next, 0);
assert_eq!(received, 0);
// After receiving packets
assert!(session.receiving_counter_mark(0).is_ok());
assert!(session.receiving_counter_mark(1).is_ok());
let (next, received) = session.current_packet_cnt();
assert_eq!(next, 2);
assert_eq!(received, 2);
}
#[test]
fn test_prepare_handshake_message_initial_state() {
let initiator_keys = generate_keypair();
let responder_keys = generate_keypair();
let psk = [3u8; 32];
let initiator_session =
create_handshake_test_session(true, &initiator_keys, &responder_keys.public, &psk);
let responder_session = create_handshake_test_session(
false,
&responder_keys,
&initiator_keys.public, // Responder also needs initiator's key for XK
&psk,
);
// Initiator should have a message to send immediately (-> e)
let initiator_msg_result = initiator_session.prepare_handshake_message();
assert!(initiator_msg_result.is_some());
let initiator_msg = initiator_msg_result
.unwrap()
.expect("Initiator msg prep failed");
assert!(!initiator_msg.is_empty());
// Responder should have nothing to send initially (waits for <- e)
let responder_msg_result = responder_session.prepare_handshake_message();
assert!(responder_msg_result.is_none());
}
#[test]
fn test_process_handshake_message_first_step() {
let initiator_keys = generate_keypair();
let responder_keys = generate_keypair();
let psk = [4u8; 32];
let initiator_session =
create_handshake_test_session(true, &initiator_keys, &responder_keys.public, &psk);
let responder_session =
create_handshake_test_session(false, &responder_keys, &initiator_keys.public, &psk);
// 1. Initiator prepares the first message (-> e)
let initiator_msg_result = initiator_session.prepare_handshake_message();
let initiator_msg = initiator_msg_result
.unwrap()
.expect("Initiator msg prep failed");
// 2. Responder processes the message (<- e)
let process_result = responder_session.process_handshake_message(&initiator_msg);
// Check the result of processing
match process_result {
Ok(ReadResult::NoOp) => {
// Expected for XK first message, responder doesn't decrypt data yet
}
Ok(other) => panic!("Unexpected process result: {:?}", other),
Err(e) => panic!("Responder processing failed: {:?}", e),
}
// 3. After processing, responder should now have a message to send (-> e, es)
let responder_response_result = responder_session.prepare_handshake_message();
assert!(responder_response_result.is_some());
let responder_response = responder_response_result
.unwrap()
.expect("Responder response prep failed");
assert!(!responder_response.is_empty());
}
#[test]
fn test_handshake_driver_simulation() {
let initiator_keys = generate_keypair();
let responder_keys = generate_keypair();
let psk = [5u8; 32];
let initiator_session =
create_handshake_test_session(true, &initiator_keys, &responder_keys.public, &psk);
let responder_session =
create_handshake_test_session(false, &responder_keys, &initiator_keys.public, &psk);
let mut initiator_to_responder_msg = None;
let mut responder_to_initiator_msg = None;
let mut rounds = 0;
const MAX_ROUNDS: usize = 10; // Safety break for the loop
// Start by priming the initiator message
initiator_to_responder_msg = initiator_session.prepare_handshake_message().unwrap().ok();
assert!(
initiator_to_responder_msg.is_some(),
"Initiator did not produce initial message"
);
while rounds < MAX_ROUNDS {
rounds += 1;
// === Initiator -> Responder ===
if let Some(msg) = initiator_to_responder_msg.take() {
// Process message
match responder_session.process_handshake_message(&msg) {
Ok(_) => {}
Err(e) => panic!("Responder processing failed: {:?}", e),
}
// Check if responder needs to send a reply
responder_to_initiator_msg = responder_session
.prepare_handshake_message()
.transpose()
.unwrap();
}
// Check completion after potentially processing responder's message below
if initiator_session.is_handshake_complete()
&& responder_session.is_handshake_complete()
{
break;
}
// === Responder -> Initiator ===
if let Some(msg) = responder_to_initiator_msg.take() {
// Process message
match initiator_session.process_handshake_message(&msg) {
Ok(_) => {}
Err(e) => panic!("Initiator processing failed: {:?}", e),
}
// Check if initiator needs to send a reply (should be last message in XK)
initiator_to_responder_msg = initiator_session
.prepare_handshake_message()
.transpose()
.unwrap();
}
// Check completion again after potentially processing initiator's message above
if initiator_session.is_handshake_complete()
&& responder_session.is_handshake_complete()
{
break;
}
}
assert!(
rounds < MAX_ROUNDS,
"Handshake did not complete within max rounds"
);
assert!(
initiator_session.is_handshake_complete(),
"Initiator handshake did not complete"
);
assert!(
responder_session.is_handshake_complete(),
"Responder handshake did not complete"
);
println!("Handshake completed in {} rounds.", rounds);
}
#[test]
fn test_encrypt_decrypt_after_handshake() {
// --- Setup Handshake ---
let initiator_keys = generate_keypair();
let responder_keys = generate_keypair();
let psk = [6u8; 32];
let initiator_session =
create_handshake_test_session(true, &initiator_keys, &responder_keys.public, &psk);
let responder_session =
create_handshake_test_session(false, &responder_keys, &initiator_keys.public, &psk);
// Drive handshake to completion (simplified loop from previous test)
let mut i_msg = initiator_session
.prepare_handshake_message()
.unwrap()
.unwrap();
responder_session.process_handshake_message(&i_msg).unwrap();
let r_msg = responder_session
.prepare_handshake_message()
.unwrap()
.unwrap();
initiator_session.process_handshake_message(&r_msg).unwrap();
i_msg = initiator_session
.prepare_handshake_message()
.unwrap()
.unwrap();
responder_session.process_handshake_message(&i_msg).unwrap();
assert!(initiator_session.is_handshake_complete());
assert!(responder_session.is_handshake_complete());
// --- Test Encryption/Decryption ---
let plaintext = b"Hello, Lewes Protocol!";
// Initiator encrypts
let ciphertext = initiator_session
.encrypt_data(plaintext)
.expect("Initiator encryption failed");
assert_ne!(ciphertext.payload(), plaintext); // Ensure it's actually encrypted
// Responder decrypts
let decrypted = responder_session
.decrypt_data(&ciphertext)
.expect("Responder decryption failed");
assert_eq!(decrypted, plaintext);
// --- Test other direction ---
let plaintext2 = b"Response from responder.";
// Responder encrypts
let ciphertext2 = responder_session
.encrypt_data(plaintext2)
.expect("Responder encryption failed");
assert_ne!(ciphertext2.payload(), plaintext2);
// Initiator decrypts
let decrypted2 = initiator_session
.decrypt_data(&ciphertext2)
.expect("Initiator decryption failed");
assert_eq!(decrypted2, plaintext2);
}
#[test]
fn test_encrypt_decrypt_before_handshake() {
let initiator_keys = generate_keypair();
let responder_keys = generate_keypair();
let psk = [7u8; 32];
let initiator_session =
create_handshake_test_session(true, &initiator_keys, &responder_keys.public, &psk);
assert!(!initiator_session.is_handshake_complete());
// Attempt to encrypt before handshake
let plaintext = b"This should fail";
let result = initiator_session.encrypt_data(plaintext);
assert!(result.is_err());
match result.unwrap_err() {
NoiseError::IncorrectStateError => {} // Expected error
e => panic!("Expected IncorrectStateError, got {:?}", e),
}
// Attempt to decrypt before handshake (using dummy ciphertext)
let dummy_ciphertext = vec![0u8; 32];
let result_decrypt =
initiator_session.decrypt_data(&LpMessage::EncryptedData(dummy_ciphertext));
assert!(result_decrypt.is_err());
match result_decrypt.unwrap_err() {
NoiseError::IncorrectStateError => {} // Expected error
e => panic!("Expected IncorrectStateError, got {:?}", e),
}
}
/*
// These tests remain commented as they rely on the old mock crypto functions
#[test]
fn test_mock_crypto() {
let session = create_test_session(true);
let data = [1, 2, 3, 4, 5];
let mut encrypted = [0; 5];
let mut decrypted = [0; 5];
// Mock encrypt should copy the data
// let encrypted_len = session.encrypt_packet(&data, &mut encrypted).unwrap(); // Removed method
// assert_eq!(encrypted_len, 5);
// assert_eq!(encrypted, data);
// Mock decrypt should copy the data
// let decrypted_len = session.decrypt_packet(&encrypted, &mut decrypted).unwrap(); // Removed method
// assert_eq!(decrypted_len, 5);
// assert_eq!(decrypted, data);
}
#[test]
fn test_mock_crypto_buffer_too_small() {
let session = create_test_session(true);
let data = [1, 2, 3, 4, 5];
let mut too_small = [0; 3];
// Should fail with buffer too small
// let result = session.encrypt_packet(&data, &mut too_small); // Removed method
// assert!(result.is_err());
// match result.unwrap_err() {
// LpError::InsufficientBufferSize => {} // Error type might change
// _ => panic!("Expected InsufficientBufferSize error"),
// }
}
*/
}
File diff suppressed because it is too large Load Diff
+296
View File
@@ -0,0 +1,296 @@
// Copyright 2025 - Nym Technologies SA <contact@nymtech.net>
// SPDX-License-Identifier: Apache-2.0
//! Session management for the Lewes Protocol.
//!
//! This module implements session lifecycle management functionality, handling
//! creation, retrieval, and storage of sessions.
use dashmap::DashMap;
use crate::keypair::{Keypair, PublicKey};
use crate::noise_protocol::ReadResult;
use crate::state_machine::{LpAction, LpInput, LpState, LpStateBare};
use crate::{LpError, LpMessage, LpSession, LpStateMachine};
/// Manages the lifecycle of Lewes Protocol sessions.
///
/// The SessionManager is responsible for creating, storing, and retrieving sessions,
/// ensuring proper thread-safety for concurrent access.
pub struct SessionManager {
/// Manages state machines directly, keyed by lp_id
state_machines: DashMap<u32, LpStateMachine>,
}
impl Default for SessionManager {
fn default() -> Self {
Self::new()
}
}
impl SessionManager {
/// Creates a new session manager with empty session storage.
pub fn new() -> Self {
Self {
state_machines: DashMap::new(),
}
}
pub fn process_input(&self, lp_id: u32, input: LpInput) -> Result<Option<LpAction>, LpError> {
self.with_state_machine_mut(lp_id, |sm| sm.process_input(input).transpose())?
}
pub fn add(&self, session: LpSession) -> Result<(), LpError> {
let sm = LpStateMachine {
state: LpState::ReadyToHandshake { session },
};
self.state_machines.insert(sm.id()?, sm);
Ok(())
}
pub fn handshaking(&self, lp_id: u32) -> Result<bool, LpError> {
Ok(self.get_state(lp_id)? == LpStateBare::Handshaking)
}
pub fn should_initiate_handshake(&self, lp_id: u32) -> Result<bool, LpError> {
Ok(self.ready_to_handshake(lp_id)? || self.closed(lp_id)?)
}
pub fn ready_to_handshake(&self, lp_id: u32) -> Result<bool, LpError> {
Ok(self.get_state(lp_id)? == LpStateBare::ReadyToHandshake)
}
pub fn closed(&self, lp_id: u32) -> Result<bool, LpError> {
Ok(self.get_state(lp_id)? == LpStateBare::Closed)
}
pub fn transport(&self, lp_id: u32) -> Result<bool, LpError> {
Ok(self.get_state(lp_id)? == LpStateBare::Transport)
}
#[cfg(test)]
fn get_state_machine_id(&self, lp_id: u32) -> Result<u32, LpError> {
self.with_state_machine(lp_id, |sm| sm.id())?
}
pub fn get_state(&self, lp_id: u32) -> Result<LpStateBare, LpError> {
self.with_state_machine(lp_id, |sm| Ok(sm.bare_state()))?
}
pub fn receiving_counter_quick_check(&self, lp_id: u32, counter: u64) -> Result<(), LpError> {
self.with_state_machine(lp_id, |sm| {
sm.session()?.receiving_counter_quick_check(counter)
})?
}
pub fn receiving_counter_mark(&self, lp_id: u32, counter: u64) -> Result<(), LpError> {
self.with_state_machine(lp_id, |sm| sm.session()?.receiving_counter_mark(counter))?
}
pub fn start_handshake(&self, lp_id: u32) -> Option<Result<LpMessage, LpError>> {
self.prepare_handshake_message(lp_id)
}
pub fn prepare_handshake_message(&self, lp_id: u32) -> Option<Result<LpMessage, LpError>> {
self.with_state_machine(lp_id, |sm| sm.session().ok()?.prepare_handshake_message())
.ok()?
}
pub fn is_handshake_complete(&self, lp_id: u32) -> Result<bool, LpError> {
self.with_state_machine(lp_id, |sm| Ok(sm.session()?.is_handshake_complete()))?
}
pub fn next_counter(&self, lp_id: u32) -> Result<u64, LpError> {
self.with_state_machine(lp_id, |sm| Ok(sm.session()?.next_counter()))?
}
pub fn decrypt_data(&self, lp_id: u32, message: &LpMessage) -> Result<Vec<u8>, LpError> {
self.with_state_machine(lp_id, |sm| {
sm.session()?
.decrypt_data(message)
.map_err(LpError::NoiseError)
})?
}
pub fn encrypt_data(&self, lp_id: u32, message: &[u8]) -> Result<LpMessage, LpError> {
self.with_state_machine(lp_id, |sm| {
sm.session()?
.encrypt_data(message)
.map_err(LpError::NoiseError)
})?
}
pub fn current_packet_cnt(&self, lp_id: u32) -> Result<(u64, u64), LpError> {
self.with_state_machine(lp_id, |sm| Ok(sm.session()?.current_packet_cnt()))?
}
pub fn process_handshake_message(
&self,
lp_id: u32,
message: &LpMessage,
) -> Result<ReadResult, LpError> {
self.with_state_machine(lp_id, |sm| {
Ok(sm.session()?.process_handshake_message(message)?)
})?
}
pub fn session_count(&self) -> usize {
self.state_machines.len()
}
pub fn state_machine_exists(&self, lp_id: u32) -> bool {
self.state_machines.contains_key(&lp_id)
}
pub fn with_state_machine<F, R>(&self, lp_id: u32, f: F) -> Result<R, LpError>
where
F: FnOnce(&LpStateMachine) -> R,
{
if let Some(sm) = self.state_machines.get(&lp_id) {
Ok(f(&sm))
} else {
Err(LpError::StateMachineNotFound { lp_id })
}
// self.state_machines.get(&lp_id).map(|sm_ref| f(&*sm_ref)) // Lock held only during closure execution
}
// For mutable access (like running process_input)
pub fn with_state_machine_mut<F, R>(&self, lp_id: u32, f: F) -> Result<R, LpError>
where
F: FnOnce(&mut LpStateMachine) -> R, // Closure takes mutable ref
{
if let Some(mut sm) = self.state_machines.get_mut(&lp_id) {
Ok(f(&mut sm))
} else {
Err(LpError::StateMachineNotFound { lp_id })
}
}
pub fn create_session_state_machine(
&self,
local_keypair: &Keypair,
remote_public_key: &PublicKey,
is_initiator: bool,
psk: &[u8],
) -> Result<u32, LpError> {
let sm = LpStateMachine::new(is_initiator, local_keypair, remote_public_key, psk)?;
let sm_id = sm.id()?;
self.state_machines.insert(sm_id, sm);
Ok(sm_id)
}
/// Method to remove a state machine
pub fn remove_state_machine(&self, lp_id: u32) -> bool {
let removed = self.state_machines.remove(&lp_id);
removed.is_some()
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_session_manager_get() {
let manager = SessionManager::new();
let sm_1_id = manager
.create_session_state_machine(
&Keypair::default(),
&PublicKey::default(),
true,
&[2u8; 32],
)
.unwrap();
let retrieved = manager.state_machine_exists(sm_1_id);
assert!(retrieved);
let not_found = manager.state_machine_exists(99);
assert!(!not_found);
}
#[test]
fn test_session_manager_remove() {
let manager = SessionManager::new();
let sm_1_id = manager
.create_session_state_machine(
&Keypair::default(),
&PublicKey::default(),
true,
&[2u8; 32],
)
.unwrap();
let removed = manager.remove_state_machine(sm_1_id);
assert!(removed);
assert_eq!(manager.session_count(), 0);
let removed_again = manager.remove_state_machine(sm_1_id);
assert!(!removed_again);
}
#[test]
fn test_multiple_sessions() {
let manager = SessionManager::new();
let sm_1 = manager
.create_session_state_machine(
&Keypair::default(),
&PublicKey::default(),
true,
&[2u8; 32],
)
.unwrap();
let sm_2 = manager
.create_session_state_machine(
&Keypair::default(),
&PublicKey::default(),
true,
&[2u8; 32],
)
.unwrap();
let sm_3 = manager
.create_session_state_machine(
&Keypair::default(),
&PublicKey::default(),
true,
&[2u8; 32],
)
.unwrap();
assert_eq!(manager.session_count(), 3);
let retrieved1 = manager.get_state_machine_id(sm_1).unwrap();
let retrieved2 = manager.get_state_machine_id(sm_2).unwrap();
let retrieved3 = manager.get_state_machine_id(sm_3).unwrap();
assert_eq!(retrieved1, sm_1);
assert_eq!(retrieved2, sm_2);
assert_eq!(retrieved3, sm_3);
}
#[test]
fn test_session_manager_create_session() {
let manager = SessionManager::new();
let sm = manager.create_session_state_machine(
&Keypair::default(),
&PublicKey::default(),
true,
&[2u8; 32],
);
assert!(sm.is_ok());
let sm = sm.unwrap();
assert_eq!(manager.session_count(), 1);
let retrieved = manager.get_state_machine_id(sm);
assert!(retrieved.is_ok());
assert_eq!(retrieved.unwrap(), sm);
}
}
+649
View File
@@ -0,0 +1,649 @@
// Copyright 2025 - Nym Technologies SA <contact@nymtech.net>
// SPDX-License-Identifier: Apache-2.0
//! Lewes Protocol State Machine for managing connection lifecycle.
use crate::{
keypair::{Keypair, PublicKey},
make_lp_id,
noise_protocol::NoiseError,
packet::LpPacket,
session::LpSession,
LpError,
};
use bytes::BytesMut;
use std::mem;
/// Represents the possible states of the Lewes Protocol connection.
#[derive(Debug, Default)]
pub enum LpState {
/// Initial state: Ready to start the handshake.
/// State machine is created with keys, lp_id is derived, session is ready.
ReadyToHandshake { session: LpSession },
/// Actively performing the Noise handshake.
/// (We might be able to merge this with ReadyToHandshake if the first step always happens)
Handshaking { session: LpSession }, // Kept for now, logic might merge later
/// Handshake complete, ready for data transport.
Transport { session: LpSession },
/// An error occurred, or the connection was intentionally closed.
Closed { reason: String },
/// Processing an input event.
#[default]
Processing,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum LpStateBare {
ReadyToHandshake,
Handshaking,
Transport,
Closed,
Processing,
}
impl From<&LpState> for LpStateBare {
fn from(state: &LpState) -> Self {
match state {
LpState::ReadyToHandshake { .. } => LpStateBare::ReadyToHandshake,
LpState::Handshaking { .. } => LpStateBare::Handshaking,
LpState::Transport { .. } => LpStateBare::Transport,
LpState::Closed { .. } => LpStateBare::Closed,
LpState::Processing => LpStateBare::Processing,
}
}
}
/// Represents inputs that drive the state machine transitions.
#[derive(Debug)]
pub enum LpInput {
/// Explicitly trigger the start of the handshake (optional, could be implicit on creation)
StartHandshake,
/// Received an LP Packet from the network.
ReceivePacket(LpPacket),
/// Application wants to send data (only valid in Transport state).
SendData(Vec<u8>), // Using Bytes for efficiency
/// Close the connection.
Close,
}
/// Represents actions the state machine requests the environment to perform.
#[derive(Debug)]
pub enum LpAction {
/// Send an LP Packet over the network.
SendPacket(LpPacket),
/// Deliver decrypted application data received from the peer.
DeliverData(BytesMut),
/// Inform the environment that the handshake is complete.
HandshakeComplete,
/// Inform the environment that the connection is closed.
ConnectionClosed,
}
/// The Lewes Protocol State Machine.
pub struct LpStateMachine {
pub state: LpState,
}
impl LpStateMachine {
pub fn bare_state(&self) -> LpStateBare {
LpStateBare::from(&self.state)
}
pub fn session(&self) -> Result<&LpSession, LpError> {
match &self.state {
LpState::ReadyToHandshake { session }
| LpState::Handshaking { session }
| LpState::Transport { session } => Ok(session),
LpState::Closed { .. } => Err(LpError::LpSessionClosed),
LpState::Processing => Err(LpError::LpSessionProcessing),
}
}
/// Consume the state machine and return the session with ownership.
/// This is useful when the handshake is complete and you want to transfer
/// ownership of the session to the caller.
pub fn into_session(self) -> Result<LpSession, LpError> {
match self.state {
LpState::ReadyToHandshake { session }
| LpState::Handshaking { session }
| LpState::Transport { session } => Ok(session),
LpState::Closed { .. } => Err(LpError::LpSessionClosed),
LpState::Processing => Err(LpError::LpSessionProcessing),
}
}
pub fn id(&self) -> Result<u32, LpError> {
Ok(self.session()?.id())
}
/// Creates a new state machine, calculates the lp_id, creates the session,
/// and sets the initial state to ReadyToHandshake.
///
/// Requires the local *full* keypair to get the public key for lp_id calculation.
pub fn new(
is_initiator: bool,
local_keypair: &Keypair, // Use Keypair
remote_public_key: &PublicKey,
psk: &[u8],
// session_manager: Arc<SessionManager> // Optional
) -> Result<Self, LpError> {
// Calculate the shared lp_id// Calculate the shared lp_id
let lp_id = make_lp_id(local_keypair.public_key(), remote_public_key);
let local_private_key = local_keypair.private_key().to_bytes();
let remote_public_key = remote_public_key.as_bytes();
// Create the session immediately
let session = LpSession::new(
lp_id,
is_initiator,
&local_private_key,
remote_public_key,
psk,
)?;
// TODO: Register the session with the SessionManager if applicable
// if let Some(manager) = session_manager {
// manager.insert_session(lp_id, session.clone())?; // Assuming insert_session exists
// }
Ok(LpStateMachine {
state: LpState::ReadyToHandshake { session },
// Store necessary info if needed for recreation, otherwise remove
// is_initiator,
// local_private_key: local_private_key.to_vec(),
// remote_public_key: remote_public_key.to_vec(),
// psk: psk.to_vec(),
})
}
/// Processes an input event and returns a list of actions to perform.
pub fn process_input(&mut self, input: LpInput) -> Option<Result<LpAction, LpError>> {
// 1. Replace current state with a placeholder, taking ownership of the real current state.
let current_state = mem::take(&mut self.state);
let mut result_action: Option<Result<LpAction, LpError>> = None;
// 2. Match on the owned current_state. Each arm calculates and returns the NEXT state.
let next_state = match (current_state, input) {
// --- ReadyToHandshake State ---
(LpState::ReadyToHandshake { session }, LpInput::StartHandshake) => {
if session.is_initiator() {
// Initiator sends the first message
match self.start_handshake(&session) {
Some(Ok(action)) => {
result_action = Some(Ok(action));
LpState::Handshaking { session } // Transition state
}
Some(Err(e)) => {
// Error occurred, move to Closed state
let reason = e.to_string();
result_action = Some(Err(e));
LpState::Closed { reason }
}
None => {
// Should not happen, treat as internal error
let err = LpError::Internal(
"start_handshake returned None unexpectedly".to_string(),
);
let reason = err.to_string();
result_action = Some(Err(err));
LpState::Closed { reason }
}
}
} else {
// Responder waits for the first message, transition to Handshaking to wait.
LpState::Handshaking { session }
// No action needed yet, result_action remains None.
}
}
// --- Handshaking State ---
(LpState::Handshaking { session }, LpInput::ReceivePacket(packet)) => {
// Check if packet lp_id matches our session
if packet.header.session_id() != session.id() {
result_action = Some(Err(LpError::UnknownSessionId(packet.header.session_id())));
// Don't change state, return the original state variant
LpState::Handshaking { session }
} else {
// --- Inline handle_handshake_packet logic ---
// 1. Check replay protection *before* processing
if let Err(e) = session.receiving_counter_quick_check(packet.header.counter) {
let _reason = e.to_string();
result_action = Some(Err(e));
LpState::Handshaking { session }
// LpState::Closed { reason }
} else {
// 2. Process the handshake message
match session.process_handshake_message(&packet.message) {
Ok(_) => {
// 3. Mark counter as received *after* successful processing
if let Err(e) = session.receiving_counter_mark(packet.header.counter) {
let _reason = e.to_string();
result_action = Some(Err(e));
// LpState::Closed { reason }
LpState::Handshaking { session }
} else {
// 4. Check if handshake is now complete
if session.is_handshake_complete() {
result_action = Some(Ok(LpAction::HandshakeComplete));
LpState::Transport { session } // Transition to Transport
} else {
// 5. Check if we need to send the next handshake message
match session.prepare_handshake_message() {
Some(Ok(message)) => {
match session.next_packet(message) {
Ok(response_packet) => {
result_action = Some(Ok(LpAction::SendPacket(response_packet)));
// Check AGAIN if handshake became complete *after preparing*
if session.is_handshake_complete() {
LpState::Transport { session } // Transition to Transport
} else {
LpState::Handshaking { session } // Remain Handshaking
}
}
Err(e) => {
let reason = e.to_string();
result_action = Some(Err(e));
LpState::Closed { reason }
}
}
}
Some(Err(e)) => {
let reason = e.to_string();
result_action = Some(Err(e));
LpState::Closed { reason }
}
None => {
// Handshake stalled unexpectedly
let err = LpError::NoiseError(NoiseError::Other(
"Handshake stalled unexpectedly".to_string(),
));
let reason = err.to_string();
result_action = Some(Err(err));
LpState::Closed { reason }
}
}
}
}
}
Err(e) => { // Error from process_handshake_message
let reason = e.to_string();
result_action = Some(Err(e.into()));
LpState::Closed { reason }
}
}
}
// --- End inline handle_handshake_packet logic ---
}
}
// Reject SendData during handshake
(LpState::Handshaking { session }, LpInput::SendData(_)) => { // Keep session if returning to this state
result_action = Some(Err(LpError::InvalidStateTransition {
state: "Handshaking".to_string(),
input: "SendData".to_string(),
}));
// Invalid input, remain in Handshaking state
LpState::Handshaking { session }
}
// Reject StartHandshake if already handshaking
(LpState::Handshaking { session }, LpInput::StartHandshake) => { // Keep session
result_action = Some(Err(LpError::InvalidStateTransition {
state: "Handshaking".to_string(),
input: "StartHandshake".to_string(),
}));
// Invalid input, remain in Handshaking state
LpState::Handshaking { session }
}
// --- Transport State ---
(LpState::Transport { session }, LpInput::ReceivePacket(packet)) => { // Needs mut session for marking counter
// Check if packet lp_id matches our session
if packet.header.session_id() != session.id() {
result_action = Some(Err(LpError::UnknownSessionId(packet.header.session_id())));
// Remain in transport state
LpState::Transport { session }
} else {
// --- Inline handle_data_packet logic ---
// 1. Check replay protection
if let Err(e) = session.receiving_counter_quick_check(packet.header.counter) {
let _reason = e.to_string();
result_action = Some(Err(e));
LpState::Transport { session }
} else {
// 2. Decrypt data
match session.decrypt_data(&packet.message) {
Ok(plaintext) => {
// 3. Mark counter as received
if let Err(e) = session.receiving_counter_mark(packet.header.counter) {
let _reason = e.to_string();
result_action = Some(Err(e));
LpState::Transport{ session }
} else {
// 4. Deliver data
result_action = Some(Ok(LpAction::DeliverData(BytesMut::from(plaintext.as_slice()))));
// Remain in transport state
LpState::Transport { session }
}
}
Err(e) => { // Error decrypting data
let reason = e.to_string();
result_action = Some(Err(e.into()));
LpState::Closed { reason }
}
}
}
// --- End inline handle_data_packet logic ---
}
}
(LpState::Transport { session }, LpInput::SendData(data)) => {
// Encrypt and send application data
match self.prepare_data_packet(&session, &data) {
Ok(packet) => result_action = Some(Ok(LpAction::SendPacket(packet))),
Err(e) => {
// If prepare fails, should we close? Let's report error and stay Transport for now.
// Alternative: transition to Closed state.
result_action = Some(Err(e.into()));
}
}
// Remain in transport state
LpState::Transport { session }
}
// Reject StartHandshake if already in transport
(LpState::Transport { session }, LpInput::StartHandshake) => { // Keep session
result_action = Some(Err(LpError::InvalidStateTransition {
state: "Transport".to_string(),
input: "StartHandshake".to_string(),
}));
// Invalid input, remain in Transport state
LpState::Transport { session }
}
// --- Close Transition (applies to ReadyToHandshake, Handshaking, Transport) ---
(
LpState::ReadyToHandshake { .. } // We consume the session here
| LpState::Handshaking { .. }
| LpState::Transport { .. },
LpInput::Close,
) => {
result_action = Some(Ok(LpAction::ConnectionClosed));
// Transition to Closed state
LpState::Closed { reason: "Closed by user".to_string() }
}
// Ignore Close if already Closed
(closed_state @ LpState::Closed { .. }, LpInput::Close) => {
// result_action remains None
// Return the original closed state
closed_state
}
// Ignore StartHandshake if Closed
// (closed_state @ LpState::Closed { .. }, LpInput::StartHandshake) => {
// result_action = Some(Err(LpError::LpSessionClosed));
// closed_state
// }
// Ignore ReceivePacket if Closed
(closed_state @ LpState::Closed { .. }, LpInput::ReceivePacket(_)) => {
result_action = Some(Err(LpError::LpSessionClosed));
closed_state
}
// Ignore SendData if Closed
(closed_state @ LpState::Closed { .. }, LpInput::SendData(_)) => {
result_action = Some(Err(LpError::LpSessionClosed));
closed_state
}
// Processing state should not be matched directly if using replace
(LpState::Processing, _) => {
// This case should ideally be unreachable if placeholder logic is correct
let err = LpError::Internal("Reached Processing state unexpectedly".to_string());
let reason = err.to_string();
result_action = Some(Err(err));
LpState::Closed { reason }
}
// --- Default: Invalid input for current state (if any combinations missed) ---
// Consider if this should transition to Closed state. For now, just report error
// and transition to Closed as a safety measure.
(invalid_state, input) => {
let err = LpError::InvalidStateTransition {
state: format!("{:?}", invalid_state), // Use owned state for debug info
input: format!("{:?}", input),
};
let reason = err.to_string();
result_action = Some(Err(err));
LpState::Closed { reason }
}
};
// 3. Put the calculated next state back into the machine.
self.state = next_state;
result_action // Return the determined action (or None)
}
// Helper to start the handshake (sends first message if initiator)
// Kept as it doesn't mutate self.state
fn start_handshake(&self, session: &LpSession) -> Option<Result<LpAction, LpError>> {
session
.prepare_handshake_message()
.map(|result| match result {
Ok(message) => match session.next_packet(message) {
Ok(packet) => Ok(LpAction::SendPacket(packet)),
Err(e) => Err(e),
},
Err(e) => Err(e),
})
}
// Helper to prepare an outgoing data packet
// Kept as it doesn't mutate self.state
fn prepare_data_packet(
&self,
session: &LpSession,
data: &[u8],
) -> Result<LpPacket, NoiseError> {
let encrypted_message = session.encrypt_data(data)?;
session
.next_packet(encrypted_message)
.map_err(|e| NoiseError::Other(e.to_string())) // Improve error conversion?
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::keypair::Keypair;
use bytes::Bytes;
#[test]
fn test_state_machine_init() {
let init_key = Keypair::new();
let resp_key = Keypair::new();
let psk = vec![0u8; 32];
let remote_pub_key = resp_key.public_key();
let initiator_sm = LpStateMachine::new(true, &init_key, &remote_pub_key, &psk);
assert!(initiator_sm.is_ok());
let initiator_sm = initiator_sm.unwrap();
assert!(matches!(
initiator_sm.state,
LpState::ReadyToHandshake { .. }
));
let init_session = initiator_sm.session().unwrap();
assert!(init_session.is_initiator());
let responder_sm = LpStateMachine::new(false, &resp_key, &init_key.public_key(), &psk);
assert!(responder_sm.is_ok());
let responder_sm = responder_sm.unwrap();
assert!(matches!(
responder_sm.state,
LpState::ReadyToHandshake { .. }
));
let resp_session = responder_sm.session().unwrap();
assert!(!resp_session.is_initiator());
// Check lp_id is the same
let expected_lp_id = make_lp_id(&init_key.public_key(), remote_pub_key);
assert_eq!(init_session.id(), expected_lp_id);
assert_eq!(resp_session.id(), expected_lp_id);
}
#[test]
fn test_state_machine_simplified_flow() {
// Create test keys
let init_key = Keypair::new();
let resp_key = Keypair::new();
let psk = vec![0u8; 32];
// Create state machines (already in ReadyToHandshake)
let mut initiator = LpStateMachine::new(
true, // is_initiator
&init_key,
&resp_key.public_key(),
&psk.clone(),
)
.unwrap();
let mut responder = LpStateMachine::new(
false, // is_initiator
&resp_key,
&init_key.public_key(),
&psk,
)
.unwrap();
let lp_id = initiator.id().unwrap();
assert_eq!(lp_id, responder.id().unwrap());
// --- Start Handshake --- (No index exchange needed)
println!("--- Step 1: Initiator starts handshake ---");
let init_actions_1 = initiator.process_input(LpInput::StartHandshake);
let init_packet_1 = if let Some(Ok(LpAction::SendPacket(packet))) = init_actions_1 {
packet.clone()
} else {
panic!("Initiator should produce 1 action");
};
assert!(
matches!(initiator.state, LpState::Handshaking { .. }),
"Initiator should be Handshaking"
);
assert_eq!(
init_packet_1.header.session_id(),
lp_id,
"Packet 1 has wrong lp_id"
);
println!("--- Step 2: Responder starts handshake (waits) ---");
let resp_actions_1 = responder.process_input(LpInput::StartHandshake);
assert!(
resp_actions_1.is_none(),
"Responder should produce 0 actions initially"
);
assert!(
matches!(responder.state, LpState::Handshaking { .. }),
"Responder should be Handshaking"
);
// --- Handshake Message Exchange ---
println!("--- Step 3: Responder receives packet 1, sends packet 2 ---");
let resp_actions_2 = responder.process_input(LpInput::ReceivePacket(init_packet_1));
let resp_packet_2 = if let Some(Ok(LpAction::SendPacket(packet))) = resp_actions_2 {
packet.clone()
} else {
panic!("Responder should send packet 2");
};
assert!(
matches!(responder.state, LpState::Handshaking { .. }),
"Responder still Handshaking"
);
assert_eq!(
resp_packet_2.header.session_id(),
lp_id,
"Packet 2 has wrong lp_id"
);
println!("--- Step 4: Initiator receives packet 2, sends packet 3 ---");
let init_actions_2 = initiator.process_input(LpInput::ReceivePacket(resp_packet_2));
let init_packet_3 = if let Some(Ok(LpAction::SendPacket(packet))) = init_actions_2 {
packet.clone()
} else {
panic!("Initiator should send packet 3");
};
assert!(
matches!(initiator.state, LpState::Transport { .. }),
"Initiator should be Transport"
);
assert_eq!(
init_packet_3.header.session_id(),
lp_id,
"Packet 3 has wrong lp_id"
);
println!("--- Step 5: Responder receives packet 3, completes handshake ---");
let resp_actions_3 = responder.process_input(LpInput::ReceivePacket(init_packet_3));
assert!(
matches!(resp_actions_3, Some(Ok(LpAction::HandshakeComplete))),
"Responder should complete handshake"
);
assert!(
matches!(responder.state, LpState::Transport { .. }),
"Responder should be Transport"
);
// --- Transport Phase ---
println!("--- Step 6: Initiator sends data ---");
let data_to_send_1 = b"hello responder";
let init_actions_3 = initiator.process_input(LpInput::SendData(data_to_send_1.to_vec()));
let data_packet_1 = if let Some(Ok(LpAction::SendPacket(packet))) = init_actions_3 {
packet.clone()
} else {
panic!("Initiator should send data packet");
};
assert_eq!(data_packet_1.header.session_id(), lp_id);
println!("--- Step 7: Responder receives data ---");
let resp_actions_4 = responder.process_input(LpInput::ReceivePacket(data_packet_1));
let resp_data_1 = if let Some(Ok(LpAction::DeliverData(data))) = resp_actions_4 {
data
} else {
panic!("Responder should deliver data");
};
assert_eq!(resp_data_1, Bytes::copy_from_slice(data_to_send_1));
println!("--- Step 8: Responder sends data ---");
let data_to_send_2 = b"hello initiator";
let resp_actions_5 = responder.process_input(LpInput::SendData(data_to_send_2.to_vec()));
let data_packet_2 = if let Some(Ok(LpAction::SendPacket(packet))) = resp_actions_5 {
packet.clone()
} else {
panic!("Responder should send data packet");
};
assert_eq!(data_packet_2.header.session_id(), lp_id);
println!("--- Step 9: Initiator receives data ---");
let init_actions_4 = initiator.process_input(LpInput::ReceivePacket(data_packet_2));
if let Some(Ok(LpAction::DeliverData(data))) = init_actions_4 {
assert_eq!(data, Bytes::copy_from_slice(data_to_send_2));
} else {
panic!("Initiator should deliver data");
}
// --- Close ---
println!("--- Step 10: Initiator closes ---");
let init_actions_5 = initiator.process_input(LpInput::Close);
assert!(matches!(
init_actions_5,
Some(Ok(LpAction::ConnectionClosed))
));
assert!(matches!(initiator.state, LpState::Closed { .. }));
println!("--- Step 11: Responder closes ---");
let resp_actions_6 = responder.process_input(LpInput::Close);
assert!(matches!(
resp_actions_6,
Some(Ok(LpAction::ConnectionClosed))
));
assert!(matches!(responder.state, LpState::Closed { .. }));
}
}
+1
View File
@@ -12,6 +12,7 @@ license.workspace = true
workspace = true
[dependencies]
serde = { workspace = true, features = ["derive"] }
tokio-util.workspace = true
nym-authenticator-requests = { path = "../authenticator-requests" }
+2 -1
View File
@@ -7,6 +7,7 @@ use nym_authenticator_requests::AuthenticatorVersion;
use nym_crypto::asymmetric::x25519::PublicKey;
use nym_ip_packet_requests::IpPair;
use nym_sphinx::addressing::{NodeIdentity, Recipient};
use serde::{Deserialize, Serialize};
#[derive(Debug, Clone, Copy, PartialEq)]
pub struct NymNode {
@@ -17,7 +18,7 @@ pub struct NymNode {
pub version: AuthenticatorVersion,
}
#[derive(Clone, Debug)]
#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct GatewayData {
pub public_key: PublicKey,
pub endpoint: SocketAddr,
+1 -1
View File
@@ -9,7 +9,6 @@
use defguard_wireguard_rs::{WGApi, WireguardInterfaceApi, host::Peer, key::Key, net::IpAddrMask};
use nym_crypto::asymmetric::x25519::KeyPair;
use nym_wireguard_types::Config;
use peer_controller::PeerControlRequest;
use std::sync::Arc;
use tokio::sync::mpsc::{self, Receiver, Sender};
use tracing::error;
@@ -26,6 +25,7 @@ pub mod peer_handle;
pub mod peer_storage_manager;
pub use error::Error;
pub use peer_controller::PeerControlRequest;
pub const CONTROL_CHANNEL_SIZE: usize = 256;
+6
View File
@@ -75,6 +75,12 @@ nym-client-core = { path = "../common/client-core", features = ["cli"] }
nym-id = { path = "../common/nym-id" }
nym-service-provider-requests-common = { path = "../common/service-provider-requests-common" }
# LP dependencies
nym-lp = { path = "../common/nym-lp" }
nym-kcp = { path = "../common/nym-kcp" }
nym-registration-common = { path = "../common/registration" }
bytes = { workspace = true }
defguard_wireguard_rs = { workspace = true }
[dev-dependencies]
+4
View File
@@ -15,6 +15,8 @@ pub struct Config {
pub upgrade_mode_watcher: UpgradeModeWatcher,
pub lp: crate::node::lp_listener::LpConfig,
pub debug: Debug,
}
@@ -24,6 +26,7 @@ impl Config {
network_requester: impl Into<NetworkRequester>,
ip_packet_router: impl Into<IpPacketRouter>,
upgrade_mode_watcher: impl Into<UpgradeModeWatcher>,
lp: impl Into<crate::node::lp_listener::LpConfig>,
debug: impl Into<Debug>,
) -> Self {
Config {
@@ -31,6 +34,7 @@ impl Config {
network_requester: network_requester.into(),
ip_packet_router: ip_packet_router.into(),
upgrade_mode_watcher: upgrade_mode_watcher.into(),
lp: lp.into(),
debug: debug.into(),
}
}
+21
View File
@@ -125,6 +125,27 @@ pub enum GatewayError {
#[error("{0}")]
CredentialVefiricationError(#[from] nym_credential_verification::Error),
#[error("LP connection error: {0}")]
LpConnectionError(String),
#[error("LP protocol error: {0}")]
LpProtocolError(String),
#[error("LP handshake error: {0}")]
LpHandshakeError(String),
#[error("Service provider {service} is not running")]
ServiceProviderNotRunning { service: String },
#[error("Internal error: {0}")]
InternalError(String),
#[error("Failed to bind listener to {address}: {source}")]
ListenerBindFailure {
address: String,
source: Box<dyn std::error::Error + Send + Sync>,
},
}
impl From<ClientCoreError> for GatewayError {
+266
View File
@@ -0,0 +1,266 @@
// Copyright 2025 - Nym Technologies SA <contact@nymtech.net>
// SPDX-License-Identifier: GPL-3.0-only
use super::handshake::LpGatewayHandshake;
use super::messages::{LpRegistrationRequest, LpRegistrationResponse};
use super::registration::process_registration;
use super::LpHandlerState;
use crate::error::GatewayError;
use nym_lp::{
keypair::{Keypair, PublicKey},
LpMessage, LpPacket, LpSession,
};
use std::net::SocketAddr;
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tokio::net::TcpStream;
use tracing::*;
pub struct LpConnectionHandler {
stream: TcpStream,
remote_addr: SocketAddr,
state: LpHandlerState,
}
impl LpConnectionHandler {
pub fn new(stream: TcpStream, remote_addr: SocketAddr, state: LpHandlerState) -> Self {
Self {
stream,
remote_addr,
state,
}
}
pub async fn handle(mut self) -> Result<(), GatewayError> {
debug!("Handling LP connection from {}", self.remote_addr);
// For LP, we need:
// 1. Gateway's keypair (from local_identity)
// 2. Client's public key (will be received during handshake)
// 3. PSK (pre-shared key) - for now use a placeholder
// Generate fresh LP keypair (x25519) for this connection
// Using Keypair::default() which generates a new random x25519 keypair
// This is secure and simple - each connection gets its own keypair
let gateway_keypair = Keypair::default();
// Receive client's public key via ClientHello message
// The client initiates by sending ClientHello as first packet
let client_pubkey = self.receive_client_hello().await?;
// Generate or retrieve PSK for this session
// TODO(nym-16): Implement proper PSK management
// Temporary solution: use gateway's identity public key as PSK
let psk = self.state.local_identity.public_key().to_bytes();
// Create LP handshake as responder
let handshake = LpGatewayHandshake::new_responder(
&gateway_keypair,
&client_pubkey,
&psk,
)?;
// Complete the LP handshake
let session = handshake.complete(&mut self.stream).await?;
info!("LP handshake completed for {} (session {})",
self.remote_addr, session.id());
// After handshake, receive registration request
let request = self.receive_registration_request(&session).await?;
debug!("LP registration request from {}: mode={:?}",
self.remote_addr, request.mode);
// Process registration (verify credentials, add peer, etc.)
let response = process_registration(request, &self.state).await;
// Send response
if let Err(e) = self.send_registration_response(&session, response.clone()).await {
warn!("Failed to send LP response to {}: {}", self.remote_addr, e);
return Err(e);
}
if response.success {
info!("LP registration successful for {} (session {})",
self.remote_addr, response.session_id);
} else {
warn!("LP registration failed for {}: {:?}",
self.remote_addr, response.error);
}
Ok(())
}
/// Receive client's public key via ClientHello message
async fn receive_client_hello(&mut self) -> Result<PublicKey, GatewayError> {
// Receive first packet which should be ClientHello
let packet = self.receive_lp_packet().await?;
// Verify it's a ClientHello message
match packet.message() {
LpMessage::ClientHello(hello_data) => {
// Validate protocol version (currently only v1)
if hello_data.protocol_version != 1 {
return Err(GatewayError::LpProtocolError(
format!("Unsupported protocol version: {}", hello_data.protocol_version)
));
}
// Convert bytes to PublicKey
PublicKey::from_bytes(&hello_data.client_lp_public_key)
.map_err(|e| GatewayError::LpProtocolError(
format!("Invalid client public key: {}", e)
))
}
other => {
Err(GatewayError::LpProtocolError(
format!("Expected ClientHello, got {}", other)
))
}
}
}
/// Receive registration request after handshake
async fn receive_registration_request(
&mut self,
session: &LpSession,
) -> Result<LpRegistrationRequest, GatewayError> {
// Read LP packet containing the registration request
let packet = self.receive_lp_packet().await?;
// Verify it's from the correct session
if packet.header().session_id != session.id() {
return Err(GatewayError::LpProtocolError(
format!("Session ID mismatch: expected {}, got {}",
session.id(), packet.header().session_id)
));
}
// Extract registration request from LP message
match packet.message() {
LpMessage::EncryptedData(data) => {
// Deserialize registration request
bincode::deserialize(&data)
.map_err(|e| GatewayError::LpProtocolError(
format!("Failed to deserialize registration request: {}", e)
))
}
other => {
Err(GatewayError::LpProtocolError(
format!("Expected EncryptedData message, got {:?}", other)
))
}
}
}
/// Send registration response after processing
async fn send_registration_response(
&mut self,
session: &LpSession,
response: LpRegistrationResponse,
) -> Result<(), GatewayError> {
// Serialize response
let data = bincode::serialize(&response)
.map_err(|e| GatewayError::LpProtocolError(
format!("Failed to serialize response: {}", e)
))?;
// Create LP packet with response
let packet = session.create_data_packet(data)
.map_err(|e| GatewayError::LpProtocolError(
format!("Failed to create data packet: {}", e)
))?;
// Send the packet
self.send_lp_packet(&packet).await
}
/// Receive an LP packet from the stream with proper length-prefixed framing
async fn receive_lp_packet(&mut self) -> Result<LpPacket, GatewayError> {
use nym_lp::codec::parse_lp_packet;
// Read 4-byte length prefix (u32 big-endian)
let mut len_buf = [0u8; 4];
self.stream.read_exact(&mut len_buf).await
.map_err(|e| GatewayError::LpConnectionError(
format!("Failed to read packet length: {}", e)
))?;
let packet_len = u32::from_be_bytes(len_buf) as usize;
// Sanity check to prevent huge allocations
const MAX_PACKET_SIZE: usize = 65536; // 64KB max
if packet_len > MAX_PACKET_SIZE {
return Err(GatewayError::LpProtocolError(
format!("Packet size {} exceeds maximum {}", packet_len, MAX_PACKET_SIZE)
));
}
// Read the actual packet data
let mut packet_buf = vec![0u8; packet_len];
self.stream.read_exact(&mut packet_buf).await
.map_err(|e| GatewayError::LpConnectionError(
format!("Failed to read packet data: {}", e)
))?;
parse_lp_packet(&packet_buf)
.map_err(|e| GatewayError::LpProtocolError(
format!("Failed to parse LP packet: {}", e)
))
}
/// Send an LP packet over the stream with proper length-prefixed framing
async fn send_lp_packet(&mut self, packet: &LpPacket) -> Result<(), GatewayError> {
use nym_lp::codec::serialize_lp_packet;
use bytes::BytesMut;
// Serialize the packet first
let mut packet_buf = BytesMut::new();
serialize_lp_packet(packet, &mut packet_buf)
.map_err(|e| GatewayError::LpProtocolError(
format!("Failed to serialize packet: {}", e)
))?;
// Send 4-byte length prefix (u32 big-endian)
let len = packet_buf.len() as u32;
self.stream.write_all(&len.to_be_bytes()).await
.map_err(|e| GatewayError::LpConnectionError(
format!("Failed to send packet length: {}", e)
))?;
// Send the actual packet data
self.stream.write_all(&packet_buf).await
.map_err(|e| GatewayError::LpConnectionError(
format!("Failed to send packet data: {}", e)
))?;
self.stream.flush().await
.map_err(|e| GatewayError::LpConnectionError(
format!("Failed to flush stream: {}", e)
))?;
Ok(())
}
}
// Extension trait for LpSession to create packets
// This would ideally be part of nym-lp
trait LpSessionExt {
fn create_data_packet(&self, data: Vec<u8>) -> Result<LpPacket, nym_lp::LpError>;
}
impl LpSessionExt for LpSession {
fn create_data_packet(&self, data: Vec<u8>) -> Result<LpPacket, nym_lp::LpError> {
use nym_lp::packet::LpHeader;
let header = LpHeader {
protocol_version: 1,
session_id: self.id(),
counter: 0, // TODO: Use actual counter from session
};
let message = LpMessage::EncryptedData(data);
Ok(LpPacket::new(header, message))
}
}
+160
View File
@@ -0,0 +1,160 @@
// Copyright 2025 - Nym Technologies SA <contact@nymtech.net>
// SPDX-License-Identifier: GPL-3.0-only
use crate::error::GatewayError;
use nym_lp::{
keypair::{Keypair, PublicKey},
state_machine::{LpAction, LpInput, LpStateMachine},
LpPacket, LpSession,
};
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tokio::net::TcpStream;
use tracing::*;
/// Wrapper around the nym-lp state machine for gateway-side LP connections
pub struct LpGatewayHandshake {
state_machine: LpStateMachine,
}
impl LpGatewayHandshake {
/// Create a new responder (gateway side) handshake
pub fn new_responder(
local_keypair: &Keypair,
remote_public_key: &PublicKey,
psk: &[u8; 32],
) -> Result<Self, GatewayError> {
let state_machine = LpStateMachine::new(
false, // responder
local_keypair,
remote_public_key,
psk,
).map_err(|e| GatewayError::LpHandshakeError(format!("Failed to create state machine: {}", e)))?;
Ok(Self { state_machine })
}
/// Complete the handshake and return the established session
pub async fn complete(
mut self,
stream: &mut TcpStream,
) -> Result<LpSession, GatewayError> {
debug!("Starting LP handshake as responder");
// Start the handshake
if let Some(action) = self.state_machine.process_input(LpInput::StartHandshake) {
match action {
Ok(LpAction::SendPacket(packet)) => {
self.send_packet(stream, &packet).await?;
}
Ok(_) => {
// Unexpected action at this stage
return Err(GatewayError::LpHandshakeError(
"Unexpected action at handshake start".to_string()
));
}
Err(e) => {
return Err(GatewayError::LpHandshakeError(
format!("Failed to start handshake: {}", e)
));
}
}
}
// Continue handshake until complete
loop {
// Read incoming packet
let packet = self.receive_packet(stream).await?;
// Process the received packet
if let Some(action) = self.state_machine.process_input(LpInput::ReceivePacket(packet)) {
match action {
Ok(LpAction::SendPacket(response_packet)) => {
self.send_packet(stream, &response_packet).await?;
}
Ok(LpAction::HandshakeComplete) => {
info!("LP handshake completed successfully");
break;
}
Ok(other) => {
debug!("Received action during handshake: {:?}", other);
}
Err(e) => {
return Err(GatewayError::LpHandshakeError(
format!("Handshake error: {}", e)
));
}
}
}
}
// Extract the session from the state machine
self.state_machine.into_session()
.map_err(|e| GatewayError::LpHandshakeError(
format!("Failed to get session after handshake: {}", e)
))
}
/// Send an LP packet over the stream with proper length-prefixed framing
async fn send_packet(
&self,
stream: &mut TcpStream,
packet: &LpPacket,
) -> Result<(), GatewayError> {
use nym_lp::codec::serialize_lp_packet;
use bytes::BytesMut;
// Serialize the packet first
let mut packet_buf = BytesMut::new();
serialize_lp_packet(packet, &mut packet_buf)
.map_err(|e| GatewayError::LpProtocolError(format!("Failed to serialize packet: {}", e)))?;
// Send 4-byte length prefix (u32 big-endian)
let len = packet_buf.len() as u32;
stream.write_all(&len.to_be_bytes()).await
.map_err(|e| GatewayError::LpConnectionError(format!("Failed to send packet length: {}", e)))?;
// Send the actual packet data
stream.write_all(&packet_buf).await
.map_err(|e| GatewayError::LpConnectionError(format!("Failed to send packet data: {}", e)))?;
stream.flush().await
.map_err(|e| GatewayError::LpConnectionError(format!("Failed to flush stream: {}", e)))?;
debug!("Sent LP packet ({} bytes + 4 byte header)", packet_buf.len());
Ok(())
}
/// Receive an LP packet from the stream with proper length-prefixed framing
async fn receive_packet(
&self,
stream: &mut TcpStream,
) -> Result<LpPacket, GatewayError> {
use nym_lp::codec::parse_lp_packet;
// Read 4-byte length prefix (u32 big-endian)
let mut len_buf = [0u8; 4];
stream.read_exact(&mut len_buf).await
.map_err(|e| GatewayError::LpConnectionError(format!("Failed to read packet length: {}", e)))?;
let packet_len = u32::from_be_bytes(len_buf) as usize;
// Sanity check to prevent huge allocations
const MAX_PACKET_SIZE: usize = 65536; // 64KB max
if packet_len > MAX_PACKET_SIZE {
return Err(GatewayError::LpProtocolError(
format!("Packet size {} exceeds maximum {}", packet_len, MAX_PACKET_SIZE)
));
}
// Read the actual packet data
let mut packet_buf = vec![0u8; packet_len];
stream.read_exact(&mut packet_buf).await
.map_err(|e| GatewayError::LpConnectionError(format!("Failed to read packet data: {}", e)))?;
let packet = parse_lp_packet(&packet_buf)
.map_err(|e| GatewayError::LpProtocolError(format!("Failed to parse packet: {}", e)))?;
debug!("Received LP packet ({} bytes + 4 byte header)", packet_len);
Ok(packet)
}
}
+124
View File
@@ -0,0 +1,124 @@
// Copyright 2025 - Nym Technologies SA <contact@nymtech.net>
// SPDX-License-Identifier: GPL-3.0-only
use nym_credentials_interface::{CredentialSpendingData, TicketType};
use nym_registration_common::GatewayData;
use nym_wireguard_types::PeerPublicKey;
use serde::{Deserialize, Serialize};
use std::net::IpAddr;
/// Registration request sent by client after LP handshake
/// Aligned with existing authenticator registration flow
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct LpRegistrationRequest {
/// Client's WireGuard public key (for dVPN mode)
pub wg_public_key: PeerPublicKey,
/// Bandwidth credential for payment
pub credential: CredentialSpendingData,
/// Ticket type for bandwidth allocation
pub ticket_type: TicketType,
/// Registration mode
pub mode: RegistrationMode,
/// Client's IP address (for tracking/metrics)
pub client_ip: IpAddr,
/// Unix timestamp for replay protection
pub timestamp: u64,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum RegistrationMode {
/// dVPN mode - register as WireGuard peer (most common)
Dvpn,
/// Mixnet mode - register for mixnet usage (future)
Mixnet {
/// Client identifier for mixnet mode
client_id: [u8; 32]
},
}
/// Registration response from gateway
/// Contains GatewayData for compatibility with existing client code
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct LpRegistrationResponse {
/// Whether registration succeeded
pub success: bool,
/// Error message if registration failed
pub error: Option<String>,
/// Gateway configuration data (same as returned by authenticator)
/// This matches what WireguardRegistrationResult expects
pub gateway_data: Option<GatewayData>,
/// Allocated bandwidth in bytes
pub allocated_bandwidth: i64,
/// Session identifier for future reference
pub session_id: u32,
}
impl LpRegistrationRequest {
/// Create a new dVPN registration request
pub fn new_dvpn(
wg_public_key: PeerPublicKey,
credential: CredentialSpendingData,
ticket_type: TicketType,
client_ip: IpAddr,
) -> Self {
Self {
wg_public_key,
credential,
ticket_type,
mode: RegistrationMode::Dvpn,
client_ip,
timestamp: std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.unwrap()
.as_secs(),
}
}
/// Validate the request timestamp is within acceptable bounds
pub fn validate_timestamp(&self, max_skew_secs: u64) -> bool {
let now = std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.unwrap()
.as_secs();
(now as i64 - self.timestamp as i64).abs() <= max_skew_secs as i64
}
}
impl LpRegistrationResponse {
/// Create a success response with GatewayData
pub fn success(
session_id: u32,
allocated_bandwidth: i64,
gateway_data: GatewayData,
) -> Self {
Self {
success: true,
error: None,
gateway_data: Some(gateway_data),
allocated_bandwidth,
session_id,
}
}
/// Create an error response
pub fn error(session_id: u32, error: String) -> Self {
Self {
success: false,
error: Some(error),
gateway_data: None,
allocated_bandwidth: 0,
session_id,
}
}
}
+217
View File
@@ -0,0 +1,217 @@
// Copyright 2025 - Nym Technologies SA <contact@nymtech.net>
// SPDX-License-Identifier: GPL-3.0-only
use crate::error::GatewayError;
use crate::node::ActiveClientsStore;
use nym_credential_verification::ecash::EcashManager;
use nym_crypto::asymmetric::ed25519;
use nym_gateway_storage::GatewayStorage;
use nym_node_metrics::NymNodeMetrics;
use nym_task::ShutdownTracker;
use nym_wireguard::{PeerControlRequest, WireguardGatewayData};
use std::net::SocketAddr;
use std::sync::Arc;
use tokio::net::TcpListener;
use tokio::sync::mpsc;
use tracing::*;
mod handler;
mod handshake;
mod messages;
mod registration;
/// Configuration for LP listener
#[derive(Debug, Clone, serde::Deserialize, serde::Serialize)]
#[serde(default)]
pub struct LpConfig {
/// Enable/disable LP listener
pub enabled: bool,
/// Bind address for control port
#[serde(default = "default_bind_address")]
pub bind_address: String,
/// Control port (default: 41264)
#[serde(default = "default_control_port")]
pub control_port: u16,
/// Data port (default: 51264)
#[serde(default = "default_data_port")]
pub data_port: u16,
/// Maximum concurrent connections
#[serde(default = "default_max_connections")]
pub max_connections: usize,
}
impl Default for LpConfig {
fn default() -> Self {
Self {
enabled: false,
bind_address: default_bind_address(),
control_port: default_control_port(),
data_port: default_data_port(),
max_connections: default_max_connections(),
}
}
}
fn default_bind_address() -> String {
"0.0.0.0".to_string()
}
fn default_control_port() -> u16 {
41264
}
fn default_data_port() -> u16 {
51264
}
fn default_max_connections() -> usize {
10000
}
/// Shared state for LP connection handlers
#[derive(Clone)]
pub struct LpHandlerState {
/// Ecash verifier for bandwidth credentials
pub ecash_verifier: Arc<EcashManager>,
/// Storage backend for persistence
pub storage: GatewayStorage,
/// Gateway's identity keypair
pub local_identity: Arc<ed25519::KeyPair>,
/// Metrics collection
pub metrics: NymNodeMetrics,
/// Active clients tracking
pub active_clients_store: ActiveClientsStore,
/// WireGuard peer controller channel (for dVPN registrations)
pub wg_peer_controller: Option<mpsc::Sender<PeerControlRequest>>,
/// WireGuard gateway data (contains keypair and config)
pub wireguard_data: Option<WireguardGatewayData>,
}
/// LP listener that accepts TCP connections on port 41264
pub struct LpListener {
/// Address to bind the LP control port (41264)
control_address: SocketAddr,
/// Port for data plane (51264) - reserved for future use
data_port: u16,
/// Shared state for connection handlers
handler_state: LpHandlerState,
/// Maximum concurrent connections
max_connections: usize,
/// Shutdown coordination
shutdown: ShutdownTracker,
}
impl LpListener {
pub fn new(
bind_address: SocketAddr,
data_port: u16,
handler_state: LpHandlerState,
max_connections: usize,
shutdown: ShutdownTracker,
) -> Self {
Self {
control_address: bind_address,
data_port,
handler_state,
max_connections,
shutdown,
}
}
pub async fn run(&mut self) -> Result<(), GatewayError> {
let listener = TcpListener::bind(self.control_address)
.await
.map_err(|e| {
error!("Failed to bind LP listener to {}: {}", self.control_address, e);
GatewayError::ListenerBindFailure {
address: self.control_address.to_string(),
source: Box::new(e),
}
})?;
info!("LP listener started on {} (data port reserved: {})",
self.control_address, self.data_port);
let shutdown_token = self.shutdown.clone_shutdown_token();
loop {
tokio::select! {
biased;
_ = shutdown_token.cancelled() => {
trace!("LP listener: received shutdown signal");
break;
}
result = listener.accept() => {
match result {
Ok((stream, addr)) => {
self.handle_connection(stream, addr);
}
Err(e) => {
warn!("Failed to accept LP connection: {}", e);
}
}
}
}
}
info!("LP listener shutdown complete");
Ok(())
}
fn handle_connection(&self, stream: tokio::net::TcpStream, remote_addr: SocketAddr) {
// Check connection limit
let active_connections = self.active_lp_connections();
if active_connections >= self.max_connections {
warn!(
"LP connection limit exceeded ({}/{}), rejecting connection from {}",
active_connections, self.max_connections, remote_addr
);
return;
}
debug!("Accepting LP connection from {} ({} active connections)",
remote_addr, active_connections);
// Increment connection counter
self.handler_state.metrics.network.new_lp_connection();
// Spawn handler task
let handler = handler::LpConnectionHandler::new(
stream,
remote_addr,
self.handler_state.clone(),
);
let metrics = self.handler_state.metrics.clone();
self.shutdown.try_spawn_named(
async move {
if let Err(e) = handler.handle().await {
warn!("LP handler error for {}: {}", remote_addr, e);
}
// Decrement connection counter on exit
metrics.network.lp_connection_closed();
},
&format!("LP::{}", remote_addr),
);
}
fn active_lp_connections(&self) -> usize {
self.handler_state.metrics.network.active_lp_connections_count()
}
}
@@ -0,0 +1,276 @@
// Copyright 2025 - Nym Technologies SA <contact@nymtech.net>
// SPDX-License-Identifier: GPL-3.0-only
use super::messages::{LpRegistrationRequest, LpRegistrationResponse, RegistrationMode};
use super::LpHandlerState;
use crate::error::GatewayError;
use defguard_wireguard_rs::host::Peer;
use defguard_wireguard_rs::key::Key;
use futures::channel::oneshot;
use nym_credential_verification::ecash::traits::EcashManager;
use nym_credential_verification::{
bandwidth_storage_manager::BandwidthStorageManager, BandwidthFlushingBehaviourConfig,
ClientBandwidth, CredentialVerifier,
};
use nym_credentials_interface::CredentialSpendingData;
use nym_gateway_requests::models::CredentialSpendingRequest;
use nym_gateway_storage::models::PersistedBandwidth;
use nym_gateway_storage::traits::BandwidthGatewayStorage;
use nym_registration_common::GatewayData;
use nym_wireguard::PeerControlRequest;
use rand::RngCore;
use std::net::{IpAddr, Ipv4Addr, Ipv6Addr, SocketAddr};
use std::str::FromStr;
use std::sync::Arc;
use tracing::*;
/// Prepare bandwidth storage for a client
async fn credential_storage_preparation(
ecash_verifier: Arc<dyn EcashManager + Send + Sync>,
client_id: i64,
) -> Result<PersistedBandwidth, GatewayError> {
ecash_verifier
.storage()
.create_bandwidth_entry(client_id)
.await?;
let bandwidth = ecash_verifier
.storage()
.get_available_bandwidth(client_id)
.await?
.ok_or_else(|| {
GatewayError::InternalError(
"bandwidth entry should have just been created".to_string(),
)
})?;
Ok(bandwidth)
}
/// Verify credential and allocate bandwidth using CredentialVerifier
async fn credential_verification(
ecash_verifier: Arc<dyn EcashManager + Send + Sync>,
credential: CredentialSpendingData,
client_id: i64,
) -> Result<i64, GatewayError> {
let bandwidth = credential_storage_preparation(ecash_verifier.clone(), client_id).await?;
let client_bandwidth = ClientBandwidth::new(bandwidth.into());
let mut verifier = CredentialVerifier::new(
CredentialSpendingRequest::new(credential),
ecash_verifier.clone(),
BandwidthStorageManager::new(
ecash_verifier.storage(),
client_bandwidth,
client_id,
BandwidthFlushingBehaviourConfig::default(),
true,
),
);
Ok(verifier.verify().await?)
}
/// Process an LP registration request
pub async fn process_registration(
request: LpRegistrationRequest,
state: &LpHandlerState,
) -> LpRegistrationResponse {
let session_id = rand::random::<u32>();
// 1. Validate timestamp for replay protection
if !request.validate_timestamp(30) {
warn!("LP registration failed: timestamp too old or too far in future");
return LpRegistrationResponse::error(
session_id,
"Invalid timestamp".to_string(),
);
}
// 2. Process based on mode
match request.mode {
RegistrationMode::Dvpn => {
// Register as WireGuard peer first to get client_id
let (gateway_data, client_id) = match register_wg_peer(
request.wg_public_key.inner().as_ref(),
request.client_ip,
request.ticket_type,
state,
).await {
Ok(result) => result,
Err(e) => {
error!("LP WireGuard peer registration failed: {}", e);
return LpRegistrationResponse::error(
session_id,
format!("WireGuard peer registration failed: {}", e),
);
}
};
// Verify credential with CredentialVerifier (handles double-spend, storage, etc.)
let allocated_bandwidth = match credential_verification(
state.ecash_verifier.clone(),
request.credential,
client_id,
).await {
Ok(bandwidth) => bandwidth,
Err(e) => {
// Credential verification failed, remove the peer
warn!("LP credential verification failed for client {}: {}", client_id, e);
if let Err(remove_err) = state.storage
.remove_wireguard_peer(&request.wg_public_key.to_string())
.await
{
error!("Failed to remove peer after credential verification failure: {}", remove_err);
}
return LpRegistrationResponse::error(
session_id,
format!("Credential verification failed: {}", e),
);
}
};
info!("LP dVPN registration successful for session {} (client_id: {})", session_id, client_id);
LpRegistrationResponse::success(
session_id,
allocated_bandwidth,
gateway_data,
)
}
RegistrationMode::Mixnet { client_id: client_id_bytes } => {
// Generate i64 client_id from the [u8; 32] in the request
let client_id = i64::from_be_bytes(client_id_bytes[0..8].try_into().unwrap());
info!("LP Mixnet registration for client_id {}, session {}", client_id, session_id);
// Verify credential with CredentialVerifier
let allocated_bandwidth = match credential_verification(
state.ecash_verifier.clone(),
request.credential,
client_id,
).await {
Ok(bandwidth) => bandwidth,
Err(e) => {
warn!("LP Mixnet credential verification failed for client {}: {}", client_id, e);
return LpRegistrationResponse::error(
session_id,
format!("Credential verification failed: {}", e),
);
}
};
// For mixnet mode, we don't have WireGuard data
// In the future, this would set up mixnet-specific state
info!("LP Mixnet registration successful for session {} (client_id: {})", session_id, client_id);
LpRegistrationResponse {
success: true,
error: None,
gateway_data: None,
allocated_bandwidth,
session_id,
}
}
}
}
/// Register a WireGuard peer and return gateway data along with the client_id
async fn register_wg_peer(
public_key_bytes: &[u8],
client_ip: IpAddr,
ticket_type: nym_credentials_interface::TicketType,
state: &LpHandlerState,
) -> Result<(GatewayData, i64), GatewayError> {
let Some(wg_controller) = &state.wg_peer_controller else {
return Err(GatewayError::ServiceProviderNotRunning {
service: "WireGuard".to_string(),
});
};
let Some(wg_data) = &state.wireguard_data else {
return Err(GatewayError::ServiceProviderNotRunning {
service: "WireGuard".to_string(),
});
};
// Convert public key bytes to WireGuard Key
let mut key_bytes = [0u8; 32];
if public_key_bytes.len() != 32 {
return Err(GatewayError::LpProtocolError(
"Invalid WireGuard public key length".to_string()
));
}
key_bytes.copy_from_slice(public_key_bytes);
let peer_key = Key::new(key_bytes);
// Allocate IP addresses for the client
// TODO: Proper IP pool management - for now use random in private range
let last_octet = {
let mut rng = rand::thread_rng();
(rng.next_u32() % 254 + 1) as u8
};
let client_ipv4 = Ipv4Addr::new(10, 1, 0, last_octet);
let client_ipv6 = Ipv6Addr::new(0xfd00, 0, 0, 0, 0, 0, 0, last_octet as u16);
// Create WireGuard peer
let mut peer = Peer::new(peer_key.clone());
peer.preshared_key = Some(Key::new(state.local_identity.public_key().to_bytes()));
peer.endpoint = Some(format!("{}:51820", client_ip).parse().unwrap_or_else(|_| {
SocketAddr::from_str("0.0.0.0:51820").unwrap()
}));
peer.allowed_ips = vec![
format!("{}/32", client_ipv4).parse().unwrap(),
format!("{}/128", client_ipv6).parse().unwrap(),
];
peer.persistent_keepalive_interval = Some(25);
// Send to WireGuard peer controller
let (tx, rx) = oneshot::channel();
wg_controller
.send(PeerControlRequest::AddPeer {
peer: peer.clone(),
response_tx: tx,
})
.await
.map_err(|e| GatewayError::InternalError(format!("Failed to send peer request: {}", e)))?;
rx.await
.map_err(|e| GatewayError::InternalError(format!("Failed to receive peer response: {}", e)))?
.map_err(|e| GatewayError::InternalError(format!("Failed to add peer: {:?}", e)))?;
// Store bandwidth allocation and get client_id
let client_id = state.storage
.insert_wireguard_peer(&peer, ticket_type.into())
.await
.map_err(|e| {
error!("Failed to store WireGuard peer in database: {}", e);
GatewayError::InternalError(format!("Failed to store peer: {}", e))
})?;
// Get gateway's actual WireGuard public key
let gateway_pubkey = *wg_data.keypair().public_key();
// Get gateway's WireGuard endpoint from config
let gateway_endpoint = wg_data.config().bind_address;
// Create GatewayData response (matching authenticator response format)
Ok((
GatewayData {
public_key: gateway_pubkey,
endpoint: gateway_endpoint,
private_ipv4: client_ipv4,
private_ipv6: client_ipv6,
},
client_id,
))
}
// Helper function to convert bandwidth to ClientBandwidth if needed
// This would integrate with the actual bandwidth controller
#[allow(dead_code)]
async fn store_client_bandwidth(
client_id: String,
bandwidth: i64,
storage: &nym_gateway_storage::GatewayStorage,
) -> Result<(), GatewayError> {
// This would integrate with the actual bandwidth storage
// For now, just log it
info!("Storing bandwidth {} for client {}", bandwidth, client_id);
Ok(())
}
+39
View File
@@ -37,6 +37,7 @@ use zeroize::Zeroizing;
pub use crate::node::upgrade_mode::watcher::UpgradeModeWatcher;
pub use client_handling::active_clients::ActiveClientsStore;
pub use lp_listener::LpConfig;
pub use nym_credential_verification::upgrade_mode::UpgradeModeCheckRequestSender;
pub use nym_gateway_stats_storage::PersistentStatsStorage;
pub use nym_gateway_storage::{
@@ -48,6 +49,7 @@ pub use nym_sdk::{NymApiTopologyProvider, NymApiTopologyProviderConfig, UserAgen
pub(crate) mod client_handling;
pub(crate) mod internal_service_providers;
pub(crate) mod lp_listener;
mod stale_data_cleaner;
pub mod upgrade_mode;
@@ -287,6 +289,43 @@ impl GatewayTasksBuilder {
))
}
pub async fn build_lp_listener(
&mut self,
active_clients_store: ActiveClientsStore,
) -> Result<lp_listener::LpListener, GatewayError> {
// Get WireGuard peer controller if available
let wg_peer_controller = if let Some(wg_data) = &self.wireguard_data {
Some(wg_data.inner.peer_tx().clone())
} else {
None
};
let handler_state = lp_listener::LpHandlerState {
ecash_verifier: self.ecash_manager().await?,
storage: self.storage.clone(),
local_identity: Arc::clone(&self.identity_keypair),
metrics: self.metrics.clone(),
active_clients_store,
wg_peer_controller,
wireguard_data: self.wireguard_data.as_ref().map(|wd| wd.inner.clone()),
};
// Parse bind address from config
let bind_addr = format!("{}:{}",
self.config.lp.bind_address,
self.config.lp.control_port
).parse()
.map_err(|e| GatewayError::InternalError(format!("Invalid LP bind address: {}", e)))?;
Ok(lp_listener::LpListener::new(
bind_addr,
self.config.lp.data_port,
handler_state,
self.config.lp.max_connections,
self.shutdown_tracker.clone(),
))
}
fn build_network_requester(
&mut self,
topology_provider: Box<dyn TopologyProvider + Send + Sync>,
+17
View File
@@ -15,6 +15,8 @@ pub struct NetworkStats {
// designed with metrics in mind and this single counter has been woven through
// the call stack
active_egress_mixnet_connections: Arc<AtomicUsize>,
active_lp_connections: AtomicUsize,
}
impl NetworkStats {
@@ -56,4 +58,19 @@ impl NetworkStats {
self.active_egress_mixnet_connections
.load(Ordering::Relaxed)
}
pub fn new_lp_connection(&self) {
self.active_lp_connections
.fetch_add(1, Ordering::Relaxed);
}
pub fn lp_connection_closed(&self) {
self.active_lp_connections
.fetch_sub(1, Ordering::Relaxed);
}
pub fn active_lp_connections_count(&self) -> usize {
self.active_lp_connections
.load(Ordering::Relaxed)
}
}
+4
View File
@@ -46,6 +46,9 @@ pub struct GatewayTasksConfig {
pub upgrade_mode: UpgradeModeWatcher,
#[serde(default)]
pub lp: nym_gateway::node::LpConfig,
#[serde(default)]
pub debug: Debug,
}
@@ -225,6 +228,7 @@ impl GatewayTasksConfig {
announce_ws_port: None,
announce_wss_port: None,
upgrade_mode: UpgradeModeWatcher::new()?,
lp: Default::default(),
debug: Default::default(),
})
}
+3
View File
@@ -27,6 +27,7 @@ fn ephemeral_gateway_config(config: &Config) -> nym_gateway::config::Config {
enabled: config.service_providers.network_requester.debug.enabled,
},
config.gateway_tasks.upgrade_mode.clone(),
config.gateway_tasks.lp.clone(),
nym_gateway::config::Debug {
client_bandwidth_max_flushing_rate: config
.gateway_tasks
@@ -91,6 +92,7 @@ pub struct GatewayTasksConfig {
pub auth_opts: Option<LocalAuthenticatorOpts>,
#[allow(dead_code)]
pub wg_opts: LocalWireguardOpts,
pub lp: nym_gateway::node::LpConfig,
}
// that function is rather disgusting, but I hope it's not going to live for too long
@@ -223,6 +225,7 @@ pub fn gateway_tasks_config(config: &Config) -> GatewayTasksConfig {
ipr_opts: Some(ipr_opts),
auth_opts: Some(auth_opts),
wg_opts,
lp: config.gateway_tasks.lp.clone(),
}
}
@@ -1353,6 +1353,7 @@ pub async fn try_upgrade_config_v10<P: AsRef<Path>>(
)
})
.unwrap_or(UpgradeModeWatcher::new_mainnet()),
lp: Default::default(),
debug: gateway_tasks::Debug {
message_retrieval_limit: old_cfg.gateway_tasks.debug.message_retrieval_limit,
maximum_open_connections: old_cfg.gateway_tasks.debug.maximum_open_connections,
+17
View File
@@ -665,6 +665,23 @@ impl NymNode {
.await?;
self.shutdown_tracker()
.try_spawn_named(async move { websocket.run().await }, "EntryWebsocket");
// Start LP listener if enabled
if self.config.gateway_tasks.lp.enabled {
info!(
"starting the LP listener on {}:{} (data port: {})",
self.config.gateway_tasks.lp.bind_address,
self.config.gateway_tasks.lp.control_port,
self.config.gateway_tasks.lp.data_port
);
let mut lp_listener = gateway_tasks_builder
.build_lp_listener(active_clients_store.clone())
.await?;
self.shutdown_tracker()
.try_spawn_named(async move { lp_listener.run().await }, "LpListener");
} else {
info!("LP listener is disabled");
}
} else {
info!("node not running in entry mode: the websocket will remain closed");
}