├── .cargo └── config.toml ├── .config └── nextest.toml ├── .github └── workflows │ ├── ci.yml │ └── release.yml ├── .gitignore ├── CHANGELOG.md ├── Cargo.toml ├── LICENSE ├── README.md ├── benches ├── data │ ├── large_size.txt │ ├── medium_size.txt │ └── small_size.txt └── micro.rs ├── codecov.yml ├── examples ├── controller.rs ├── proxy.rs └── tracing.rs ├── justfile ├── rust-toolchain ├── rustfmt.toml ├── src ├── client │ ├── conn │ │ ├── mod.rs │ │ └── tokio.rs │ ├── handler │ │ ├── mod.rs │ │ ├── offline.rs │ │ └── online.rs │ └── mod.rs ├── codec │ ├── decoder │ │ ├── body.rs │ │ ├── dedup.rs │ │ ├── fragment.rs │ │ ├── mod.rs │ │ └── ordered.rs │ ├── encoder │ │ ├── body.rs │ │ ├── fragment.rs │ │ └── mod.rs │ ├── frame.rs │ ├── mod.rs │ └── tokio.rs ├── congestion │ ├── bbr.rs │ ├── cubic.rs │ ├── legacy.rs │ └── mod.rs ├── errors.rs ├── estimator.rs ├── lib.rs ├── link.rs ├── opts.rs ├── packet │ ├── connected │ │ ├── ack.rs │ │ ├── frame_set.rs │ │ └── mod.rs │ ├── mod.rs │ └── unconnected.rs ├── reliable.rs ├── server │ ├── handler │ │ ├── mod.rs │ │ ├── offline.rs │ │ └── online.rs │ ├── incoming │ │ ├── mod.rs │ │ └── tokio.rs │ └── mod.rs ├── state.rs ├── tests.rs └── utils │ ├── bit_queue.rs │ ├── fastrace.rs │ ├── hash.rs │ ├── log.rs │ ├── mod.rs │ ├── reactor.rs │ ├── seq_num.rs │ └── tests │ ├── flusher.rs │ ├── mod.rs │ ├── sim_net.rs │ └── tracing.rs └── xdp ├── ebpf ├── .cargo │ └── config.toml ├── Cargo.toml └── src │ ├── buffer.rs │ ├── inet │ ├── ecn.rs │ ├── ethernet.rs │ ├── ip.rs │ ├── ipv4.rs │ ├── ipv6.rs │ ├── mod.rs │ ├── udp.rs │ └── zerocopy.rs │ └── main.rs └── raknet-xdp ├── Cargo.toml ├── bindings └── input.h ├── build.rs └── src ├── bindings.rs ├── if_xdp.rs ├── io ├── mod.rs ├── rx.rs └── tx.rs ├── lib.rs ├── mmap.rs ├── socket.rs ├── syscall.rs └── umem.rs /.cargo/config.toml: -------------------------------------------------------------------------------- 1 | [target.'cfg(all())'] 2 | rustflags = ["--cfg", "tokio_unstable"] 3 | -------------------------------------------------------------------------------- /.config/nextest.toml: -------------------------------------------------------------------------------- 1 | [profile.default] 2 | retries = 0 3 | slow-timeout = { period = "5s", terminate-after = 6 } 4 | status-level = "all" 5 | final-status-level = "slow" 6 | -------------------------------------------------------------------------------- /.github/workflows/ci.yml: -------------------------------------------------------------------------------- 1 | name: CI 2 | 3 | on: 4 | push: 5 | branches: [ "main" ] 6 | pull_request: 7 | branches: [ "main" ] 8 | 9 | env: 10 | CARGO_TERM_COLOR: always 11 | 12 | jobs: 13 | check: 14 | runs-on: ubuntu-latest 15 | steps: 16 | - uses: actions/checkout@v3 17 | - uses: taiki-e/install-action@cargo-binstall 18 | 19 | - name: Install Dependencies 20 | run: | 21 | rustup component add rustfmt llvm-tools-preview clippy 22 | cargo binstall -y --no-symlinks cargo-llvm-cov cargo-nextest cargo-sort cargo-audit cargo-machete 23 | wget https://github.com/mozilla/sccache/releases/download/v0.7.5/sccache-v0.7.5-x86_64-unknown-linux-musl.tar.gz 24 | mkdir sccache-install && tar -xzf sccache-v0.7.5-x86_64-unknown-linux-musl.tar.gz -C sccache-install --strip-components=1 25 | mv sccache-install/sccache /usr/local/bin && rm -rf sccache-install/ 26 | 27 | - name: Configure sccache 28 | uses: actions/github-script@v6 29 | with: 30 | script: | 31 | core.exportVariable('RUSTC_WRAPPER', 'sccache') 32 | core.exportVariable('SCCACHE_GHA_ENABLED', 'on') 33 | core.exportVariable('ACTIONS_CACHE_URL', process.env.ACTIONS_CACHE_URL || ''); 34 | core.exportVariable('ACTIONS_RUNTIME_TOKEN', process.env.ACTIONS_RUNTIME_TOKEN || ''); 35 | 36 | - name: Run quick check 37 | run: | 38 | cargo audit 39 | cargo fmt --all -- --check 40 | cargo sort --check 41 | cargo machete 42 | 43 | - name: Sccache clear 44 | run: sccache --zero-stats > /dev/null 45 | 46 | - name: Build in release mode 47 | run: cargo build --release --all-features --features fastrace/enable 48 | 49 | - name: Test with codecov 50 | run: cargo llvm-cov --no-report nextest --all-features --features fastrace/enable 51 | 52 | - name: Test example with codecov 53 | run: | 54 | cargo llvm-cov --no-report run --example proxy 55 | cargo llvm-cov --no-report run --example controller 56 | cargo llvm-cov --no-report run --example tracing --features fastrace/enable 57 | 58 | - name: Test bench with codecov 59 | run: | 60 | cargo llvm-cov --no-report --bench micro --all-features 61 | 62 | - name: Generate codecov report 63 | run: | 64 | cargo llvm-cov report --lcov --output-path lcov.info 65 | 66 | - name: Sccache state 67 | run: sccache --show-stats 68 | 69 | - uses: codecov/codecov-action@v4 70 | env: 71 | CODECOV_TOKEN: ${{ secrets.CODECOV_TOKEN }} 72 | with: 73 | verbose: true 74 | -------------------------------------------------------------------------------- /.github/workflows/release.yml: -------------------------------------------------------------------------------- 1 | name: Release 2 | 3 | on: 4 | push: 5 | tags: 6 | - "v*" 7 | 8 | jobs: 9 | create_release: 10 | name: Create Release 11 | runs-on: ubuntu-latest 12 | steps: 13 | - uses: actions/checkout@v3 14 | 15 | - name: Generate Release Notes 16 | run: | 17 | awk '/---/{i++; next} i==1' CHANGELOG.md > RELEASE_NOTES.md 18 | 19 | - name: Create Release 20 | uses: actions/create-release@v1.1.4 21 | env: 22 | GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} 23 | with: 24 | tag_name: ${{ github.ref }} 25 | release_name: Release ${{ github.ref }} 26 | body_path: RELEASE_NOTES.md 27 | draft: false 28 | prerelease: false 29 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | .* 2 | !.github 3 | !.gitignore 4 | !.cargo 5 | !.config/ 6 | target/ 7 | Cargo.lock 8 | **/*.rs.bk 9 | profile.json 10 | -------------------------------------------------------------------------------- /CHANGELOG.md: -------------------------------------------------------------------------------- 1 | # Changelog 2 | 3 | ## Unreleased 4 | 5 | - Add `serde` feature to support config serializing and deserializing 6 | - Add `ConnectionInfo` 7 | 8 | --- 9 | ## 0.1.4 10 | 11 | - Split `IO` into `Stream` and `Sink` 12 | - Fix some fragmentation bugs 13 | - Add `FlushStrategy` 14 | - Add `Estimator` based on RFC6298 15 | 16 | --- 17 | ## 0.1.3 18 | 19 | - Preliminary implementing RFC6298 RttEstimator 20 | - Zero copy packets decoding 21 | - Support offline packets lifetime span tracing 22 | - Add experimental `Ping` for clients 23 | - Reduced the overhead when calling `poll_close` 24 | - By introducing a simple time reactor 25 | - Add examples for E2E tracing 26 | - Fixed some API call issues 27 | 28 | --- 29 | ## 0.1.2 30 | 31 | - Fix fragmented packets boundary 32 | - Refactor lots of codes 33 | - Ensure all packets are delivered after `poll_close` 34 | - Add packet tracing powered by `minitrace` 35 | 36 | --- 37 | ## 0.1.1 38 | 39 | - Add `tokio-udp` feature 40 | -------------------------------------------------------------------------------- /Cargo.toml: -------------------------------------------------------------------------------- 1 | [package] 2 | name = "raknet-rs" 3 | version = "0.1.4" 4 | readme = "README.md" 5 | resolver = "2" 6 | description = "Raknet protocol implementation by rust" 7 | edition = { workspace = true } 8 | license = { workspace = true } 9 | authors = { workspace = true } 10 | homepage = { workspace = true } 11 | repository = { workspace = true } 12 | categories = { workspace = true } 13 | keywords = { workspace = true } 14 | 15 | [dependencies] 16 | async-channel = "2.3.1" 17 | bytes = "1" 18 | concurrent-queue = "2.5.0" 19 | fastrace = "0.7" 20 | futures = { version = "0.3.5", default-features = false } 21 | futures-async-stream = "0.2" 22 | log = "0.4" 23 | lru = "0.13" 24 | parking_lot = "0.12" 25 | pin-project-lite = "0.2" 26 | rand = "0.8" 27 | rustc-hash = { version = "2.1.0", optional = true } 28 | serde = { version = "1", features = ["derive"], optional = true } 29 | thiserror = "1" 30 | tokio = { version = "1", features = ["net", "rt"], optional = true } 31 | 32 | [dev-dependencies] 33 | criterion = { version = "0.5", features = ["async_futures", "async_tokio"] } 34 | env_logger = "0.11" 35 | indexmap = "2" 36 | reqwest = "0.12" 37 | tokio = { version = "1", features = ["full"] } 38 | 39 | [features] 40 | default = ["tokio-rt", "serde"] 41 | tokio-rt = ["dep:tokio"] 42 | serde = ["dep:serde"] 43 | rustc-hash = ["dep:rustc-hash"] # enable faster hash if the network is trusted 44 | micro-bench = [] # for benchmark, do not enable it in normal use 45 | 46 | [[bench]] 47 | name = "micro" 48 | harness = false 49 | required-features = ["micro-bench", "rustc-hash"] 50 | 51 | [profile.bench] 52 | opt-level = 3 53 | lto = true 54 | debug = true 55 | 56 | [lints] 57 | workspace = true 58 | # 59 | ############################# 60 | ## workspace configuration ## 61 | ############################# 62 | # 63 | [workspace] 64 | members = ["xdp/ebpf", "xdp/raknet-xdp"] 65 | 66 | [workspace.package] 67 | edition = "2021" 68 | license = "Apache-2.0" 69 | authors = ["iGxnon"] 70 | homepage = "https://github.com/MemoriesOfTime/raknet-rs" 71 | repository = "https://github.com/MemoriesOfTime/raknet-rs" 72 | categories = ["network-programming"] 73 | keywords = ["raknet", "network", "udp"] 74 | 75 | [workspace.lints.rust] 76 | keyword_idents = "warn" 77 | missing_copy_implementations = "warn" 78 | missing_debug_implementations = "warn" 79 | non_ascii_idents = "warn" 80 | noop_method_call = "warn" 81 | unused_extern_crates = "warn" 82 | unused_import_braces = "warn" 83 | 84 | [workspace.lints.clippy] 85 | cognitive_complexity = "warn" 86 | dbg_macro = "warn" 87 | debug_assert_with_mut_call = "warn" 88 | doc_link_with_quotes = "warn" 89 | doc_markdown = "warn" 90 | empty_line_after_outer_attr = "warn" 91 | empty_structs_with_brackets = "warn" 92 | float_cmp = "warn" 93 | float_cmp_const = "warn" 94 | float_equality_without_abs = "warn" 95 | missing_errors_doc = "warn" 96 | missing_panics_doc = "warn" 97 | option_if_let_else = "warn" 98 | print_stderr = "warn" 99 | print_stdout = "warn" 100 | semicolon_if_nothing_returned = "warn" 101 | unseparated_literal_suffix = "warn" 102 | shadow_unrelated = "warn" 103 | similar_names = "warn" 104 | suspicious_operation_groupings = "warn" 105 | unused_self = "warn" 106 | use_debug = "warn" 107 | used_underscore_binding = "warn" 108 | useless_let_if_seq = "warn" 109 | wildcard_dependencies = "warn" 110 | wildcard_imports = "warn" 111 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | ![Header](https://capsule-render.vercel.app/api?type=Waving&color=timeGradient&height=300&animation=fadeIn§ion=header&text=raknet-rs&fontSize=90&fontAlignY=45) 2 | 3 | [![Crates.io](https://img.shields.io/crates/v/raknet-rs.svg?style=flat-square&logo=rust)](https://crates.io/crates/raknet-rs) 4 | [![CI Status](https://img.shields.io/github/actions/workflow/status/MemoriesOfTime/raknet-rs/ci.yml?style=flat-square&logo=github)](https://github.com/MemoriesOfTime/raknet-rs/actions) 5 | [![Coverage](https://img.shields.io/codecov/c/github/MemoriesOfTime/raknet-rs?style=flat-square&logo=codecov)](https://app.codecov.io/github/MemoriesOfTime/raknet-rs) 6 | [![License](https://img.shields.io/crates/l/raknet-rs?style=flat-square)](https://github.com/MemoriesOfTime/raknet-rs/blob/master/LICENSE) 7 | 8 | Yet another project rewritten in Rust. 9 | 10 | ## Features 11 | 12 | - `Stream`/`Sink`/`Future` based async API. 13 | - Low level API but easy to use. 14 | - RakNet features: 15 | - Support `Unreliable`, `Reliable` and `ReliableOrdered` packets. 16 | - Support multiple order channels. 17 | - Support `ACK`/`NACK` mechanism. 18 | - Support message priority for unordered reliability. 19 | - Full tracing: 20 | - You can track a packet's span during deduplication, fragmentation, ... 21 | 22 | ## Roadmap 23 | 24 | > Ordered by priority 25 | 26 | - Documentation 27 | - Simulation testing 28 | - Bulk benchmark 29 | - AF_XDP socket & XDP redirect from UDP 30 | - Optimize performance for single thread runtime (IO without Send) 31 | - Robust client implementation 32 | - Use stable rust toolchain (~~I like nightly~~) 33 | 34 | ## Getting Started 35 | 36 | See [examples](examples/) or [integration testing](src/tests.rs) for basic usage. 37 | 38 | ### Server 39 | 40 | Most operations are performed on `Stream` and `Sink`. There will be some options in [opts](src/opts.rs). 41 | 42 | The implementation details are obscured, and you can only see a very high level of abstraction, including the `Error` type, which is just `std::io::Error`. 43 | 44 | Keep polling `incoming` because it also serves as the router to every connections. 45 | 46 | Apply `Sink::poll_flush` to IO will trigger to flush all pending packets, `ACK`/`NACK`, and stale packets. So you have to call `poll_flush` periodically. You can configure the [flush strategy](src/opts.rs) you want. 47 | 48 | Apply `Sink::poll_close` to IO will ensure that all data is received by the peer before returning. It may keep resending infinitely unless you cancel the task. So you'd better set a timeout for each `poll_close`. 49 | 50 | > [!NOTE] 51 | > All calculations are lazy. The state will not update if you do not poll it. 52 | 53 | ```rust 54 | use bytes::Bytes; 55 | use futures::{SinkExt, StreamExt}; 56 | use raknet_rs::server::{self, MakeIncoming}; 57 | 58 | let socket = tokio::net::UdpSocket::bind("127.0.0.1:0").await?; 59 | let config = server::Config::new() 60 | .sever_guid(114514) 61 | .advertisement(&b"Hello, I am server"[..]) 62 | ... 63 | let mut incoming = socket.make_incoming(config); 64 | let (reader, _) = incoming.next().await.unwrap(); 65 | tokio::pin!(reader); 66 | let data: Bytes = reader.next().await.unwrap(); 67 | ``` 68 | 69 | ### Client 70 | 71 | > [!WARNING] 72 | > The current version of the client only has the most basic handshake implementation, and it is not recommended to use it directly. 73 | 74 | ```rust 75 | use bytes::Bytes; 76 | use futures::{SinkExt, StreamExt}; 77 | use raknet_rs::client::{self, ConnectTo}; 78 | use raknet_rs::Reliability; 79 | 80 | let socket = tokio::net::UdpSocket::bind("0.0.0.0:0").await?; 81 | let config = client::Config::new() 82 | .client_guid(1919810) 83 | ... 84 | let (_, writer) = socket.connect_to(, config).await?; 85 | tokio::pin!(writer); 86 | writer.send(Message::new(Bytes::from_static(b"Hello, Anyone there?"))) 87 | .await?; 88 | ``` 89 | -------------------------------------------------------------------------------- /benches/data/large_size.txt: -------------------------------------------------------------------------------- 1 | "They're made out of meat." 2 | 3 | "Meat?" 4 | 5 | "Meat. They're made out of meat." 6 | 7 | "Meat?" 8 | 9 | "There's no doubt about it. We picked several from different parts of the planet, took them aboard our recon vessels, probed them all the way through. They're completely meat." 10 | 11 | "That's impossible. What about the radio signals? The messages to the stars." 12 | 13 | "They use the radio waves to talk, but the signals don't come from them. The signals come from machines." 14 | 15 | "So who made the machines? That's who we want to contact." 16 | 17 | "They made the machines. That's what I'm trying to tell you. Meat made the machines." 18 | 19 | "That's ridiculous. How can meat make a machine? You're asking me to believe in sentient meat." 20 | 21 | "I'm not asking you, I'm telling you. These creatures are the only sentient race in the sector and they're made out of meat." 22 | 23 | "Maybe they're like the Orfolei. You know, a carbon-based intelligence that goes through a meat stage." 24 | 25 | "Nope. They're born meat and they die meat. We studied them for several of their life spans, which didn't take too long. Do you have any idea the life span of meat?" 26 | 27 | "Spare me. Okay, maybe they're only part meat. You know, like the Weddilei. A meat head with an electron plasma brain inside." 28 | 29 | "Nope. We thought of that, since they do have meat heads like the Weddilei. But I told you, we probed them. They're meat all the way through." 30 | 31 | "No brain?" 32 | 33 | "Oh, there is a brain all right. It's just that the brain is made out of meat!" 34 | 35 | "So... what does the thinking?" 36 | 37 | "You're not understanding, are you? The brain does the thinking. The meat." 38 | 39 | "Thinking meat! You're asking me to believe in thinking meat!" 40 | 41 | "Yes, thinking meat! Conscious meat! Loving meat. Dreaming meat. The meat is the whole deal! Are you getting the picture?" 42 | 43 | "Omigod. You're serious then. They're made out of meat." 44 | 45 | "Finally, Yes. They are indeed made out meat. And they've been trying to get in touch with us for almost a hundred of their years." 46 | 47 | "So what does the meat have in mind." 48 | 49 | "First it wants to talk to us. Then I imagine it wants to explore the universe, contact other sentients, swap ideas and information. The usual." 50 | 51 | "We're supposed to talk to meat?" 52 | 53 | "That's the idea. That's the message they're sending out by radio. 'Hello. Anyone out there? Anyone home?' That sort of thing." 54 | 55 | "They actually do talk, then. They use words, ideas, concepts?" 56 | 57 | "Oh, yes. Except they do it with meat." 58 | 59 | "I thought you just told me they used radio." 60 | 61 | "They do, but what do you think is on the radio? Meat sounds. You know how when you slap or flap meat it makes a noise? They talk by flapping their meat at each other. They can even sing by squirting air through their meat." 62 | 63 | "Omigod. Singing meat. This is altogether too much. So what do you advise?" 64 | 65 | "Officially or unofficially?" 66 | 67 | "Both." 68 | 69 | "Officially, we are required to contact, welcome, and log in any and all sentient races or multibeings in the quadrant, without prejudice, fear, or favor. Unofficially, I advise that we erase the records and forget the whole thing." 70 | 71 | "I was hoping you would say that." 72 | 73 | "It seems harsh, but there is a limit. Do we really want to make contact with meat?" 74 | 75 | "I agree one hundred percent. What's there to say?" `Hello, meat. How's it going?' But will this work? How many planets are we dealing with here?" 76 | 77 | "Just one. They can travel to other planets in special meat containers, but they can't live on them. And being meat, they only travel through C space. Which limits them to the speed of light and makes the possibility of their ever making contact pretty slim. Infinitesimal, in fact." 78 | 79 | "So we just pretend there's no one home in the universe." 80 | 81 | "That's it." 82 | 83 | "Cruel. But you said it yourself, who wants to meet meat? And the ones who have been aboard our vessels, the ones you have probed? You're sure they won't remember?" 84 | 85 | "They'll be considered crackpots if they do. We went into their heads and smoothed out their meat so that we're just a dream to them." 86 | 87 | "A dream to meat! How strangely appropriate, that we should be meat's dream." 88 | 89 | "And we can marked this sector unoccupied." 90 | 91 | "Good. Agreed, officially and unofficially. Case closed. Any others? Anyone interesting on that side of the galaxy?" 92 | 93 | "Yes, a rather shy but sweet hydrogen core cluster intelligence in a class nine star in G445 zone. Was in contact two galactic rotation ago, wants to be friendly again." 94 | 95 | "They always come around." 96 | 97 | "And why not? Imagine how unbearably, how unutterably cold the universe would be if one were all alone." 98 | -------------------------------------------------------------------------------- /benches/data/medium_size.txt: -------------------------------------------------------------------------------- 1 | Strange about learning; the farther I go the more I see that I never knew even existed. 2 | A short while ago I foolishly thought I could learn everything - all the knowledge in the world. 3 | Now I hope only to be able to know of its existence, and to understand one grain of it. 4 | Is there time? 5 | -------------------------------------------------------------------------------- /benches/data/small_size.txt: -------------------------------------------------------------------------------- 1 | In me, the tiger sniffs the rose 2 | -------------------------------------------------------------------------------- /benches/micro.rs: -------------------------------------------------------------------------------- 1 | //! Micro benches 2 | 3 | use std::iter::repeat; 4 | 5 | use bytes::Bytes; 6 | use criterion::async_executor::FuturesExecutor; 7 | use criterion::measurement::WallTime; 8 | use criterion::{ 9 | black_box, criterion_group, criterion_main, BatchSize, BenchmarkGroup, Criterion, Throughput, 10 | }; 11 | use raknet_rs::micro_bench; 12 | use raknet_rs::micro_bench::codec::BenchOpts; 13 | 14 | pub fn codec_benchmark(c: &mut Criterion) { 15 | let mut group = c.benchmark_group("codec"); 16 | 17 | fn decode( 18 | group: &mut BenchmarkGroup, 19 | datagram: &'static [u8], 20 | cnt: usize, 21 | throughput: impl Fn(&BenchOpts) -> Throughput, 22 | ) { 23 | let opts = micro_bench::codec::BenchOpts { 24 | datagrams: repeat(Bytes::from_static(datagram)).take(cnt).collect(), 25 | seed: 114514, 26 | dup_ratio: 0., 27 | shuffle_ratio: 0., 28 | mtu: 1480, 29 | }; 30 | group.throughput(throughput(&opts)); 31 | group.bench_function( 32 | format!("decode_cnt-{cnt}_size-{}", datagram.len()), 33 | |bencher| { 34 | bencher.to_async(FuturesExecutor).iter_batched( 35 | || black_box(opts.gen_inputs()), 36 | micro_bench::codec::run_bench, 37 | BatchSize::SmallInput, 38 | ); 39 | }, 40 | ); 41 | } 42 | 43 | let el = |opts: &BenchOpts| Throughput::Elements(opts.elements()); 44 | let by = |opts: &BenchOpts| Throughput::Bytes(opts.bytes()); 45 | 46 | let small_size = include_bytes!("data/small_size.txt"); 47 | let medium_size = include_bytes!("data/medium_size.txt"); 48 | let large_size = include_bytes!("data/large_size.txt"); 49 | 50 | decode(&mut group, small_size, 1, el); 51 | decode(&mut group, small_size, 5, el); 52 | decode(&mut group, small_size, 10, el); 53 | decode(&mut group, small_size, 50, el); 54 | decode(&mut group, small_size, 100, el); 55 | decode(&mut group, small_size, 1000, el); 56 | 57 | decode(&mut group, small_size, 1, by); 58 | decode(&mut group, small_size, 5, by); 59 | decode(&mut group, small_size, 10, by); 60 | decode(&mut group, small_size, 50, by); 61 | decode(&mut group, small_size, 100, by); 62 | decode(&mut group, small_size, 1000, by); 63 | 64 | decode(&mut group, medium_size, 1, el); 65 | decode(&mut group, medium_size, 5, el); 66 | decode(&mut group, medium_size, 10, el); 67 | decode(&mut group, medium_size, 50, el); 68 | decode(&mut group, medium_size, 100, el); 69 | decode(&mut group, medium_size, 1000, el); 70 | 71 | decode(&mut group, medium_size, 1, by); 72 | decode(&mut group, medium_size, 5, by); 73 | decode(&mut group, medium_size, 10, by); 74 | decode(&mut group, medium_size, 50, by); 75 | decode(&mut group, medium_size, 100, by); 76 | decode(&mut group, medium_size, 1000, by); 77 | 78 | decode(&mut group, large_size, 1, el); 79 | decode(&mut group, large_size, 5, el); 80 | decode(&mut group, large_size, 10, el); 81 | decode(&mut group, large_size, 50, el); 82 | decode(&mut group, large_size, 100, el); 83 | decode(&mut group, large_size, 1000, el); 84 | 85 | decode(&mut group, large_size, 1, by); 86 | decode(&mut group, large_size, 5, by); 87 | decode(&mut group, large_size, 10, by); 88 | decode(&mut group, large_size, 50, by); 89 | decode(&mut group, large_size, 100, by); 90 | decode(&mut group, large_size, 1000, by); 91 | 92 | group.finish(); 93 | } 94 | 95 | criterion_group!(benches, codec_benchmark); 96 | criterion_main!(benches); 97 | -------------------------------------------------------------------------------- /codecov.yml: -------------------------------------------------------------------------------- 1 | # codecov config 2 | # Reference: https://docs.codecov.com/docs/codecovyml-reference 3 | # Tips. You may run following command to validate before committing any changes 4 | # curl --data-binary @codecov.yml https://codecov.io/validate 5 | coverage: 6 | range: "70..95" 7 | precision: 2 8 | round: down 9 | status: 10 | patch: 11 | default: 12 | informational: true 13 | project: 14 | default: 15 | target: auto 16 | threshold: "1%" # allow drop by 1% 17 | changes: 18 | default: 19 | informational: true 20 | codecov: 21 | allow_coverage_offsets: true 22 | -------------------------------------------------------------------------------- /examples/controller.rs: -------------------------------------------------------------------------------- 1 | #![feature(local_waker)] 2 | #![feature(context_ext)] 3 | #![allow(clippy::print_stdout)] 4 | 5 | use std::error::Error; 6 | use std::net::SocketAddr; 7 | use std::pin::Pin; 8 | use std::sync::Arc; 9 | use std::task::ContextBuilder; 10 | use std::{cmp, io}; 11 | 12 | use bytes::Bytes; 13 | use concurrent_queue::ConcurrentQueue; 14 | use futures::future::poll_fn; 15 | use futures::{Sink, SinkExt, StreamExt}; 16 | use raknet_rs::client::{self, ConnectTo}; 17 | use raknet_rs::opts::FlushStrategy; 18 | use raknet_rs::server::MakeIncoming; 19 | use raknet_rs::{server, Message}; 20 | use tokio::net::UdpSocket; 21 | 22 | /// Self-balancing flush controller 23 | struct FlushController { 24 | write: Pin + Send + Sync + 'static>>, 25 | next_flush: Option, 26 | buffer: Arc>, 27 | delay: u64, // us 28 | } 29 | 30 | impl FlushController { 31 | fn new(write: Pin + Send + Sync + 'static>>) -> Self { 32 | Self { 33 | write, 34 | next_flush: None, 35 | buffer: Arc::new(ConcurrentQueue::unbounded()), 36 | delay: 5_000, // 5ms 37 | } 38 | } 39 | 40 | fn shared_sender(&self) -> Arc> { 41 | self.buffer.clone() 42 | } 43 | 44 | async fn wait(&self) { 45 | if let Some(next_flush) = self.next_flush { 46 | tokio::time::sleep_until(next_flush).await; 47 | } 48 | } 49 | 50 | async fn flush(&mut self) -> io::Result<()> { 51 | // Drain buffer 52 | while let Ok(msg) = self.buffer.pop() { 53 | self.write.feed(msg).await?; 54 | } 55 | 56 | // Flush 57 | let mut strategy = FlushStrategy::new(true, true, true); 58 | poll_fn(|cx| { 59 | let mut cx = ContextBuilder::from(cx).ext(&mut strategy).build(); 60 | self.write.as_mut().poll_flush(&mut cx) 61 | }) 62 | .await?; 63 | 64 | // A naive exponential backoff algorithm 65 | if strategy.flushed_ack() + strategy.flushed_nack() + strategy.flushed_pack() > 0 { 66 | self.delay = cmp::max(self.delay / 2, 5_000); 67 | } else { 68 | self.delay = cmp::min(self.delay * 2, 100_000); 69 | } 70 | self.next_flush = 71 | Some(tokio::time::Instant::now() + tokio::time::Duration::from_micros(self.delay)); 72 | Ok(()) 73 | } 74 | } 75 | 76 | #[tokio::main] 77 | async fn main() -> Result<(), Box> { 78 | env_logger::init(); 79 | 80 | let socket = UdpSocket::bind("127.0.0.1:0").await?; 81 | let local_addr = socket.local_addr()?; 82 | println!("[server] server listening on {local_addr} with flush controller"); 83 | let mut incoming = socket.make_incoming( 84 | server::Config::new() 85 | .sever_guid(114514) 86 | .advertisement("Hello, I am proxy server") 87 | .min_mtu(500) 88 | .max_mtu(1400) 89 | .support_version(vec![9, 11, 13]) 90 | .max_pending(64), 91 | ); 92 | 93 | tokio::spawn(async move { 94 | loop { 95 | let (src, dst) = incoming.next().await.unwrap(); 96 | tokio::spawn(async move { 97 | tokio::pin!(src); 98 | let mut controller = FlushController::new(Box::pin(dst)); 99 | { 100 | let sender1 = controller.shared_sender(); 101 | sender1 102 | .push(Message::new(Bytes::from_static(b"Welcome to the server"))) 103 | .unwrap(); 104 | } 105 | let sender2 = controller.shared_sender(); 106 | loop { 107 | tokio::select! { 108 | Some(data) = src.next() => { 109 | // echo back 110 | sender2.push(Message::new( 111 | data, 112 | )).unwrap(); 113 | } 114 | _ = controller.wait() => { 115 | controller.flush().await.unwrap(); 116 | } 117 | } 118 | } 119 | }); 120 | } 121 | }); 122 | client(local_addr).await?; 123 | Ok(()) 124 | } 125 | 126 | async fn client(addr: SocketAddr) -> Result<(), Box> { 127 | let socket = UdpSocket::bind("0.0.0.0:0").await?; 128 | let (src, dst) = socket 129 | .connect_to( 130 | addr, 131 | client::Config::new() 132 | .mtu(1000) 133 | .client_guid(1919810) 134 | .protocol_version(11), 135 | ) 136 | .await?; 137 | tokio::pin!(src); 138 | tokio::pin!(dst); 139 | dst.send(Message::new(Bytes::from_static(b"User pack1"))) 140 | .await?; 141 | dst.send(Message::new(Bytes::from_static(b"User pack2"))) 142 | .await?; 143 | let pack1 = src.next().await.unwrap(); 144 | let pack2 = src.next().await.unwrap(); 145 | let pack3 = src.next().await.unwrap(); 146 | assert_eq!(pack1, Bytes::from_static(b"Welcome to the server")); 147 | assert_eq!(pack2, Bytes::from_static(b"User pack1")); 148 | assert_eq!(pack3, Bytes::from_static(b"User pack2")); 149 | Ok(()) 150 | } 151 | -------------------------------------------------------------------------------- /examples/proxy.rs: -------------------------------------------------------------------------------- 1 | #![allow(clippy::print_stdout)] 2 | 3 | use std::error::Error; 4 | use std::net::SocketAddr; 5 | 6 | use bytes::Bytes; 7 | use futures::{SinkExt, StreamExt}; 8 | use raknet_rs::client::{self, ConnectTo}; 9 | use raknet_rs::server::{self, MakeIncoming}; 10 | use raknet_rs::Message; 11 | use tokio::net::UdpSocket; 12 | 13 | #[tokio::main] 14 | async fn main() -> Result<(), Box> { 15 | let socket = UdpSocket::bind("127.0.0.1:0").await?; 16 | let local_addr = socket.local_addr()?; 17 | println!("[server] proxy server listening on {local_addr}"); 18 | let mut incoming = socket.make_incoming( 19 | server::Config::new() 20 | .sever_guid(114514) 21 | .advertisement("Hello, I am proxy server") 22 | .min_mtu(500) 23 | .max_mtu(1400) 24 | .support_version(vec![9, 11, 13]) 25 | .max_pending(64), 26 | ); 27 | 28 | tokio::spawn(async move { 29 | loop { 30 | let (src, dst) = incoming.next().await.unwrap(); 31 | tokio::spawn(async move { 32 | tokio::pin!(src); 33 | tokio::pin!(dst); 34 | loop { 35 | if let Some(data) = src.next().await { 36 | println!( 37 | "[server] got proxy data: '{}'", 38 | String::from_utf8_lossy(&data) 39 | ); 40 | let client = reqwest::Client::new(); 41 | let res = client 42 | .post("http://httpbin.org/post") 43 | .body(data) 44 | .send() 45 | .await 46 | .unwrap(); 47 | dst.send(Message::new(res.bytes().await.unwrap())) 48 | .await 49 | .unwrap(); 50 | continue; 51 | } 52 | break; 53 | } 54 | }); 55 | } 56 | }); 57 | 58 | client(local_addr, "paopao").await?; 59 | client(local_addr, "yui").await?; 60 | Ok(()) 61 | } 62 | 63 | async fn client(addr: SocketAddr, name: &str) -> Result<(), Box> { 64 | let socket = UdpSocket::bind("0.0.0.0:0").await?; 65 | println!("[{name}] I am listening on {}", socket.local_addr()?); 66 | let (src, dst) = socket 67 | .connect_to( 68 | addr, 69 | client::Config::new() 70 | .mtu(1000) 71 | .client_guid(1919810) 72 | .protocol_version(11), 73 | ) 74 | .await?; 75 | tokio::pin!(src); 76 | tokio::pin!(dst); 77 | dst.send(Message::new(Bytes::from_static(b"Hello, Anyone there?"))) 78 | .await?; 79 | let res = src.next().await.unwrap(); 80 | println!( 81 | "[{name}] got server response: {}", 82 | String::from_utf8_lossy(&res) 83 | ); 84 | Ok(()) 85 | } 86 | -------------------------------------------------------------------------------- /examples/tracing.rs: -------------------------------------------------------------------------------- 1 | #![allow(clippy::print_stdout)] 2 | 3 | use std::collections::HashMap; 4 | use std::error::Error; 5 | use std::net::SocketAddr; 6 | use std::process::exit; 7 | use std::time::Duration; 8 | 9 | use bytes::Bytes; 10 | use fastrace::collector::{SpanContext, SpanId, SpanRecord, TraceId}; 11 | use fastrace::Span; 12 | use futures::{SinkExt, StreamExt}; 13 | use raknet_rs::client::{self, ConnectTo}; 14 | use raknet_rs::opts::TraceInfo; 15 | use raknet_rs::server::{self, MakeIncoming}; 16 | use raknet_rs::Message; 17 | use tokio::net::UdpSocket; 18 | 19 | #[tokio::main] 20 | async fn main() -> Result<(), Box> { 21 | let (reporter, spans) = fastrace::collector::TestReporter::new(); 22 | fastrace::set_reporter( 23 | reporter, 24 | fastrace::collector::Config::default().report_before_root_finish(true), 25 | ); 26 | 27 | let socket = UdpSocket::bind("127.0.0.1:0").await?; 28 | let local_addr = socket.local_addr()?; 29 | let mut incoming = socket.make_incoming( 30 | server::Config::new() 31 | .sever_guid(114514) 32 | .advertisement("Hello, I am proxy server") 33 | .min_mtu(500) 34 | .max_mtu(1400) 35 | .support_version(vec![9, 11, 13]) 36 | .max_pending(64), 37 | ); 38 | 39 | tokio::spawn(async move { 40 | loop { 41 | let (reader, writer) = incoming.next().await.unwrap(); 42 | tokio::spawn(async move { 43 | tokio::pin!(reader); 44 | tokio::pin!(writer); 45 | loop { 46 | if let Some(data) = reader.next().await { 47 | let trace_id = reader.last_trace_id().unwrap_or_else(|| { 48 | println!("Please run with `--features fastrace/enable` and try again"); 49 | exit(0) 50 | }); 51 | let root_span = Span::root( 52 | "user root span", 53 | SpanContext::new(trace_id, SpanId::default()), 54 | ); 55 | // do something with data 56 | tokio::time::sleep(Duration::from_millis(10)).await; 57 | let _span = Span::enter_with_parent("user child span", &root_span); 58 | writer.send(Message::new(data)).await.unwrap(); 59 | continue; 60 | } 61 | break; 62 | } 63 | }); 64 | } 65 | }); 66 | 67 | client(local_addr).await?; 68 | 69 | fastrace::flush(); 70 | display(spans.lock().clone()); 71 | Ok(()) 72 | } 73 | 74 | async fn client(addr: SocketAddr) -> Result<(), Box> { 75 | let socket = UdpSocket::bind("0.0.0.0:0").await?; 76 | let (src, dst) = socket 77 | .connect_to( 78 | addr, 79 | client::Config::new() 80 | .mtu(1000) 81 | .client_guid(1919810) 82 | .protocol_version(11), 83 | ) 84 | .await?; 85 | tokio::pin!(src); 86 | tokio::pin!(dst); 87 | dst.send(Message::new(Bytes::from_static(b"User pack1"))) 88 | .await?; 89 | dst.send(Message::new(Bytes::from_static(b"User pack2"))) 90 | .await?; 91 | let pack1 = src.next().await.unwrap(); 92 | let pack2 = src.next().await.unwrap(); 93 | assert_eq!(pack1, Bytes::from_static(b"User pack1")); 94 | assert_eq!(pack2, Bytes::from_static(b"User pack2")); 95 | Ok(()) 96 | } 97 | 98 | fn display(spans: Vec) { 99 | let spans_map: HashMap = spans 100 | .iter() 101 | .map(|span| (span.span_id, span.clone())) 102 | .collect(); 103 | let adjacency_lists: HashMap>> = spans.iter().fold( 104 | HashMap::new(), 105 | |mut map, 106 | SpanRecord { 107 | trace_id, 108 | span_id, 109 | parent_id, 110 | .. 111 | }| { 112 | map.entry(*trace_id) 113 | .or_default() 114 | .entry(*parent_id) 115 | .or_default() 116 | .push(*span_id); 117 | map 118 | }, 119 | ); 120 | fn dfs( 121 | adjacency_list: &HashMap>, 122 | spans: &HashMap, 123 | span_id: SpanId, 124 | depth: usize, 125 | last: bool, 126 | ) { 127 | let span = &spans[&span_id]; 128 | let mut properties = String::new(); 129 | for (key, value) in &span.properties { 130 | properties.push_str(&format!("{}: {}, ", key, value)); 131 | } 132 | let mut events = String::new(); 133 | for ev in &span.events { 134 | events.push_str(&format!("'{}'", ev.name)); 135 | } 136 | let prefix = if depth == 0 { 137 | String::new() 138 | } else if last { 139 | "╰".to_owned() + &"─".repeat(depth) + " " 140 | } else { 141 | "├".to_owned() + &"─".repeat(depth) + " " 142 | }; 143 | println!( 144 | "{}{}({}{{{}}}) [{}us]", 145 | prefix, 146 | span.name, 147 | properties, 148 | events, 149 | span.duration_ns as f64 / 1_000.0, 150 | ); 151 | if let Some(children) = adjacency_list.get(&span_id) { 152 | for (i, child) in children.iter().enumerate() { 153 | dfs( 154 | adjacency_list, 155 | spans, 156 | *child, 157 | depth + 1, 158 | i == children.len() - 1 && last, 159 | ); 160 | } 161 | } 162 | } 163 | for (trace_id, list) in adjacency_lists { 164 | if list.is_empty() { 165 | continue; 166 | } 167 | println!("trace_id: {}", trace_id.0); 168 | let l = &list[&SpanId::default()]; 169 | for (i, root) in l.iter().enumerate() { 170 | dfs(&list, &spans_map, *root, 0, i == l.len() - 1); 171 | } 172 | println!(); 173 | } 174 | } 175 | -------------------------------------------------------------------------------- /justfile: -------------------------------------------------------------------------------- 1 | default: 2 | just --list 3 | 4 | fix: 5 | cargo fmt --all 6 | cargo sort --workspace 7 | cargo machete --fix 8 | 9 | check: 10 | cargo fmt --all -- --check 11 | cargo sort --check 12 | cargo machete 13 | 14 | test: 15 | cargo nextest run --all-features --features fastrace/enable 16 | 17 | install-pre-commit: 18 | (! (grep -q 'just check' .git/hooks/pre-commit 2> /dev/null && echo 'You have installed.') && echo -n 'just check' >> .git/hooks/pre-commit && chmod +x .git/hooks/pre-commit && echo 'Installed!') || true 19 | 20 | remove-pre-commit: 21 | (! (sed -i.bak 's/just check//g' .git/hooks/pre-commit 2> /dev/null && echo 'Removed') && echo 'No pre-commit file found') || true 22 | -------------------------------------------------------------------------------- /rust-toolchain: -------------------------------------------------------------------------------- 1 | nightly-2025-01-01 2 | -------------------------------------------------------------------------------- /rustfmt.toml: -------------------------------------------------------------------------------- 1 | comment_width = 120 2 | format_code_in_doc_comments = true 3 | format_macro_bodies = true 4 | format_macro_matchers = true 5 | normalize_comments = true 6 | normalize_doc_attributes = true 7 | imports_granularity = "Module" 8 | group_imports = "StdExternalCrate" 9 | reorder_impl_items = true 10 | reorder_imports = true 11 | tab_spaces = 4 12 | wrap_comments = true 13 | -------------------------------------------------------------------------------- /src/client/conn/mod.rs: -------------------------------------------------------------------------------- 1 | use std::future::Future; 2 | use std::io; 3 | use std::net::ToSocketAddrs; 4 | use std::pin::Pin; 5 | use std::sync::Arc; 6 | 7 | use bytes::Bytes; 8 | use futures::{Sink, SinkExt, Stream, StreamExt}; 9 | use log::{error, trace}; 10 | 11 | use super::handler::offline; 12 | use super::handler::online::HandleOnline; 13 | use crate::codec::frame::Framed; 14 | use crate::codec::{AsyncSocket, Decoded, Encoded}; 15 | use crate::link::{Route, TransferLink}; 16 | use crate::opts::{ConnectionInfo, Ping, WrapConnectionInfo}; 17 | use crate::reliable::WrapReliable; 18 | use crate::state::{IncomingStateManage, OutgoingStateManage}; 19 | use crate::utils::Logged; 20 | use crate::{codec, Message, Role}; 21 | 22 | /// Connection implementation by using tokio's UDP framework 23 | #[cfg(feature = "tokio-rt")] 24 | mod tokio; 25 | 26 | #[derive(Debug, Clone, Copy)] 27 | #[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))] 28 | pub struct Config { 29 | /// The given mtu, the default value is 1400 30 | mtu: u16, 31 | /// The client guid, used to identify the client, initialized by random 32 | client_guid: u64, 33 | /// Raknet protocol version, default is 9 34 | protocol_version: u8, 35 | /// Limit the max size of a parted frames set, 0 means no limit 36 | /// It will abort the split frame if the `parted_size` reaches limit. 37 | /// The maximum number of inflight parted frames is `max_parted_size` * `max_parted_count` 38 | max_parted_size: u32, 39 | /// Limit the max count of **all** parted frames sets from an address. 40 | /// It might cause client resending frames if the limit is reached. 41 | /// The maximum number of inflight parted frames is `max_parted_size` * `max_parted_count` 42 | max_parted_count: usize, 43 | /// Maximum ordered channel, the value should be less than 256 44 | max_channels: usize, 45 | } 46 | 47 | impl Default for Config { 48 | fn default() -> Self { 49 | Self::new() 50 | } 51 | } 52 | 53 | impl Config { 54 | pub fn new() -> Self { 55 | Self { 56 | mtu: 1400, 57 | client_guid: rand::random(), 58 | protocol_version: 9, 59 | max_parted_size: 256, 60 | max_parted_count: 256, 61 | max_channels: 1, 62 | } 63 | } 64 | 65 | /// Give the mtu of the connection 66 | pub fn mtu(mut self, mtu: u16) -> Self { 67 | self.mtu = mtu; 68 | self 69 | } 70 | 71 | /// Set the client guid 72 | pub fn client_guid(mut self, client_guid: u64) -> Self { 73 | self.client_guid = client_guid; 74 | self 75 | } 76 | 77 | /// Set the protocol version 78 | pub fn protocol_version(mut self, protocol_version: u8) -> Self { 79 | self.protocol_version = protocol_version; 80 | self 81 | } 82 | 83 | /// Set the maximum parted size 84 | /// The default value is 256 85 | /// The maximum number of inflight parted frames is `max_parted_size`*`max_parted_count`nt 86 | pub fn max_parted_size(mut self, size: u32) -> Self { 87 | self.max_parted_size = size; 88 | self 89 | } 90 | 91 | /// Set the maximum parted count 92 | /// The default value is 256 93 | /// The maximum number of inflight parted frames is `max_parted_size`*`max_parted_count`nt 94 | pub fn max_parted_count(mut self, count: usize) -> Self { 95 | self.max_parted_count = count; 96 | self 97 | } 98 | 99 | /// Set the maximum channels 100 | /// The default value is 1 101 | /// The maximum value should be less than 256 102 | /// # Panics 103 | /// Panics if the channels is greater than 256 104 | pub fn max_channels(mut self, channels: usize) -> Self { 105 | assert!(channels < 256, "max_channels should be less than 256"); 106 | self.max_channels = channels; 107 | self 108 | } 109 | 110 | fn offline_config(&self) -> offline::Config { 111 | offline::Config { 112 | client_guid: self.client_guid, 113 | mtu: self.mtu, 114 | protocol_version: self.protocol_version, 115 | } 116 | } 117 | 118 | fn codec_config(&self) -> codec::Config { 119 | codec::Config { 120 | max_parted_count: self.max_parted_count, 121 | max_parted_size: self.max_parted_size, 122 | max_channels: self.max_channels, 123 | } 124 | } 125 | 126 | fn client_role(&self) -> Role { 127 | Role::Client { 128 | guid: self.client_guid, 129 | } 130 | } 131 | } 132 | 133 | pub trait ConnectTo: Sized { 134 | #[allow(async_fn_in_trait)] // No need to consider the auto trait for now. 135 | async fn connect_to( 136 | self, 137 | addrs: impl ToSocketAddrs, 138 | config: Config, 139 | ) -> io::Result<( 140 | impl Stream, 141 | impl Sink + Ping + ConnectionInfo, 142 | )>; 143 | } 144 | 145 | pub(crate) async fn connect_to( 146 | socket: impl AsyncSocket, 147 | addrs: impl ToSocketAddrs, 148 | config: super::Config, 149 | runtime: impl Fn(Pin + Send + Sync + 'static>>) -> H, 150 | ) -> io::Result<( 151 | impl Stream, 152 | impl Sink + Ping + ConnectionInfo, 153 | )> { 154 | let mut lookups = addrs.to_socket_addrs()?; 155 | let addr = lookups 156 | .next() 157 | .ok_or_else(|| io::Error::new(io::ErrorKind::AddrNotAvailable, "invalid address"))?; 158 | 159 | let (mut incoming, peer) = offline::OfflineHandler::new( 160 | Framed::new(socket.clone(), config.mtu as usize), // TODO: discover MTU 161 | addr, 162 | config.offline_config(), 163 | ) 164 | .await?; 165 | let role = config.client_role(); 166 | 167 | let link = TransferLink::new_arc(role, peer); 168 | let mut dst = Framed::new(socket.clone(), peer.mtu as usize) 169 | .wrap_reliable(Arc::clone(&link), peer, role) 170 | .frame_encoded(peer.mtu, config.codec_config(), Arc::clone(&link)) 171 | .manage_outgoing_state(None) 172 | .wrap_connection_info(peer); 173 | 174 | let (mut router, route) = Route::new(Arc::clone(&link)); 175 | 176 | runtime(Box::pin(async move { 177 | while let Some(pack) = incoming.next().await { 178 | // deliver the packet actively so that we do not miss ACK/NACK packets to advance 179 | // the outgoing state 180 | router.deliver(pack); 181 | } 182 | })); 183 | 184 | let src = route 185 | .frame_decoded(config.codec_config()) 186 | .logged( 187 | move |frame| trace!("[{role}] received {frame:?} from {peer}"), 188 | move |err| error!("[{role}] decode error: {err} from {peer}"), 189 | ) 190 | .manage_incoming_state() 191 | .handle_online(addr, config.client_guid, Arc::clone(&link)); 192 | 193 | // make all internal packets flushed 194 | SinkExt::::flush(&mut dst).await?; 195 | Ok((src, dst)) 196 | } 197 | -------------------------------------------------------------------------------- /src/client/conn/tokio.rs: -------------------------------------------------------------------------------- 1 | use std::io; 2 | use std::net::ToSocketAddrs; 3 | use std::sync::Arc; 4 | 5 | use bytes::Bytes; 6 | use futures::{Sink, Stream}; 7 | use tokio::net::UdpSocket as TokioUdpSocket; 8 | 9 | use super::ConnectTo; 10 | use crate::opts::{ConnectionInfo, Ping}; 11 | use crate::Message; 12 | 13 | impl ConnectTo for TokioUdpSocket { 14 | async fn connect_to( 15 | self, 16 | addrs: impl ToSocketAddrs, 17 | config: super::Config, 18 | ) -> io::Result<( 19 | impl Stream, 20 | impl Sink + Ping + ConnectionInfo, 21 | )> { 22 | let socket = Arc::new(self); 23 | super::connect_to(socket, addrs, config, tokio::spawn).await 24 | } 25 | } 26 | -------------------------------------------------------------------------------- /src/client/handler/mod.rs: -------------------------------------------------------------------------------- 1 | pub(crate) mod offline; 2 | pub(crate) mod online; 3 | -------------------------------------------------------------------------------- /src/client/handler/offline.rs: -------------------------------------------------------------------------------- 1 | use std::io; 2 | use std::net::SocketAddr; 3 | use std::pin::Pin; 4 | use std::task::{ready, Context, Poll}; 5 | 6 | use futures::{Future, Sink, SinkExt, Stream, StreamExt}; 7 | use pin_project_lite::pin_project; 8 | 9 | use crate::packet::connected::{self, FramesMut}; 10 | use crate::packet::{unconnected, Packet}; 11 | use crate::{Peer, Role}; 12 | 13 | #[derive(Debug, Clone, Copy)] 14 | pub(crate) struct Config { 15 | pub(crate) mtu: u16, 16 | pub(crate) client_guid: u64, 17 | pub(crate) protocol_version: u8, 18 | } 19 | 20 | pin_project! { 21 | pub(crate) struct OfflineHandler { 22 | frame: Option, 23 | state: State, 24 | server_addr: SocketAddr, 25 | config: Config, 26 | role: Role, 27 | } 28 | } 29 | 30 | impl OfflineHandler 31 | where 32 | F: Stream, SocketAddr)> 33 | + Sink<(unconnected::Packet, SocketAddr), Error = io::Error> 34 | + Unpin, 35 | { 36 | pub(crate) fn new(frame: F, server_addr: SocketAddr, config: Config) -> Self { 37 | Self { 38 | frame: Some(frame), 39 | state: State::SendOpenConnReq1(unconnected::Packet::OpenConnectionRequest1 { 40 | magic: (), 41 | protocol_version: config.protocol_version, 42 | mtu: config.mtu, 43 | }), 44 | server_addr, 45 | role: Role::Client { 46 | guid: config.client_guid, 47 | }, 48 | config, 49 | } 50 | } 51 | } 52 | 53 | enum State { 54 | SendOpenConnReq1(unconnected::Packet), 55 | SendOpenConnReq1Flush, 56 | WaitOpenConnReply1, 57 | SendOpenConnReq2(unconnected::Packet), 58 | SendOpenConnReq2Flush, 59 | WaitOpenConnReply2, 60 | } 61 | 62 | impl Future for OfflineHandler 63 | where 64 | F: Stream, SocketAddr)> 65 | + Sink<(unconnected::Packet, SocketAddr), Error = io::Error> 66 | + Unpin, 67 | { 68 | type Output = Result<(impl Stream>, Peer), io::Error>; 69 | 70 | fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { 71 | let this = self.project(); 72 | let frame = this.frame.as_mut().unwrap(); 73 | loop { 74 | match this.state { 75 | State::SendOpenConnReq1(pack) => { 76 | ready!(frame.poll_ready_unpin(cx))?; 77 | frame.start_send_unpin((pack.clone(), *this.server_addr))?; 78 | *this.state = State::SendOpenConnReq1Flush; 79 | } 80 | State::SendOpenConnReq1Flush => { 81 | ready!(frame.poll_flush_unpin(cx))?; 82 | *this.state = State::WaitOpenConnReply1; 83 | } 84 | State::WaitOpenConnReply1 => { 85 | // TODO: Add timeout 86 | let (pack, addr) = ready!(frame.poll_next_unpin(cx)).ok_or(io::Error::new( 87 | io::ErrorKind::ConnectionReset, 88 | "connection reset by peer", 89 | ))?; 90 | if addr != *this.server_addr { 91 | continue; 92 | } 93 | let next = match pack { 94 | Packet::Unconnected(unconnected::Packet::OpenConnectionReply1 { 95 | mtu, 96 | .. 97 | }) => unconnected::Packet::OpenConnectionRequest2 { 98 | magic: (), 99 | server_address: *this.server_addr, 100 | mtu, 101 | client_guid: this.config.client_guid, 102 | }, 103 | _ => continue, 104 | }; 105 | *this.state = State::SendOpenConnReq2(next); 106 | } 107 | State::SendOpenConnReq2(pack) => { 108 | ready!(frame.poll_ready_unpin(cx))?; 109 | frame.start_send_unpin((pack.clone(), *this.server_addr))?; 110 | *this.state = State::SendOpenConnReq2Flush; 111 | } 112 | State::SendOpenConnReq2Flush => { 113 | ready!(frame.poll_flush_unpin(cx))?; 114 | *this.state = State::WaitOpenConnReply2; 115 | } 116 | State::WaitOpenConnReply2 => { 117 | // TODO: Add timeout 118 | let (pack, addr) = ready!(frame.poll_next_unpin(cx)).ok_or(io::Error::new( 119 | io::ErrorKind::ConnectionReset, 120 | "connection reset by peer", 121 | ))?; 122 | if addr != *this.server_addr { 123 | continue; 124 | } 125 | match pack { 126 | Packet::Unconnected(unconnected::Packet::OpenConnectionReply2 { 127 | server_guid: guid, 128 | .. 129 | }) => { 130 | return Poll::Ready(Ok(( 131 | FilterConnected { 132 | frame: this.frame.take().unwrap(), 133 | server_addr: *this.server_addr, 134 | }, 135 | Peer { 136 | addr: *this.server_addr, 137 | mtu: this.config.mtu, 138 | guid, 139 | }, 140 | ))) 141 | } 142 | _ => continue, 143 | }; 144 | } 145 | } 146 | } 147 | } 148 | } 149 | 150 | pin_project! { 151 | struct FilterConnected { 152 | frame: F, 153 | server_addr: SocketAddr, 154 | } 155 | } 156 | 157 | impl Stream for FilterConnected 158 | where 159 | F: Stream, SocketAddr)> + Unpin, 160 | { 161 | type Item = connected::Packet; 162 | 163 | fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { 164 | let this = self.project(); 165 | loop { 166 | let Some((pack, addr)) = ready!(this.frame.poll_next_unpin(cx)) else { 167 | return Poll::Ready(None); 168 | }; 169 | if addr != *this.server_addr { 170 | continue; 171 | } 172 | match pack { 173 | Packet::Connected(pack) => return Poll::Ready(Some(pack)), 174 | _ => continue, 175 | }; 176 | } 177 | } 178 | } 179 | -------------------------------------------------------------------------------- /src/client/handler/online.rs: -------------------------------------------------------------------------------- 1 | use std::net::SocketAddr; 2 | use std::pin::Pin; 3 | use std::task::{ready, Context, Poll}; 4 | 5 | use bytes::Bytes; 6 | use futures::Stream; 7 | use log::debug; 8 | use pin_project_lite::pin_project; 9 | 10 | use crate::link::SharedLink; 11 | use crate::packet::connected::FrameBody; 12 | use crate::utils::timestamp; 13 | use crate::Role; 14 | 15 | pub(crate) trait HandleOnline: Sized { 16 | fn handle_online( 17 | self, 18 | addr: SocketAddr, 19 | client_guid: u64, 20 | link: SharedLink, 21 | ) -> OnlineHandler; 22 | } 23 | 24 | impl HandleOnline for F 25 | where 26 | F: Stream, 27 | { 28 | fn handle_online( 29 | self, 30 | addr: SocketAddr, 31 | client_guid: u64, 32 | link: SharedLink, 33 | ) -> OnlineHandler { 34 | link.send_frame_body(FrameBody::ConnectionRequest { 35 | client_guid, 36 | request_timestamp: timestamp(), 37 | use_encryption: false, 38 | }); 39 | OnlineHandler { 40 | frame: self, 41 | state: State::WaitConnRes, 42 | addr, 43 | link, 44 | role: Role::Client { guid: client_guid }, 45 | } 46 | } 47 | } 48 | 49 | pin_project! { 50 | pub(crate) struct OnlineHandler { 51 | #[pin] 52 | frame: F, 53 | state: State, 54 | addr: SocketAddr, 55 | link: SharedLink, 56 | role: Role, 57 | } 58 | } 59 | 60 | enum State { 61 | WaitConnRes, 62 | Connected, 63 | } 64 | 65 | impl Stream for OnlineHandler 66 | where 67 | F: Stream, 68 | { 69 | type Item = Bytes; 70 | 71 | fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { 72 | let mut this = self.project(); 73 | loop { 74 | match this.state { 75 | State::WaitConnRes => { 76 | let Some(body) = ready!(this.frame.as_mut().poll_next(cx)) else { 77 | return Poll::Ready(None); 78 | }; 79 | match body { 80 | FrameBody::ConnectionRequestAccepted { 81 | system_addresses, 82 | accepted_timestamp, 83 | .. 84 | } => { 85 | this.link.send_frame_body(FrameBody::NewIncomingConnection { 86 | server_address: *this.addr, 87 | system_addresses, 88 | request_timestamp: timestamp(), 89 | accepted_timestamp, 90 | }); 91 | *this.state = State::Connected; 92 | debug!( 93 | "[{}] connected to server {addr:?}", 94 | this.role, 95 | addr = this.addr 96 | ); 97 | continue; 98 | } 99 | FrameBody::User(data) => return Poll::Ready(Some(data)), /* FIXME: we should ignore this packet */ 100 | _ => { 101 | debug!("[{}] ignore packet {body:?} on WaitConnRes", this.role); 102 | } 103 | } 104 | } 105 | State::Connected => { 106 | let Some(body) = ready!(this.frame.as_mut().poll_next(cx)) else { 107 | return Poll::Ready(None); 108 | }; 109 | match body { 110 | FrameBody::DetectLostConnections => { 111 | this.link.send_frame_body(FrameBody::ConnectedPing { 112 | client_timestamp: timestamp(), 113 | }); 114 | } 115 | FrameBody::User(data) => return Poll::Ready(Some(data)), 116 | _ => { 117 | debug!("[{}] ignore packet {body:?} on Connected", this.role); 118 | } 119 | } 120 | } 121 | } 122 | } 123 | } 124 | } 125 | -------------------------------------------------------------------------------- /src/client/mod.rs: -------------------------------------------------------------------------------- 1 | //! This is a basic client implementation, without any retry, do NOT use it directly 2 | 3 | mod conn; 4 | mod handler; 5 | 6 | pub use self::conn::*; 7 | -------------------------------------------------------------------------------- /src/codec/decoder/body.rs: -------------------------------------------------------------------------------- 1 | use std::pin::Pin; 2 | use std::task::{ready, Context, Poll}; 3 | 4 | use fastrace::local::LocalSpan; 5 | use fastrace::Event; 6 | use futures::Stream; 7 | use pin_project_lite::pin_project; 8 | 9 | use crate::errors::CodecError; 10 | use crate::packet::connected::{Frame, FrameBody, FrameSet}; 11 | 12 | pin_project! { 13 | pub(crate) struct BodyDecoder { 14 | #[pin] 15 | frame: F 16 | } 17 | } 18 | 19 | pub(crate) trait BodyDecoded: Sized { 20 | fn body_decoded(self) -> BodyDecoder; 21 | } 22 | 23 | impl BodyDecoded for F 24 | where 25 | F: Stream, CodecError>>, 26 | { 27 | fn body_decoded(self) -> BodyDecoder { 28 | BodyDecoder { frame: self } 29 | } 30 | } 31 | 32 | impl Stream for BodyDecoder 33 | where 34 | F: Stream, CodecError>>, 35 | { 36 | type Item = Result; 37 | 38 | fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { 39 | let mut this = self.project(); 40 | 41 | let Some(frame_set) = ready!(this.frame.as_mut().poll_next(cx)?) else { 42 | return Poll::Ready(None); 43 | }; 44 | 45 | let span = LocalSpan::enter_with_local_parent("codec.body_decoder") 46 | .with_properties(|| [("seq_num", frame_set.seq_num.to_string())]); 47 | 48 | match FrameBody::read(frame_set.set.body) { 49 | Ok(body) => { 50 | let _ = span.with_property(|| ("frame_type", format!("{:?}", body))); 51 | Poll::Ready(Some(Ok(body))) 52 | } 53 | Err(err) => { 54 | LocalSpan::add_event(Event::new(err.to_string())); 55 | Poll::Ready(Some(Err(err))) 56 | } 57 | } 58 | } 59 | } 60 | -------------------------------------------------------------------------------- /src/codec/decoder/mod.rs: -------------------------------------------------------------------------------- 1 | mod body; 2 | mod dedup; 3 | mod fragment; 4 | mod ordered; 5 | 6 | use std::pin::Pin; 7 | use std::task::{Context, Poll}; 8 | 9 | use fastrace::Span; 10 | use futures::Stream; 11 | use pin_project_lite::pin_project; 12 | 13 | pub(super) use self::body::*; 14 | pub(super) use self::dedup::*; 15 | pub(super) use self::fragment::*; 16 | pub(super) use self::ordered::*; 17 | 18 | pin_project! { 19 | // Trace stream pending span 20 | pub(super) struct PendingTracer { 21 | #[pin] 22 | frame: F, 23 | span: Option, 24 | } 25 | } 26 | 27 | pub(super) trait TracePending: Sized { 28 | fn trace_pending(self) -> PendingTracer; 29 | } 30 | 31 | impl TracePending for F 32 | where 33 | F: Stream, 34 | { 35 | fn trace_pending(self) -> PendingTracer { 36 | PendingTracer { 37 | frame: self, 38 | span: None, 39 | } 40 | } 41 | } 42 | 43 | impl Stream for PendingTracer { 44 | type Item = F::Item; 45 | 46 | fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { 47 | let this = self.project(); 48 | this.span 49 | .get_or_insert_with(|| Span::enter_with_local_parent("codec.pending")); 50 | match this.frame.poll_next(cx) { 51 | r @ Poll::Ready(_) => { 52 | this.span.take(); 53 | r 54 | } 55 | p @ Poll::Pending => p, 56 | } 57 | } 58 | } 59 | -------------------------------------------------------------------------------- /src/codec/decoder/ordered.rs: -------------------------------------------------------------------------------- 1 | use std::pin::Pin; 2 | use std::task::{ready, Context, Poll}; 3 | 4 | use fastrace::local::LocalSpan; 5 | use fastrace::{Event, Span}; 6 | use futures::Stream; 7 | use log::warn; 8 | use pin_project_lite::pin_project; 9 | 10 | use crate::errors::CodecError; 11 | use crate::packet::connected::{self, Frame, FrameSet}; 12 | use crate::utils::u24; 13 | use crate::HashMap; 14 | 15 | struct Ordering { 16 | map: HashMap>, 17 | read: u24, 18 | } 19 | 20 | impl Default for Ordering { 21 | fn default() -> Self { 22 | Self { 23 | map: HashMap::default(), 24 | read: 0.into(), 25 | } 26 | } 27 | } 28 | 29 | pin_project! { 30 | // Ordering layer, ordered the packets based on ordering_frame_index. 31 | pub(crate) struct Order { 32 | #[pin] 33 | frame: F, 34 | // Max ordered channel that will be used in detailed protocol 35 | max_channels: usize, 36 | ordering: Vec, 37 | span: Option, 38 | } 39 | } 40 | 41 | pub(crate) trait Ordered: Sized { 42 | fn ordered(self, max_channels: usize) -> Order; 43 | } 44 | 45 | impl Ordered for F 46 | where 47 | F: Stream, CodecError>>, 48 | { 49 | fn ordered(self, max_channels: usize) -> Order { 50 | assert!( 51 | max_channels < 256, 52 | "max channels should not be larger than u8::MAX" 53 | ); 54 | assert!(max_channels > 0, "max_channels > 0"); 55 | 56 | Order { 57 | frame: self, 58 | max_channels, 59 | ordering: std::iter::repeat_with(Ordering::default) 60 | .take(max_channels) 61 | .collect(), 62 | span: None, 63 | } 64 | } 65 | } 66 | 67 | impl Stream for Order 68 | where 69 | F: Stream, CodecError>>, 70 | { 71 | type Item = Result, CodecError>; 72 | 73 | fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { 74 | let mut this = self.project(); 75 | loop { 76 | // empty each channel in order 77 | for channel in 0..*this.max_channels { 78 | let ordering = this 79 | .ordering 80 | .get_mut(channel) 81 | .expect("channel < max_channels"); 82 | // check if we could read next 83 | if let Some(next) = ordering.map.remove(&ordering.read) { 84 | ordering.read += 1; 85 | this.span.take(); 86 | return Poll::Ready(Some(Ok(next))); 87 | } 88 | } 89 | 90 | let Some(frame_set) = ready!(this.frame.as_mut().poll_next(cx)?) else { 91 | return Poll::Ready(None); 92 | }; 93 | this.span.get_or_insert_with(|| { 94 | Span::enter_with_local_parent("codec.reorder").with_properties(|| { 95 | [( 96 | "pending", 97 | this.ordering 98 | .iter() 99 | .fold(0, |acc, o| acc + o.map.len()) 100 | .to_string(), 101 | )] 102 | }) 103 | }); 104 | if let Some(connected::Ordered { 105 | frame_index, 106 | channel, 107 | }) = frame_set.set.ordered 108 | { 109 | let channel = usize::from(channel); 110 | if channel >= *this.max_channels { 111 | let err = format!("channel {} >= max_channels {}", channel, *this.max_channels); 112 | LocalSpan::add_event(Event::new(err.clone())); 113 | return Poll::Ready(Some(Err(CodecError::OrderedFrame(err)))); 114 | } 115 | let ordering = this 116 | .ordering 117 | .get_mut(channel) 118 | .expect("channel < max_channels"); 119 | if frame_index < ordering.read { 120 | warn!("ignore old ordered frame index {frame_index}"); 121 | continue; 122 | } 123 | ordering.map.insert(frame_index, frame_set); 124 | // we cannot read anymore 125 | continue; 126 | } 127 | // the frame set which does not require ordered 128 | this.span.take(); 129 | return Poll::Ready(Some(Ok(frame_set))); 130 | } 131 | } 132 | } 133 | 134 | #[cfg(test)] 135 | mod test { 136 | use bytes::Bytes; 137 | use futures::StreamExt; 138 | use futures_async_stream::stream; 139 | 140 | use super::Ordered; 141 | use crate::errors::CodecError; 142 | use crate::packet::connected::{Flags, Frame, FrameSet, Ordered as OrderedFlag}; 143 | 144 | fn frame_sets(idx: impl IntoIterator) -> Vec> { 145 | idx.into_iter() 146 | .map(|(channel, frame_index)| FrameSet { 147 | seq_num: 0.into(), 148 | set: Frame { 149 | flags: Flags::parse(0b011_11100), 150 | reliable_frame_index: None, 151 | seq_frame_index: None, 152 | ordered: Some(OrderedFlag { 153 | frame_index: frame_index.into(), 154 | channel, 155 | }), 156 | fragment: None, 157 | body: Bytes::new(), 158 | }, 159 | }) 160 | .collect() 161 | } 162 | 163 | #[tokio::test] 164 | async fn test_ordered_works() { 165 | let frame = { 166 | #[stream] 167 | async { 168 | for frame_set in 169 | frame_sets([(0, 1), (0, 0), (0, 2), (0, 0), (0, 4), (0, 3), (1, 1)]) 170 | { 171 | yield frame_set; 172 | } 173 | } 174 | }; 175 | tokio::pin!(frame); 176 | let mut ordered = frame.map(Ok).ordered(10); 177 | let cmp_sets = frame_sets([(0, 0), (0, 1), (0, 2), (0, 3), (0, 4)]).into_iter(); 178 | for next in cmp_sets { 179 | assert_eq!(ordered.next().await.unwrap().unwrap(), next); 180 | } 181 | 182 | assert!(ordered.next().await.is_none()); 183 | } 184 | 185 | #[tokio::test] 186 | async fn test_ordered_channel_exceed() { 187 | let frame = { 188 | #[stream] 189 | async { 190 | for frame_set in frame_sets([(10, 1)]) { 191 | yield frame_set; 192 | } 193 | } 194 | }; 195 | tokio::pin!(frame); 196 | let mut ordered = frame.map(Ok).ordered(10); 197 | assert!(matches!( 198 | ordered.next().await.unwrap().unwrap_err(), 199 | CodecError::OrderedFrame(_) 200 | )); 201 | } 202 | } 203 | -------------------------------------------------------------------------------- /src/codec/encoder/body.rs: -------------------------------------------------------------------------------- 1 | use std::io; 2 | use std::pin::Pin; 3 | use std::task::{ready, Context, Poll}; 4 | 5 | use bytes::BytesMut; 6 | use futures::Sink; 7 | use pin_project_lite::pin_project; 8 | 9 | use crate::link::SharedLink; 10 | use crate::packet::connected::FrameBody; 11 | use crate::{Message, Priority, Reliability}; 12 | 13 | pin_project! { 14 | // BodyEncoder encodes internal frame body into Message 15 | pub(crate) struct BodyEncoder { 16 | #[pin] 17 | frame: F, 18 | link: SharedLink, 19 | } 20 | } 21 | 22 | pub(crate) trait BodyEncoded: Sized { 23 | fn body_encoded(self, link: SharedLink) -> BodyEncoder; 24 | } 25 | 26 | impl BodyEncoded for F 27 | where 28 | F: Sink, 29 | { 30 | fn body_encoded(self, link: SharedLink) -> BodyEncoder { 31 | BodyEncoder { frame: self, link } 32 | } 33 | } 34 | 35 | #[inline(always)] 36 | fn encode(body: FrameBody) -> Message { 37 | const DEFAULT_FRAME_BODY_ORDERED_CHANNEL: u8 = 0; 38 | const DEFAULT_FRAME_BODY_PRIORITY: Priority = Priority::Medium; 39 | 40 | let reliability = match body { 41 | FrameBody::ConnectedPing { .. } => Reliability::Unreliable, 42 | FrameBody::ConnectedPong { .. } => Reliability::Unreliable, 43 | FrameBody::ConnectionRequest { .. } => Reliability::ReliableOrdered, 44 | FrameBody::ConnectionRequestAccepted { .. } => Reliability::Reliable, 45 | FrameBody::NewIncomingConnection { .. } => Reliability::ReliableOrdered, 46 | FrameBody::DisconnectNotification => Reliability::Reliable, 47 | FrameBody::DetectLostConnections => Reliability::Reliable, 48 | FrameBody::User(_) => { 49 | panic!("you should not send user packet into BodyEncoder, please send `Message`") 50 | } 51 | }; 52 | let mut data = BytesMut::new(); 53 | body.write(&mut data); 54 | Message::new(data.freeze()) 55 | .reliability(reliability) 56 | .order_channel(DEFAULT_FRAME_BODY_ORDERED_CHANNEL) 57 | .priority(DEFAULT_FRAME_BODY_PRIORITY) 58 | } 59 | 60 | impl BodyEncoder 61 | where 62 | F: Sink, 63 | { 64 | /// Empty the link buffer all the frame body, insure the frame is ready to send 65 | pub(crate) fn poll_empty( 66 | self: Pin<&mut Self>, 67 | cx: &mut Context<'_>, 68 | ) -> Poll> { 69 | if self.link.frame_body_empty() { 70 | return Poll::Ready(Ok(())); 71 | } 72 | let mut this = self.project(); 73 | for body in this.link.process_frame_body() { 74 | ready!(this.frame.as_mut().poll_ready(cx))?; 75 | this.frame.as_mut().start_send(encode(body))?; 76 | } 77 | Poll::Ready(Ok(())) 78 | } 79 | } 80 | 81 | impl Sink for BodyEncoder 82 | where 83 | F: Sink, 84 | { 85 | type Error = io::Error; 86 | 87 | fn poll_ready(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { 88 | Sink::::poll_ready(self, cx) 89 | } 90 | 91 | fn start_send(self: Pin<&mut Self>, item: Message) -> Result<(), Self::Error> { 92 | // skip encode 93 | self.project().frame.start_send(item) 94 | } 95 | 96 | fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { 97 | Sink::::poll_flush(self, cx) 98 | } 99 | 100 | fn poll_close(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { 101 | Sink::::poll_close(self, cx) 102 | } 103 | } 104 | 105 | impl Sink for BodyEncoder 106 | where 107 | F: Sink, 108 | { 109 | type Error = io::Error; 110 | 111 | fn poll_ready(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { 112 | ready!(self.project().frame.poll_ready(cx))?; 113 | Poll::Ready(Ok(())) 114 | } 115 | 116 | fn start_send(self: Pin<&mut Self>, body: FrameBody) -> Result<(), Self::Error> { 117 | self.project().frame.start_send(encode(body)) 118 | } 119 | 120 | fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { 121 | ready!(self.as_mut().poll_empty(cx))?; 122 | self.project().frame.poll_flush(cx) 123 | } 124 | 125 | fn poll_close(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { 126 | ready!(self.as_mut().poll_empty(cx))?; 127 | self.project().frame.poll_close(cx) 128 | } 129 | } 130 | -------------------------------------------------------------------------------- /src/codec/encoder/mod.rs: -------------------------------------------------------------------------------- 1 | mod body; 2 | mod fragment; 3 | 4 | pub(super) use self::body::*; 5 | pub(super) use self::fragment::*; 6 | -------------------------------------------------------------------------------- /src/codec/frame.rs: -------------------------------------------------------------------------------- 1 | use std::io; 2 | use std::net::SocketAddr; 3 | use std::pin::Pin; 4 | use std::task::{ready, Context, Poll}; 5 | 6 | use bytes::BytesMut; 7 | use fastrace::local::LocalSpan; 8 | use fastrace::{Event, Span}; 9 | use futures::{Sink, Stream}; 10 | use log::error; 11 | 12 | use super::AsyncSocket; 13 | use crate::packet::connected::{FramesMut, FramesRef}; 14 | use crate::packet::{unconnected, Packet}; 15 | 16 | /// `Framed` is a base structure for socket communication. 17 | /// In this project, it wraps an asynchronous UDP socket and implements the 18 | /// [`Stream`](futures::stream::Stream) and [`Sink`](futures::sink::Sink) traits. 19 | /// This allows for reading and writing `RakNet` frames over the socket. 20 | /// It supports both receiving and sending unconnected and connected packets. 21 | pub(crate) struct Framed { 22 | /// The underlying socket that is being wrapped. 23 | socket: T, 24 | /// [maximum transmission unit](https://en.wikipedia.org/wiki/Maximum_transmission_unit), 25 | /// used to pre-allocate the capacity for both the `rd` and `wr` byte buffers 26 | mtu: usize, 27 | /// a buffer for storing incoming data read from the socket. 28 | rd: BytesMut, 29 | /// a buffer for writing bytes 30 | wr: BytesMut, 31 | /// the socket address to which the data will be sent 32 | out_addr: SocketAddr, 33 | /// indicates whether the data has been sent to the peer, and the 34 | /// `wr` buffer has been cleared 35 | flushed: bool, 36 | /// indicates whether data has been read from the peer and written 37 | /// into the `rd` buffer. When `true`, it signifies that the data is ready for frame packet 38 | /// decoding and will be passed to the upper layer for further processing. 39 | is_readable: bool, 40 | /// the address of the current peer 41 | current_addr: Option, 42 | decode_span: Option, 43 | read_span: Option, 44 | } 45 | 46 | impl Framed { 47 | pub(crate) fn new(socket: T, mtu: usize) -> Self { 48 | Self { 49 | socket, 50 | mtu, 51 | rd: BytesMut::new(), // we may never use rd at all 52 | wr: BytesMut::with_capacity(mtu), 53 | out_addr: SocketAddr::from(([0, 0, 0, 0], 0)), 54 | flushed: true, 55 | is_readable: false, 56 | current_addr: None, 57 | decode_span: None, 58 | read_span: None, 59 | } 60 | } 61 | 62 | #[inline] 63 | fn poll_flush_0(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { 64 | if self.flushed { 65 | return Poll::Ready(Ok(())); 66 | } 67 | 68 | let Self { 69 | ref socket, 70 | ref mut out_addr, 71 | ref mut wr, 72 | .. 73 | } = *self; 74 | 75 | let n = ready!(socket.poll_send_to(cx, wr, *out_addr))?; 76 | 77 | let wrote_all = n == self.wr.len(); 78 | self.wr.clear(); 79 | self.flushed = true; 80 | 81 | let res = if wrote_all { 82 | Ok(()) 83 | } else { 84 | Err(io::Error::new( 85 | io::ErrorKind::Other, 86 | "failed to write entire datagram to socket", 87 | )) 88 | }; 89 | 90 | Poll::Ready(res) 91 | } 92 | } 93 | 94 | impl Stream for Framed { 95 | type Item = (Packet, SocketAddr); 96 | 97 | fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { 98 | let pin = self.get_mut(); 99 | 100 | pin.rd.reserve(pin.mtu); 101 | 102 | loop { 103 | // Are there still bytes left in the read buffer to decode? 104 | if pin.is_readable { 105 | pin.is_readable = false; 106 | 107 | // decode one packet at a time 108 | match Packet::read(&mut pin.rd) { 109 | Ok(frame) => { 110 | let current_addr = pin 111 | .current_addr 112 | .expect("will always be set before this line is called"); 113 | LocalSpan::add_event(Event::new(format!( 114 | "{:?} decoded", 115 | frame.pack_type() 116 | ))); 117 | pin.decode_span.take(); 118 | pin.rd.clear(); 119 | 120 | return Poll::Ready(Some((frame, current_addr))); 121 | } 122 | Err(err) => { 123 | error!("failed to decode packet: {}", err); 124 | LocalSpan::add_event(Event::new(err.to_string())); 125 | pin.decode_span.take(); 126 | pin.rd.clear(); 127 | } 128 | } 129 | } 130 | 131 | // We're out of data. Try and fetch more data to decode 132 | pin.read_span 133 | .get_or_insert_with(|| Span::enter_with_local_parent("codec.frame.read")); 134 | let addr = match ready!(pin.socket.poll_recv_from(cx, &mut pin.rd)) { 135 | Ok(addr) => addr, 136 | Err(err) => { 137 | error!("failed to receive data: {:?}", err); 138 | LocalSpan::add_event(Event::new(err.to_string())); 139 | pin.rd.clear(); 140 | continue; 141 | } 142 | }; 143 | // finish the read span 144 | pin.read_span.take(); 145 | // start a new decode span 146 | pin.decode_span.get_or_insert_with(|| { 147 | Span::enter_with_local_parent("codec.frame.decode").with_properties(|| { 148 | [ 149 | ("addr", addr.to_string()), 150 | ("datagram_size", pin.rd.len().to_string()), 151 | ] 152 | }) 153 | }); 154 | pin.current_addr = Some(addr); 155 | pin.is_readable = true; 156 | } 157 | } 158 | } 159 | 160 | /// The `Sink` implementation for cheap buffer cloning (i.e. `bytes::Bytes`). 161 | impl<'a, T: AsyncSocket> Sink<(Packet>, SocketAddr)> for Framed { 162 | type Error = io::Error; 163 | 164 | fn poll_ready(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { 165 | self.poll_flush_0(cx) 166 | } 167 | 168 | fn start_send( 169 | self: Pin<&mut Self>, 170 | item: (Packet>, SocketAddr), 171 | ) -> Result<(), Self::Error> { 172 | let (frame, out_addr) = item; 173 | 174 | let pin = self.get_mut(); 175 | 176 | frame.write(&mut pin.wr); 177 | pin.out_addr = out_addr; 178 | pin.flushed = false; 179 | 180 | Ok(()) 181 | } 182 | 183 | fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { 184 | self.poll_flush_0(cx) 185 | } 186 | 187 | fn poll_close(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { 188 | self.poll_flush_0(cx) 189 | } 190 | } 191 | 192 | // Separate the unconnected packets for offline and online states. 193 | // So the offline handler could take account of the unconnected packets and ignore the generic from 194 | // the connected packets. 195 | impl Sink<(unconnected::Packet, SocketAddr)> for Framed { 196 | type Error = io::Error; 197 | 198 | fn poll_ready(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { 199 | self.poll_flush_0(cx) 200 | } 201 | 202 | fn start_send( 203 | self: Pin<&mut Self>, 204 | item: (unconnected::Packet, SocketAddr), 205 | ) -> Result<(), Self::Error> { 206 | let (frame, out_addr) = item; 207 | 208 | let pin = self.get_mut(); 209 | 210 | frame.write(&mut pin.wr); 211 | pin.out_addr = out_addr; 212 | pin.flushed = false; 213 | 214 | Ok(()) 215 | } 216 | 217 | fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { 218 | self.poll_flush_0(cx) 219 | } 220 | 221 | fn poll_close(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { 222 | self.poll_flush_0(cx) 223 | } 224 | } 225 | -------------------------------------------------------------------------------- /src/codec/tokio.rs: -------------------------------------------------------------------------------- 1 | use std::io; 2 | use std::mem::MaybeUninit; 3 | use std::net::SocketAddr; 4 | use std::sync::Arc; 5 | use std::task::{ready, Context, Poll}; 6 | 7 | use bytes::{BufMut, BytesMut}; 8 | use tokio::net::UdpSocket as TokioUdpSocket; 9 | 10 | use super::AsyncSocket; 11 | 12 | impl AsyncSocket for Arc { 13 | fn poll_recv_from( 14 | &self, 15 | cx: &mut Context<'_>, 16 | rd: &mut BytesMut, 17 | ) -> Poll> { 18 | // Safety: `chunk_mut()` returns a `&mut UninitSlice`, and `UninitSlice` is a 19 | // transparent wrapper around `[MaybeUninit]`. 20 | let buf = unsafe { &mut *(rd.chunk_mut() as *mut _ as *mut [MaybeUninit]) }; 21 | let mut read = tokio::io::ReadBuf::uninit(buf); 22 | let ptr = read.filled().as_ptr(); 23 | let res = ready!(self.as_ref().poll_recv_from(cx, &mut read)); 24 | 25 | assert_eq!(ptr, read.filled().as_ptr()); 26 | let addr = res?; 27 | 28 | // Safety: This is guaranteed to be the number of initialized (and read) bytes due 29 | // to the invariants provided by `ReadBuf::filled`. 30 | unsafe { rd.advance_mut(read.filled().len()) }; 31 | 32 | Poll::Ready(Ok(addr)) 33 | } 34 | 35 | fn poll_send_to( 36 | &self, 37 | cx: &mut Context<'_>, 38 | buf: &[u8], 39 | target: SocketAddr, 40 | ) -> Poll> { 41 | self.as_ref().poll_send_to(cx, buf, target) 42 | } 43 | } 44 | -------------------------------------------------------------------------------- /src/congestion/bbr.rs: -------------------------------------------------------------------------------- 1 | 2 | -------------------------------------------------------------------------------- /src/congestion/cubic.rs: -------------------------------------------------------------------------------- 1 | 2 | -------------------------------------------------------------------------------- /src/congestion/legacy.rs: -------------------------------------------------------------------------------- 1 | pub(crate) struct LegacyCongester { 2 | state: State, 3 | cwnd: f32, 4 | ssthresh: f32, 5 | } 6 | 7 | enum State { 8 | SlowStart, 9 | CongestionAvoidance, 10 | } 11 | 12 | impl LegacyCongester { 13 | const INITIAL_CWND: f32 = 2.0; 14 | 15 | pub(crate) fn new() -> LegacyCongester { 16 | LegacyCongester { 17 | state: State::SlowStart, 18 | cwnd: Self::INITIAL_CWND, 19 | ssthresh: f32::INFINITY, 20 | } 21 | } 22 | } 23 | 24 | impl super::CongestionController for LegacyCongester { 25 | fn congestion_window(&self) -> usize { 26 | self.cwnd as usize 27 | } 28 | 29 | fn on_ack(&mut self, cnt: usize) { 30 | match self.state { 31 | State::SlowStart => { 32 | self.cwnd += cnt as f32; 33 | if self.cwnd >= self.ssthresh { 34 | self.state = State::CongestionAvoidance; 35 | } 36 | } 37 | State::CongestionAvoidance => { 38 | self.cwnd += cnt as f32 / self.cwnd; 39 | } 40 | } 41 | } 42 | 43 | fn on_nack(&mut self, cnt: usize) { 44 | // A simple implementation of fast recovery 45 | // Because we can only receive new ack packets later, which will trigger to quit 46 | // fast recovery into congestion avoidance. 47 | self.state = State::CongestionAvoidance; 48 | self.ssthresh = self.cwnd / 2.0; 49 | self.cwnd = self.ssthresh + cnt as f32; 50 | } 51 | 52 | fn on_timeout(&mut self) { 53 | self.state = State::SlowStart; 54 | self.ssthresh = self.cwnd / 2.0; 55 | self.cwnd = Self::INITIAL_CWND; 56 | } 57 | } 58 | -------------------------------------------------------------------------------- /src/congestion/mod.rs: -------------------------------------------------------------------------------- 1 | /// Legacy congestion control algorithms 2 | pub(crate) mod legacy; 3 | 4 | pub(crate) mod cubic; 5 | 6 | pub(crate) mod bbr; 7 | 8 | pub(crate) trait CongestionController: Send + Sync + 'static { 9 | /// Returns the size of the current congestion window in frames 10 | fn congestion_window(&self) -> usize; 11 | 12 | fn on_ack(&mut self, cnt: usize); 13 | 14 | fn on_nack(&mut self, cnt: usize); 15 | 16 | fn on_timeout(&mut self); 17 | } 18 | -------------------------------------------------------------------------------- /src/errors.rs: -------------------------------------------------------------------------------- 1 | #[derive(thiserror::Error, Debug)] 2 | pub(crate) enum CodecError { 3 | #[error("io error {0}")] 4 | IO(#[from] std::io::Error), 5 | #[error("invalid ip version {0}")] 6 | InvalidIPVer(u8), 7 | #[error("expect IPv6 family 0x17, got {0}")] 8 | InvalidIPV6Family(u16), 9 | #[error("invalid packet length when decode {0}")] 10 | InvalidPacketLength(&'static str), 11 | #[error("invalid record type {0}")] 12 | InvalidRecordType(u8), 13 | #[error("invalid packet type {0}, maybe it is a user packet")] 14 | InvalidPacketType(u8), 15 | #[error("parted frame error, reason: {0}")] 16 | PartedFrame(String), 17 | #[error("ordered frame error, reason: {0}")] 18 | OrderedFrame(String), 19 | #[error("ack count {0} exceeded")] 20 | AckCountExceed(usize), 21 | #[error("magic number not matched, pos {0}, byte {1}")] 22 | MagicNotMatched(usize, u8), 23 | } 24 | -------------------------------------------------------------------------------- /src/estimator.rs: -------------------------------------------------------------------------------- 1 | use std::cmp; 2 | use std::time::Duration; 3 | 4 | pub(crate) trait Estimator: Send + Sync + 'static { 5 | /// The current RTO estimation without penalty. 6 | fn rto(&self) -> Duration; 7 | 8 | /// The current RTT estimation. 9 | fn rtt(&self) -> Duration; 10 | 11 | /// Update the RTT estimator with a new RTT sample. 12 | fn update(&mut self, rtt: Duration); 13 | 14 | /// Clear the estimator's state. 15 | fn clear(&mut self); 16 | } 17 | 18 | /// RTT estimation based on RFC6298 19 | #[derive(Copy, Clone)] 20 | pub struct RFC6298Impl { 21 | /// The most recent RTT measurement made when receiving an ack for a previously unacked packet 22 | latest: Duration, 23 | /// The smoothed RTT of the connection, computed as described in RFC6298 24 | smoothed: Option, 25 | /// The RTT variance, computed as described in RFC6298 26 | var: Duration, 27 | } 28 | 29 | impl RFC6298Impl { 30 | pub(crate) fn new() -> Self { 31 | // Initial RTO value as suggested in RFC6298 2.1 will be 1s. 32 | // We will use 50ms as the initial value for the latest RTT as well. 33 | // https://www.rfc-editor.org/rfc/rfc6298.html 34 | 35 | const INITIAL_RTT: Duration = Duration::from_millis(500); 36 | Self { 37 | latest: INITIAL_RTT, 38 | smoothed: None, 39 | var: Duration::from_secs(0), 40 | } 41 | } 42 | 43 | /// The current best RTT estimation. 44 | fn rtt(&self) -> Duration { 45 | self.smoothed.unwrap_or(self.latest) 46 | } 47 | 48 | /// The current RTO estimation. 49 | pub(crate) fn rto(&self) -> Duration { 50 | // The granularity of the timer 51 | // https://www.rfc-editor.org/rfc/rfc9002#section-6.1.2 52 | // # The RECOMMENDED value of the timer granularity is 1 millisecond. 53 | const TIMER_GRANULARITY: Duration = Duration::from_millis(1); 54 | 55 | // RFC6298 2.4 suggests a minimum of 1 second, which may be 56 | // a conservative choice. 57 | // We disregard the minimum of 1 second for now 58 | self.rtt() + cmp::max(TIMER_GRANULARITY, 4 * self.var) 59 | } 60 | 61 | /// Once smoothed and var are cleared, they should be initialized with the next RTT sample 62 | pub(crate) fn clear(&mut self) { 63 | self.smoothed = None; 64 | } 65 | 66 | pub(crate) fn update(&mut self, rtt: Duration) { 67 | self.latest = rtt; 68 | if let Some(smoothed) = self.smoothed { 69 | let var_sample = if smoothed > rtt { 70 | smoothed - rtt 71 | } else { 72 | rtt - smoothed 73 | }; 74 | self.var = (3 * self.var + var_sample) / 4; 75 | self.smoothed = Some((7 * smoothed + rtt) / 8); 76 | } else { 77 | self.smoothed = Some(rtt); 78 | self.var = rtt / 2; 79 | } 80 | } 81 | } 82 | 83 | impl Estimator for RFC6298Impl { 84 | fn rto(&self) -> Duration { 85 | self.rto() 86 | } 87 | 88 | fn rtt(&self) -> Duration { 89 | self.rtt() 90 | } 91 | 92 | fn update(&mut self, rtt: Duration) { 93 | self.update(rtt); 94 | } 95 | 96 | fn clear(&mut self) { 97 | self.clear(); 98 | } 99 | } 100 | -------------------------------------------------------------------------------- /src/lib.rs: -------------------------------------------------------------------------------- 1 | //! Raknet implementation by rust 2 | 3 | #![feature(impl_trait_in_assoc_type)] 4 | #![feature(coroutines, proc_macro_hygiene, stmt_expr_attributes)] 5 | #![feature(binary_heap_into_iter_sorted)] 6 | #![feature(let_chains)] 7 | #![feature(context_ext)] 8 | #![feature(local_waker)] 9 | 10 | /// Protocol codec 11 | mod codec; 12 | 13 | /// Errors 14 | mod errors; 15 | 16 | /// Protocol packet 17 | mod packet; 18 | 19 | /// Utils 20 | mod utils; 21 | 22 | /// Reliable layer 23 | mod reliable; 24 | 25 | /// Sink & Stream state 26 | mod state; 27 | 28 | /// Transfer link 29 | mod link; 30 | 31 | /// Estimators 32 | mod estimator; 33 | 34 | /// Congestion control 35 | mod congestion; 36 | 37 | /// Raknet server 38 | pub mod server; 39 | 40 | /// Raknet client 41 | pub mod client; 42 | 43 | /// Connection optional settings 44 | pub mod opts; 45 | 46 | #[cfg(feature = "micro-bench")] 47 | pub mod micro_bench { 48 | pub mod codec { 49 | pub use crate::codec::micro_bench::*; 50 | } 51 | } 52 | 53 | #[cfg(feature = "rustc-hash")] 54 | pub type HashMap = std::collections::HashMap; 55 | 56 | #[cfg(not(feature = "rustc-hash"))] 57 | pub type HashMap = std::collections::HashMap; 58 | 59 | /// Unit tests 60 | #[cfg(test)] 61 | mod tests; 62 | 63 | use std::net::SocketAddr; 64 | 65 | use bytes::Bytes; 66 | 67 | /// The `Role` enum is used to identify the `Client` and `Server`, and it stores their GUID. 68 | /// The GUID is a globally unique identifier that is not affected by changes to IP address or port. 69 | /// It is application-defined and ensures unique identification. 70 | #[derive(Debug, Clone, Copy)] 71 | enum Role { 72 | Client { guid: u64 }, 73 | Server { guid: u64 }, 74 | } 75 | 76 | impl Role { 77 | #[cfg(test)] 78 | fn test_server() -> Self { 79 | Role::Server { guid: 114514 } 80 | } 81 | 82 | fn guid(&self) -> u64 { 83 | match self { 84 | Role::Client { guid } => *guid, 85 | Role::Server { guid } => *guid, 86 | } 87 | } 88 | } 89 | 90 | impl std::fmt::Display for Role { 91 | fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { 92 | match self { 93 | Role::Client { guid } => write!(f, "client({guid})"), 94 | Role::Server { guid } => write!(f, "server({guid})"), 95 | } 96 | } 97 | } 98 | 99 | #[derive(Debug, Clone, Copy)] 100 | struct Peer { 101 | guid: u64, 102 | addr: SocketAddr, 103 | mtu: u16, 104 | } 105 | 106 | impl Peer { 107 | #[cfg(test)] 108 | fn test() -> Self { 109 | Self { 110 | guid: 114514, 111 | addr: SocketAddr::from(([11, 45, 14, 19], 19810)), 112 | mtu: 1919, 113 | } 114 | } 115 | } 116 | 117 | impl std::fmt::Display for Peer { 118 | fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { 119 | write!(f, "{guid}@{addr}", guid = self.guid, addr = self.addr,) 120 | } 121 | } 122 | 123 | #[derive(Debug, Clone, Eq, PartialEq, Copy)] 124 | #[repr(u8)] 125 | pub enum Reliability { 126 | /// Same as regular UDP, except that it will also discard duplicate datagrams. `RakNet` adds 127 | /// (6 to 17) + 21 bits of overhead, 16 of which is used to detect duplicate packets and 6 128 | /// to 17 of which is used for message length. 129 | Unreliable = 0b000, 130 | 131 | /// Regular UDP with a sequence counter. Out of order messages will be discarded. 132 | /// Sequenced and ordered messages sent on the same channel will arrive in the order sent. 133 | UnreliableSequenced = 0b001, 134 | 135 | /// The message is sent reliably, but not necessarily in any order. Same overhead as 136 | /// UNRELIABLE. 137 | Reliable = 0b010, 138 | 139 | /// This message is reliable and will arrive in the order you sent it. Messages will be 140 | /// delayed while waiting for out of order messages. Same overhead as `UnreliableSequenced`. 141 | /// Sequenced and ordered messages sent on the same channel will arrive in the order sent. 142 | ReliableOrdered = 0b011, 143 | 144 | /// This message is reliable and will arrive in the sequence you sent it. Out of order 145 | /// messages will be dropped. Same overhead as `UnreliableSequenced`. Sequenced and ordered 146 | /// messages sent on the same channel will arrive in the order sent. 147 | ReliableSequenced = 0b100, 148 | 149 | /// Same as Unreliable, however the peer will get either ACK or 150 | /// NACK based on the result of sending this message when calling. 151 | UnreliableWithAckReceipt = 0b101, 152 | 153 | /// Same as Reliable, however the peer will get either ACK or 154 | /// NACK based on the result of sending this message when calling. 155 | ReliableWithAckReceipt = 0b110, 156 | 157 | /// Same as `ReliableOrdered`, however the peer will get either ACK or 158 | /// NACK based on the result of sending this message when calling. 159 | ReliableOrderedWithAckReceipt = 0b111, 160 | } 161 | 162 | impl Reliability { 163 | /// Reliable ensures that the packet is not duplicated. 164 | pub(crate) fn is_reliable(&self) -> bool { 165 | matches!( 166 | self, 167 | Reliability::Reliable 168 | | Reliability::ReliableSequenced 169 | | Reliability::ReliableOrdered 170 | | Reliability::ReliableWithAckReceipt 171 | | Reliability::ReliableOrderedWithAckReceipt 172 | ) 173 | } 174 | 175 | /// Sequenced or Ordered ensures that packets should be received in order at their 176 | /// `order_channels` as they are sent. 177 | pub(crate) fn is_sequenced_or_ordered(&self) -> bool { 178 | matches!( 179 | self, 180 | Reliability::ReliableSequenced 181 | | Reliability::ReliableOrdered 182 | | Reliability::UnreliableSequenced 183 | | Reliability::ReliableOrderedWithAckReceipt 184 | ) 185 | } 186 | 187 | /// TODO: implement sequenced 188 | pub(crate) fn is_sequenced(&self) -> bool { 189 | matches!( 190 | self, 191 | Reliability::UnreliableSequenced | Reliability::ReliableSequenced 192 | ) 193 | } 194 | 195 | /// The header size (without fragment part) implied from reliability 196 | pub(crate) fn size(&self) -> usize { 197 | // flag(1B) + length(2B) 198 | let mut size = 3; 199 | if self.is_reliable() { 200 | size += 3; 201 | } 202 | if self.is_sequenced() { 203 | size += 3; 204 | } 205 | if self.is_sequenced_or_ordered() { 206 | size += 4; 207 | } 208 | size 209 | } 210 | } 211 | 212 | // Message priority, messages with higher priority will be transmitted first. 213 | // 214 | // There are three tiers: High, Medium, Low 215 | // The High and Low tiers come with a numeric penalty. 216 | // The lowest priority one is Low(255), and the highest priority one is High(0). 217 | // 218 | // Low(255) ~ Low(0) | Medium | High(255) ~ High(0) 219 | // >>>-------- priority from low to high ------->>> 220 | // 221 | // The message will be sent at a Medium tier by default. 222 | #[derive(Debug, Clone, Copy, Eq, PartialEq)] 223 | pub enum Priority { 224 | High(u8), 225 | Medium, 226 | Low(u8), 227 | } 228 | 229 | impl Default for Priority { 230 | fn default() -> Self { 231 | Self::Medium 232 | } 233 | } 234 | 235 | /// Raknet message 236 | #[derive(Debug, Clone)] 237 | pub struct Message { 238 | pub reliability: Reliability, 239 | pub order_channel: u8, 240 | pub priority: Priority, 241 | pub data: Bytes, 242 | } 243 | 244 | impl Message { 245 | pub fn new(data: Bytes) -> Self { 246 | Self { 247 | reliability: Reliability::ReliableOrdered, 248 | order_channel: 0, 249 | priority: Priority::default(), 250 | data, 251 | } 252 | } 253 | 254 | pub fn reliability(mut self, reliability: Reliability) -> Self { 255 | self.reliability = reliability; 256 | self 257 | } 258 | 259 | pub fn order_channel(mut self, channel: u8) -> Self { 260 | self.order_channel = channel; 261 | self 262 | } 263 | 264 | pub fn priority(mut self, priority: Priority) -> Self { 265 | self.priority = priority; 266 | self 267 | } 268 | } 269 | -------------------------------------------------------------------------------- /src/opts.rs: -------------------------------------------------------------------------------- 1 | use std::future::Future; 2 | use std::io; 3 | use std::net::SocketAddr; 4 | use std::pin::Pin; 5 | use std::task::{Context, Poll}; 6 | 7 | use fastrace::collector::TraceId; 8 | use futures::{Sink, SinkExt}; 9 | use pin_project_lite::pin_project; 10 | 11 | use crate::link::SharedLink; 12 | use crate::packet::connected::FrameBody; 13 | use crate::reliable::SendBuffer; 14 | use crate::utils::timestamp; 15 | use crate::{Message, Peer}; 16 | 17 | /// Trace info extension for server 18 | pub trait TraceInfo { 19 | fn last_trace_id(&self) -> Option; 20 | } 21 | 22 | /// Obtain the connection information 23 | pub trait ConnectionInfo { 24 | fn mtu(&self) -> u16; 25 | fn remote_addr(&self) -> SocketAddr; 26 | fn guid(&self) -> u64; 27 | } 28 | 29 | pub(crate) trait WrapConnectionInfo: Sized { 30 | fn wrap_connection_info(self, peer: Peer) -> ConnectionInfoWrapper; 31 | } 32 | 33 | pin_project! { 34 | pub(crate) struct ConnectionInfoWrapper { 35 | #[pin] 36 | inner: I, 37 | peer: Peer, 38 | } 39 | } 40 | 41 | impl ConnectionInfo for ConnectionInfoWrapper { 42 | fn mtu(&self) -> u16 { 43 | self.peer.mtu 44 | } 45 | 46 | fn remote_addr(&self) -> SocketAddr { 47 | self.peer.addr 48 | } 49 | 50 | fn guid(&self) -> u64 { 51 | self.peer.guid 52 | } 53 | } 54 | 55 | impl> Sink for ConnectionInfoWrapper { 56 | type Error = I::Error; 57 | 58 | fn poll_ready(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { 59 | self.project().inner.poll_ready(cx) 60 | } 61 | 62 | fn start_send(self: Pin<&mut Self>, item: T) -> Result<(), Self::Error> { 63 | self.project().inner.start_send(item) 64 | } 65 | 66 | fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { 67 | self.project().inner.poll_flush(cx) 68 | } 69 | 70 | fn poll_close(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { 71 | self.project().inner.poll_close(cx) 72 | } 73 | } 74 | 75 | impl> WrapConnectionInfo for S { 76 | fn wrap_connection_info(self, peer: Peer) -> ConnectionInfoWrapper { 77 | ConnectionInfoWrapper { inner: self, peer } 78 | } 79 | } 80 | 81 | /// Ping extension for client, experimental 82 | pub trait Ping { 83 | fn ping(self: Pin<&mut Self>) -> impl Future> + Send; 84 | } 85 | 86 | impl Ping for S 87 | where 88 | S: Sink + Send, 89 | { 90 | async fn ping(mut self: Pin<&mut Self>) -> Result<(), io::Error> { 91 | self.send(FrameBody::ConnectedPing { 92 | client_timestamp: timestamp(), 93 | }) 94 | .await 95 | } 96 | } 97 | 98 | /// Flush strategy can be used as ext data of [`std::task::Context`] to guide how 99 | /// [`Sink::poll_flush`] perform flush. And the results after flush will be stored here. 100 | /// The default strategy will flush all buffers. 101 | /// 102 | /// Customizing your own strategy can achieve many features: 103 | /// 104 | /// 1. [**Delayed ack**](https://en.wikipedia.org/wiki/TCP_delayed_acknowledgment) based on timing, 105 | /// thereby reducing the number of ack packets and improving bandwidth utilization. At the same 106 | /// time, sending based on timing can avoid deadlocks or regressions caused by delaying based on 107 | /// the number of packets. 108 | /// 109 | /// 2. More aggressive nack/pack flush strategy which would be more beneficial for retransmitting 110 | /// packets. 111 | /// 112 | /// After the flush is completed, the strategy will store the number of frames that have been 113 | /// flushed. You can use this number to determine when to take the next flush. 114 | /// 115 | /// Note that it can only be used in [`Sink::poll_flush`]. 116 | #[derive(Debug, Default, Clone, Copy)] 117 | pub struct FlushStrategy { 118 | ack_tag: isize, 119 | nack_tag: isize, 120 | pack_tag: isize, 121 | } 122 | 123 | impl FlushStrategy { 124 | /// Create a new flush strategy with specified flush options. 125 | pub fn new(ack: bool, nack: bool, pack: bool) -> Self { 126 | FlushStrategy { 127 | ack_tag: if ack { 0 } else { -1 }, 128 | nack_tag: if nack { 0 } else { -1 }, 129 | pack_tag: if pack { 0 } else { -1 }, 130 | } 131 | } 132 | 133 | /// Get how many ack frames have been flushed. 134 | /// 135 | /// # Panics 136 | /// It will panic if ack flush is not enabled. 137 | pub fn flushed_ack(&self) -> usize { 138 | assert!( 139 | self.ack_tag != -1, 140 | "you should enable flush ack before checking result of flushed ack" 141 | ); 142 | self.ack_tag as usize 143 | } 144 | 145 | /// Get how many nack frames have been flushed. 146 | /// 147 | /// # Panics 148 | /// It will panic if nack flush is not enabled. 149 | pub fn flushed_nack(&self) -> usize { 150 | assert!( 151 | self.nack_tag != -1, 152 | "you should enable flush nack before checking result of flushed nack" 153 | ); 154 | self.nack_tag as usize 155 | } 156 | 157 | /// Get how many pack frames have been flushed. 158 | /// 159 | /// # Panics 160 | /// It will panic if pack flush is not enabled. 161 | pub fn flushed_pack(&self) -> usize { 162 | assert!( 163 | self.pack_tag != -1, 164 | "you should enable flush pack before checking result of flushed pack" 165 | ); 166 | self.pack_tag as usize 167 | } 168 | 169 | pub(crate) fn check_flushed(&self, link: &SharedLink, buf: &SendBuffer) -> bool { 170 | let mut ret = true; 171 | if self.ack_tag != -1 { 172 | ret &= link.outgoing_ack_empty(); 173 | } 174 | if self.nack_tag != -1 { 175 | ret &= link.outgoing_nack_empty(); 176 | } 177 | if self.pack_tag != -1 { 178 | ret &= link.unconnected_empty() && buf.is_empty(); 179 | } 180 | ret 181 | } 182 | 183 | pub(crate) fn flush_ack(&self) -> bool { 184 | self.ack_tag != -1 185 | } 186 | 187 | pub(crate) fn flush_nack(&self) -> bool { 188 | self.nack_tag != -1 189 | } 190 | 191 | pub(crate) fn flush_pack(&self) -> bool { 192 | self.pack_tag != -1 193 | } 194 | 195 | pub(crate) fn mark_flushed_ack(&mut self, cnt: usize) { 196 | if self.ack_tag == -1 { 197 | return; 198 | } 199 | self.ack_tag += cnt as isize; 200 | } 201 | 202 | pub(crate) fn mark_flushed_nack(&mut self, cnt: usize) { 203 | if self.nack_tag == -1 { 204 | return; 205 | } 206 | self.nack_tag += cnt as isize; 207 | } 208 | 209 | pub(crate) fn mark_flushed_pack(&mut self, cnt: usize) { 210 | if self.pack_tag == -1 { 211 | return; 212 | } 213 | self.pack_tag += cnt as isize; 214 | } 215 | } 216 | -------------------------------------------------------------------------------- /src/packet/connected/ack.rs: -------------------------------------------------------------------------------- 1 | use bytes::{Buf, BufMut, BytesMut}; 2 | 3 | use crate::errors::CodecError; 4 | use crate::packet::read_buf; 5 | use crate::utils::{u24, BufExt, BufMutExt}; 6 | 7 | #[derive(Clone)] 8 | #[cfg_attr(test, derive(PartialEq, Eq))] 9 | pub(crate) struct AckOrNack { 10 | pub(crate) records: Vec, 11 | } 12 | 13 | impl std::fmt::Debug for AckOrNack { 14 | fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { 15 | let mut records = String::new(); 16 | for record in &self.records { 17 | match record { 18 | Record::Range(start, end) => { 19 | records += &format!("{}-{},", start, end); 20 | } 21 | Record::Single(idx) => { 22 | records += &format!("{},", idx); 23 | } 24 | } 25 | } 26 | if !records.is_empty() { 27 | records.pop(); 28 | } 29 | write!(f, "AckOrNack({})", records) 30 | } 31 | } 32 | 33 | impl AckOrNack { 34 | /// This function implements **Sequence Number Compression** on an `Iterator`. Consecutive 35 | /// sequence numbers are grouped into `Record::Range`, while non-consecutive sequence 36 | /// numbers are stored as `Record::Single`. This approach reduces the data payload size. 37 | /// 38 | /// - A `Record::Range` consumes 7 bytes. 39 | /// - A `Record::Single` consumes 4 bytes. 40 | pub(crate) fn extend_from>( 41 | mut sorted_seq_nums: I, 42 | mut mtu: u16, 43 | ) -> Option { 44 | // pack_type(1) + length(2) + single record(4) = 7 45 | debug_assert!(mtu >= 7, "7 is the least size of mtu"); 46 | 47 | let mut first = sorted_seq_nums.next()?; 48 | 49 | let mut records = vec![]; 50 | let mut last = first; 51 | let mut upgrade_flag = true; 52 | // first byte is pack_type, next 2 bytes are length, the first seq_num takes at least 4 53 | // bytes 54 | mtu -= 7; 55 | loop { 56 | // we cannot poll sorted_seq_nums because 4 is the least size of a record 57 | if mtu < 4 { 58 | break; 59 | } 60 | let Some(seq_num) = sorted_seq_nums.next() else { 61 | break; 62 | }; 63 | if seq_num == last + 1 { 64 | if upgrade_flag { 65 | mtu -= 3; 66 | upgrade_flag = false; 67 | } 68 | last = seq_num; 69 | continue; 70 | } 71 | mtu -= 4; 72 | upgrade_flag = true; 73 | if first != last { 74 | records.push(Record::Range(first, last)); 75 | } else { 76 | records.push(Record::Single(first)); 77 | } 78 | first = seq_num; 79 | last = seq_num; 80 | } 81 | 82 | if first != last { 83 | records.push(Record::Range(first, last)); 84 | } else { 85 | records.push(Record::Single(first)); 86 | } 87 | 88 | Some(Self { records }) 89 | } 90 | 91 | pub(super) fn read(buf: &mut BytesMut) -> Result { 92 | const MAX_ACKNOWLEDGEMENT_PACKETS: usize = 16384; 93 | 94 | let mut ack_cnt = 0; 95 | let record_cnt = buf.get_u16(); 96 | let mut records = Vec::with_capacity(record_cnt as usize); 97 | for _ in 0..record_cnt { 98 | let record = Record::read(buf)?; 99 | ack_cnt += record.ack_cnt(); 100 | if ack_cnt > MAX_ACKNOWLEDGEMENT_PACKETS { 101 | return Err(CodecError::AckCountExceed(ack_cnt)); 102 | } 103 | records.push(record); 104 | } 105 | Ok(Self { records }) 106 | } 107 | 108 | pub(super) fn write(self, buf: &mut BytesMut) { 109 | debug_assert!( 110 | self.records.len() < u16::MAX as usize, 111 | "self.records should be constructed based on mtu" 112 | ); 113 | buf.put_u16(self.records.len() as u16); 114 | for record in self.records { 115 | record.write(buf); 116 | } 117 | } 118 | 119 | pub(crate) fn total_cnt(&self) -> usize { 120 | self.records.iter().map(|record| record.ack_cnt()).sum() 121 | } 122 | } 123 | 124 | const RECORD_RANGE: u8 = 0; 125 | const RECORD_SINGLE: u8 = 1; 126 | 127 | #[derive(Clone)] 128 | #[cfg_attr(test, derive(PartialEq, Eq))] 129 | pub(crate) enum Record { 130 | Range(u24, u24), 131 | Single(u24), 132 | } 133 | 134 | impl Record { 135 | fn read(buf: &mut BytesMut) -> Result { 136 | let record_type = read_buf!(buf, 1, buf.get_u8()); 137 | match record_type { 138 | RECORD_RANGE => read_buf!( 139 | buf, 140 | 6, 141 | Ok(Record::Range(buf.get_u24_le(), buf.get_u24_le())) 142 | ), 143 | RECORD_SINGLE => read_buf!(buf, 3, Ok(Record::Single(buf.get_u24_le()))), 144 | _ => Err(CodecError::InvalidRecordType(record_type)), 145 | } 146 | } 147 | 148 | fn write(self, buf: &mut BytesMut) { 149 | match self { 150 | Record::Range(start, end) => { 151 | buf.put_u8(RECORD_RANGE); 152 | buf.put_u24_le(start); 153 | buf.put_u24_le(end); 154 | } 155 | Record::Single(idx) => { 156 | buf.put_u8(RECORD_SINGLE); 157 | buf.put_u24_le(idx); 158 | } 159 | } 160 | } 161 | 162 | fn ack_cnt(&self) -> usize { 163 | match self { 164 | Record::Range(start, end) => usize::from(end - start + 1), 165 | Record::Single(_) => 1, 166 | } 167 | } 168 | } 169 | 170 | #[cfg(test)] 171 | mod test { 172 | use super::*; 173 | 174 | #[test] 175 | fn test_ack_should_not_overflow_mtu() { 176 | let mtu: u16 = 21; 177 | let mut buf = BytesMut::with_capacity(mtu as usize); 178 | 179 | let test_cases = [ 180 | // 3 + 0-2(7) + 4-5(7) + 7(4) = 21, remain 8 181 | (vec![0, 1, 2, 4, 5, 7, 8], 21, 1), 182 | // 3 + 0-1(7) + 3-4(7) + 6(4) = 21, remain 7, 9 183 | (vec![0, 1, 3, 4, 6, 7, 9], 21, 2), 184 | // 3 + 0(4) + 2(4) + 4(4) + 6(4) = 19, remain 8, 10, 12 185 | (vec![0, 2, 4, 6, 8, 10, 12], 19, 3), 186 | // 3 + 0(4) + 2(4) + 5-6(7) = 18, remain 8, 9, 12 187 | (vec![0, 2, 5, 6, 8, 9, 12], 18, 3), 188 | // 3 + 0-1(7) = 10, no remain 189 | (vec![0, 1], 10, 0), 190 | // 3 + 0(4) + 2-3(7) = 14, no remain 191 | (vec![0, 2, 3], 14, 0), 192 | // 3 + 0(4) + 2(4) + 4(4) = 15, no remain 193 | (vec![0, 2, 4], 15, 0), 194 | // 3 + 1(4) + 1(4) + 1(4) = 15, no remain 195 | (vec![1, 1, 1], 15, 0), 196 | // 3 + 0-999(7) = 10, no remain 197 | (Vec::from_iter(0..1000), 10, 0), 198 | ]; 199 | for (seq_nums, len, remain) in test_cases { 200 | buf.clear(); 201 | // pack type 202 | buf.put_u8(0); 203 | let mut seq_nums = seq_nums.into_iter().map(u24::from); 204 | let ack = AckOrNack::extend_from(&mut seq_nums, mtu).unwrap(); 205 | ack.write(&mut buf); 206 | assert_eq!(buf.len(), len); 207 | assert_eq!(seq_nums.len(), remain); 208 | } 209 | } 210 | } 211 | -------------------------------------------------------------------------------- /src/packet/connected/mod.rs: -------------------------------------------------------------------------------- 1 | use bytes::{BufMut, BytesMut}; 2 | 3 | use crate::errors::CodecError; 4 | use crate::packet::PackType; 5 | 6 | mod ack; 7 | mod frame_set; 8 | 9 | pub(crate) use ack::*; 10 | pub(crate) use frame_set::*; 11 | 12 | use super::{ACK_FLAG, CONTINUOUS_SEND_FLAG, NACK_FLAG, NEEDS_B_AND_AS_FLAG, VALID_FLAG}; 13 | 14 | // Packet when RakNet has established a connection 15 | #[derive(Debug, Clone)] 16 | #[cfg_attr(test, derive(PartialEq, Eq))] 17 | pub(crate) enum Packet { 18 | FrameSet(FrameSet), 19 | Ack(AckOrNack), 20 | Nack(AckOrNack), 21 | } 22 | 23 | impl Packet { 24 | pub(crate) fn pack_type(&self) -> PackType { 25 | match self { 26 | Packet::FrameSet(_) => PackType::FrameSet, 27 | Packet::Ack(_) => PackType::Ack, 28 | Packet::Nack(_) => PackType::Nack, 29 | } 30 | } 31 | 32 | pub(super) fn read_ack(buf: &mut BytesMut) -> Result { 33 | Ok(Packet::Ack(AckOrNack::read(buf)?)) 34 | } 35 | 36 | pub(super) fn read_nack(buf: &mut BytesMut) -> Result { 37 | Ok(Packet::Nack(AckOrNack::read(buf)?)) 38 | } 39 | } 40 | 41 | /// For cheap buffers cloning (i.e. `bytes::Bytes`) 42 | impl Packet> { 43 | pub(crate) fn write(self, buf: &mut BytesMut) { 44 | match self { 45 | Packet::FrameSet(frame) => { 46 | let mut flag = VALID_FLAG | NEEDS_B_AND_AS_FLAG; 47 | if frame.set[0].flags.parted { 48 | flag |= CONTINUOUS_SEND_FLAG; 49 | } 50 | buf.put_u8(flag); 51 | frame.write(buf); 52 | } 53 | Packet::Ack(ack) => { 54 | buf.put_u8(ACK_FLAG); 55 | ack.write(buf); 56 | } 57 | Packet::Nack(ack) => { 58 | buf.put_u8(NACK_FLAG); 59 | ack.write(buf); 60 | } 61 | } 62 | } 63 | } 64 | 65 | impl Packet { 66 | pub(super) fn read_frame_set(buf: &mut BytesMut) -> Result { 67 | Ok(Packet::FrameSet(FrameSet::read(buf)?)) 68 | } 69 | } 70 | -------------------------------------------------------------------------------- /src/server/handler/mod.rs: -------------------------------------------------------------------------------- 1 | pub(crate) mod offline; 2 | pub(crate) mod online; 3 | -------------------------------------------------------------------------------- /src/server/handler/online.rs: -------------------------------------------------------------------------------- 1 | use std::net::{Ipv4Addr, Ipv6Addr, SocketAddr}; 2 | use std::pin::Pin; 3 | use std::task::{ready, Context, Poll}; 4 | 5 | use bytes::Bytes; 6 | use futures::Stream; 7 | use log::{debug, trace}; 8 | use pin_project_lite::pin_project; 9 | 10 | use crate::link::SharedLink; 11 | use crate::packet::connected::FrameBody; 12 | use crate::packet::unconnected; 13 | use crate::utils::timestamp; 14 | use crate::{Peer, Role}; 15 | 16 | pub(crate) trait HandleOnline: Sized { 17 | fn handle_online(self, role: Role, peer: Peer, link: SharedLink) -> OnlineHandler; 18 | } 19 | 20 | impl HandleOnline for F 21 | where 22 | F: Stream, 23 | { 24 | fn handle_online(self, role: Role, peer: Peer, link: SharedLink) -> OnlineHandler { 25 | OnlineHandler { 26 | frame: self, 27 | role, 28 | peer, 29 | state: HandshakeState::WaitConnRequest, 30 | link, 31 | } 32 | } 33 | } 34 | 35 | pin_project! { 36 | pub(crate) struct OnlineHandler { 37 | #[pin] 38 | frame: F, 39 | role: Role, 40 | peer: Peer, 41 | state: HandshakeState, 42 | link: SharedLink, 43 | } 44 | } 45 | 46 | enum HandshakeState { 47 | WaitConnRequest, 48 | WaitNewIncomingConn, 49 | Connected, 50 | } 51 | 52 | impl Stream for OnlineHandler 53 | where 54 | F: Stream, 55 | { 56 | type Item = Bytes; 57 | 58 | fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { 59 | let mut this = self.project(); 60 | loop { 61 | match this.state { 62 | HandshakeState::WaitConnRequest => { 63 | let Some(body) = ready!(this.frame.as_mut().poll_next(cx)) else { 64 | return Poll::Ready(None); 65 | }; 66 | if let FrameBody::ConnectionRequest { 67 | request_timestamp, 68 | use_encryption, 69 | .. 70 | } = body 71 | { 72 | if use_encryption { 73 | this.link.send_unconnected( 74 | unconnected::Packet::ConnectionRequestFailed { 75 | magic: (), 76 | server_guid: this.role.guid(), 77 | }, 78 | ); 79 | continue; 80 | } 81 | let system_addr = if this.peer.addr.is_ipv6() { 82 | SocketAddr::new( 83 | std::net::IpAddr::V6(Ipv6Addr::new(0, 0, 0, 0, 0, 0, 0, 0)), 84 | 0, 85 | ) 86 | } else { 87 | SocketAddr::new(std::net::IpAddr::V4(Ipv4Addr::new(0, 0, 0, 0)), 0) 88 | }; 89 | this.link 90 | .send_frame_body(FrameBody::ConnectionRequestAccepted { 91 | client_address: this.peer.addr, 92 | system_index: 0, 93 | system_addresses: [system_addr; 20], 94 | request_timestamp, 95 | accepted_timestamp: timestamp(), 96 | }); 97 | *this.state = HandshakeState::WaitNewIncomingConn; 98 | continue; 99 | } 100 | debug!("[{}] ignore packet {body:?} on WaitConnRequest", this.role); 101 | } 102 | HandshakeState::WaitNewIncomingConn => { 103 | let Some(body) = ready!(this.frame.as_mut().poll_next(cx)) else { 104 | return Poll::Ready(None); 105 | }; 106 | if let FrameBody::NewIncomingConnection { .. } = body { 107 | debug!("[{}] accept new incoming connection", this.role); 108 | *this.state = HandshakeState::Connected; 109 | continue; 110 | } 111 | // FIXME: it is wrong, client should finish handshake before sending user data 112 | // currently, the client's handshake is lazy, so user data may be sent before 113 | match body { 114 | FrameBody::User(data) => return Poll::Ready(Some(data)), 115 | _ => { 116 | debug!( 117 | "[{}] ignore packet {body:?} on WaitNewIncomingConn", 118 | this.role 119 | ); 120 | } 121 | } 122 | } 123 | HandshakeState::Connected => { 124 | let Some(body) = ready!(this.frame.as_mut().poll_next(cx)) else { 125 | return Poll::Ready(None); 126 | }; 127 | match body { 128 | FrameBody::ConnectedPing { client_timestamp } => { 129 | trace!( 130 | "[{}] receive ConnectedPing from {}, send ConnectedPong back", 131 | this.role, 132 | this.peer 133 | ); 134 | this.link.send_frame_body(FrameBody::ConnectedPong { 135 | client_timestamp, 136 | server_timestamp: timestamp(), 137 | }); 138 | } 139 | FrameBody::User(data) => return Poll::Ready(Some(data)), 140 | _ => { 141 | debug!("[{}] ignore packet {body:?} on Connected", this.role); 142 | } 143 | } 144 | } 145 | } 146 | } 147 | } 148 | } 149 | -------------------------------------------------------------------------------- /src/server/incoming/tokio.rs: -------------------------------------------------------------------------------- 1 | use std::io; 2 | use std::sync::Arc; 3 | 4 | use bytes::Bytes; 5 | use futures::{Sink, Stream}; 6 | use tokio::net::UdpSocket as TokioUdpSocket; 7 | 8 | use super::{Config, Incoming, MakeIncoming}; 9 | use crate::opts::{ConnectionInfo, TraceInfo}; 10 | use crate::Message; 11 | 12 | impl MakeIncoming for TokioUdpSocket { 13 | fn make_incoming( 14 | self, 15 | config: Config, 16 | ) -> impl Stream< 17 | Item = ( 18 | impl Stream + TraceInfo, 19 | impl Sink + ConnectionInfo, 20 | ), 21 | > { 22 | let socket = Arc::new(self); 23 | Incoming::new(socket, config) 24 | } 25 | } 26 | 27 | impl MakeIncoming for Arc { 28 | fn make_incoming( 29 | self, 30 | config: Config, 31 | ) -> impl Stream< 32 | Item = ( 33 | impl Stream + TraceInfo, 34 | impl Sink + ConnectionInfo, 35 | ), 36 | > { 37 | Incoming::new(self, config) 38 | } 39 | } 40 | -------------------------------------------------------------------------------- /src/server/mod.rs: -------------------------------------------------------------------------------- 1 | mod handler; 2 | mod incoming; 3 | 4 | pub use self::incoming::*; 5 | -------------------------------------------------------------------------------- /src/utils/bit_queue.rs: -------------------------------------------------------------------------------- 1 | use std::collections::VecDeque; 2 | 3 | const DEFAULT_BIT_VEC_QUEUE_CAP: usize = 256; 4 | 5 | /// A one-direction bit vector queue 6 | /// It use a ring buffer `VecDeque` to store bits 7 | #[derive(Debug, Clone)] 8 | pub(crate) struct BitVecQueue { 9 | store: VecDeque, 10 | head: usize, 11 | tail: usize, 12 | } 13 | 14 | impl Default for BitVecQueue { 15 | #[inline] 16 | fn default() -> Self { 17 | Self::with_capacity(DEFAULT_BIT_VEC_QUEUE_CAP) 18 | } 19 | } 20 | 21 | impl BitVecQueue { 22 | pub(crate) fn with_capacity(cap: usize) -> Self { 23 | Self { 24 | store: VecDeque::with_capacity(cap), 25 | head: 0, 26 | tail: 0, 27 | } 28 | } 29 | 30 | #[inline] 31 | pub(crate) fn len(&self) -> usize { 32 | self.tail - self.head 33 | } 34 | 35 | #[inline] 36 | fn to_physical_slot(&self, idx: usize) -> (usize, usize) { 37 | let physical_idx = idx + self.head; 38 | (physical_idx / 128, physical_idx & 127) 39 | } 40 | 41 | #[inline] 42 | pub(crate) fn get(&self, idx: usize) -> Option { 43 | if idx >= self.len() { 44 | return None; 45 | } 46 | let (idx, slot) = self.to_physical_slot(idx); 47 | Some(self.store[idx] & (1 << slot) != 0) 48 | } 49 | 50 | #[inline] 51 | pub(crate) fn set(&mut self, idx: usize, v: bool) { 52 | if idx >= self.len() { 53 | return; 54 | } 55 | let (idx, slot) = self.to_physical_slot(idx); 56 | let mask = 1 << slot; 57 | if v { 58 | self.store[idx] |= mask; 59 | } else { 60 | self.store[idx] &= !mask; 61 | } 62 | } 63 | 64 | pub(crate) fn push_back(&mut self, v: bool) { 65 | if let Some(back) = self.store.back_mut() 66 | && self.tail & 127 != 0 67 | { 68 | let mask = 1 << (self.tail & 127); 69 | if v { 70 | *back |= mask; 71 | } else { 72 | *back &= !mask; 73 | } 74 | } else { 75 | self.store.push_back(u128::from(v)); 76 | } 77 | self.tail += 1; 78 | } 79 | 80 | pub(crate) fn front(&self) -> Option { 81 | self.store.front().map(|front| { 82 | let mask = 1 << self.head; 83 | front & mask != 0 84 | }) 85 | } 86 | 87 | pub(crate) fn pop_front(&mut self) { 88 | if self.head == self.tail { 89 | self.clear(); 90 | return; 91 | } 92 | self.head += 1; 93 | if self.head == 128 { 94 | self.store.pop_front(); 95 | self.head = 0; 96 | self.tail -= 128; 97 | } 98 | } 99 | 100 | /// Clear the bit queue 101 | pub(crate) fn clear(&mut self) { 102 | self.head = 0; 103 | self.tail = 0; 104 | self.store.clear(); 105 | } 106 | } 107 | 108 | #[cfg(test)] 109 | mod test { 110 | use super::BitVecQueue; 111 | 112 | #[test] 113 | fn test_it_works() { 114 | let mut bit = BitVecQueue::default(); 115 | assert_eq!(bit.front(), None); 116 | bit.push_back(true); 117 | bit.push_back(false); 118 | bit.push_back(true); 119 | assert_eq!(bit.front(), Some(true)); 120 | bit.pop_front(); 121 | assert_eq!(bit.get(0), Some(false)); 122 | bit.set(1, false); 123 | assert_eq!(bit.get(1), Some(false)); 124 | bit.pop_front(); 125 | bit.pop_front(); 126 | assert_eq!(bit.len(), 0); 127 | 128 | assert_eq!(bit.head, 3); 129 | assert_eq!(bit.tail, 3); 130 | assert!(!bit.store.is_empty()); 131 | 132 | bit.pop_front(); 133 | assert_eq!(bit.head, 0); 134 | assert_eq!(bit.tail, 0); 135 | assert!(bit.store.is_empty()); 136 | } 137 | 138 | #[test] 139 | fn test_large_store() { 140 | let mut bit = BitVecQueue::default(); 141 | for i in 0..1_000_000 { 142 | bit.push_back(i % 2 == 0); 143 | } 144 | for i in 0..1_000_000 { 145 | assert_eq!(bit.get(i), Some(i % 2 == 0)); 146 | } 147 | for i in 0..1_000_000 { 148 | bit.set(i, i % 2 != 0); 149 | } 150 | for i in 0..1_000_000 { 151 | assert_eq!(bit.get(i), Some(i % 2 != 0)); 152 | } 153 | assert_eq!(bit.len(), 1_000_000); 154 | for _ in 0..1_000_000 { 155 | bit.pop_front(); 156 | } 157 | assert_eq!(bit.len(), 0); 158 | bit.set(2, true); 159 | assert_eq!(bit.get(2), None); 160 | } 161 | } 162 | -------------------------------------------------------------------------------- /src/utils/fastrace.rs: -------------------------------------------------------------------------------- 1 | use std::pin::Pin; 2 | use std::task::{Context, Poll}; 3 | 4 | use fastrace::collector::{SpanContext, TraceId}; 5 | use fastrace::Span; 6 | use futures::Stream; 7 | use pin_project_lite::pin_project; 8 | 9 | use crate::opts::TraceInfo; 10 | 11 | pub(crate) trait TraceStreamExt: Stream + Sized { 12 | /// It starts a span at every time an item is generating from the stream, and the span will end 13 | /// when an option yield from the stream. So it could be used to track the span from last 14 | /// reception to the current reception of each packet . 15 | /// 16 | /// [--------------**SPAN**-------------------] 17 | /// v v 18 | /// [---packet1---] [-----------packet2-----------] 19 | /// ^ ^ 20 | /// [----codec children spans----] 21 | /// 22 | /// ------------------------- timeline ------------------------------>>> 23 | fn enter_on_item Span>(self, span_fn: O) -> EnterOnItem { 24 | EnterOnItem { 25 | inner: self, 26 | span: None, 27 | last_trace_id: None, 28 | span_fn, 29 | } 30 | } 31 | } 32 | 33 | impl TraceStreamExt for S {} 34 | 35 | pin_project! { 36 | pub(crate) struct EnterOnItem { 37 | #[pin] 38 | inner: T, 39 | span: Option, 40 | last_trace_id: Option, 41 | span_fn: O, 42 | } 43 | } 44 | 45 | impl Stream for EnterOnItem 46 | where 47 | T: Stream, 48 | O: Fn() -> Span, 49 | { 50 | type Item = T::Item; 51 | 52 | fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { 53 | let this = self.project(); 54 | let span = this.span.get_or_insert_with(this.span_fn); 55 | *this.last_trace_id = SpanContext::from_span(span).map(|ctx| ctx.trace_id); 56 | let guard = span.set_local_parent(); // set the span as the local thread parent for every poll_next call 57 | let res = this.inner.poll_next(cx); 58 | match res { 59 | r @ Poll::Pending => r, // guard is dropped here before the task is moved 60 | other => { 61 | drop(guard); 62 | // ready for produce a result 63 | this.span.take(); 64 | other 65 | } 66 | } 67 | } 68 | } 69 | 70 | impl TraceInfo for EnterOnItem { 71 | fn last_trace_id(&self) -> Option { 72 | self.last_trace_id 73 | } 74 | } 75 | -------------------------------------------------------------------------------- /src/utils/hash.rs: -------------------------------------------------------------------------------- 1 | use std::hash::{BuildHasher, Hasher}; 2 | 3 | // Only support integer key with maximum 64 bits width. 4 | #[derive(Debug, Clone, Copy, Default)] 5 | pub(crate) struct NoHashBuilder; 6 | 7 | impl BuildHasher for NoHashBuilder { 8 | type Hasher = NoHashHasher; 9 | 10 | fn build_hasher(&self) -> Self::Hasher { 11 | NoHashHasher(0) 12 | } 13 | } 14 | 15 | #[derive(Debug, Clone, Copy)] 16 | pub(crate) struct NoHashHasher(u64); 17 | 18 | impl Hasher for NoHashHasher { 19 | fn finish(&self) -> u64 { 20 | self.0 21 | } 22 | 23 | fn write_u8(&mut self, i: u8) { 24 | self.0 = u64::from(i); 25 | } 26 | 27 | fn write_u16(&mut self, i: u16) { 28 | self.0 = u64::from(i); 29 | } 30 | 31 | fn write_u32(&mut self, i: u32) { 32 | self.0 = u64::from(i); 33 | } 34 | 35 | fn write_u64(&mut self, i: u64) { 36 | self.0 = i; 37 | } 38 | 39 | fn write_usize(&mut self, i: usize) { 40 | self.0 = i as u64; 41 | } 42 | 43 | fn write_i8(&mut self, i: i8) { 44 | self.0 = i as u64; 45 | } 46 | 47 | fn write_i16(&mut self, i: i16) { 48 | self.0 = i as u64; 49 | } 50 | 51 | fn write_i32(&mut self, i: i32) { 52 | self.0 = i as u64; 53 | } 54 | 55 | fn write_i64(&mut self, i: i64) { 56 | self.0 = i as u64; 57 | } 58 | 59 | fn write_isize(&mut self, i: isize) { 60 | self.0 = i as u64; 61 | } 62 | 63 | fn write_u128(&mut self, _: u128) { 64 | unimplemented!("unsupported") 65 | } 66 | 67 | fn write_i128(&mut self, _: i128) { 68 | unimplemented!("unsupported") 69 | } 70 | 71 | fn write(&mut self, _: &[u8]) { 72 | unimplemented!("unsupported") 73 | } 74 | } 75 | 76 | // From Google's city hash. 77 | #[inline(always)] 78 | pub(crate) fn combine_hashes(upper: u64, lower: u64) -> u64 { 79 | const MUL: u64 = 0x9ddfea08eb382d69; 80 | 81 | let mut a = (lower ^ upper).wrapping_mul(MUL); 82 | a ^= a >> 47; 83 | let mut b = (upper ^ a).wrapping_mul(MUL); 84 | b ^= b >> 47; 85 | b = b.wrapping_mul(MUL); 86 | b 87 | } 88 | -------------------------------------------------------------------------------- /src/utils/log.rs: -------------------------------------------------------------------------------- 1 | use std::pin::Pin; 2 | use std::task::{ready, Context, Poll}; 3 | 4 | use futures::{Sink, Stream}; 5 | use pin_project_lite::pin_project; 6 | 7 | pub(crate) trait Logged: Sized { 8 | fn logged( 9 | self, 10 | ok_f: impl Fn(&T) + Send + Sync + 'static, 11 | err_f: impl Fn(&E) + Send + Sync + 'static, 12 | ) -> Log; 13 | } 14 | 15 | impl Logged for F 16 | where 17 | F: Stream>, 18 | { 19 | fn logged( 20 | self, 21 | ok_f: impl Fn(&T) + Send + Sync + 'static, 22 | err_f: impl Fn(&E) + Send + Sync + 'static, 23 | ) -> Log { 24 | Log { 25 | source: self, 26 | err_f: Box::new(err_f), 27 | ok_f: Box::new(ok_f), 28 | } 29 | } 30 | } 31 | 32 | pin_project! { 33 | /// Log the error of the stream while reading. 34 | pub(crate) struct Log { 35 | #[pin] 36 | source: F, 37 | ok_f: Box, 38 | err_f: Box, 39 | } 40 | } 41 | 42 | impl Stream for Log 43 | where 44 | F: Stream>, 45 | E: std::error::Error, 46 | { 47 | type Item = T; 48 | 49 | fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { 50 | let mut this = self.project(); 51 | loop { 52 | let Some(res) = ready!(this.source.as_mut().poll_next(cx)) else { 53 | return Poll::Ready(None); 54 | }; 55 | let v = match res { 56 | Ok(v) => v, 57 | Err(ref err) => { 58 | (*this.err_f)(err); 59 | continue; 60 | } 61 | }; 62 | (*this.ok_f)(&v); 63 | return Poll::Ready(Some(v)); 64 | } 65 | } 66 | } 67 | 68 | // Propagate sink for logged wrapper 69 | // Do not log error for sink as it already throws error to the user 70 | impl Sink for Log 71 | where 72 | F: Sink, 73 | { 74 | type Error = E; 75 | 76 | fn poll_ready(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { 77 | self.project().source.poll_ready(cx) 78 | } 79 | 80 | fn start_send(self: Pin<&mut Self>, item: Si) -> Result<(), Self::Error> { 81 | self.project().source.start_send(item) 82 | } 83 | 84 | fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { 85 | self.project().source.poll_flush(cx) 86 | } 87 | 88 | fn poll_close(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { 89 | self.project().source.poll_close(cx) 90 | } 91 | } 92 | -------------------------------------------------------------------------------- /src/utils/mod.rs: -------------------------------------------------------------------------------- 1 | mod bit_queue; 2 | mod fastrace; 3 | mod hash; 4 | mod log; 5 | mod reactor; 6 | mod seq_num; 7 | 8 | pub(crate) use self::bit_queue::*; 9 | pub(crate) use self::fastrace::*; 10 | pub(crate) use self::hash::*; 11 | pub(crate) use self::log::*; 12 | pub(crate) use self::reactor::*; 13 | pub(crate) use self::seq_num::*; 14 | 15 | /// Test utils. 16 | #[cfg(test)] 17 | pub(crate) mod tests; 18 | 19 | #[inline] 20 | pub(crate) fn timestamp() -> i64 { 21 | std::time::SystemTime::now() 22 | .duration_since(std::time::UNIX_EPOCH) 23 | .unwrap() 24 | .as_millis() as i64 25 | } 26 | -------------------------------------------------------------------------------- /src/utils/reactor.rs: -------------------------------------------------------------------------------- 1 | use std::collections::BTreeMap; 2 | use std::sync::OnceLock; 3 | use std::task::Waker; 4 | use std::time::{Duration, Instant}; 5 | use std::{mem, panic, thread}; 6 | 7 | use super::NoHashBuilder; 8 | use crate::HashMap; 9 | 10 | /// Timers are in the order in which they fire. The `usize` in this type is a timer ID used to 11 | /// distinguish timers that fire at the same time. The `Waker` represents the task awaiting 12 | /// the timer. 13 | type Timers = BTreeMap<(Instant, usize), Waker>; 14 | 15 | /// A reactor that manages timers. 16 | pub(crate) struct Reactor { 17 | /// Map of registered timers, distinguished by their unique timer key. 18 | conn_timers: parking_lot::Mutex>, 19 | /// A condvar used to wake up the reactor when timers changed. 20 | cond: parking_lot::Condvar, 21 | } 22 | 23 | impl Reactor { 24 | pub(crate) fn get() -> &'static Reactor { 25 | static REACTOR: OnceLock = OnceLock::new(); 26 | 27 | fn main_loop() { 28 | let reactor = Reactor::get(); 29 | loop { 30 | reactor.process_timers(); 31 | } 32 | } 33 | 34 | REACTOR.get_or_init(|| { 35 | // Spawn the daemon thread to motivate the reactor. 36 | thread::Builder::new() 37 | .name("raknet-timer-reactor".to_string()) 38 | .spawn(main_loop) 39 | .expect("cannot spawn timer-reactor thread"); 40 | 41 | Reactor { 42 | conn_timers: parking_lot::Mutex::new(HashMap::default()), 43 | cond: parking_lot::Condvar::new(), 44 | } 45 | }) 46 | } 47 | 48 | pub(crate) fn insert_timer(&self, key: u64, when: Instant, waker: &Waker) { 49 | let mut timers = self.conn_timers.lock(); 50 | let timers = timers.entry(key).or_default(); 51 | timers.insert((when, timers.len()), waker.clone()); 52 | self.cond.notify_one(); 53 | } 54 | 55 | pub(crate) fn cancel_all_timers(&self, key: u64) -> impl Iterator { 56 | let mut timers = self.conn_timers.lock(); 57 | let res = timers 58 | .remove(&key) 59 | .into_iter() 60 | .flat_map(BTreeMap::into_values); 61 | self.cond.notify_one(); 62 | res 63 | } 64 | 65 | /// Processes ready timers and waits for the next timer changed. 66 | fn process_timers(&self) { 67 | let mut region_timers = self.conn_timers.lock(); 68 | let now = Instant::now(); 69 | 70 | let mut dur: Option = None; 71 | 72 | for timers in region_timers.values_mut() { 73 | // Split timers into ready and pending timers. 74 | // 75 | // Careful to split just *after* `now`, so that a timer set for exactly `now` is 76 | // considered ready. 77 | let pending = timers.split_off(&(now + Duration::from_nanos(1), 0)); 78 | let ready = mem::replace(timers, pending); 79 | let next_time = timers 80 | .keys() 81 | .next() 82 | .map(|&(when, _)| when.saturating_duration_since(now)); 83 | for (_, waker) in ready { 84 | // Don't let a panicking waker blow everything up. 85 | panic::catch_unwind(|| waker.wake()).ok(); 86 | } 87 | match (dur, next_time) { 88 | (Some(d), Some(n)) => dur = Some(d.min(n)), 89 | (None, Some(n)) => dur = Some(n), 90 | _ => {} 91 | } 92 | } 93 | 94 | if let Some(dur) = dur { 95 | self.cond.wait_for(&mut region_timers, dur); 96 | } else { 97 | self.cond.wait(&mut region_timers); 98 | } 99 | } 100 | } 101 | 102 | #[cfg(test)] 103 | mod test { 104 | use super::*; 105 | use crate::utils::tests::TestWaker; 106 | 107 | #[test] 108 | fn test_it_works() { 109 | let reactor = Reactor::get(); 110 | 111 | let dur = Duration::from_millis(100); 112 | let when = Instant::now() + dur; 113 | { 114 | let (waker, test) = TestWaker::pair(); 115 | reactor.insert_timer(1, when, &waker); 116 | assert_eq!(reactor.cancel_all_timers(1).count(), 1); 117 | assert!(!test.woken.load(std::sync::atomic::Ordering::Relaxed)); 118 | } 119 | 120 | { 121 | let (waker, test) = TestWaker::pair(); 122 | reactor.insert_timer(2, when, &waker); 123 | std::thread::sleep(dur + Duration::from_millis(10)); 124 | assert_eq!(reactor.cancel_all_timers(2).count(), 0); 125 | assert!(test.woken.load(std::sync::atomic::Ordering::Relaxed)); 126 | } 127 | } 128 | } 129 | -------------------------------------------------------------------------------- /src/utils/seq_num.rs: -------------------------------------------------------------------------------- 1 | use std::{fmt, ops}; 2 | 3 | use bytes::{Buf, BufMut}; 4 | 5 | /// Unsigned 24bits integer (actually occupied 32 bits) with litter endian and wrapping checking 6 | #[allow(non_camel_case_types)] 7 | #[derive(Copy, Clone, Ord, PartialOrd, Eq, PartialEq, Hash, Default)] 8 | pub(crate) struct u24(u32); 9 | 10 | impl fmt::Debug for u24 { 11 | fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { 12 | write!(f, "{}", self.0) 13 | } 14 | } 15 | 16 | impl fmt::Display for u24 { 17 | fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { 18 | write!(f, "{}", self.0) 19 | } 20 | } 21 | 22 | impl fmt::Binary for u24 { 23 | fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { 24 | write!(f, "{:b}", self.0) 25 | } 26 | } 27 | 28 | impl fmt::Octal for u24 { 29 | fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { 30 | write!(f, "{:o}", self.0) 31 | } 32 | } 33 | 34 | impl fmt::LowerHex for u24 { 35 | fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { 36 | write!(f, "{:x}", self.0) 37 | } 38 | } 39 | 40 | impl fmt::UpperHex for u24 { 41 | fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { 42 | write!(f, "{:X}", self.0) 43 | } 44 | } 45 | 46 | impl u24 { 47 | pub(crate) fn to_u32(self) -> u32 { 48 | self.0 49 | } 50 | 51 | pub(crate) fn to_usize(self) -> usize { 52 | self.0 as usize 53 | } 54 | } 55 | 56 | macro_rules! for_all_primitives { 57 | ($macro:ident) => { 58 | $macro! { u8, u16, u32, u64, usize, i8, i16, i32, i64, isize } 59 | }; 60 | } 61 | 62 | macro_rules! impl_to_for_u24 { 63 | ($($t:ty),*) => { 64 | $( 65 | impl From for $t { 66 | fn from(value: u24) -> Self { 67 | value.0 as $t 68 | } 69 | } 70 | )* 71 | }; 72 | } 73 | 74 | for_all_primitives! { impl_to_for_u24 } 75 | 76 | macro_rules! impl_from_for_u24 { 77 | ($($t:ty),*) => { 78 | $( 79 | impl From<$t> for u24 { 80 | fn from(value: $t) -> Self { 81 | debug_assert!((value as u32 >> 24) == 0, "{value} exceed the maximum of u24"); 82 | Self(value as u32) 83 | } 84 | } 85 | )* 86 | }; 87 | } 88 | 89 | for_all_primitives! { impl_from_for_u24 } 90 | 91 | pub(crate) trait BufExt { 92 | fn get_u24_le(&mut self) -> u24; 93 | } 94 | 95 | pub(crate) trait BufMutExt { 96 | fn put_u24_le(&mut self, v: u24); 97 | } 98 | 99 | impl BufExt for B { 100 | fn get_u24_le(&mut self) -> u24 { 101 | u24(self.get_uint_le(3) as u32) 102 | } 103 | } 104 | 105 | impl BufMutExt for B { 106 | fn put_u24_le(&mut self, v: u24) { 107 | self.put_uint_le(v.0 as u64, 3); 108 | } 109 | } 110 | 111 | ///// Checked operations for u24 112 | 113 | fn checked_add(a: u32, b: u32) -> u24 { 114 | let value = a + b; 115 | assert!((value >> 24) == 0, "{value} exceed the maximum of u24"); 116 | u24(value) 117 | } 118 | 119 | impl ops::Add<&u24> for &u24 { 120 | type Output = u24; 121 | 122 | fn add(self, rhs: &u24) -> Self::Output { 123 | checked_add(self.0, rhs.0) 124 | } 125 | } 126 | 127 | impl ops::Add<&u24> for u24 { 128 | type Output = u24; 129 | 130 | fn add(self, rhs: &u24) -> Self::Output { 131 | checked_add(self.0, rhs.0) 132 | } 133 | } 134 | 135 | impl ops::Add for &u24 { 136 | type Output = u24; 137 | 138 | fn add(self, rhs: u24) -> Self::Output { 139 | checked_add(self.0, rhs.0) 140 | } 141 | } 142 | 143 | impl ops::Add for u24 { 144 | type Output = u24; 145 | 146 | fn add(self, rhs: u24) -> Self::Output { 147 | checked_add(self.0, rhs.0) 148 | } 149 | } 150 | 151 | // for primitives 152 | 153 | macro_rules! impl_add_for_u24 { 154 | ($($t:ty),*) => { 155 | $( 156 | impl ops::Add<$t> for u24 { 157 | type Output = u24; 158 | 159 | fn add(self, rhs: $t) -> Self::Output { 160 | checked_add(self.0, rhs as u32) 161 | } 162 | } 163 | 164 | impl ops::Add<$t> for &u24 { 165 | type Output = u24; 166 | 167 | fn add(self, rhs: $t) -> Self::Output { 168 | checked_add(self.0, rhs as u32) 169 | } 170 | } 171 | )* 172 | }; 173 | } 174 | 175 | for_all_primitives! { impl_add_for_u24 } 176 | 177 | impl ops::AddAssign<&u24> for u24 { 178 | fn add_assign(&mut self, rhs: &u24) { 179 | self.0 += rhs.0; 180 | assert!((self.0 >> 24) == 0, "{self} exceed the maximum of u24"); 181 | } 182 | } 183 | 184 | impl ops::AddAssign for u24 { 185 | fn add_assign(&mut self, rhs: u24) { 186 | self.0 += rhs.0; 187 | assert!((self.0 >> 24) == 0, "{self} exceed the maximum of u24"); 188 | } 189 | } 190 | 191 | // for primitives 192 | 193 | macro_rules! impl_add_assign_for_u24 { 194 | ($($t:ty),*) => { 195 | $( 196 | impl ops::AddAssign<$t> for u24 { 197 | fn add_assign(&mut self, rhs: $t) { 198 | self.0 += rhs as u32; 199 | assert!((self.0 >> 24) == 0, "{self} exceed the maximum of u24"); 200 | } 201 | } 202 | )* 203 | }; 204 | } 205 | 206 | for_all_primitives! { impl_add_assign_for_u24 } 207 | 208 | impl ops::Sub<&u24> for &u24 { 209 | type Output = u24; 210 | 211 | fn sub(self, rhs: &u24) -> Self::Output { 212 | u24(self.0 - rhs.0) 213 | } 214 | } 215 | 216 | impl ops::Sub for &u24 { 217 | type Output = u24; 218 | 219 | fn sub(self, rhs: u24) -> Self::Output { 220 | u24(self.0 - rhs.0) 221 | } 222 | } 223 | 224 | impl ops::Sub<&u24> for u24 { 225 | type Output = u24; 226 | 227 | fn sub(self, rhs: &u24) -> Self::Output { 228 | u24(self.0 - rhs.0) 229 | } 230 | } 231 | 232 | impl ops::Sub for u24 { 233 | type Output = u24; 234 | 235 | fn sub(self, rhs: u24) -> Self::Output { 236 | u24(self.0 - rhs.0) 237 | } 238 | } 239 | 240 | // for primitives 241 | 242 | macro_rules! impl_sub_for_u24 { 243 | ($($t:ty),*) => { 244 | $( 245 | impl ops::Sub<$t> for &u24 { 246 | type Output = u24; 247 | 248 | fn sub(self, rhs: $t) -> Self::Output { 249 | u24(self.0 - rhs as u32) 250 | } 251 | } 252 | 253 | impl ops::Sub<$t> for u24 { 254 | type Output = u24; 255 | 256 | fn sub(self, rhs: $t) -> Self::Output { 257 | u24(self.0 - rhs as u32) 258 | } 259 | } 260 | )* 261 | }; 262 | } 263 | 264 | for_all_primitives! { impl_sub_for_u24 } 265 | 266 | impl ops::SubAssign<&u24> for u24 { 267 | fn sub_assign(&mut self, rhs: &u24) { 268 | self.0 -= rhs.0; 269 | } 270 | } 271 | 272 | impl ops::SubAssign for u24 { 273 | fn sub_assign(&mut self, rhs: u24) { 274 | self.0 -= rhs.0; 275 | } 276 | } 277 | 278 | macro_rules! impl_sub_assign_for_u24 { 279 | ($($t:ty),*) => { 280 | $( 281 | impl ops::SubAssign<$t> for u24 { 282 | fn sub_assign(&mut self, rhs: $t) { 283 | self.0 -= rhs as u32; 284 | } 285 | } 286 | )* 287 | }; 288 | } 289 | 290 | for_all_primitives! { impl_sub_assign_for_u24 } 291 | 292 | // almost no multiplication, division or other operations on u24 293 | 294 | #[cfg(test)] 295 | mod test { 296 | use super::*; 297 | 298 | #[test] 299 | #[should_panic] 300 | fn test_u24_overflow_1() { 301 | let mut a1: u24 = 0.into(); 302 | a1 -= 1; 303 | } 304 | 305 | #[test] 306 | #[should_panic] 307 | fn test_u24_overflow_2() { 308 | let _a1: u24 = (1 << 24).into(); 309 | } 310 | 311 | #[test] 312 | #[should_panic] 313 | fn test_u24_overflow_3() { 314 | let mut a1: u24 = ((1 << 24) - 1).into(); 315 | a1 += 1; 316 | } 317 | 318 | #[test] 319 | #[should_panic] 320 | fn test_u24_overflow_4() { 321 | let a1: u24 = ((1 << 24) - 1).into(); 322 | let _b1 = a1 + 1; 323 | } 324 | 325 | #[test] 326 | #[should_panic] 327 | fn test_u24_overflow_5() { 328 | let a1: u24 = 0.into(); 329 | let _b1 = a1 - 1; 330 | } 331 | 332 | #[test] 333 | fn test_u24_works() { 334 | let a1: u24 = 1.into(); 335 | let mut a2: u24 = 2.into(); 336 | a2 += 1; 337 | let mut b1 = a1 + a2; 338 | b1 -= 1; 339 | assert_eq!(b1.to_u32(), 3); 340 | } 341 | } 342 | -------------------------------------------------------------------------------- /src/utils/tests/flusher.rs: -------------------------------------------------------------------------------- 1 | use std::io; 2 | use std::sync::Arc; 3 | use std::time::Duration; 4 | 5 | use async_channel::{Receiver, Sender}; 6 | use futures::{Sink, SinkExt}; 7 | use tokio::sync::Notify; 8 | use tokio::time; 9 | 10 | use crate::Message; 11 | 12 | pub(crate) struct Flusher { 13 | writer: W, 14 | ticker: time::Interval, 15 | notifier: Arc, 16 | rx: Receiver, 17 | } 18 | 19 | impl + Unpin + Send + Sync + 'static> Flusher { 20 | const DEFAULT_FLUSH_DELAY: Duration = Duration::from_millis(1); 21 | 22 | pub(crate) fn new(writer: W) -> (Self, Arc, Sender) { 23 | let mut ticker = time::interval(Self::DEFAULT_FLUSH_DELAY); 24 | ticker.set_missed_tick_behavior(time::MissedTickBehavior::Delay); 25 | let notifier = Arc::new(Notify::new()); 26 | let (tx, rx) = async_channel::unbounded(); 27 | let me = Self { 28 | writer, 29 | ticker, 30 | notifier: notifier.clone(), 31 | rx, 32 | }; 33 | (me, notifier, tx) 34 | } 35 | 36 | pub(crate) async fn run(&mut self) -> io::Result<()> { 37 | loop { 38 | tokio::select! { 39 | _ = self.ticker.tick() => { 40 | self.writer.flush().await?; 41 | } 42 | _ = self.notifier.notified() => { 43 | self.writer.flush().await?; 44 | } 45 | Ok(msg) = self.rx.recv() => { 46 | self.writer.feed(msg).await?; 47 | } 48 | } 49 | } 50 | } 51 | } 52 | -------------------------------------------------------------------------------- /src/utils/tests/mod.rs: -------------------------------------------------------------------------------- 1 | use std::io; 2 | use std::sync::atomic::AtomicBool; 3 | use std::sync::Arc; 4 | use std::task::Waker; 5 | 6 | use bytes::Bytes; 7 | use futures::{Sink, Stream, StreamExt}; 8 | 9 | mod flusher; 10 | mod sim_net; 11 | mod tracing; 12 | 13 | pub(crate) use flusher::*; 14 | pub(crate) use sim_net::*; 15 | pub(crate) use tracing::*; 16 | 17 | use crate::Message; 18 | 19 | pub(crate) struct TestWaker { 20 | pub(crate) woken: AtomicBool, 21 | } 22 | 23 | impl std::task::Wake for TestWaker { 24 | fn wake_by_ref(self: &Arc) { 25 | self.woken.store(true, std::sync::atomic::Ordering::Relaxed); 26 | } 27 | 28 | fn wake(self: Arc) { 29 | self.wake_by_ref(); 30 | } 31 | } 32 | 33 | impl TestWaker { 34 | pub(crate) fn create() -> Waker { 35 | Arc::new(TestWaker { 36 | woken: AtomicBool::new(false), 37 | }) 38 | .into() 39 | } 40 | 41 | pub(crate) fn pair() -> (Waker, Arc) { 42 | let arc = Arc::new(TestWaker { 43 | woken: AtomicBool::new(false), 44 | }); 45 | (Waker::from(Arc::clone(&arc)), arc) 46 | } 47 | } 48 | 49 | pub(crate) fn spawn_echo_server( 50 | incoming: impl Stream< 51 | Item = ( 52 | impl Stream + Send + Sync + 'static, 53 | impl Sink + Send + Sync + 'static, 54 | ), 55 | > + Send 56 | + Sync 57 | + 'static, 58 | ) { 59 | tokio::spawn(async move { 60 | tokio::pin!(incoming); 61 | loop { 62 | let (reader, sender) = incoming.next().await.unwrap(); 63 | tokio::spawn(async move { 64 | tokio::pin!(reader); 65 | let (mut flusher, _notify, tx) = Flusher::new(Box::pin(sender)); 66 | tokio::spawn(async move { 67 | flusher.run().await.unwrap(); 68 | }); 69 | while let Some(data) = reader.next().await { 70 | tx.send(Message::new(data)).await.unwrap(); 71 | } 72 | }); 73 | } 74 | }); 75 | } 76 | -------------------------------------------------------------------------------- /src/utils/tests/sim_net.rs: -------------------------------------------------------------------------------- 1 | use std::collections::{HashMap, HashSet, VecDeque}; 2 | use std::io; 3 | use std::net::{IpAddr, SocketAddr, ToSocketAddrs}; 4 | use std::sync::Arc; 5 | use std::task::{Context, Poll}; 6 | use std::time::{Duration, Instant}; 7 | 8 | use bytes::{BufMut, Bytes, BytesMut}; 9 | use futures::{Sink, Stream}; 10 | use parking_lot::Mutex; 11 | use rand::seq::SliceRandom; 12 | use rand::Rng; 13 | 14 | use crate::client::{self, ConnectTo}; 15 | use crate::codec::AsyncSocket; 16 | use crate::opts::{ConnectionInfo, Ping, TraceInfo}; 17 | use crate::server::{Incoming, MakeIncoming}; 18 | use crate::utils::Reactor; 19 | use crate::Message; 20 | 21 | // packets will be received in order by default 22 | type NetChan = VecDeque<(Vec /* data */, Instant /* receive at */)>; 23 | 24 | pub(crate) struct SimNet { 25 | // fake ip, used to distinguish different sim net 26 | ip: IpAddr, 27 | // fake port 28 | allocated_port: u16, 29 | // sim net topology 30 | adjacency: HashMap>, 31 | } 32 | 33 | impl SimNet { 34 | pub(crate) fn bind(ip: IpAddr) -> Self { 35 | Self { 36 | ip, 37 | allocated_port: 0, 38 | adjacency: HashMap::new(), 39 | } 40 | } 41 | 42 | pub(crate) fn add_endpoint(&mut self) -> Endpoint { 43 | self.allocated_port += 1; 44 | Endpoint { 45 | addr: SocketAddr::new(self.ip, self.allocated_port), 46 | inboxes: HashMap::new(), 47 | mailboxes: HashMap::new(), 48 | nemeses: HashMap::new(), 49 | } 50 | } 51 | 52 | pub(crate) fn connect(&mut self, p1: &mut Endpoint, p2: &mut Endpoint) { 53 | assert_eq!(p1.addr.ip(), p2.addr.ip()); 54 | assert_eq!(p1.addr.ip(), self.ip); 55 | 56 | let p1top2 = Arc::new(Mutex::new(NetChan::new())); 57 | let p2top1 = Arc::new(Mutex::new(NetChan::new())); 58 | 59 | p1.mailboxes.insert(p2.addr, p1top2.clone()); 60 | p1.nemeses.insert(p2.addr, Nemeses::new()); 61 | p1.inboxes.insert(p2.addr, p2top1.clone()); 62 | 63 | p2.mailboxes.insert(p1.addr, p2top1); 64 | p2.nemeses.insert(p1.addr, Nemeses::new()); 65 | p2.inboxes.insert(p1.addr, p1top2); 66 | 67 | self.adjacency.entry(p1.addr).or_default().insert(p2.addr); 68 | self.adjacency.entry(p2.addr).or_default().insert(p1.addr); 69 | } 70 | 71 | pub(crate) fn add_nemesis(&self, from: &mut Endpoint, to: &Endpoint, nemesis: impl Nemesis) { 72 | assert_eq!(from.addr.ip(), to.addr.ip()); 73 | assert_eq!(from.addr.ip(), self.ip); 74 | 75 | from.nemeses 76 | .get_mut(&to.addr) 77 | .expect("endpoints do not have connected") 78 | .push(Box::new(nemesis)); 79 | } 80 | } 81 | 82 | pub(crate) struct Endpoint { 83 | addr: SocketAddr, 84 | inboxes: HashMap>>, 85 | mailboxes: HashMap>>, 86 | nemeses: HashMap, 87 | } 88 | 89 | impl Endpoint { 90 | pub(crate) fn addr(&self) -> SocketAddr { 91 | self.addr 92 | } 93 | } 94 | 95 | impl MakeIncoming for Endpoint { 96 | fn make_incoming( 97 | self, 98 | config: crate::server::Config, 99 | ) -> impl Stream< 100 | Item = ( 101 | impl Stream + TraceInfo, 102 | impl Sink + ConnectionInfo, 103 | ), 104 | > { 105 | let socket = Arc::new(self); 106 | Incoming::new(socket, config) 107 | } 108 | } 109 | 110 | impl ConnectTo for Endpoint { 111 | async fn connect_to( 112 | self, 113 | addrs: impl ToSocketAddrs, 114 | config: crate::client::Config, 115 | ) -> io::Result<( 116 | impl Stream, 117 | impl Sink + Ping + ConnectionInfo, 118 | )> { 119 | let socket = Arc::new(self); 120 | client::connect_to(socket, addrs, config, tokio::spawn).await 121 | } 122 | } 123 | 124 | impl AsyncSocket for Arc { 125 | fn poll_recv_from( 126 | &self, 127 | cx: &mut Context<'_>, 128 | buf: &mut BytesMut, 129 | ) -> Poll> { 130 | let now = Instant::now(); 131 | let mut next_poll = now + Duration::from_millis(10); 132 | for (from, chan) in &self.inboxes { 133 | let mut chan = chan.lock(); 134 | if chan.is_empty() { 135 | continue; 136 | } 137 | let receive_at = chan.front().map(|x| x.1).unwrap(); 138 | if receive_at > now { 139 | // net delay 140 | next_poll = next_poll.min(receive_at); 141 | continue; 142 | } 143 | let (data, _) = chan.pop_front().unwrap(); 144 | if data.is_empty() { 145 | // data loss 146 | continue; 147 | } 148 | buf.put(&data[..]); 149 | return Poll::Ready(Ok(*from)); 150 | } 151 | Reactor::get().insert_timer(0, next_poll, cx.waker()); 152 | Poll::Pending 153 | } 154 | 155 | fn poll_send_to( 156 | &self, 157 | _: &mut Context<'_>, 158 | buf: &[u8], 159 | target: SocketAddr, 160 | ) -> Poll> { 161 | let mut chan = self.mailboxes.get(&target).expect("not connected").lock(); 162 | let nemeses = self.nemeses.get(&target).expect("not connect"); 163 | let now = Instant::now(); 164 | chan.push_back((buf.to_vec(), now)); 165 | // skip for the first 2 packets 166 | if chan.len() > 2 { 167 | for nemesis in nemeses { 168 | nemesis.strike(&mut chan); 169 | } 170 | } 171 | Poll::Ready(Ok(buf.len())) 172 | } 173 | } 174 | 175 | type Nemeses = Vec>; 176 | 177 | pub(crate) trait Nemesis: Send + Sync + 'static { 178 | fn strike(&self, inflight: &mut NetChan); 179 | } 180 | 181 | // Make packet disorder 182 | pub(crate) struct RandomDataDisorder { 183 | pub(crate) odd: f32, 184 | pub(crate) rate: f32, 185 | } 186 | 187 | impl Nemesis for RandomDataDisorder { 188 | fn strike(&self, inflight: &mut NetChan) { 189 | let mut rng = rand::thread_rng(); 190 | if self.odd <= rng.r#gen() { 191 | return; 192 | } 193 | let amount = (inflight.len() as f32 * self.rate) as usize; 194 | log::error!("nemesis strike! disorder {} packets", amount); 195 | let (front, back) = inflight.as_mut_slices(); 196 | front.partial_shuffle(&mut rng, amount); 197 | if amount > front.len() { 198 | back.partial_shuffle(&mut rng, amount - front.len()); 199 | } 200 | } 201 | } 202 | 203 | // Make all packets lost 204 | pub(crate) struct RandomDataLost { 205 | pub(crate) odd: f32, 206 | } 207 | 208 | impl Nemesis for RandomDataLost { 209 | fn strike(&self, inflight: &mut NetChan) { 210 | let mut rng = rand::thread_rng(); 211 | if self.odd <= rng.r#gen() { 212 | return; 213 | } 214 | log::error!("nemesis strike! loss all({}) packets", inflight.len()); 215 | for packet in inflight.iter_mut() { 216 | packet.0.clear(); 217 | } 218 | } 219 | } 220 | 221 | // Delay all packets in a sim net connection 222 | pub(crate) struct RandomNetDelay { 223 | pub(crate) odd: f32, 224 | pub(crate) delay: Duration, 225 | } 226 | 227 | impl Nemesis for RandomNetDelay { 228 | fn strike(&self, inflight: &mut NetChan) { 229 | if self.odd <= rand::random() { 230 | return; 231 | } 232 | log::error!( 233 | "nemesis strike! delay {} packets {:?}", 234 | inflight.len(), 235 | self.delay 236 | ); 237 | inflight[0].1 += self.delay; 238 | } 239 | } 240 | -------------------------------------------------------------------------------- /src/utils/tests/tracing.rs: -------------------------------------------------------------------------------- 1 | use std::collections::HashMap; 2 | use std::sync::Arc; 3 | 4 | use fastrace::collector::{SpanId, SpanRecord, TraceId}; 5 | use parking_lot::Mutex; 6 | 7 | pub(crate) struct TestTraceLogGuard { 8 | spans: Arc>>, 9 | } 10 | 11 | impl Drop for TestTraceLogGuard { 12 | #[allow(clippy::print_stderr)] 13 | fn drop(&mut self) { 14 | fastrace::flush(); 15 | 16 | let spans = self.spans.lock().clone(); 17 | let spans_map: HashMap = spans 18 | .iter() 19 | .map(|span| (span.span_id, span.clone())) 20 | .collect(); 21 | let adjacency_lists: HashMap>> = spans.iter().fold( 22 | HashMap::new(), 23 | |mut map, 24 | SpanRecord { 25 | trace_id, 26 | span_id, 27 | parent_id, 28 | .. 29 | }| { 30 | map.entry(*trace_id) 31 | .or_default() 32 | .entry(*parent_id) 33 | .or_default() 34 | .push(*span_id); 35 | map 36 | }, 37 | ); 38 | fn dfs( 39 | adjacency_list: &HashMap>, 40 | spans: &HashMap, 41 | span_id: SpanId, 42 | depth: usize, 43 | last: bool, 44 | ) { 45 | let span = &spans[&span_id]; 46 | let mut properties = String::new(); 47 | for (key, value) in &span.properties { 48 | properties.push_str(&format!("{}: {}, ", key, value)); 49 | } 50 | let mut events = String::new(); 51 | for ev in &span.events { 52 | events.push_str(&format!("'{}'", ev.name)); 53 | } 54 | let prefix = if depth == 0 { 55 | String::new() 56 | } else if last { 57 | "╰".to_owned() + &"─".repeat(depth) + " " 58 | } else { 59 | "├".to_owned() + &"─".repeat(depth) + " " 60 | }; 61 | eprintln!( 62 | "{}{}({}{{{}}}) [{}us]", 63 | prefix, 64 | span.name, 65 | properties, 66 | events, 67 | span.duration_ns as f64 / 1_000.0, 68 | ); 69 | if let Some(children) = adjacency_list.get(&span_id) { 70 | for (i, child) in children.iter().enumerate() { 71 | dfs( 72 | adjacency_list, 73 | spans, 74 | *child, 75 | depth + 1, 76 | i == children.len() - 1 && last, 77 | ); 78 | } 79 | } 80 | } 81 | for (trace_id, list) in adjacency_lists { 82 | if list.is_empty() { 83 | continue; 84 | } 85 | eprintln!("[trace_id] {}", trace_id.0); 86 | if !list.contains_key(&SpanId::default()) { 87 | log::error!( 88 | "root span not found, trace may lost record, list: {:?}", 89 | list 90 | ); 91 | continue; 92 | } 93 | let l = &list[&SpanId::default()]; 94 | for (i, root) in l.iter().enumerate() { 95 | dfs(&list, &spans_map, *root, 0, i == l.len() - 1); 96 | } 97 | eprintln!(); 98 | } 99 | } 100 | } 101 | 102 | #[must_use = "guard should be kept alive to keep the trace log"] 103 | pub(crate) fn test_trace_log_setup() -> TestTraceLogGuard { 104 | std::env::set_var("RUST_LOG", "trace"); 105 | let (reporter, spans) = fastrace::collector::TestReporter::new(); 106 | fastrace::set_reporter( 107 | reporter, 108 | fastrace::collector::Config::default().report_before_root_finish(true), 109 | ); 110 | let _ignore = env_logger::try_init(); 111 | TestTraceLogGuard { spans } 112 | } 113 | -------------------------------------------------------------------------------- /xdp/ebpf/.cargo/config.toml: -------------------------------------------------------------------------------- 1 | [build] 2 | target = ["bpfeb-unknown-none", "bpfel-unknown-none"] 3 | 4 | [unstable] 5 | build-std = ["core"] 6 | -------------------------------------------------------------------------------- /xdp/ebpf/Cargo.toml: -------------------------------------------------------------------------------- 1 | [package] 2 | name = "ebpf" 3 | version = "0.1.0" 4 | publish = false 5 | edition = { workspace = true } 6 | license = { workspace = true } 7 | authors = { workspace = true } 8 | homepage = { workspace = true } 9 | repository = { workspace = true } 10 | categories = { workspace = true } 11 | keywords = { workspace = true } 12 | 13 | [dependencies] 14 | # These crates are not published so use a git dependency for now. See https://github.com/aya-rs/aya/issues/464 15 | aya-bpf = { git = "https://github.com/aya-rs/aya", tag = "aya-v0.12.0" } 16 | aya-log-ebpf = { git = "https://github.com/aya-rs/aya", tag = "aya-v0.12.0" } 17 | zerocopy = { version = "0.7", features = ["derive"] } 18 | 19 | [features] 20 | default = [] 21 | trace = [] 22 | 23 | [[bin]] 24 | name = "raknet-xdp-ebpf" 25 | path = "src/main.rs" 26 | 27 | [lints] 28 | workspace = true 29 | -------------------------------------------------------------------------------- /xdp/ebpf/src/buffer.rs: -------------------------------------------------------------------------------- 1 | pub struct DecodeBuffer<'a> { 2 | bytes: &'a mut [u8], 3 | end: *mut u8, 4 | } 5 | 6 | #[derive(Clone, Copy, Debug)] 7 | pub enum DecoderError { 8 | UnexpectedEof(usize), 9 | InvariantViolation(&'static str), 10 | } 11 | 12 | impl<'a> DecodeBuffer<'a> { 13 | #[inline] 14 | pub unsafe fn new(start: *mut u8, end: *mut u8) -> Self { 15 | let len = end as usize - start as usize; 16 | let bytes = core::slice::from_raw_parts_mut(start, len); 17 | Self { bytes, end } 18 | } 19 | 20 | #[inline] 21 | pub fn new_checked(bytes: &'a mut [u8], end: *mut u8) -> Result { 22 | if bytes.as_ptr() > end { 23 | return Err(DecoderError::UnexpectedEof(0)); 24 | } 25 | Ok(Self { bytes, end }) 26 | } 27 | 28 | #[inline] 29 | pub fn decode_slice( 30 | self, 31 | count: usize, 32 | ) -> Result<(DecodeBuffer<'a>, DecodeBuffer<'a>), DecoderError> { 33 | self.ensure_len(count)?; 34 | let end = self.end; 35 | let (slice, remaining) = self.bytes.split_at_mut(count); 36 | Ok(( 37 | Self::new_checked(slice, end)?, 38 | Self::new_checked(remaining, end)?, 39 | )) 40 | } 41 | 42 | #[inline] 43 | pub fn ensure_len(&self, len: usize) -> Result<(), DecoderError> { 44 | if self.len() < len { 45 | Err(DecoderError::UnexpectedEof(len)) 46 | } else { 47 | Ok(()) 48 | } 49 | } 50 | 51 | #[inline] 52 | pub fn len(&self) -> usize { 53 | self.bytes.len() 54 | } 55 | 56 | #[inline] 57 | pub fn is_empty(&self) -> bool { 58 | self.bytes.is_empty() 59 | } 60 | 61 | #[inline] 62 | pub fn into_bytes(self) -> &'a mut [u8] { 63 | self.bytes 64 | } 65 | } 66 | -------------------------------------------------------------------------------- /xdp/ebpf/src/inet/ecn.rs: -------------------------------------------------------------------------------- 1 | //# https://www.rfc-editor.org/rfc/rfc3168#section-5 2 | #[repr(u8)] 3 | #[derive(Clone, Copy, Debug, PartialEq, PartialOrd, Eq, Ord, Hash)] 4 | pub enum ExplicitCongestionNotification { 5 | /// The not-ECT codepoint '00' indicates a packet that is not using ECN. 6 | NotEct = 0b00, 7 | 8 | /// ECT(1) is set by the data sender to indicate that the end-points of the transport 9 | /// protocol are ECN-capable. 10 | Ect1 = 0b01, 11 | 12 | /// ECT(0) is set by the data sender to indicate that the end-points of the transport 13 | /// protocol are ECN-capable. 14 | /// Protocols and senders that only require a single ECT codepoint SHOULD use ECT(0). 15 | Ect0 = 0b10, 16 | 17 | /// The CE codepoint '11' is set by a router to indicate congestion to the end nodes. 18 | Ce = 0b11, 19 | } 20 | 21 | impl Default for ExplicitCongestionNotification { 22 | #[inline] 23 | fn default() -> Self { 24 | Self::NotEct 25 | } 26 | } 27 | 28 | impl ExplicitCongestionNotification { 29 | /// Create a `ExplicitCongestionNotification` from the ECN field in the IP header 30 | #[inline] 31 | pub fn new(ecn_field: u8) -> Self { 32 | match ecn_field & 0b11 { 33 | 0b00 => ExplicitCongestionNotification::NotEct, 34 | 0b01 => ExplicitCongestionNotification::Ect1, 35 | 0b10 => ExplicitCongestionNotification::Ect0, 36 | 0b11 => ExplicitCongestionNotification::Ce, 37 | _ => unreachable!(), 38 | } 39 | } 40 | 41 | /// Returns true if congestion was experienced by the peer 42 | #[inline] 43 | pub fn congestion_experienced(self) -> bool { 44 | self == Self::Ce 45 | } 46 | 47 | /// Returns true if ECN is in use 48 | #[inline] 49 | pub fn using_ecn(self) -> bool { 50 | self != Self::NotEct 51 | } 52 | } 53 | -------------------------------------------------------------------------------- /xdp/ebpf/src/inet/ethernet.rs: -------------------------------------------------------------------------------- 1 | //# https://www.rfc-editor.org/rfc/rfc826 2 | 3 | use crate::define_inet_type; 4 | 5 | define_inet_type! { 6 | pub struct Header { 7 | destination: MacAddress, 8 | source: MacAddress, 9 | ethertype: EtherType, 10 | } 11 | } 12 | 13 | impl Header { 14 | /// Swaps the direction of the header 15 | #[inline] 16 | pub fn swap(&mut self) { 17 | core::mem::swap(&mut self.source, &mut self.destination); 18 | } 19 | 20 | #[inline] 21 | pub const fn destination(&self) -> &MacAddress { 22 | &self.destination 23 | } 24 | 25 | #[inline] 26 | pub fn destination_mut(&mut self) -> &mut MacAddress { 27 | &mut self.destination 28 | } 29 | 30 | #[inline] 31 | pub const fn source(&self) -> &MacAddress { 32 | &self.source 33 | } 34 | 35 | #[inline] 36 | pub fn source_mut(&mut self) -> &mut MacAddress { 37 | &mut self.source 38 | } 39 | 40 | #[inline] 41 | pub const fn ethertype(&self) -> &EtherType { 42 | &self.ethertype 43 | } 44 | 45 | #[inline] 46 | pub fn ethertype_mut(&mut self) -> &mut EtherType { 47 | &mut self.ethertype 48 | } 49 | } 50 | 51 | const MAC_LEN: usize = 48 / 8; 52 | 53 | define_inet_type! { 54 | pub struct MacAddress { 55 | octets: [u8; MAC_LEN], 56 | } 57 | } 58 | 59 | define_inet_type! { 60 | pub struct EtherType { 61 | id: [u8; 2], 62 | } 63 | } 64 | 65 | macro_rules! impl_type { 66 | ($fun:ident, $cap:ident, $val:expr) => { 67 | pub const $cap: Self = Self { id: $val }; 68 | 69 | #[inline] 70 | pub const fn $fun(self) -> bool { 71 | matches!(self, Self::$cap) 72 | } 73 | }; 74 | } 75 | 76 | impl EtherType { 77 | impl_type!(is_ipv4, IPV4, [0x08, 0x00]); 78 | 79 | impl_type!(is_arp, ARP, [0x08, 0x06]); 80 | 81 | impl_type!(is_ipv6, IPV6, [0x86, 0xDD]); 82 | 83 | impl_type!(is_ppp, PPP, [0x88, 0x0B]); 84 | 85 | impl_type!(is_vlan, VLAN, [0x88, 0xA8]); 86 | } 87 | -------------------------------------------------------------------------------- /xdp/ebpf/src/inet/ip.rs: -------------------------------------------------------------------------------- 1 | use crate::define_inet_type; 2 | 3 | define_inet_type! { 4 | pub struct Protocol { 5 | id: u8, 6 | } 7 | } 8 | 9 | macro_rules! impl_p { 10 | ($fun:ident, $cap:ident, $val:literal) => { 11 | pub const $cap: Self = Self { id: $val }; 12 | 13 | #[inline] 14 | pub const fn $fun(self) -> bool { 15 | self.id == $val 16 | } 17 | }; 18 | } 19 | 20 | impl Protocol { 21 | // https://www.iana.org/assignments/protocol-numbers/protocol-numbers.xhtml 22 | // NOTE: these variants were added as the ones we think we'll need. feel free to add more as 23 | // needed. 24 | impl_p!(is_hop_by_hop, HOPOPT, 0); 25 | 26 | impl_p!(is_icmp, ICMP, 1); 27 | 28 | impl_p!(is_ipv4, IPV4, 4); 29 | 30 | impl_p!(is_tcp, TCP, 6); 31 | 32 | impl_p!(is_udp, UDP, 17); 33 | 34 | impl_p!(is_ipv6, IPV6, 41); 35 | 36 | impl_p!(is_ipv6_route, IPV6_ROUTE, 43); 37 | 38 | impl_p!(is_ipv6_fragment, IPV6_FRAG, 44); 39 | 40 | impl_p!(is_ipv6_icmp, IPV6_ICMP, 58); 41 | 42 | impl_p!(is_ipv6_no_next, IPV6_NO_NXT, 59); 43 | 44 | impl_p!(is_ipv6_options, IPV6_OPTS, 60); 45 | 46 | impl_p!(is_udplite, UDPLITE, 136); 47 | } 48 | -------------------------------------------------------------------------------- /xdp/ebpf/src/inet/ipv4.rs: -------------------------------------------------------------------------------- 1 | use core::fmt; 2 | 3 | use super::zerocopy::U16; 4 | use crate::define_inet_type; 5 | use crate::inet::ecn::ExplicitCongestionNotification; 6 | use crate::inet::ip; 7 | 8 | //# 0 1 2 3 9 | //# 0 1 2 3 4 5 6 7 8 9 0 1 2 3 4 5 6 7 8 9 0 1 2 3 4 5 6 7 8 9 0 1 10 | //# +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+ 11 | //# |Version| IHL |Type of Service| Total Length | 12 | //# +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+ 13 | //# | Identification |Flags| Fragment Offset | 14 | //# +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+ 15 | //# | Time to Live | Protocol | Header Checksum | 16 | //# +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+ 17 | //# | Source Address | 18 | //# +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+ 19 | //# | Destination Address | 20 | //# +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+ 21 | //# | Options | Padding | 22 | //# +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+ 23 | 24 | define_inet_type! { 25 | pub struct Header { 26 | vihl: Vihl, 27 | tos: Tos, 28 | total_len: U16, 29 | id: U16, 30 | flag_fragment: FlagFragment, 31 | ttl: u8, 32 | protocol: ip::Protocol, 33 | checksum: U16, 34 | source: IpV4Address, 35 | destination: IpV4Address, 36 | } 37 | } 38 | 39 | impl Header { 40 | /// Swaps the direction of the header 41 | #[inline] 42 | pub fn swap(&mut self) { 43 | core::mem::swap(&mut self.source, &mut self.destination) 44 | } 45 | 46 | #[inline] 47 | pub const fn vihl(&self) -> &Vihl { 48 | &self.vihl 49 | } 50 | 51 | #[inline] 52 | pub fn vihl_mut(&mut self) -> &mut Vihl { 53 | &mut self.vihl 54 | } 55 | 56 | #[inline] 57 | pub const fn tos(&self) -> &Tos { 58 | &self.tos 59 | } 60 | 61 | #[inline] 62 | pub fn tos_mut(&mut self) -> &mut Tos { 63 | &mut self.tos 64 | } 65 | 66 | #[inline] 67 | pub const fn total_len(&self) -> &U16 { 68 | &self.total_len 69 | } 70 | 71 | #[inline] 72 | pub fn total_len_mut(&mut self) -> &mut U16 { 73 | &mut self.total_len 74 | } 75 | 76 | #[inline] 77 | pub const fn id(&self) -> &U16 { 78 | &self.id 79 | } 80 | 81 | #[inline] 82 | pub fn id_mut(&mut self) -> &mut U16 { 83 | &mut self.id 84 | } 85 | 86 | #[inline] 87 | pub const fn flag_fragment(&self) -> &FlagFragment { 88 | &self.flag_fragment 89 | } 90 | 91 | #[inline] 92 | pub fn flag_fragment_mut(&mut self) -> &mut FlagFragment { 93 | &mut self.flag_fragment 94 | } 95 | 96 | #[inline] 97 | pub const fn ttl(&self) -> &u8 { 98 | &self.ttl 99 | } 100 | 101 | #[inline] 102 | pub fn ttl_mut(&mut self) -> &mut u8 { 103 | &mut self.ttl 104 | } 105 | 106 | #[inline] 107 | pub const fn protocol(&self) -> &ip::Protocol { 108 | &self.protocol 109 | } 110 | 111 | #[inline] 112 | pub fn protocol_mut(&mut self) -> &mut ip::Protocol { 113 | &mut self.protocol 114 | } 115 | 116 | #[inline] 117 | pub const fn checksum(&self) -> &U16 { 118 | &self.checksum 119 | } 120 | 121 | #[inline] 122 | pub fn checksum_mut(&mut self) -> &mut U16 { 123 | &mut self.checksum 124 | } 125 | 126 | #[inline] 127 | pub const fn source(&self) -> &IpV4Address { 128 | &self.source 129 | } 130 | 131 | #[inline] 132 | pub fn source_mut(&mut self) -> &mut IpV4Address { 133 | &mut self.source 134 | } 135 | 136 | #[inline] 137 | pub const fn destination(&self) -> &IpV4Address { 138 | &self.destination 139 | } 140 | 141 | #[inline] 142 | pub fn destination_mut(&mut self) -> &mut IpV4Address { 143 | &mut self.destination 144 | } 145 | 146 | #[inline] 147 | pub fn update_checksum(&mut self) { 148 | // TODO 149 | } 150 | } 151 | 152 | // the bits for version and IHL (header len). 153 | define_inet_type! { 154 | pub struct Vihl { 155 | value: u8, 156 | } 157 | } 158 | 159 | impl fmt::Debug for Vihl { 160 | fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { 161 | f.debug_struct("Vihl") 162 | .field("version", &self.version()) 163 | .field("header_len", &self.header_len()) 164 | .finish() 165 | } 166 | } 167 | 168 | impl Vihl { 169 | #[inline] 170 | fn version(&self) -> u8 { 171 | self.value >> 4 172 | } 173 | 174 | #[inline] 175 | pub fn set_version(&mut self, value: u8) -> &mut Self { 176 | self.value = (value << 4) | (self.value & 0x0F); 177 | self 178 | } 179 | 180 | #[inline] 181 | pub fn header_len(&self) -> u8 { 182 | self.value & 0x0F 183 | } 184 | 185 | #[inline] 186 | pub fn set_header_len(&mut self, value: u8) -> &mut Self { 187 | self.value = (self.value & 0xF0) | (value & 0x0F); 188 | self 189 | } 190 | } 191 | 192 | define_inet_type! { 193 | pub struct Tos { 194 | value: u8, 195 | } 196 | } 197 | 198 | impl Tos { 199 | /// Differentiated Services Code Point 200 | #[inline] 201 | pub fn dscp(&self) -> u8 { 202 | self.value >> 2 203 | } 204 | 205 | #[inline] 206 | pub fn set_dscp(&mut self, value: u8) -> &mut Self { 207 | self.value = (value << 2) | (self.value & 0b11); 208 | self 209 | } 210 | 211 | #[inline] 212 | pub fn ecn(&self) -> ExplicitCongestionNotification { 213 | ExplicitCongestionNotification::new(self.value & 0b11) 214 | } 215 | 216 | #[inline] 217 | pub fn set_ecn(&mut self, ecn: ExplicitCongestionNotification) -> &mut Self { 218 | self.value = (self.value & !0b11) | ecn as u8; 219 | self 220 | } 221 | } 222 | 223 | impl fmt::Debug for Tos { 224 | fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { 225 | f.debug_struct("ipv4::Tos") 226 | .field("dscp", &self.dscp()) 227 | .field("ecn", &self.ecn()) 228 | .finish() 229 | } 230 | } 231 | 232 | define_inet_type!( 233 | pub struct FlagFragment { 234 | value: U16, 235 | } 236 | ); 237 | 238 | impl FlagFragment { 239 | const FRAGMENT_MASK: u16 = 0b0001_1111_1111_1111; 240 | 241 | #[inline] 242 | pub fn reserved(&self) -> bool { 243 | self.get(1 << 15) 244 | } 245 | 246 | pub fn set_reserved(&mut self, enabled: bool) -> &mut Self { 247 | self.set(1 << 15, enabled) 248 | } 249 | 250 | #[inline] 251 | pub fn dont_fragment(&self) -> bool { 252 | self.get(1 << 14) 253 | } 254 | 255 | #[inline] 256 | pub fn set_dont_fragment(&mut self, enabled: bool) -> &mut Self { 257 | self.set(1 << 14, enabled) 258 | } 259 | 260 | #[inline] 261 | pub fn more_fragments(&self) -> bool { 262 | self.get(1 << 13) 263 | } 264 | 265 | #[inline] 266 | pub fn set_more_fragments(&mut self, enabled: bool) -> &mut Self { 267 | self.set(1 << 13, enabled) 268 | } 269 | 270 | #[inline] 271 | pub fn fragment_offset(&self) -> u16 { 272 | self.value.get() & Self::FRAGMENT_MASK 273 | } 274 | 275 | #[inline] 276 | pub fn set_fragment_offset(&mut self, offset: u16) -> &mut Self { 277 | self.value 278 | .set(self.value.get() & !Self::FRAGMENT_MASK | offset & Self::FRAGMENT_MASK); 279 | self 280 | } 281 | 282 | #[inline] 283 | fn get(&self, mask: u16) -> bool { 284 | self.value.get() & mask == mask 285 | } 286 | 287 | #[inline] 288 | fn set(&mut self, mask: u16, enabled: bool) -> &mut Self { 289 | let value = self.value.get(); 290 | let value = if enabled { value | mask } else { value & !mask }; 291 | self.value.set(value); 292 | self 293 | } 294 | } 295 | 296 | impl fmt::Debug for FlagFragment { 297 | fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { 298 | f.debug_struct("ipv4::FlagFragment") 299 | .field("reserved", &self.reserved()) 300 | .field("dont_fragment", &self.dont_fragment()) 301 | .field("more_fragments", &self.more_fragments()) 302 | .field("fragment_offset", &self.fragment_offset()) 303 | .finish() 304 | } 305 | } 306 | 307 | const IPV4_LEN: usize = 32 / 8; 308 | 309 | define_inet_type!( 310 | pub struct IpV4Address { 311 | octets: [u8; IPV4_LEN], 312 | } 313 | ); 314 | -------------------------------------------------------------------------------- /xdp/ebpf/src/inet/ipv6.rs: -------------------------------------------------------------------------------- 1 | use core::fmt; 2 | 3 | use super::ecn::ExplicitCongestionNotification; 4 | use super::zerocopy::U16; 5 | use crate::define_inet_type; 6 | use crate::inet::ip; 7 | 8 | //= https://www.rfc-editor.org/rfc/rfc8200#section-3 9 | //# +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+ 10 | //# |Version| Traffic Class | Flow Label | 11 | //# +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+ 12 | //# | Payload Length | Next Header | Hop Limit | 13 | //# +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+ 14 | //# | | 15 | //# + + 16 | //# | | 17 | //# + Source Address + 18 | //# | | 19 | //# + + 20 | //# | | 21 | //# +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+ 22 | //# | | 23 | //# + + 24 | //# | | 25 | //# + Destination Address + 26 | //# | | 27 | //# + + 28 | //# | | 29 | //# +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+ 30 | 31 | define_inet_type! { 32 | pub struct Header { 33 | vtcfl: Vtcfl, 34 | payload_len: U16, 35 | next_header: ip::Protocol, 36 | hop_limit: u8, 37 | source: IpV6Address, 38 | destination: IpV6Address, 39 | } 40 | } 41 | 42 | impl Header { 43 | /// Swaps the direction of the header 44 | #[inline] 45 | pub fn swap(&mut self) { 46 | core::mem::swap(&mut self.source, &mut self.destination); 47 | } 48 | 49 | #[inline] 50 | pub const fn vtcfl(&self) -> &Vtcfl { 51 | &self.vtcfl 52 | } 53 | 54 | #[inline] 55 | pub fn vtcfl_mut(&mut self) -> &mut Vtcfl { 56 | &mut self.vtcfl 57 | } 58 | 59 | #[inline] 60 | pub const fn payload_len(&self) -> &U16 { 61 | &self.payload_len 62 | } 63 | 64 | #[inline] 65 | pub fn payload_len_mut(&mut self) -> &mut U16 { 66 | &mut self.payload_len 67 | } 68 | 69 | #[inline] 70 | pub const fn next_header(&self) -> &ip::Protocol { 71 | &self.next_header 72 | } 73 | 74 | #[inline] 75 | pub fn next_header_mut(&mut self) -> &mut ip::Protocol { 76 | &mut self.next_header 77 | } 78 | 79 | #[inline] 80 | pub const fn hop_limit(&self) -> &u8 { 81 | &self.hop_limit 82 | } 83 | 84 | #[inline] 85 | pub fn hop_limit_mut(&mut self) -> &mut u8 { 86 | &mut self.hop_limit 87 | } 88 | 89 | #[inline] 90 | pub const fn source(&self) -> &IpV6Address { 91 | &self.source 92 | } 93 | 94 | #[inline] 95 | pub fn source_mut(&mut self) -> &mut IpV6Address { 96 | &mut self.source 97 | } 98 | 99 | #[inline] 100 | pub const fn destination(&self) -> &IpV6Address { 101 | &self.destination 102 | } 103 | 104 | #[inline] 105 | pub fn destination_mut(&mut self) -> &mut IpV6Address { 106 | &mut self.destination 107 | } 108 | } 109 | 110 | define_inet_type! { 111 | pub struct Vtcfl { 112 | octets: [u8; 4], 113 | } 114 | } 115 | 116 | impl fmt::Debug for Vtcfl { 117 | fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { 118 | f.debug_struct("ipv6::Vtf") 119 | .field("version", &self.version()) 120 | .field("dscp", &self.dscp()) 121 | .field("ecn", &self.ecn()) 122 | .field("flow_label", &format_args!("0x{:05x}", self.flow_label())) 123 | .finish() 124 | } 125 | } 126 | 127 | impl Vtcfl { 128 | #[inline] 129 | pub const fn version(&self) -> u8 { 130 | self.octets[0] >> 4 131 | } 132 | 133 | #[inline] 134 | pub fn set_version(&mut self, version: u8) -> &mut Self { 135 | self.octets[0] = (version << 4) | self.octets[0] & 0x0F; 136 | self 137 | } 138 | 139 | #[inline] 140 | pub fn dscp(&self) -> u8 { 141 | let value = (self.octets[0] << 4) | (self.octets[1] >> 4); 142 | value >> 2 143 | } 144 | 145 | #[inline] 146 | pub fn set_dscp(&mut self, value: u8) -> &mut Self { 147 | let value = value << 2; 148 | self.octets[0] = self.octets[0] & 0xF0 | (value >> 4); 149 | self.octets[1] = (value << 4) | self.octets[1] & 0b11_1111; 150 | self 151 | } 152 | 153 | #[inline] 154 | pub fn ecn(&self) -> ExplicitCongestionNotification { 155 | ExplicitCongestionNotification::new((self.octets[1] >> 4) & 0b11) 156 | } 157 | 158 | #[inline] 159 | pub fn set_ecn(&mut self, ecn: ExplicitCongestionNotification) -> &mut Self { 160 | self.octets[1] = (self.octets[1] & !(0b11 << 4)) | ((ecn as u8) << 4); 161 | self 162 | } 163 | 164 | #[inline] 165 | pub const fn flow_label(&self) -> u32 { 166 | u32::from_be_bytes([0, self.octets[1] & 0x0F, self.octets[2], self.octets[3]]) 167 | } 168 | 169 | #[inline] 170 | pub fn set_flow_label(&mut self, flow_label: u32) -> &mut Self { 171 | let bytes = flow_label.to_be_bytes(); 172 | self.octets[1] = self.octets[1] & 0xF0 | bytes[1] & 0x0F; 173 | self.octets[2] = bytes[2]; 174 | self.octets[3] = bytes[3]; 175 | self 176 | } 177 | } 178 | 179 | const IPV6_LEN: usize = 128 / 8; 180 | 181 | define_inet_type!( 182 | pub struct IpV6Address { 183 | octets: [u8; IPV6_LEN], 184 | } 185 | ); 186 | -------------------------------------------------------------------------------- /xdp/ebpf/src/inet/mod.rs: -------------------------------------------------------------------------------- 1 | pub mod ecn; 2 | pub mod ethernet; 3 | pub mod ip; 4 | pub mod ipv4; 5 | pub mod ipv6; 6 | pub mod udp; 7 | 8 | pub mod zerocopy; 9 | 10 | #[macro_export] 11 | macro_rules! define_inet_type { 12 | ($($vis:ident)? struct $name:ident { 13 | $( 14 | $field:ident: $field_ty:ty 15 | ),* 16 | $(,)* 17 | }) => { 18 | #[derive(Clone, Copy, Default, Eq, PartialEq, PartialOrd, Ord, zerocopy::FromZeroes, zerocopy::FromBytes, zerocopy::AsBytes, zerocopy::Unaligned)] 19 | #[repr(C)] 20 | $($vis)? struct $name { 21 | $( 22 | pub(crate) $field: $field_ty, 23 | )* 24 | } 25 | 26 | // By letting the compiler derive PartialEq, we can do structural matching on these 27 | // structs. But we also want hashing to be on the byte level, since we know the struct 28 | // has no allocations and has a simple layout. 29 | // 30 | // See: https://godbolt.org/z/czohnrWxK 31 | #[allow(clippy::derived_hash_with_manual_eq)] 32 | impl core::hash::Hash for $name { 33 | #[inline] 34 | fn hash(&self, hasher: &mut H) { 35 | self.as_bytes().hash(hasher); 36 | } 37 | } 38 | 39 | impl $name { 40 | #[allow(non_camel_case_types)] 41 | #[allow(clippy::too_many_arguments)] 42 | #[inline] 43 | pub fn new<$($field: Into<$field_ty>),*>($($field: $field),*) -> Self { 44 | Self { 45 | $( 46 | $field: $field.into() 47 | ),* 48 | } 49 | } 50 | 51 | #[inline] 52 | pub fn as_bytes(&self) -> &[u8] { 53 | zerocopy::AsBytes::as_bytes(self) 54 | } 55 | 56 | #[inline] 57 | pub fn as_bytes_mut(&mut self) -> &mut [u8] { 58 | zerocopy::AsBytes::as_bytes_mut(self) 59 | } 60 | 61 | #[inline] 62 | pub fn decode(buffer: $crate::buffer::DecodeBuffer<'_>) -> Result<(&Self, $crate::buffer::DecodeBuffer<'_>), $crate::buffer::DecoderError> { 63 | let (value, buffer) = buffer.decode_slice(core::mem::size_of::<$name>())?; 64 | let value = value.into_bytes(); 65 | let value = unsafe { 66 | // Safety: the type implements FromBytes 67 | &*(value as *const _ as *const $name) 68 | }; 69 | Ok((value, buffer)) 70 | } 71 | } 72 | }; 73 | } 74 | -------------------------------------------------------------------------------- /xdp/ebpf/src/inet/udp.rs: -------------------------------------------------------------------------------- 1 | use core::fmt; 2 | 3 | use super::zerocopy::U16; 4 | use crate::define_inet_type; 5 | 6 | //# ------ 7 | //# 0 7 8 15 16 23 24 31 8 | //# +--------+--------+--------+--------+ 9 | //# | Source | Destination | 10 | //# | Port | Port | 11 | //# +--------+--------+--------+--------+ 12 | //# | | | 13 | //# | Length | Checksum | 14 | //# +--------+--------+--------+--------+ 15 | //# | 16 | //# | data octets ... 17 | //# +---------------- ... 18 | 19 | define_inet_type!( 20 | pub struct Header { 21 | source: U16, 22 | destination: U16, 23 | len: U16, 24 | checksum: U16, 25 | } 26 | ); 27 | 28 | impl Header { 29 | /// Swaps the direction of the header 30 | #[inline] 31 | pub fn swap(&mut self) { 32 | core::mem::swap(&mut self.source, &mut self.destination); 33 | } 34 | 35 | #[inline] 36 | pub const fn source(&self) -> &U16 { 37 | &self.source 38 | } 39 | 40 | #[inline] 41 | pub fn source_mut(&mut self) -> &mut U16 { 42 | &mut self.source 43 | } 44 | 45 | #[inline] 46 | pub const fn destination(&self) -> &U16 { 47 | &self.destination 48 | } 49 | 50 | #[inline] 51 | pub fn destination_mut(&mut self) -> &mut U16 { 52 | &mut self.destination 53 | } 54 | 55 | #[inline] 56 | pub const fn len(&self) -> &U16 { 57 | &self.len 58 | } 59 | 60 | #[inline] 61 | pub fn len_mut(&mut self) -> &mut U16 { 62 | &mut self.len 63 | } 64 | 65 | #[inline] 66 | pub const fn checksum(&self) -> &U16 { 67 | &self.checksum 68 | } 69 | 70 | #[inline] 71 | pub fn checksum_mut(&mut self) -> &mut U16 { 72 | &mut self.checksum 73 | } 74 | } 75 | 76 | impl fmt::Debug for Header { 77 | fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { 78 | f.debug_struct("udp::Header") 79 | .field("source", &self.source) 80 | .field("destination", &self.destination) 81 | .field("len", &self.len) 82 | .field("checksum", &format_args!("0x{:04x}", &self.checksum.get())) 83 | .finish() 84 | } 85 | } 86 | -------------------------------------------------------------------------------- /xdp/ebpf/src/inet/zerocopy.rs: -------------------------------------------------------------------------------- 1 | use core::cmp::Ordering; 2 | use core::fmt; 3 | use core::hash::{Hash, Hasher}; 4 | 5 | pub use zerocopy::*; 6 | 7 | macro_rules! zerocopy_network_integer { 8 | ($native:ident, $name:ident) => { 9 | #[derive( 10 | Clone, 11 | Copy, 12 | Default, 13 | Eq, 14 | $crate::inet::zerocopy::FromBytes, 15 | $crate::inet::zerocopy::FromZeroes, 16 | $crate::inet::zerocopy::AsBytes, 17 | $crate::inet::zerocopy::Unaligned, 18 | )] 19 | #[repr(C)] 20 | pub struct $name(::zerocopy::byteorder::$name); 21 | 22 | impl $name { 23 | pub const ZERO: Self = Self(::zerocopy::byteorder::$name::ZERO); 24 | 25 | #[inline(always)] 26 | pub fn new(value: $native) -> Self { 27 | value.into() 28 | } 29 | 30 | #[inline(always)] 31 | pub fn get(&self) -> $native { 32 | self.get_be().to_be() 33 | } 34 | 35 | #[inline(always)] 36 | pub fn get_be(&self) -> $native { 37 | unsafe { 38 | $native::from_ne_bytes( 39 | *(self.0.as_bytes().as_ptr() 40 | as *const [u8; ::core::mem::size_of::<$native>()]), 41 | ) 42 | } 43 | } 44 | 45 | #[inline(always)] 46 | pub fn set(&mut self, value: $native) { 47 | self.0.as_bytes_mut().copy_from_slice(&value.to_be_bytes()); 48 | } 49 | 50 | #[inline(always)] 51 | pub fn set_be(&mut self, value: $native) { 52 | self.0.as_bytes_mut().copy_from_slice(&value.to_ne_bytes()); 53 | } 54 | } 55 | 56 | impl PartialEq for $name { 57 | #[inline] 58 | fn eq(&self, other: &Self) -> bool { 59 | self.cmp(other) == Ordering::Equal 60 | } 61 | } 62 | 63 | impl PartialEq<$native> for $name { 64 | #[inline] 65 | fn eq(&self, other: &$native) -> bool { 66 | self.partial_cmp(other) == Some(Ordering::Equal) 67 | } 68 | } 69 | 70 | impl PartialOrd for $name { 71 | #[inline] 72 | fn partial_cmp(&self, other: &Self) -> Option { 73 | Some(self.cmp(other)) 74 | } 75 | } 76 | 77 | impl PartialOrd<$native> for $name { 78 | #[inline] 79 | fn partial_cmp(&self, other: &$native) -> Option { 80 | Some(self.get().cmp(other)) 81 | } 82 | } 83 | 84 | impl Ord for $name { 85 | #[inline] 86 | fn cmp(&self, other: &Self) -> Ordering { 87 | self.get_be().cmp(&other.get_be()) 88 | } 89 | } 90 | 91 | impl Hash for $name { 92 | fn hash(&self, state: &mut H) { 93 | self.get().hash(state); 94 | } 95 | } 96 | 97 | impl fmt::Debug for $name { 98 | fn fmt(&self, formatter: &mut fmt::Formatter) -> fmt::Result { 99 | write!(formatter, "{}", self.get()) 100 | } 101 | } 102 | 103 | impl fmt::Display for $name { 104 | fn fmt(&self, formatter: &mut fmt::Formatter) -> fmt::Result { 105 | write!(formatter, "{}", self.get()) 106 | } 107 | } 108 | 109 | impl From<$native> for $name { 110 | #[inline] 111 | fn from(value: $native) -> Self { 112 | Self(::zerocopy::byteorder::$name::new(value)) 113 | } 114 | } 115 | 116 | impl From<$name> for $native { 117 | #[inline] 118 | fn from(v: $name) -> $native { 119 | v.get() 120 | } 121 | } 122 | }; 123 | } 124 | 125 | zerocopy_network_integer!(i16, I16); 126 | zerocopy_network_integer!(u16, U16); 127 | zerocopy_network_integer!(i32, I32); 128 | zerocopy_network_integer!(u32, U32); 129 | zerocopy_network_integer!(i64, I64); 130 | zerocopy_network_integer!(u64, U64); 131 | zerocopy_network_integer!(i128, I128); 132 | zerocopy_network_integer!(u128, U128); 133 | -------------------------------------------------------------------------------- /xdp/ebpf/src/main.rs: -------------------------------------------------------------------------------- 1 | #![no_std] 2 | #![no_main] 3 | #![allow(unused)] 4 | 5 | use aya_bpf::bindings::xdp_action; 6 | use aya_bpf::macros::{map, xdp}; 7 | use aya_bpf::maps::{HashMap, XskMap}; 8 | use aya_bpf::programs::XdpContext; 9 | use buffer::{DecodeBuffer, DecoderError}; 10 | use inet::ethernet::{self, EtherType}; 11 | use inet::{ip, ipv4, ipv6, udp}; 12 | 13 | mod buffer; 14 | mod inet; 15 | 16 | // Never panic 17 | #[cfg(not(test))] 18 | #[panic_handler] 19 | fn panic(_info: &core::panic::PanicInfo) -> ! { 20 | unsafe { core::hint::unreachable_unchecked() } 21 | } 22 | 23 | #[map(name = "RAKNET_XDP_SOCKETS")] 24 | static XDP_SOCKETS: XskMap = XskMap::with_max_entries(1024, 0); 25 | 26 | #[map(name = "RAKNET_XDP_PORTS")] 27 | static XDP_PORTS: HashMap = HashMap::with_max_entries(1024, 0); 28 | 29 | #[xdp] 30 | pub fn raknet_xdp(ctx: XdpContext) -> u32 { 31 | let action = redirect_udp(&ctx); 32 | 33 | #[cfg(feature = "trace")] 34 | { 35 | use aya_log_ebpf as log; 36 | 37 | match action { 38 | xdp_action::XDP_DROP => log::trace!(&ctx, "ACTION: DROP"), 39 | xdp_action::XDP_PASS => log::trace!(&ctx, "ACTION: PASS"), 40 | xdp_action::XDP_REDIRECT => log::trace!(&ctx, "ACTION: REDIRECT"), 41 | xdp_action::XDP_ABORTED => log::trace!(&ctx, "ACTION: ABORTED"), 42 | _ => (), 43 | } 44 | } 45 | 46 | action 47 | } 48 | 49 | fn redirect_udp(ctx: &XdpContext) -> u32 { 50 | let start = ctx.data() as *mut u8; 51 | let end = ctx.data_end() as *mut u8; 52 | let buffer = unsafe { 53 | // Safety: start and end come from the caller and have been validated 54 | DecodeBuffer::new(start, end) 55 | }; 56 | match decode(buffer) { 57 | Ok(Some(payload)) => { 58 | // if the payload is empty there isn't much we can do with it 59 | if payload.is_empty() { 60 | return xdp_action::XDP_DROP; 61 | } 62 | 63 | // if the packet is valid forward it on to the associated AF_XDP socket 64 | let queue_id = unsafe { (*ctx.ctx).rx_queue_index }; 65 | XDP_SOCKETS 66 | .redirect(queue_id, 0) 67 | .unwrap_or(xdp_action::XDP_PASS) 68 | } 69 | Ok(None) => xdp_action::XDP_PASS, 70 | Err(_) => xdp_action::XDP_ABORTED, 71 | } 72 | } 73 | 74 | #[inline] 75 | fn decode(buffer: DecodeBuffer<'_>) -> Result>, DecoderError> { 76 | let (header, buffer) = ethernet::Header::decode(buffer)?; 77 | match *header.ethertype() { 78 | EtherType::IPV4 => decode_ipv4(buffer), 79 | EtherType::IPV6 => decode_ipv6(buffer), 80 | // pass the packet on to the OS network stack if we don't understand it 81 | _ => Ok(None), 82 | } 83 | } 84 | 85 | #[inline] 86 | fn decode_ipv4(buffer: DecodeBuffer<'_>) -> Result>, DecoderError> { 87 | let (header, buffer) = ipv4::Header::decode(buffer)?; 88 | let protocol = header.protocol(); 89 | 90 | //= https://www.rfc-editor.org/rfc/rfc791#section-3.1 91 | //# IHL: 4 bits 92 | //# 93 | //# Internet Header Length is the length of the internet header in 32 94 | //# bit words, and thus points to the beginning of the data. Note that 95 | //# the minimum value for a correct header is 5. 96 | 97 | // subtract the fixed header size 98 | let count_without_header = header 99 | .vihl() 100 | .header_len() 101 | .checked_sub(5) 102 | .ok_or(DecoderError::InvariantViolation("invalid IPv4 IHL value"))?; 103 | 104 | // skip the options and go to the actual payload 105 | let options_len = count_without_header as usize * (32 / 8); 106 | let (_options, buffer) = buffer.decode_slice(options_len)?; 107 | 108 | parse_ip_protocol(*protocol, buffer) 109 | } 110 | 111 | #[inline] 112 | fn decode_ipv6(buffer: DecodeBuffer<'_>) -> Result>, DecoderError> { 113 | let (header, buffer) = ipv6::Header::decode(buffer)?; 114 | let protocol = header.next_header(); 115 | 116 | // TODO parse Hop-by-hop/Options headers, for now we'll just forward the packet on to the OS 117 | 118 | parse_ip_protocol(*protocol, buffer) 119 | } 120 | 121 | #[inline] 122 | fn parse_ip_protocol( 123 | protocol: ip::Protocol, 124 | buffer: DecodeBuffer<'_>, 125 | ) -> Result>, DecoderError> { 126 | match protocol { 127 | ip::Protocol::UDP | ip::Protocol::UDPLITE => parse_udp(buffer), 128 | // pass the packet on to the OS network stack if we don't understand it 129 | _ => Ok(None), 130 | } 131 | } 132 | 133 | #[inline] 134 | fn parse_udp(buffer: DecodeBuffer<'_>) -> Result>, DecoderError> { 135 | let (header, buffer) = udp::Header::decode(buffer)?; 136 | 137 | // Make sure the port is in the port map. Otherwise, forward the packet to the OS. 138 | if XDP_PORTS.get_ptr(&header.destination().get()).is_none() { 139 | return Ok(None); 140 | } 141 | 142 | // NOTE: duvet doesn't know how to parse this RFC since it doesn't follow more modern formatting 143 | //# https://www.rfc-editor.org/rfc/rfc768 144 | //# Length is the length in octets of this user datagram including this 145 | //# header and the data. (This means the minimum value of the length is 146 | //# eight.) 147 | let total_len = header.len().get(); 148 | let payload_len = total_len 149 | .checked_sub(8) 150 | .ok_or(DecoderError::InvariantViolation("invalid UDP length"))?; 151 | let (udp_payload, _remaining) = buffer.decode_slice(payload_len as usize)?; 152 | 153 | Ok(Some(udp_payload)) 154 | } 155 | -------------------------------------------------------------------------------- /xdp/raknet-xdp/Cargo.toml: -------------------------------------------------------------------------------- 1 | [package] 2 | name = "raknet-xdp" 3 | version = "0.1.0" 4 | edition = { workspace = true } 5 | license = { workspace = true } 6 | authors = { workspace = true } 7 | homepage = { workspace = true } 8 | repository = { workspace = true } 9 | categories = { workspace = true } 10 | keywords = { workspace = true } 11 | 12 | [dependencies] 13 | aya = { version = "0.13", default-features = false } 14 | bitflags = "2" 15 | errno = "0.3" 16 | libc = "0.2" 17 | 18 | [dev-dependencies] 19 | env_logger = "0.11" 20 | 21 | [build-dependencies] 22 | bindgen = "0.71" 23 | cargo_metadata = "0.19" 24 | 25 | [lints] 26 | workspace = true 27 | -------------------------------------------------------------------------------- /xdp/raknet-xdp/bindings/input.h: -------------------------------------------------------------------------------- 1 | #include 2 | -------------------------------------------------------------------------------- /xdp/raknet-xdp/build.rs: -------------------------------------------------------------------------------- 1 | use std::io::{BufRead, BufReader}; 2 | use std::path::{Path, PathBuf}; 3 | use std::process::{Child, Command, Stdio}; 4 | use std::{env, fs}; 5 | 6 | use cargo_metadata::{ 7 | Artifact, CompilerMessage, Message, Metadata, MetadataCommand, Package, Target, 8 | }; 9 | 10 | const RECOMPILE_OPT: &str = "RAK_XDP_EBPF_COMPILE"; 11 | const TRACE_OPT: &str = "RAK_XDP_EBFP_TRACE"; 12 | 13 | fn main() { 14 | println!("cargo:rerun-if-env-changed={}", RECOMPILE_OPT); 15 | println!("cargo:rerun-if-env-changed={}", TRACE_OPT); 16 | 17 | let enable_trace = env::var(TRACE_OPT).unwrap_or_default(); 18 | let out_dir = env::var_os("OUT_DIR").unwrap(); 19 | let out_dir = PathBuf::from(out_dir); 20 | let endian = env::var_os("CARGO_CFG_TARGET_ENDIAN").unwrap(); 21 | let target = if endian == "big" { 22 | "bpfeb-unknown-none" 23 | } else if endian == "little" { 24 | "bpfel-unknown-none" 25 | } else { 26 | panic!("unsupported endian={:?}", endian) 27 | }; 28 | let Metadata { packages, .. } = MetadataCommand::new().no_deps().exec().unwrap(); 29 | let ebpf_package = packages 30 | .into_iter() 31 | .find(|Package { name, .. }| name == "ebpf") 32 | .unwrap(); 33 | let Package { manifest_path, .. } = ebpf_package; 34 | let ebpf_dir = manifest_path.parent().unwrap(); 35 | 36 | println!("cargo:rerun-if-changed={}", ebpf_dir.as_str()); 37 | 38 | // build all bins in ebpf 39 | let mut cmd = Command::new("cargo"); 40 | cmd.current_dir(ebpf_dir); 41 | cmd.args([ 42 | "build", 43 | "-Z", 44 | "build-std=core", 45 | "--bins", 46 | "--message-format=json", 47 | "--release", 48 | "--target", 49 | target, 50 | ]); 51 | if matches!(enable_trace.to_lowercase().as_str(), "1" | "on" | "true") { 52 | cmd.args(["--features", "trace"]); 53 | } 54 | 55 | let mut child = cmd 56 | .stdout(Stdio::piped()) 57 | .stderr(Stdio::piped()) 58 | .spawn() 59 | .unwrap_or_else(|err| panic!("failed to spawn {cmd:?}: {err}")); 60 | let Child { stdout, stderr, .. } = &mut child; 61 | 62 | // Trampoline stdout to cargo warnings. 63 | let stderr = BufReader::new(stderr.take().unwrap()); 64 | let stderr = std::thread::spawn(move || { 65 | for line in stderr.lines() { 66 | let line = line.unwrap(); 67 | println!("cargo:warning={line}"); 68 | } 69 | }); 70 | 71 | let stdout = BufReader::new(stdout.take().unwrap()); 72 | let mut executables = Vec::new(); 73 | for message in Message::parse_stream(stdout) { 74 | #[allow(clippy::collapsible_match)] 75 | match message.expect("valid JSON") { 76 | Message::CompilerArtifact(Artifact { 77 | executable, 78 | target: Target { name, .. }, 79 | .. 80 | }) => { 81 | if let Some(executable) = executable { 82 | executables.push((name, executable.into_std_path_buf())); 83 | } 84 | } 85 | Message::CompilerMessage(CompilerMessage { message: msg, .. }) => { 86 | for line in msg.rendered.unwrap_or_default().split('\n') { 87 | println!("cargo:warning={line}"); 88 | } 89 | } 90 | Message::TextLine(line) => { 91 | println!("cargo:warning={line}"); 92 | } 93 | _ => {} 94 | } 95 | } 96 | 97 | let status = child 98 | .wait() 99 | .unwrap_or_else(|err| panic!("failed to wait for {cmd:?}: {err}")); 100 | assert_eq!(status.code(), Some(0), "{cmd:?} failed: {status:?}"); 101 | 102 | stderr.join().map_err(std::panic::resume_unwind).unwrap(); 103 | 104 | // copy to $OUT_DIR 105 | for (name, binary) in executables { 106 | let dst = out_dir.join(name); 107 | let _: u64 = fs::copy(&binary, &dst) 108 | .unwrap_or_else(|err| panic!("failed to copy {binary:?} to {dst:?}: {err}")); 109 | } 110 | 111 | make_binding(); 112 | } 113 | 114 | fn make_binding() { 115 | let root = Path::new(env!("CARGO_MANIFEST_DIR")); 116 | 117 | let bindings = bindgen::Builder::default() 118 | .header(root.join("bindings/input.h").display().to_string()) 119 | .allowlist_var("ETHTOOL_GCHANNELS") 120 | .allowlist_type("ethtool_channels") 121 | .rust_target(bindgen::RustTarget::Stable_1_47) 122 | .layout_tests(false) 123 | .raw_line( 124 | r#" 125 | #![allow(non_camel_case_types)] 126 | "# 127 | .trim(), 128 | ) 129 | .generate() 130 | .expect("gen binding failed"); 131 | 132 | let out = root.join("src/bindings.rs"); 133 | bindings 134 | .write_to_file(out) 135 | .expect("write binding.rs failed"); 136 | } 137 | -------------------------------------------------------------------------------- /xdp/raknet-xdp/src/bindings.rs: -------------------------------------------------------------------------------- 1 | /* automatically generated by rust-bindgen 0.71.1 */ 2 | 3 | #![allow(non_camel_case_types)] 4 | 5 | pub const ETHTOOL_GCHANNELS: u32 = 60; 6 | pub type __u32 = ::std::os::raw::c_uint; 7 | #[repr(C)] 8 | #[derive(Debug, Copy, Clone)] 9 | pub struct ethtool_channels { 10 | pub cmd: __u32, 11 | pub max_rx: __u32, 12 | pub max_tx: __u32, 13 | pub max_other: __u32, 14 | pub max_combined: __u32, 15 | pub rx_count: __u32, 16 | pub tx_count: __u32, 17 | pub other_count: __u32, 18 | pub combined_count: __u32, 19 | } 20 | -------------------------------------------------------------------------------- /xdp/raknet-xdp/src/io/mod.rs: -------------------------------------------------------------------------------- 1 | mod rx; 2 | mod tx; 3 | -------------------------------------------------------------------------------- /xdp/raknet-xdp/src/io/rx.rs: -------------------------------------------------------------------------------- 1 | 2 | -------------------------------------------------------------------------------- /xdp/raknet-xdp/src/io/tx.rs: -------------------------------------------------------------------------------- 1 | 2 | -------------------------------------------------------------------------------- /xdp/raknet-xdp/src/lib.rs: -------------------------------------------------------------------------------- 1 | use aya::include_bytes_aligned; 2 | 3 | pub const DEFAULT_PROGRAM: &[u8] = 4 | include_bytes_aligned!(concat!(env!("OUT_DIR"), "/raknet-xdp-ebpf")); 5 | pub const DEFAULT_PROGRAM_NAME: &str = "raknet_xdp"; 6 | 7 | /// The name of the `AF_XDP` socket map 8 | pub static XSK_MAP_NAME: &str = "RAKNET_XDP_SOCKETS"; 9 | 10 | /// The name of the port map 11 | pub static PORT_MAP_NAME: &str = "RAKNET_XDP_PORTS"; 12 | 13 | type Result = core::result::Result; 14 | 15 | #[rustfmt::skip] 16 | mod bindings; 17 | 18 | mod if_xdp; 19 | mod io; 20 | mod mmap; 21 | mod socket; 22 | mod syscall; 23 | mod umem; 24 | -------------------------------------------------------------------------------- /xdp/raknet-xdp/src/mmap.rs: -------------------------------------------------------------------------------- 1 | // Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. 2 | // SPDX-License-Identifier: Apache-2.0 3 | 4 | use core::ffi::c_void; 5 | use core::ops::{Deref, DerefMut}; 6 | use core::ptr::NonNull; 7 | use std::os::unix::io::RawFd; 8 | 9 | use crate::syscall::{mmap, munmap}; 10 | use crate::Result; 11 | 12 | /// A mmap'd region in memory 13 | #[derive(Debug)] 14 | pub struct Mmap { 15 | addr: NonNull, 16 | len: usize, 17 | } 18 | 19 | #[derive(Debug)] 20 | pub enum Options { 21 | Huge, 22 | Fd(RawFd), 23 | } 24 | 25 | /// Safety: Mmap pointer can be sent between threads 26 | unsafe impl Send for Mmap {} 27 | 28 | /// Safety: Mmap pointer can be shared between threads 29 | unsafe impl Sync for Mmap {} 30 | 31 | impl Mmap { 32 | /// Creates a new mmap'd region, with an optional file descriptor. 33 | #[inline] 34 | pub fn new(len: usize, offset: usize, flags: Option) -> Result { 35 | let addr = match flags { 36 | Some(Options::Huge) => mmap(len, offset, None, true), 37 | Some(Options::Fd(fd)) => mmap(len, offset, Some(fd), false), 38 | None => mmap(len, offset, None, false), 39 | }?; 40 | Ok(Self { addr, len }) 41 | } 42 | 43 | /// Returns the raw address for the mmap region 44 | #[inline] 45 | pub fn addr(&self) -> NonNull { 46 | self.addr 47 | } 48 | } 49 | 50 | impl Deref for Mmap { 51 | type Target = [u8]; 52 | 53 | #[inline] 54 | fn deref(&self) -> &Self::Target { 55 | unsafe { core::slice::from_raw_parts(self.addr.as_ptr() as _, self.len) } 56 | } 57 | } 58 | 59 | impl DerefMut for Mmap { 60 | #[inline] 61 | fn deref_mut(&mut self) -> &mut Self::Target { 62 | unsafe { core::slice::from_raw_parts_mut(self.addr.as_ptr() as _, self.len) } 63 | } 64 | } 65 | 66 | impl Drop for Mmap { 67 | #[inline] 68 | fn drop(&mut self) { 69 | let _ = unsafe { 70 | // Safety: the len is the same value as on creation 71 | munmap(self.addr, self.len) 72 | }; 73 | } 74 | } 75 | 76 | #[cfg(test)] 77 | mod tests { 78 | use super::*; 79 | 80 | #[test] 81 | fn mmap_test() { 82 | let mut area = Mmap::new(32, 0, None).unwrap(); 83 | assert_eq!(area.len(), 32); 84 | let _ = &area[..]; 85 | let _ = &mut area[..]; 86 | } 87 | } 88 | -------------------------------------------------------------------------------- /xdp/raknet-xdp/src/socket.rs: -------------------------------------------------------------------------------- 1 | // Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. 2 | // SPDX-License-Identifier: Apache-2.0 3 | 4 | use core::fmt; 5 | use std::os::unix::io::{AsRawFd, RawFd}; 6 | use std::sync::Arc; 7 | 8 | use crate::{syscall, Result}; 9 | 10 | /// A structure for reference counting an AF-XDP socket 11 | #[derive(Clone, PartialEq, Eq)] 12 | pub struct Fd(Arc); 13 | 14 | impl fmt::Debug for Fd { 15 | fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { 16 | f.debug_tuple("Fd").field(&(self.0).0).finish() 17 | } 18 | } 19 | 20 | impl Fd { 21 | /// Opens an AF_XDP socket 22 | /// 23 | /// This call requires `CAP_NET_RAW` capabilities to succeed. 24 | #[inline] 25 | pub fn open() -> Result { 26 | let fd = syscall::open()?; 27 | let fd = Arc::new(Inner(fd)); 28 | Ok(Self(fd)) 29 | } 30 | 31 | pub fn attach_umem(&self, umem: &crate::umem::Umem) -> Result<()> { 32 | umem.attach(self)?; 33 | // TODO store the umem 34 | Ok(()) 35 | } 36 | 37 | /// Creates a socket from a raw file descriptor 38 | /// 39 | /// This can be useful for automatically cleaning up a socket on drop 40 | pub(crate) fn from_raw(value: RawFd) -> Self { 41 | Self(Arc::new(Inner(value))) 42 | } 43 | } 44 | 45 | impl AsRawFd for Fd { 46 | #[inline] 47 | fn as_raw_fd(&self) -> RawFd { 48 | (self.0).0 49 | } 50 | } 51 | 52 | /// Wrap the RawFd in a structure that automatically closes the socket on drop 53 | #[derive(PartialEq, Eq)] 54 | struct Inner(RawFd); 55 | 56 | impl Drop for Inner { 57 | fn drop(&mut self) { 58 | unsafe { 59 | let _ = libc::close(self.0); 60 | } 61 | } 62 | } 63 | -------------------------------------------------------------------------------- /xdp/raknet-xdp/src/umem.rs: -------------------------------------------------------------------------------- 1 | // Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. 2 | // SPDX-License-Identifier: Apache-2.0 3 | 4 | use core::ptr::NonNull; 5 | use std::os::unix::io::AsRawFd; 6 | use std::sync::Arc; 7 | 8 | use crate::if_xdp::{RxTxDescriptor, UmemDescriptor, UmemFlags, UmemReg}; 9 | use crate::mmap::{self, Mmap}; 10 | use crate::{syscall, Result}; 11 | 12 | /// The default value for frame sizes 13 | pub const DEFAULT_FRAME_SIZE: u32 = 4096; 14 | 15 | #[derive(Clone, Copy, Debug)] 16 | pub struct Builder { 17 | /// The maximum number of bytes a frame can hold (MTU) 18 | pub frame_size: u32, 19 | /// The number of frames that should be allocated 20 | pub frame_count: u32, 21 | /// The headroom size for each frame 22 | pub frame_headroom: u32, 23 | /// The flags for the Umem 24 | pub flags: UmemFlags, 25 | /// Back the umem with a hugepage 26 | pub hugepage: bool, 27 | } 28 | 29 | impl Default for Builder { 30 | fn default() -> Self { 31 | Self { 32 | frame_size: DEFAULT_FRAME_SIZE, 33 | frame_count: 1024, 34 | frame_headroom: 0, 35 | flags: Default::default(), 36 | hugepage: false, 37 | } 38 | } 39 | } 40 | 41 | impl Builder { 42 | pub fn build(self) -> Result { 43 | let len = self.frame_size as usize * self.frame_count as usize; 44 | let options = if self.hugepage { 45 | Some(mmap::Options::Huge) 46 | } else { 47 | None 48 | }; 49 | let area = Mmap::new(len, 0, options)?; 50 | let area = Arc::new(area); 51 | let mem = area.addr().cast(); 52 | 53 | Ok(Umem { 54 | area, 55 | mem, 56 | frame_size: self.frame_size, 57 | frame_count: self.frame_count, 58 | flags: self.flags, 59 | frame_headroom: self.frame_headroom, 60 | }) 61 | } 62 | } 63 | 64 | /// A shared region of memory for holding frame (packet) data 65 | /// 66 | /// Callers are responsible for correct descriptor allocation. This means only one descriptor 67 | /// should be alive for each frame. If this invariant is not held, the borrowing rules will be 68 | /// violated and potentially result in UB. 69 | #[derive(Clone)] 70 | pub struct Umem { 71 | area: Arc, 72 | mem: NonNull, 73 | frame_size: u32, 74 | frame_count: u32, 75 | frame_headroom: u32, 76 | flags: UmemFlags, 77 | } 78 | 79 | /// Safety: The umem mmap region can be sent to other threads 80 | unsafe impl Send for Umem {} 81 | 82 | /// Safety: The umem mmap region is synchronized by Rings and the data getters are marked as 83 | /// `unsafe`. 84 | unsafe impl Sync for Umem {} 85 | 86 | impl Umem { 87 | /// Creates a Umem builder with defaults 88 | pub fn builder() -> Builder { 89 | Builder::default() 90 | } 91 | 92 | /// Returns the configured size for each frame 93 | #[inline] 94 | pub fn frame_size(&self) -> u32 { 95 | self.frame_size 96 | } 97 | 98 | /// Returns the total number of frames in the Umem 99 | #[inline] 100 | pub fn frame_count(&self) -> u32 { 101 | self.frame_count 102 | } 103 | 104 | /// Returns the configured headroom for each frame 105 | #[inline] 106 | pub fn frame_headroom(&self) -> u32 { 107 | self.frame_headroom 108 | } 109 | 110 | /// Returns the flags for the Umem 111 | #[inline] 112 | pub fn flags(&self) -> UmemFlags { 113 | self.flags 114 | } 115 | 116 | /// Returns an iterator over all of the frame descriptors 117 | /// 118 | /// This can be used to initialize a frame allocator 119 | pub fn frames(&self) -> impl Iterator { 120 | let size = self.frame_size as u64; 121 | (0..self.frame_count as u64).map(move |idx| UmemDescriptor { 122 | address: idx * size, 123 | }) 124 | } 125 | 126 | /// Returns the number of bytes in the Umem 127 | #[inline] 128 | pub fn len(&self) -> usize { 129 | self.area.len() 130 | } 131 | 132 | /// Returns the pointer to the umem memory region 133 | #[inline] 134 | pub fn as_ptr(&self) -> *mut u8 { 135 | self.mem.as_ptr() 136 | } 137 | 138 | /// Returns `true` if the Umem is empty 139 | #[inline] 140 | pub fn is_empty(&self) -> bool { 141 | self.area.is_empty() 142 | } 143 | 144 | /// Returns the region of memory as specified by the index type 145 | /// 146 | /// # Safety 147 | /// 148 | /// The caller MUST ensure that this index is not already mutably borrowed and that the index 149 | /// is in bounds for this Umem. 150 | #[inline] 151 | pub unsafe fn get(&self, idx: T) -> &[u8] 152 | where 153 | Self: UnsafeIndex, 154 | { 155 | self.index(idx) 156 | } 157 | 158 | /// Returns the mutable region of memory as specified by the index type 159 | /// 160 | /// # Safety 161 | /// 162 | /// The caller MUST ensure that this index is not already mutably borrowed and that the index 163 | /// is in bounds for this Umem. 164 | #[inline] 165 | #[allow(clippy::mut_from_ref)] // interior mutability safety is enforced by the caller 166 | pub unsafe fn get_mut(&self, idx: T) -> &mut [u8] 167 | where 168 | Self: UnsafeIndex, 169 | { 170 | self.index_mut(idx) 171 | } 172 | 173 | /// Attaches the Umem to the specified socket 174 | pub(crate) fn attach(&self, socket: &Fd) -> Result<()> { 175 | let umem_conf = UmemReg { 176 | addr: self.area.addr().as_ptr() as _, 177 | chunk_size: self.frame_size, 178 | flags: self.flags, 179 | headroom: self.frame_headroom, 180 | len: self.area.len() as _, 181 | }; 182 | 183 | syscall::set_umem(socket, &umem_conf)?; 184 | 185 | Ok(()) 186 | } 187 | 188 | #[inline] 189 | fn validate_rx_tx_descriptor(&self, desc: RxTxDescriptor) -> *mut u8 { 190 | debug_assert!(desc.len <= self.frame_size, "frame too large"); 191 | debug_assert!( 192 | desc.address + desc.len as u64 <= self.area.len() as u64, 193 | "pointer out of bounds" 194 | ); 195 | unsafe { self.as_ptr().add(desc.address as _) } 196 | } 197 | 198 | #[inline] 199 | fn validate_umem_descriptor(&self, desc: UmemDescriptor) -> *mut u8 { 200 | debug_assert!( 201 | desc.address + self.frame_size as u64 <= self.area.len() as u64, 202 | "pointer out of bounds" 203 | ); 204 | unsafe { self.as_ptr().add(desc.address as _) } 205 | } 206 | } 207 | 208 | /// Specifies an indexable value, which relies on the caller to guarantee borrowing rules are not 209 | /// violated. 210 | pub trait UnsafeIndex { 211 | /// # Safety 212 | /// 213 | /// Callers need to guarantee the reference is not already exclusively borrowed 214 | unsafe fn index(&self, idx: T) -> &[u8]; 215 | 216 | /// # Safety 217 | /// 218 | /// Callers need to guarantee the reference is not already exclusively borrowed 219 | #[allow(clippy::mut_from_ref)] // interior mutability safety is enforced by the caller 220 | unsafe fn index_mut(&self, idx: T) -> &mut [u8]; 221 | } 222 | 223 | impl UnsafeIndex for Umem { 224 | #[inline] 225 | unsafe fn index(&self, idx: RxTxDescriptor) -> &[u8] { 226 | let ptr = self.validate_rx_tx_descriptor(idx); 227 | core::slice::from_raw_parts(ptr, idx.len as _) 228 | } 229 | 230 | #[inline] 231 | unsafe fn index_mut(&self, idx: RxTxDescriptor) -> &mut [u8] { 232 | let ptr = self.validate_rx_tx_descriptor(idx); 233 | core::slice::from_raw_parts_mut(ptr, idx.len as _) 234 | } 235 | } 236 | 237 | impl UnsafeIndex for Umem { 238 | #[inline] 239 | unsafe fn index(&self, idx: UmemDescriptor) -> &[u8] { 240 | let ptr = self.validate_umem_descriptor(idx); 241 | core::slice::from_raw_parts(ptr, self.frame_size as _) 242 | } 243 | 244 | #[inline] 245 | unsafe fn index_mut(&self, idx: UmemDescriptor) -> &mut [u8] { 246 | let ptr = self.validate_umem_descriptor(idx); 247 | core::slice::from_raw_parts_mut(ptr, self.frame_size as _) 248 | } 249 | } 250 | 251 | #[cfg(test)] 252 | mod tests { 253 | use super::*; 254 | 255 | #[test] 256 | fn index_test() { 257 | let umem = Umem::builder().build().unwrap(); 258 | 259 | for descriptor in umem.frames() { 260 | unsafe { 261 | let _ = umem.get(descriptor); 262 | let _ = umem.get_mut(descriptor); 263 | } 264 | 265 | let rx = RxTxDescriptor { 266 | address: descriptor.address, 267 | len: 1, 268 | options: Default::default(), 269 | }; 270 | 271 | unsafe { 272 | let _ = umem.get(rx); 273 | let _ = umem.get_mut(rx); 274 | } 275 | } 276 | } 277 | } 278 | --------------------------------------------------------------------------------