mirror of
https://github.com/YGGverse/aquatic.git
synced 2026-03-31 17:55:36 +00:00
aquatic_ws: mio: replace native_tls with rustls, rewrite connection state logic completely (#38)
* ws: mio: use rustls, rewrite Connection logic * ws: mio: improve poll register/deregister handling * ws: mio: work on type-level poll registry safety * ws: mio: use stronger poll registry type-level guarantees * ws: mio: fix stream reading * ws: clean up, run fmt * ws: mio: don't require registered connection for writing, improve docs * ws: mio: add Connection::get_meta(), make Connection::meta private * ws: mio: add ConnectionMap struct; remove utils.rs * ws: mio: move token counter into ConnectionMap, improve docs * ws: mio: connection: move Connection struct above state structs * Update TODO * ws: fix build errors * ws: upgrade to tungstenite 0.16 * ws load test: don't panic on Close message; print shorter errors * ws: fix socket worker bugs, add log statements * ws: mio: wait for write availability if would block for ws messages * Update README * ws: mio: limit channels & queues; read 1 message only; other fixes * ws: mio: send local responses each event; decrease channel size * Update TODO * ws: mio: limit ws send queue, fixing memory leak; limit pending messages Also change some log output levels and run rustfmt * Update TODO * Update TODO
This commit is contained in:
Joakim Frostegård
GitHub
parent
222fac0e09
commit
67c4c02bbd
13 changed files with 893 additions and 776 deletions
147
Cargo.lock
generated
147
Cargo.lock
generated
|
|
@ -245,12 +245,12 @@ dependencies = [
|
||||||
"log",
|
"log",
|
||||||
"mimalloc",
|
"mimalloc",
|
||||||
"mio",
|
"mio",
|
||||||
"native-tls",
|
|
||||||
"parking_lot",
|
"parking_lot",
|
||||||
"privdrop",
|
"privdrop",
|
||||||
"quickcheck",
|
"quickcheck",
|
||||||
"quickcheck_macros",
|
"quickcheck_macros",
|
||||||
"rand",
|
"rand",
|
||||||
|
"rustls",
|
||||||
"rustls-pemfile",
|
"rustls-pemfile",
|
||||||
"serde",
|
"serde",
|
||||||
"signal-hook",
|
"signal-hook",
|
||||||
|
|
@ -272,6 +272,7 @@ dependencies = [
|
||||||
"futures-rustls",
|
"futures-rustls",
|
||||||
"glommio 0.6.0 (git+https://github.com/DataDog/glommio.git?rev=2efe2f2a08f54394a435b674e8e0125057cbff03)",
|
"glommio 0.6.0 (git+https://github.com/DataDog/glommio.git?rev=2efe2f2a08f54394a435b674e8e0125057cbff03)",
|
||||||
"hashbrown 0.11.2",
|
"hashbrown 0.11.2",
|
||||||
|
"log",
|
||||||
"mimalloc",
|
"mimalloc",
|
||||||
"quickcheck",
|
"quickcheck",
|
||||||
"quickcheck_macros",
|
"quickcheck_macros",
|
||||||
|
|
@ -517,22 +518,6 @@ dependencies = [
|
||||||
"winapi 0.3.9",
|
"winapi 0.3.9",
|
||||||
]
|
]
|
||||||
|
|
||||||
[[package]]
|
|
||||||
name = "core-foundation"
|
|
||||||
version = "0.9.2"
|
|
||||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
|
||||||
checksum = "6888e10551bb93e424d8df1d07f1a8b4fceb0001a3a4b048bfc47554946f47b3"
|
|
||||||
dependencies = [
|
|
||||||
"core-foundation-sys",
|
|
||||||
"libc",
|
|
||||||
]
|
|
||||||
|
|
||||||
[[package]]
|
|
||||||
name = "core-foundation-sys"
|
|
||||||
version = "0.8.3"
|
|
||||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
|
||||||
checksum = "5827cebf4670468b8772dd191856768aedcb1b0278a04f989f7766351917b9dc"
|
|
||||||
|
|
||||||
[[package]]
|
[[package]]
|
||||||
name = "cpufeatures"
|
name = "cpufeatures"
|
||||||
version = "0.2.1"
|
version = "0.2.1"
|
||||||
|
|
@ -772,21 +757,6 @@ version = "1.0.7"
|
||||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||||
checksum = "3f9eec918d3f24069decb9af1554cad7c880e2da24a9afd88aca000531ab82c1"
|
checksum = "3f9eec918d3f24069decb9af1554cad7c880e2da24a9afd88aca000531ab82c1"
|
||||||
|
|
||||||
[[package]]
|
|
||||||
name = "foreign-types"
|
|
||||||
version = "0.3.2"
|
|
||||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
|
||||||
checksum = "f6f339eb8adc052cd2ca78910fda869aefa38d22d5cb648e6485e4d3fc06f3b1"
|
|
||||||
dependencies = [
|
|
||||||
"foreign-types-shared",
|
|
||||||
]
|
|
||||||
|
|
||||||
[[package]]
|
|
||||||
name = "foreign-types-shared"
|
|
||||||
version = "0.1.1"
|
|
||||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
|
||||||
checksum = "00b0228411908ca8685dba7fc2cdd70ec9990a6e753e89b6ac91a84c40fbaf4b"
|
|
||||||
|
|
||||||
[[package]]
|
[[package]]
|
||||||
name = "form_urlencoded"
|
name = "form_urlencoded"
|
||||||
version = "1.0.1"
|
version = "1.0.1"
|
||||||
|
|
@ -1323,24 +1293,6 @@ dependencies = [
|
||||||
"winapi 0.3.9",
|
"winapi 0.3.9",
|
||||||
]
|
]
|
||||||
|
|
||||||
[[package]]
|
|
||||||
name = "native-tls"
|
|
||||||
version = "0.2.8"
|
|
||||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
|
||||||
checksum = "48ba9f7719b5a0f42f338907614285fb5fd70e53858141f69898a1fb7203b24d"
|
|
||||||
dependencies = [
|
|
||||||
"lazy_static",
|
|
||||||
"libc",
|
|
||||||
"log",
|
|
||||||
"openssl",
|
|
||||||
"openssl-probe",
|
|
||||||
"openssl-sys",
|
|
||||||
"schannel",
|
|
||||||
"security-framework",
|
|
||||||
"security-framework-sys",
|
|
||||||
"tempfile",
|
|
||||||
]
|
|
||||||
|
|
||||||
[[package]]
|
[[package]]
|
||||||
name = "nix"
|
name = "nix"
|
||||||
version = "0.23.0"
|
version = "0.23.0"
|
||||||
|
|
@ -1464,39 +1416,6 @@ version = "0.3.0"
|
||||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||||
checksum = "624a8340c38c1b80fd549087862da4ba43e08858af025b236e509b6649fc13d5"
|
checksum = "624a8340c38c1b80fd549087862da4ba43e08858af025b236e509b6649fc13d5"
|
||||||
|
|
||||||
[[package]]
|
|
||||||
name = "openssl"
|
|
||||||
version = "0.10.38"
|
|
||||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
|
||||||
checksum = "0c7ae222234c30df141154f159066c5093ff73b63204dcda7121eb082fc56a95"
|
|
||||||
dependencies = [
|
|
||||||
"bitflags 1.3.2",
|
|
||||||
"cfg-if",
|
|
||||||
"foreign-types",
|
|
||||||
"libc",
|
|
||||||
"once_cell",
|
|
||||||
"openssl-sys",
|
|
||||||
]
|
|
||||||
|
|
||||||
[[package]]
|
|
||||||
name = "openssl-probe"
|
|
||||||
version = "0.1.4"
|
|
||||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
|
||||||
checksum = "28988d872ab76095a6e6ac88d99b54fd267702734fd7ffe610ca27f533ddb95a"
|
|
||||||
|
|
||||||
[[package]]
|
|
||||||
name = "openssl-sys"
|
|
||||||
version = "0.9.71"
|
|
||||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
|
||||||
checksum = "7df13d165e607909b363a4757a6f133f8a818a74e9d3a98d09c6128e15fa4c73"
|
|
||||||
dependencies = [
|
|
||||||
"autocfg",
|
|
||||||
"cc",
|
|
||||||
"libc",
|
|
||||||
"pkg-config",
|
|
||||||
"vcpkg",
|
|
||||||
]
|
|
||||||
|
|
||||||
[[package]]
|
[[package]]
|
||||||
name = "owned-alloc"
|
name = "owned-alloc"
|
||||||
version = "0.2.0"
|
version = "0.2.0"
|
||||||
|
|
@ -1761,15 +1680,6 @@ version = "0.6.25"
|
||||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||||
checksum = "f497285884f3fcff424ffc933e56d7cbca511def0c9831a7f9b5f6153e3cc89b"
|
checksum = "f497285884f3fcff424ffc933e56d7cbca511def0c9831a7f9b5f6153e3cc89b"
|
||||||
|
|
||||||
[[package]]
|
|
||||||
name = "remove_dir_all"
|
|
||||||
version = "0.5.3"
|
|
||||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
|
||||||
checksum = "3acd125665422973a33ac9d3dd2df85edad0f4ae9b00dafb1a05e43a9f5ef8e7"
|
|
||||||
dependencies = [
|
|
||||||
"winapi 0.3.9",
|
|
||||||
]
|
|
||||||
|
|
||||||
[[package]]
|
[[package]]
|
||||||
name = "ring"
|
name = "ring"
|
||||||
version = "0.16.20"
|
version = "0.16.20"
|
||||||
|
|
@ -1845,16 +1755,6 @@ dependencies = [
|
||||||
"winapi-util",
|
"winapi-util",
|
||||||
]
|
]
|
||||||
|
|
||||||
[[package]]
|
|
||||||
name = "schannel"
|
|
||||||
version = "0.1.19"
|
|
||||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
|
||||||
checksum = "8f05ba609c234e60bee0d547fe94a4c7e9da733d1c962cf6e59efa4cd9c8bc75"
|
|
||||||
dependencies = [
|
|
||||||
"lazy_static",
|
|
||||||
"winapi 0.3.9",
|
|
||||||
]
|
|
||||||
|
|
||||||
[[package]]
|
[[package]]
|
||||||
name = "scoped-tls"
|
name = "scoped-tls"
|
||||||
version = "1.0.0"
|
version = "1.0.0"
|
||||||
|
|
@ -1877,29 +1777,6 @@ dependencies = [
|
||||||
"untrusted",
|
"untrusted",
|
||||||
]
|
]
|
||||||
|
|
||||||
[[package]]
|
|
||||||
name = "security-framework"
|
|
||||||
version = "2.4.2"
|
|
||||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
|
||||||
checksum = "525bc1abfda2e1998d152c45cf13e696f76d0a4972310b22fac1658b05df7c87"
|
|
||||||
dependencies = [
|
|
||||||
"bitflags 1.3.2",
|
|
||||||
"core-foundation",
|
|
||||||
"core-foundation-sys",
|
|
||||||
"libc",
|
|
||||||
"security-framework-sys",
|
|
||||||
]
|
|
||||||
|
|
||||||
[[package]]
|
|
||||||
name = "security-framework-sys"
|
|
||||||
version = "2.4.2"
|
|
||||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
|
||||||
checksum = "a9dd14d83160b528b7bfd66439110573efcfbe281b17fc2ca9f39f550d619c7e"
|
|
||||||
dependencies = [
|
|
||||||
"core-foundation-sys",
|
|
||||||
"libc",
|
|
||||||
]
|
|
||||||
|
|
||||||
[[package]]
|
[[package]]
|
||||||
name = "semver"
|
name = "semver"
|
||||||
version = "1.0.4"
|
version = "1.0.4"
|
||||||
|
|
@ -2105,20 +1982,6 @@ dependencies = [
|
||||||
"unicode-xid",
|
"unicode-xid",
|
||||||
]
|
]
|
||||||
|
|
||||||
[[package]]
|
|
||||||
name = "tempfile"
|
|
||||||
version = "3.2.0"
|
|
||||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
|
||||||
checksum = "dac1c663cfc93810f88aed9b8941d48cabf856a1b111c29a40439018d870eb22"
|
|
||||||
dependencies = [
|
|
||||||
"cfg-if",
|
|
||||||
"libc",
|
|
||||||
"rand",
|
|
||||||
"redox_syscall",
|
|
||||||
"remove_dir_all",
|
|
||||||
"winapi 0.3.9",
|
|
||||||
]
|
|
||||||
|
|
||||||
[[package]]
|
[[package]]
|
||||||
name = "termcolor"
|
name = "termcolor"
|
||||||
version = "1.1.2"
|
version = "1.1.2"
|
||||||
|
|
@ -2337,12 +2200,6 @@ dependencies = [
|
||||||
"ryu",
|
"ryu",
|
||||||
]
|
]
|
||||||
|
|
||||||
[[package]]
|
|
||||||
name = "vcpkg"
|
|
||||||
version = "0.2.15"
|
|
||||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
|
||||||
checksum = "accd4ea62f7bb7a82fe23066fb0957d48ef677f6eeb8215f372f52e48bb32426"
|
|
||||||
|
|
||||||
[[package]]
|
[[package]]
|
||||||
name = "version_check"
|
name = "version_check"
|
||||||
version = "0.9.3"
|
version = "0.9.3"
|
||||||
|
|
|
||||||
23
README.md
23
README.md
|
|
@ -17,7 +17,7 @@ of sub-implementations for different protocols:
|
||||||
|--------------|--------------------------------------------|------------------------------------------------------------|
|
|--------------|--------------------------------------------|------------------------------------------------------------|
|
||||||
| aquatic_udp | [BitTorrent over UDP] | Unix-like |
|
| aquatic_udp | [BitTorrent over UDP] | Unix-like |
|
||||||
| aquatic_http | [BitTorrent over HTTP] with TLS ([rustls]) | Linux 5.8+ |
|
| aquatic_http | [BitTorrent over HTTP] with TLS ([rustls]) | Linux 5.8+ |
|
||||||
| aquatic_ws | [WebTorrent] | Unix-like with [mio] (default) / Linux 5.8+ with [glommio] |
|
| aquatic_ws | [WebTorrent] over TLS ([rustls]) | Unix-like with [mio] (default) / Linux 5.8+ with [glommio] |
|
||||||
|
|
||||||
## Usage
|
## Usage
|
||||||
|
|
||||||
|
|
@ -166,28 +166,11 @@ tls_private_key_path = './key.pem'
|
||||||
Aims for compatibility with [WebTorrent](https://github.com/webtorrent)
|
Aims for compatibility with [WebTorrent](https://github.com/webtorrent)
|
||||||
clients, with some exceptions:
|
clients, with some exceptions:
|
||||||
|
|
||||||
|
* Only runs over TLS
|
||||||
* Doesn't track of the number of torrent downloads (0 is always sent).
|
* Doesn't track of the number of torrent downloads (0 is always sent).
|
||||||
* Doesn't allow full scrapes, i.e. of all registered info hashes
|
* Doesn't allow full scrapes, i.e. of all registered info hashes
|
||||||
|
|
||||||
|
For TLS setup instructions, please see `aquatic_http` TLS section above.
|
||||||
#### TLS: mio version
|
|
||||||
|
|
||||||
To run over TLS, a pkcs12 file (`.pkx`) is needed. It can be generated from
|
|
||||||
Let's Encrypt certificates as follows, assuming you are in the directory where
|
|
||||||
they are stored:
|
|
||||||
|
|
||||||
```sh
|
|
||||||
openssl pkcs12 -export -out identity.pfx -inkey privkey.pem -in cert.pem -certfile fullchain.pem
|
|
||||||
```
|
|
||||||
|
|
||||||
Enter a password when prompted. Then move `identity.pfx` somewhere suitable,
|
|
||||||
and enter the path into the tracker configuration field `tls_pkcs12_path`. Set
|
|
||||||
the password in the field `tls_pkcs12_password` and set `use_tls` to true.
|
|
||||||
|
|
||||||
#### TLS: glommio version
|
|
||||||
|
|
||||||
The glommio version only runs over TLS. For setup instructions, please see
|
|
||||||
`aquatic_http` TLS section above.
|
|
||||||
|
|
||||||
#### Benchmarks
|
#### Benchmarks
|
||||||
|
|
||||||
|
|
|
||||||
9
TODO.md
9
TODO.md
|
|
@ -40,6 +40,15 @@
|
||||||
messages can be sent back (e.g., "full scrapes are not supported")
|
messages can be sent back (e.g., "full scrapes are not supported")
|
||||||
|
|
||||||
* aquatic_ws
|
* aquatic_ws
|
||||||
|
* mio
|
||||||
|
* shard torrent state. this could decrease dropped messages too, since
|
||||||
|
request handlers won't send large batches of them
|
||||||
|
* connection cleaning interval
|
||||||
|
* use access list cache
|
||||||
|
* use write event interest for handshakes too
|
||||||
|
* deregistering before closing is required by mio, but it hurts performance
|
||||||
|
* blocked on https://github.com/snapview/tungstenite-rs/issues/51
|
||||||
|
* connection closing: send tls close message etc?
|
||||||
* glommio
|
* glommio
|
||||||
* proper cpu set pinning
|
* proper cpu set pinning
|
||||||
* RES memory still high after traffic stops, even if torrent maps and connection slabs go down to 0 len and capacity
|
* RES memory still high after traffic stops, even if torrent maps and connection slabs go down to 0 len and capacity
|
||||||
|
|
|
||||||
|
|
@ -16,8 +16,8 @@ name = "aquatic_ws"
|
||||||
[features]
|
[features]
|
||||||
default = ["with-mio"]
|
default = ["with-mio"]
|
||||||
cpu-pinning = ["aquatic_common/cpu-pinning"]
|
cpu-pinning = ["aquatic_common/cpu-pinning"]
|
||||||
with-glommio = ["cpu-pinning", "async-tungstenite", "futures-lite", "futures", "futures-rustls", "glommio", "rustls-pemfile"]
|
with-glommio = ["cpu-pinning", "async-tungstenite", "futures-lite", "futures", "futures-rustls", "glommio"]
|
||||||
with-mio = ["crossbeam-channel", "histogram", "mio", "native-tls", "parking_lot", "socket2"]
|
with-mio = ["crossbeam-channel", "histogram", "mio", "parking_lot", "socket2"]
|
||||||
|
|
||||||
[dependencies]
|
[dependencies]
|
||||||
anyhow = "1"
|
anyhow = "1"
|
||||||
|
|
@ -31,6 +31,8 @@ log = "0.4"
|
||||||
mimalloc = { version = "0.1", default-features = false }
|
mimalloc = { version = "0.1", default-features = false }
|
||||||
privdrop = "0.5"
|
privdrop = "0.5"
|
||||||
rand = { version = "0.8", features = ["small_rng"] }
|
rand = { version = "0.8", features = ["small_rng"] }
|
||||||
|
rustls = "0.20"
|
||||||
|
rustls-pemfile = "0.2"
|
||||||
serde = { version = "1", features = ["derive"] }
|
serde = { version = "1", features = ["derive"] }
|
||||||
signal-hook = { version = "0.3" }
|
signal-hook = { version = "0.3" }
|
||||||
slab = "0.4"
|
slab = "0.4"
|
||||||
|
|
@ -40,7 +42,6 @@ tungstenite = "0.16"
|
||||||
crossbeam-channel = { version = "0.5", optional = true }
|
crossbeam-channel = { version = "0.5", optional = true }
|
||||||
histogram = { version = "0.6", optional = true }
|
histogram = { version = "0.6", optional = true }
|
||||||
mio = { version = "0.8", features = ["net", "os-poll"], optional = true }
|
mio = { version = "0.8", features = ["net", "os-poll"], optional = true }
|
||||||
native-tls = { version = "0.2", optional = true }
|
|
||||||
parking_lot = { version = "0.11", optional = true }
|
parking_lot = { version = "0.11", optional = true }
|
||||||
socket2 = { version = "0.4", features = ["all"], optional = true }
|
socket2 = { version = "0.4", features = ["all"], optional = true }
|
||||||
|
|
||||||
|
|
@ -50,7 +51,6 @@ futures-lite = { version = "1", optional = true }
|
||||||
futures = { version = "0.3", optional = true }
|
futures = { version = "0.3", optional = true }
|
||||||
futures-rustls = { version = "0.22", optional = true }
|
futures-rustls = { version = "0.22", optional = true }
|
||||||
glommio = { git = "https://github.com/DataDog/glommio.git", rev = "2efe2f2a08f54394a435b674e8e0125057cbff03", optional = true }
|
glommio = { git = "https://github.com/DataDog/glommio.git", rev = "2efe2f2a08f54394a435b674e8e0125057cbff03", optional = true }
|
||||||
rustls-pemfile = { version = "0.2", optional = true }
|
|
||||||
|
|
||||||
[dev-dependencies]
|
[dev-dependencies]
|
||||||
quickcheck = "1"
|
quickcheck = "1"
|
||||||
|
|
|
||||||
|
|
@ -1,5 +1,7 @@
|
||||||
pub mod handlers;
|
pub mod handlers;
|
||||||
|
|
||||||
|
use std::fs::File;
|
||||||
|
use std::io::BufReader;
|
||||||
use std::net::{IpAddr, SocketAddr};
|
use std::net::{IpAddr, SocketAddr};
|
||||||
use std::sync::Arc;
|
use std::sync::Arc;
|
||||||
use std::time::Instant;
|
use std::time::Instant;
|
||||||
|
|
@ -142,3 +144,32 @@ impl TorrentMaps {
|
||||||
torrent_map.shrink_to_fit();
|
torrent_map.shrink_to_fit();
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
pub fn create_tls_config(config: &Config) -> anyhow::Result<rustls::ServerConfig> {
|
||||||
|
let certs = {
|
||||||
|
let f = File::open(&config.network.tls_certificate_path)?;
|
||||||
|
let mut f = BufReader::new(f);
|
||||||
|
|
||||||
|
rustls_pemfile::certs(&mut f)?
|
||||||
|
.into_iter()
|
||||||
|
.map(|bytes| rustls::Certificate(bytes))
|
||||||
|
.collect()
|
||||||
|
};
|
||||||
|
|
||||||
|
let private_key = {
|
||||||
|
let f = File::open(&config.network.tls_private_key_path)?;
|
||||||
|
let mut f = BufReader::new(f);
|
||||||
|
|
||||||
|
rustls_pemfile::pkcs8_private_keys(&mut f)?
|
||||||
|
.first()
|
||||||
|
.map(|bytes| rustls::PrivateKey(bytes.clone()))
|
||||||
|
.ok_or(anyhow::anyhow!("No private keys in file"))?
|
||||||
|
};
|
||||||
|
|
||||||
|
let tls_config = rustls::ServerConfig::builder()
|
||||||
|
.with_safe_defaults()
|
||||||
|
.with_no_client_auth()
|
||||||
|
.with_single_cert(certs, private_key)?;
|
||||||
|
|
||||||
|
Ok(tls_config)
|
||||||
|
}
|
||||||
|
|
|
||||||
|
|
@ -1,5 +1,4 @@
|
||||||
use std::net::SocketAddr;
|
use std::net::SocketAddr;
|
||||||
#[cfg(feature = "with-glommio")]
|
|
||||||
use std::path::PathBuf;
|
use std::path::PathBuf;
|
||||||
|
|
||||||
#[cfg(feature = "cpu-pinning")]
|
#[cfg(feature = "cpu-pinning")]
|
||||||
|
|
@ -48,17 +47,9 @@ pub struct NetworkConfig {
|
||||||
pub websocket_max_message_size: usize,
|
pub websocket_max_message_size: usize,
|
||||||
pub websocket_max_frame_size: usize,
|
pub websocket_max_frame_size: usize,
|
||||||
|
|
||||||
#[cfg(feature = "with-glommio")]
|
|
||||||
pub tls_certificate_path: PathBuf,
|
pub tls_certificate_path: PathBuf,
|
||||||
#[cfg(feature = "with-glommio")]
|
|
||||||
pub tls_private_key_path: PathBuf,
|
pub tls_private_key_path: PathBuf,
|
||||||
|
|
||||||
#[cfg(feature = "with-mio")]
|
|
||||||
pub use_tls: bool,
|
|
||||||
#[cfg(feature = "with-mio")]
|
|
||||||
pub tls_pkcs12_path: String,
|
|
||||||
#[cfg(feature = "with-mio")]
|
|
||||||
pub tls_pkcs12_password: String,
|
|
||||||
#[cfg(feature = "with-mio")]
|
#[cfg(feature = "with-mio")]
|
||||||
pub poll_event_capacity: usize,
|
pub poll_event_capacity: usize,
|
||||||
#[cfg(feature = "with-mio")]
|
#[cfg(feature = "with-mio")]
|
||||||
|
|
@ -143,17 +134,9 @@ impl Default for NetworkConfig {
|
||||||
websocket_max_message_size: 64 * 1024,
|
websocket_max_message_size: 64 * 1024,
|
||||||
websocket_max_frame_size: 16 * 1024,
|
websocket_max_frame_size: 16 * 1024,
|
||||||
|
|
||||||
#[cfg(feature = "with-glommio")]
|
|
||||||
tls_certificate_path: "".into(),
|
tls_certificate_path: "".into(),
|
||||||
#[cfg(feature = "with-glommio")]
|
|
||||||
tls_private_key_path: "".into(),
|
tls_private_key_path: "".into(),
|
||||||
|
|
||||||
#[cfg(feature = "with-mio")]
|
|
||||||
use_tls: false,
|
|
||||||
#[cfg(feature = "with-mio")]
|
|
||||||
tls_pkcs12_path: "".into(),
|
|
||||||
#[cfg(feature = "with-mio")]
|
|
||||||
tls_pkcs12_password: "".into(),
|
|
||||||
#[cfg(feature = "with-mio")]
|
#[cfg(feature = "with-mio")]
|
||||||
poll_event_capacity: 4096,
|
poll_event_capacity: 4096,
|
||||||
#[cfg(feature = "with-mio")]
|
#[cfg(feature = "with-mio")]
|
||||||
|
|
@ -176,7 +159,7 @@ impl Default for ProtocolConfig {
|
||||||
impl Default for HandlerConfig {
|
impl Default for HandlerConfig {
|
||||||
fn default() -> Self {
|
fn default() -> Self {
|
||||||
Self {
|
Self {
|
||||||
max_requests_per_iter: 10000,
|
max_requests_per_iter: 256,
|
||||||
channel_recv_timeout_microseconds: 200,
|
channel_recv_timeout_microseconds: 200,
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
|
||||||
|
|
@ -2,13 +2,9 @@ pub mod common;
|
||||||
pub mod request;
|
pub mod request;
|
||||||
pub mod socket;
|
pub mod socket;
|
||||||
|
|
||||||
use std::{
|
use std::sync::{atomic::AtomicUsize, Arc};
|
||||||
fs::File,
|
|
||||||
io::BufReader,
|
|
||||||
sync::{atomic::AtomicUsize, Arc},
|
|
||||||
};
|
|
||||||
|
|
||||||
use crate::config::Config;
|
use crate::{common::create_tls_config, config::Config};
|
||||||
#[cfg(feature = "cpu-pinning")]
|
#[cfg(feature = "cpu-pinning")]
|
||||||
use aquatic_common::cpu_pinning::{pin_current_if_configured_to, WorkerIndex};
|
use aquatic_common::cpu_pinning::{pin_current_if_configured_to, WorkerIndex};
|
||||||
use aquatic_common::privileges::drop_privileges_after_socket_binding;
|
use aquatic_common::privileges::drop_privileges_after_socket_binding;
|
||||||
|
|
@ -109,32 +105,3 @@ pub fn run(config: Config, state: State) -> anyhow::Result<()> {
|
||||||
|
|
||||||
Ok(())
|
Ok(())
|
||||||
}
|
}
|
||||||
|
|
||||||
fn create_tls_config(config: &Config) -> anyhow::Result<TlsConfig> {
|
|
||||||
let certs = {
|
|
||||||
let f = File::open(&config.network.tls_certificate_path)?;
|
|
||||||
let mut f = BufReader::new(f);
|
|
||||||
|
|
||||||
rustls_pemfile::certs(&mut f)?
|
|
||||||
.into_iter()
|
|
||||||
.map(|bytes| futures_rustls::rustls::Certificate(bytes))
|
|
||||||
.collect()
|
|
||||||
};
|
|
||||||
|
|
||||||
let private_key = {
|
|
||||||
let f = File::open(&config.network.tls_private_key_path)?;
|
|
||||||
let mut f = BufReader::new(f);
|
|
||||||
|
|
||||||
rustls_pemfile::pkcs8_private_keys(&mut f)?
|
|
||||||
.first()
|
|
||||||
.map(|bytes| futures_rustls::rustls::PrivateKey(bytes.clone()))
|
|
||||||
.ok_or(anyhow::anyhow!("No private keys in file"))?
|
|
||||||
};
|
|
||||||
|
|
||||||
let tls_config = futures_rustls::rustls::ServerConfig::builder()
|
|
||||||
.with_safe_defaults()
|
|
||||||
.with_no_client_auth()
|
|
||||||
.with_single_cert(certs, private_key)?;
|
|
||||||
|
|
||||||
Ok(tls_config)
|
|
||||||
}
|
|
||||||
|
|
|
||||||
|
|
@ -1,5 +1,3 @@
|
||||||
use std::fs::File;
|
|
||||||
use std::io::Read;
|
|
||||||
use std::sync::Arc;
|
use std::sync::Arc;
|
||||||
use std::thread::Builder;
|
use std::thread::Builder;
|
||||||
use std::time::Duration;
|
use std::time::Duration;
|
||||||
|
|
@ -9,7 +7,6 @@ use anyhow::Context;
|
||||||
use aquatic_common::cpu_pinning::{pin_current_if_configured_to, WorkerIndex};
|
use aquatic_common::cpu_pinning::{pin_current_if_configured_to, WorkerIndex};
|
||||||
use histogram::Histogram;
|
use histogram::Histogram;
|
||||||
use mio::{Poll, Waker};
|
use mio::{Poll, Waker};
|
||||||
use native_tls::{Identity, TlsAcceptor};
|
|
||||||
use parking_lot::Mutex;
|
use parking_lot::Mutex;
|
||||||
use privdrop::PrivDrop;
|
use privdrop::PrivDrop;
|
||||||
|
|
||||||
|
|
@ -17,11 +14,13 @@ pub mod common;
|
||||||
pub mod request;
|
pub mod request;
|
||||||
pub mod socket;
|
pub mod socket;
|
||||||
|
|
||||||
use crate::config::Config;
|
use crate::{common::create_tls_config, config::Config};
|
||||||
use common::*;
|
use common::*;
|
||||||
|
|
||||||
pub const APP_NAME: &str = "aquatic_ws: WebTorrent tracker";
|
pub const APP_NAME: &str = "aquatic_ws: WebTorrent tracker";
|
||||||
|
|
||||||
|
const SHARED_IN_CHANNEL_SIZE: usize = 1024;
|
||||||
|
|
||||||
pub fn run(config: Config, state: State) -> anyhow::Result<()> {
|
pub fn run(config: Config, state: State) -> anyhow::Result<()> {
|
||||||
start_workers(config.clone(), state.clone()).expect("couldn't start workers");
|
start_workers(config.clone(), state.clone()).expect("couldn't start workers");
|
||||||
|
|
||||||
|
|
@ -44,9 +43,10 @@ pub fn run(config: Config, state: State) -> anyhow::Result<()> {
|
||||||
}
|
}
|
||||||
|
|
||||||
pub fn start_workers(config: Config, state: State) -> anyhow::Result<()> {
|
pub fn start_workers(config: Config, state: State) -> anyhow::Result<()> {
|
||||||
let opt_tls_acceptor = create_tls_acceptor(&config)?;
|
let tls_config = Arc::new(create_tls_config(&config)?);
|
||||||
|
|
||||||
let (in_message_sender, in_message_receiver) = ::crossbeam_channel::unbounded();
|
let (in_message_sender, in_message_receiver) =
|
||||||
|
::crossbeam_channel::bounded(SHARED_IN_CHANNEL_SIZE);
|
||||||
|
|
||||||
let mut out_message_senders = Vec::new();
|
let mut out_message_senders = Vec::new();
|
||||||
let mut wakers = Vec::new();
|
let mut wakers = Vec::new();
|
||||||
|
|
@ -66,11 +66,12 @@ pub fn start_workers(config: Config, state: State) -> anyhow::Result<()> {
|
||||||
let state = state.clone();
|
let state = state.clone();
|
||||||
let socket_worker_statuses = socket_worker_statuses.clone();
|
let socket_worker_statuses = socket_worker_statuses.clone();
|
||||||
let in_message_sender = in_message_sender.clone();
|
let in_message_sender = in_message_sender.clone();
|
||||||
let opt_tls_acceptor = opt_tls_acceptor.clone();
|
let tls_config = tls_config.clone();
|
||||||
let poll = Poll::new()?;
|
let poll = Poll::new()?;
|
||||||
let waker = Arc::new(Waker::new(poll.registry(), CHANNEL_TOKEN)?);
|
let waker = Arc::new(Waker::new(poll.registry(), CHANNEL_TOKEN)?);
|
||||||
|
|
||||||
let (out_message_sender, out_message_receiver) = ::crossbeam_channel::unbounded();
|
let (out_message_sender, out_message_receiver) =
|
||||||
|
::crossbeam_channel::bounded(SHARED_IN_CHANNEL_SIZE * 16);
|
||||||
|
|
||||||
out_message_senders.push(out_message_sender);
|
out_message_senders.push(out_message_sender);
|
||||||
wakers.push(waker);
|
wakers.push(waker);
|
||||||
|
|
@ -93,7 +94,7 @@ pub fn start_workers(config: Config, state: State) -> anyhow::Result<()> {
|
||||||
poll,
|
poll,
|
||||||
in_message_sender,
|
in_message_sender,
|
||||||
out_message_receiver,
|
out_message_receiver,
|
||||||
opt_tls_acceptor,
|
tls_config,
|
||||||
);
|
);
|
||||||
})?;
|
})?;
|
||||||
}
|
}
|
||||||
|
|
@ -180,27 +181,6 @@ pub fn start_workers(config: Config, state: State) -> anyhow::Result<()> {
|
||||||
Ok(())
|
Ok(())
|
||||||
}
|
}
|
||||||
|
|
||||||
pub fn create_tls_acceptor(config: &Config) -> anyhow::Result<Option<TlsAcceptor>> {
|
|
||||||
if config.network.use_tls {
|
|
||||||
let mut identity_bytes = Vec::new();
|
|
||||||
let mut file = File::open(&config.network.tls_pkcs12_path)
|
|
||||||
.context("Couldn't open pkcs12 identity file")?;
|
|
||||||
|
|
||||||
file.read_to_end(&mut identity_bytes)
|
|
||||||
.context("Couldn't read pkcs12 identity file")?;
|
|
||||||
|
|
||||||
let identity = Identity::from_pkcs12(&identity_bytes, &config.network.tls_pkcs12_password)
|
|
||||||
.context("Couldn't parse pkcs12 identity file")?;
|
|
||||||
|
|
||||||
let acceptor = TlsAcceptor::new(identity)
|
|
||||||
.context("Couldn't create TlsAcceptor from pkcs12 identity")?;
|
|
||||||
|
|
||||||
Ok(Some(acceptor))
|
|
||||||
} else {
|
|
||||||
Ok(None)
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
fn print_statistics(state: &State) {
|
fn print_statistics(state: &State) {
|
||||||
let mut peers_per_torrent = Histogram::new();
|
let mut peers_per_torrent = Histogram::new();
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -1,298 +1,577 @@
|
||||||
use std::io::{Read, Write};
|
use std::{collections::VecDeque, io::ErrorKind, marker::PhantomData, net::Shutdown, sync::Arc};
|
||||||
use std::net::SocketAddr;
|
|
||||||
|
|
||||||
use either::Either;
|
use aquatic_common::ValidUntil;
|
||||||
use hashbrown::HashMap;
|
use aquatic_ws_protocol::{InMessage, OutMessage};
|
||||||
use log::info;
|
use mio::{net::TcpStream, Interest, Poll, Token};
|
||||||
use mio::net::TcpStream;
|
use rustls::{ServerConfig, ServerConnection};
|
||||||
use mio::{Poll, Token};
|
use tungstenite::{
|
||||||
use native_tls::{MidHandshakeTlsStream, TlsAcceptor, TlsStream};
|
handshake::{server::NoCallback, MidHandshake},
|
||||||
use tungstenite::handshake::{server::NoCallback, HandshakeError, MidHandshake};
|
protocol::WebSocketConfig,
|
||||||
use tungstenite::protocol::WebSocketConfig;
|
HandshakeError, ServerHandshake,
|
||||||
use tungstenite::ServerHandshake;
|
};
|
||||||
use tungstenite::WebSocket;
|
|
||||||
|
|
||||||
use crate::common::*;
|
use crate::common::ConnectionMeta;
|
||||||
|
|
||||||
pub enum Stream {
|
const MAX_PENDING_MESSAGES: usize = 16;
|
||||||
TcpStream(TcpStream),
|
|
||||||
TlsStream(TlsStream<TcpStream>),
|
type TlsStream = rustls::StreamOwned<ServerConnection, TcpStream>;
|
||||||
|
|
||||||
|
type WsHandshakeResult<S> =
|
||||||
|
Result<tungstenite::WebSocket<S>, HandshakeError<ServerHandshake<S, NoCallback>>>;
|
||||||
|
|
||||||
|
type ConnectionReadResult<T> = ::std::io::Result<ConnectionReadStatus<T>>;
|
||||||
|
|
||||||
|
pub trait RegistryStatus {}
|
||||||
|
|
||||||
|
pub struct Registered;
|
||||||
|
|
||||||
|
impl RegistryStatus for Registered {}
|
||||||
|
|
||||||
|
pub struct NotRegistered;
|
||||||
|
|
||||||
|
impl RegistryStatus for NotRegistered {}
|
||||||
|
|
||||||
|
enum ConnectionReadStatus<T> {
|
||||||
|
Message(T, InMessage),
|
||||||
|
Ok(T),
|
||||||
|
WouldBlock(T),
|
||||||
}
|
}
|
||||||
|
|
||||||
impl Stream {
|
enum ConnectionState<R: RegistryStatus> {
|
||||||
#[inline]
|
TlsHandshaking(TlsHandshaking<R>),
|
||||||
pub fn get_peer_addr(&self) -> ::std::io::Result<SocketAddr> {
|
WsHandshaking(WsHandshaking<R>),
|
||||||
match self {
|
WsConnection(WsConnection<R>),
|
||||||
Self::TcpStream(stream) => stream.peer_addr(),
|
}
|
||||||
Self::TlsStream(stream) => stream.get_ref().peer_addr(),
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
#[inline]
|
pub struct Connection<R: RegistryStatus> {
|
||||||
pub fn deregister(&mut self, poll: &mut Poll) -> ::std::io::Result<()> {
|
pub valid_until: ValidUntil,
|
||||||
match self {
|
meta: ConnectionMeta,
|
||||||
Self::TcpStream(stream) => poll.registry().deregister(stream),
|
state: ConnectionState<R>,
|
||||||
Self::TlsStream(stream) => poll.registry().deregister(stream.get_mut()),
|
pub message_queue: VecDeque<OutMessage>,
|
||||||
}
|
pub interest: Interest,
|
||||||
|
phantom_data: PhantomData<R>,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl<R: RegistryStatus> Connection<R> {
|
||||||
|
pub fn get_meta(&self) -> ConnectionMeta {
|
||||||
|
self.meta
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
impl Read for Stream {
|
impl Connection<NotRegistered> {
|
||||||
#[inline]
|
pub fn new(
|
||||||
fn read(&mut self, buf: &mut [u8]) -> Result<usize, ::std::io::Error> {
|
tls_config: Arc<ServerConfig>,
|
||||||
match self {
|
|
||||||
Self::TcpStream(stream) => stream.read(buf),
|
|
||||||
Self::TlsStream(stream) => stream.read(buf),
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
/// Not used but provided for completeness
|
|
||||||
#[inline]
|
|
||||||
fn read_vectored(
|
|
||||||
&mut self,
|
|
||||||
bufs: &mut [::std::io::IoSliceMut<'_>],
|
|
||||||
) -> ::std::io::Result<usize> {
|
|
||||||
match self {
|
|
||||||
Self::TcpStream(stream) => stream.read_vectored(bufs),
|
|
||||||
Self::TlsStream(stream) => stream.read_vectored(bufs),
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
impl Write for Stream {
|
|
||||||
#[inline]
|
|
||||||
fn write(&mut self, buf: &[u8]) -> ::std::io::Result<usize> {
|
|
||||||
match self {
|
|
||||||
Self::TcpStream(stream) => stream.write(buf),
|
|
||||||
Self::TlsStream(stream) => stream.write(buf),
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
/// Not used but provided for completeness
|
|
||||||
#[inline]
|
|
||||||
fn write_vectored(&mut self, bufs: &[::std::io::IoSlice<'_>]) -> ::std::io::Result<usize> {
|
|
||||||
match self {
|
|
||||||
Self::TcpStream(stream) => stream.write_vectored(bufs),
|
|
||||||
Self::TlsStream(stream) => stream.write_vectored(bufs),
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
#[inline]
|
|
||||||
fn flush(&mut self) -> ::std::io::Result<()> {
|
|
||||||
match self {
|
|
||||||
Self::TcpStream(stream) => stream.flush(),
|
|
||||||
Self::TlsStream(stream) => stream.flush(),
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
enum HandshakeMachine {
|
|
||||||
TcpStream(TcpStream),
|
|
||||||
TlsStream(TlsStream<TcpStream>),
|
|
||||||
TlsMidHandshake(MidHandshakeTlsStream<TcpStream>),
|
|
||||||
WsMidHandshake(MidHandshake<ServerHandshake<Stream, NoCallback>>),
|
|
||||||
}
|
|
||||||
|
|
||||||
impl HandshakeMachine {
|
|
||||||
#[inline]
|
|
||||||
fn new(tcp_stream: TcpStream) -> Self {
|
|
||||||
Self::TcpStream(tcp_stream)
|
|
||||||
}
|
|
||||||
|
|
||||||
#[inline]
|
|
||||||
fn advance(
|
|
||||||
self,
|
|
||||||
ws_config: WebSocketConfig,
|
ws_config: WebSocketConfig,
|
||||||
opt_tls_acceptor: &Option<TlsAcceptor>, // If set, run TLS
|
tcp_stream: TcpStream,
|
||||||
) -> (Option<Either<EstablishedWs, Self>>, bool) {
|
valid_until: ValidUntil,
|
||||||
// bool = stop looping
|
meta: ConnectionMeta,
|
||||||
match self {
|
) -> Self {
|
||||||
HandshakeMachine::TcpStream(stream) => {
|
let state =
|
||||||
if let Some(tls_acceptor) = opt_tls_acceptor {
|
ConnectionState::TlsHandshaking(TlsHandshaking::new(tls_config, ws_config, tcp_stream));
|
||||||
Self::handle_tls_handshake_result(tls_acceptor.accept(stream))
|
|
||||||
} else {
|
|
||||||
let handshake_result = ::tungstenite::accept_with_config(
|
|
||||||
Stream::TcpStream(stream),
|
|
||||||
Some(ws_config),
|
|
||||||
);
|
|
||||||
|
|
||||||
Self::handle_ws_handshake_result(handshake_result)
|
Self {
|
||||||
|
valid_until,
|
||||||
|
meta,
|
||||||
|
state,
|
||||||
|
message_queue: Default::default(),
|
||||||
|
interest: Interest::READABLE,
|
||||||
|
phantom_data: PhantomData::default(),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Read until stream blocks (or error occurs)
|
||||||
|
///
|
||||||
|
/// Requires Connection not to be registered, since it might be dropped on errors
|
||||||
|
pub fn read<F>(
|
||||||
|
mut self,
|
||||||
|
message_handler: &mut F,
|
||||||
|
) -> ::std::io::Result<Connection<NotRegistered>>
|
||||||
|
where
|
||||||
|
F: FnMut(ConnectionMeta, InMessage),
|
||||||
|
{
|
||||||
|
loop {
|
||||||
|
let result = match self.state {
|
||||||
|
ConnectionState::TlsHandshaking(inner) => inner.read(),
|
||||||
|
ConnectionState::WsHandshaking(inner) => inner.read(),
|
||||||
|
ConnectionState::WsConnection(inner) => inner.read(),
|
||||||
|
};
|
||||||
|
|
||||||
|
match result {
|
||||||
|
Ok(ConnectionReadStatus::Message(state, message)) => {
|
||||||
|
self.state = state;
|
||||||
|
|
||||||
|
message_handler(self.meta, message);
|
||||||
|
|
||||||
|
// Stop looping even if WouldBlock wasn't necessarily reached. Otherwise,
|
||||||
|
// we might get stuck reading from this connection only. Since we register
|
||||||
|
// the connection again upon reinsertion into the ConnectionMap, we should
|
||||||
|
// be getting new events anyway.
|
||||||
|
return Ok(self);
|
||||||
}
|
}
|
||||||
}
|
Ok(ConnectionReadStatus::Ok(state)) => {
|
||||||
HandshakeMachine::TlsStream(stream) => {
|
self.state = state;
|
||||||
let handshake_result = ::tungstenite::accept(Stream::TlsStream(stream));
|
|
||||||
|
|
||||||
Self::handle_ws_handshake_result(handshake_result)
|
::log::debug!("read connection");
|
||||||
}
|
}
|
||||||
HandshakeMachine::TlsMidHandshake(handshake) => {
|
Ok(ConnectionReadStatus::WouldBlock(state)) => {
|
||||||
Self::handle_tls_handshake_result(handshake.handshake())
|
self.state = state;
|
||||||
}
|
|
||||||
HandshakeMachine::WsMidHandshake(handshake) => {
|
|
||||||
Self::handle_ws_handshake_result(handshake.handshake())
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
#[inline]
|
::log::debug!("reading connection would block");
|
||||||
fn handle_tls_handshake_result(
|
|
||||||
result: Result<TlsStream<TcpStream>, ::native_tls::HandshakeError<TcpStream>>,
|
|
||||||
) -> (Option<Either<EstablishedWs, Self>>, bool) {
|
|
||||||
match result {
|
|
||||||
Ok(stream) => {
|
|
||||||
::log::trace!(
|
|
||||||
"established tls handshake with peer with addr: {:?}",
|
|
||||||
stream.get_ref().peer_addr()
|
|
||||||
);
|
|
||||||
|
|
||||||
(Some(Either::Right(Self::TlsStream(stream))), false)
|
return Ok(self);
|
||||||
}
|
|
||||||
Err(native_tls::HandshakeError::WouldBlock(handshake)) => {
|
|
||||||
(Some(Either::Right(Self::TlsMidHandshake(handshake))), true)
|
|
||||||
}
|
|
||||||
Err(native_tls::HandshakeError::Failure(err)) => {
|
|
||||||
info!("tls handshake error: {}", err);
|
|
||||||
|
|
||||||
(None, false)
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
#[inline]
|
|
||||||
fn handle_ws_handshake_result(
|
|
||||||
result: Result<WebSocket<Stream>, HandshakeError<ServerHandshake<Stream, NoCallback>>>,
|
|
||||||
) -> (Option<Either<EstablishedWs, Self>>, bool) {
|
|
||||||
match result {
|
|
||||||
Ok(mut ws) => match ws.get_mut().get_peer_addr() {
|
|
||||||
Ok(peer_addr) => {
|
|
||||||
::log::trace!(
|
|
||||||
"established ws handshake with peer with addr: {:?}",
|
|
||||||
peer_addr
|
|
||||||
);
|
|
||||||
|
|
||||||
let established_ws = EstablishedWs { ws, peer_addr };
|
|
||||||
|
|
||||||
(Some(Either::Left(established_ws)), false)
|
|
||||||
}
|
}
|
||||||
Err(err) => {
|
Err(err) => {
|
||||||
::log::info!(
|
::log::debug!("Connection::read error: {}", err);
|
||||||
"get_peer_addr failed during handshake, removing connection: {:?}",
|
|
||||||
err
|
|
||||||
);
|
|
||||||
|
|
||||||
(None, false)
|
return Err(err);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn register(self, poll: &mut Poll, token: Token) -> Connection<Registered> {
|
||||||
|
let state = match self.state {
|
||||||
|
ConnectionState::TlsHandshaking(inner) => {
|
||||||
|
ConnectionState::TlsHandshaking(inner.register(poll, token, self.interest))
|
||||||
|
}
|
||||||
|
ConnectionState::WsHandshaking(inner) => {
|
||||||
|
ConnectionState::WsHandshaking(inner.register(poll, token, self.interest))
|
||||||
|
}
|
||||||
|
ConnectionState::WsConnection(inner) => {
|
||||||
|
ConnectionState::WsConnection(inner.register(poll, token, self.interest))
|
||||||
|
}
|
||||||
|
};
|
||||||
|
|
||||||
|
Connection {
|
||||||
|
valid_until: self.valid_until,
|
||||||
|
meta: self.meta,
|
||||||
|
state,
|
||||||
|
message_queue: self.message_queue,
|
||||||
|
interest: self.interest,
|
||||||
|
phantom_data: PhantomData::default(),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn close(self) {
|
||||||
|
::log::debug!("will close connection to {}", self.meta.naive_peer_addr);
|
||||||
|
|
||||||
|
match self.state {
|
||||||
|
ConnectionState::TlsHandshaking(inner) => inner.close(),
|
||||||
|
ConnectionState::WsHandshaking(inner) => inner.close(),
|
||||||
|
ConnectionState::WsConnection(inner) => inner.close(),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
impl Connection<Registered> {
|
||||||
|
pub fn write_or_queue_message(
|
||||||
|
&mut self,
|
||||||
|
poll: &mut Poll,
|
||||||
|
message: OutMessage,
|
||||||
|
) -> ::std::io::Result<()> {
|
||||||
|
let message_clone = message.clone();
|
||||||
|
|
||||||
|
match self.write_message(message) {
|
||||||
|
Ok(()) => Ok(()),
|
||||||
|
Err(err) if err.kind() == ErrorKind::WouldBlock => {
|
||||||
|
if self.message_queue.len() < MAX_PENDING_MESSAGES {
|
||||||
|
self.message_queue.push_back(message_clone);
|
||||||
|
|
||||||
|
if !self.interest.is_writable() {
|
||||||
|
self.interest = Interest::WRITABLE;
|
||||||
|
|
||||||
|
self.reregister(poll)?;
|
||||||
|
}
|
||||||
|
} else {
|
||||||
|
::log::info!("Connection::message_queue is full, dropping message");
|
||||||
|
}
|
||||||
|
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
Err(err) => Err(err),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn write(&mut self, poll: &mut Poll) -> ::std::io::Result<()> {
|
||||||
|
if let ConnectionState::WsConnection(_) = self.state {
|
||||||
|
while let Some(message) = self.message_queue.pop_front() {
|
||||||
|
let message_clone = message.clone();
|
||||||
|
|
||||||
|
match self.write_message(message) {
|
||||||
|
Ok(()) => {}
|
||||||
|
Err(err) if err.kind() == ErrorKind::WouldBlock => {
|
||||||
|
// Can't make message queue longer than it was before pop_front
|
||||||
|
self.message_queue.push_front(message_clone);
|
||||||
|
|
||||||
|
return Ok(());
|
||||||
|
}
|
||||||
|
Err(err) => {
|
||||||
|
return Err(err);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
if self.message_queue.is_empty() {
|
||||||
|
self.interest = Interest::READABLE;
|
||||||
|
}
|
||||||
|
|
||||||
|
self.reregister(poll)?;
|
||||||
|
|
||||||
|
Ok(())
|
||||||
|
} else {
|
||||||
|
Err(std::io::Error::new(
|
||||||
|
ErrorKind::NotConnected,
|
||||||
|
"WebSocket connection not established",
|
||||||
|
))
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn write_message(&mut self, message: OutMessage) -> ::std::io::Result<()> {
|
||||||
|
if let ConnectionState::WsConnection(WsConnection {
|
||||||
|
ref mut web_socket, ..
|
||||||
|
}) = self.state
|
||||||
|
{
|
||||||
|
match web_socket.write_message(message.to_ws_message()) {
|
||||||
|
Ok(_) => {}
|
||||||
|
Err(tungstenite::Error::SendQueueFull(_message)) => {
|
||||||
|
return Err(std::io::Error::new(
|
||||||
|
ErrorKind::WouldBlock,
|
||||||
|
"Send queue full",
|
||||||
|
))
|
||||||
|
}
|
||||||
|
Err(tungstenite::Error::Io(err)) => return Err(err),
|
||||||
|
Err(err) => return Err(std::io::Error::new(ErrorKind::Other, err))?,
|
||||||
|
}
|
||||||
|
|
||||||
|
match web_socket.write_pending() {
|
||||||
|
Ok(()) => Ok(()),
|
||||||
|
Err(tungstenite::Error::Io(err)) => Err(err),
|
||||||
|
Err(err) => Err(std::io::Error::new(ErrorKind::Other, err))?,
|
||||||
|
}
|
||||||
|
} else {
|
||||||
|
Err(std::io::Error::new(
|
||||||
|
ErrorKind::NotConnected,
|
||||||
|
"WebSocket connection not established",
|
||||||
|
))
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn reregister(&mut self, poll: &mut Poll) -> ::std::io::Result<()> {
|
||||||
|
let token = Token(self.meta.connection_id.0);
|
||||||
|
|
||||||
|
match self.state {
|
||||||
|
ConnectionState::TlsHandshaking(ref mut inner) => {
|
||||||
|
inner.reregister(poll, token, self.interest)
|
||||||
|
}
|
||||||
|
ConnectionState::WsHandshaking(ref mut inner) => {
|
||||||
|
inner.reregister(poll, token, self.interest)
|
||||||
|
}
|
||||||
|
ConnectionState::WsConnection(ref mut inner) => {
|
||||||
|
inner.reregister(poll, token, self.interest)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn deregister(self, poll: &mut Poll) -> Connection<NotRegistered> {
|
||||||
|
let state = match self.state {
|
||||||
|
ConnectionState::TlsHandshaking(inner) => {
|
||||||
|
ConnectionState::TlsHandshaking(inner.deregister(poll))
|
||||||
|
}
|
||||||
|
ConnectionState::WsHandshaking(inner) => {
|
||||||
|
ConnectionState::WsHandshaking(inner.deregister(poll))
|
||||||
|
}
|
||||||
|
ConnectionState::WsConnection(inner) => {
|
||||||
|
ConnectionState::WsConnection(inner.deregister(poll))
|
||||||
|
}
|
||||||
|
};
|
||||||
|
|
||||||
|
Connection {
|
||||||
|
valid_until: self.valid_until,
|
||||||
|
meta: self.meta,
|
||||||
|
state,
|
||||||
|
message_queue: self.message_queue,
|
||||||
|
interest: self.interest,
|
||||||
|
phantom_data: PhantomData::default(),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
struct TlsHandshaking<R: RegistryStatus> {
|
||||||
|
tls_conn: ServerConnection,
|
||||||
|
ws_config: WebSocketConfig,
|
||||||
|
tcp_stream: TcpStream,
|
||||||
|
phantom_data: PhantomData<R>,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl TlsHandshaking<NotRegistered> {
|
||||||
|
fn new(tls_config: Arc<ServerConfig>, ws_config: WebSocketConfig, stream: TcpStream) -> Self {
|
||||||
|
Self {
|
||||||
|
tls_conn: ServerConnection::new(tls_config).unwrap(),
|
||||||
|
ws_config,
|
||||||
|
tcp_stream: stream,
|
||||||
|
phantom_data: PhantomData::default(),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn read(mut self) -> ConnectionReadResult<ConnectionState<NotRegistered>> {
|
||||||
|
match self.tls_conn.read_tls(&mut self.tcp_stream) {
|
||||||
|
Ok(0) => {
|
||||||
|
return Err(::std::io::Error::new(
|
||||||
|
ErrorKind::ConnectionReset,
|
||||||
|
"Connection closed",
|
||||||
|
))
|
||||||
|
}
|
||||||
|
Ok(_) => match self.tls_conn.process_new_packets() {
|
||||||
|
Ok(_) => {
|
||||||
|
while self.tls_conn.wants_write() {
|
||||||
|
self.tls_conn.write_tls(&mut self.tcp_stream)?;
|
||||||
|
}
|
||||||
|
|
||||||
|
if self.tls_conn.is_handshaking() {
|
||||||
|
Ok(ConnectionReadStatus::WouldBlock(
|
||||||
|
ConnectionState::TlsHandshaking(self),
|
||||||
|
))
|
||||||
|
} else {
|
||||||
|
let tls_stream = TlsStream::new(self.tls_conn, self.tcp_stream);
|
||||||
|
|
||||||
|
WsHandshaking::handle_handshake_result(tungstenite::accept_with_config(
|
||||||
|
tls_stream,
|
||||||
|
Some(self.ws_config),
|
||||||
|
))
|
||||||
|
}
|
||||||
|
}
|
||||||
|
Err(err) => {
|
||||||
|
let _ = self.tls_conn.write_tls(&mut self.tcp_stream);
|
||||||
|
|
||||||
|
Err(::std::io::Error::new(ErrorKind::InvalidData, err))
|
||||||
}
|
}
|
||||||
},
|
},
|
||||||
Err(HandshakeError::Interrupted(handshake)) => (
|
Err(err) if err.kind() == ErrorKind::WouldBlock => {
|
||||||
Some(Either::Right(HandshakeMachine::WsMidHandshake(handshake))),
|
return Ok(ConnectionReadStatus::WouldBlock(
|
||||||
true,
|
ConnectionState::TlsHandshaking(self),
|
||||||
),
|
))
|
||||||
Err(HandshakeError::Failure(err)) => {
|
|
||||||
info!("ws handshake error: {}", err);
|
|
||||||
|
|
||||||
(None, false)
|
|
||||||
}
|
}
|
||||||
}
|
Err(err) => return Err(err),
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
pub struct EstablishedWs {
|
|
||||||
pub ws: WebSocket<Stream>,
|
|
||||||
pub peer_addr: SocketAddr,
|
|
||||||
}
|
|
||||||
|
|
||||||
pub struct Connection {
|
|
||||||
ws_config: WebSocketConfig,
|
|
||||||
pub valid_until: ValidUntil,
|
|
||||||
inner: Either<EstablishedWs, HandshakeMachine>,
|
|
||||||
}
|
|
||||||
|
|
||||||
/// Create from TcpStream. Run `advance_handshakes` until `get_established_ws`
|
|
||||||
/// returns Some(EstablishedWs).
|
|
||||||
///
|
|
||||||
/// advance_handshakes takes ownership of self because the TLS and WebSocket
|
|
||||||
/// handshake methods do. get_established_ws doesn't, since work can be done
|
|
||||||
/// on a mutable reference to a tungstenite websocket, and this way, the whole
|
|
||||||
/// Connection doesn't have to be removed from and reinserted into the
|
|
||||||
/// TorrentMap. This is also the reason for wrapping Container.inner in an
|
|
||||||
/// Either instead of combining all states into one structure just having a
|
|
||||||
/// single method for advancing handshakes and maybe returning a websocket.
|
|
||||||
impl Connection {
|
|
||||||
#[inline]
|
|
||||||
pub fn new(ws_config: WebSocketConfig, valid_until: ValidUntil, tcp_stream: TcpStream) -> Self {
|
|
||||||
Self {
|
|
||||||
ws_config,
|
|
||||||
valid_until,
|
|
||||||
inner: Either::Right(HandshakeMachine::new(tcp_stream)),
|
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
#[inline]
|
fn register(
|
||||||
pub fn get_established_ws(&mut self) -> Option<&mut EstablishedWs> {
|
mut self,
|
||||||
match self.inner {
|
poll: &mut Poll,
|
||||||
Either::Left(ref mut ews) => Some(ews),
|
token: Token,
|
||||||
Either::Right(_) => None,
|
interest: Interest,
|
||||||
|
) -> TlsHandshaking<Registered> {
|
||||||
|
poll.registry()
|
||||||
|
.register(&mut self.tcp_stream, token, interest)
|
||||||
|
.unwrap();
|
||||||
|
|
||||||
|
TlsHandshaking {
|
||||||
|
tls_conn: self.tls_conn,
|
||||||
|
ws_config: self.ws_config,
|
||||||
|
tcp_stream: self.tcp_stream,
|
||||||
|
phantom_data: PhantomData::default(),
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
#[inline]
|
fn close(self) {
|
||||||
pub fn advance_handshakes(
|
::log::debug!("closing connection (TlsHandshaking state)");
|
||||||
self,
|
|
||||||
opt_tls_acceptor: &Option<TlsAcceptor>,
|
|
||||||
valid_until: ValidUntil,
|
|
||||||
) -> (Option<Self>, bool) {
|
|
||||||
match self.inner {
|
|
||||||
Either::Left(_) => (Some(self), false),
|
|
||||||
Either::Right(machine) => {
|
|
||||||
let ws_config = self.ws_config;
|
|
||||||
|
|
||||||
let (opt_inner, stop_loop) = machine.advance(ws_config, opt_tls_acceptor);
|
let _ = self.tcp_stream.shutdown(Shutdown::Both);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
let opt_new_self = opt_inner.map(|inner| Self {
|
impl TlsHandshaking<Registered> {
|
||||||
ws_config,
|
fn deregister(mut self, poll: &mut Poll) -> TlsHandshaking<NotRegistered> {
|
||||||
valid_until,
|
poll.registry().deregister(&mut self.tcp_stream).unwrap();
|
||||||
inner,
|
|
||||||
|
TlsHandshaking {
|
||||||
|
tls_conn: self.tls_conn,
|
||||||
|
ws_config: self.ws_config,
|
||||||
|
tcp_stream: self.tcp_stream,
|
||||||
|
phantom_data: PhantomData::default(),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn reregister(
|
||||||
|
&mut self,
|
||||||
|
poll: &mut Poll,
|
||||||
|
token: Token,
|
||||||
|
interest: Interest,
|
||||||
|
) -> std::io::Result<()> {
|
||||||
|
poll.registry()
|
||||||
|
.reregister(&mut self.tcp_stream, token, interest)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
struct WsHandshaking<R: RegistryStatus> {
|
||||||
|
mid_handshake: MidHandshake<ServerHandshake<TlsStream, NoCallback>>,
|
||||||
|
phantom_data: PhantomData<R>,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl WsHandshaking<NotRegistered> {
|
||||||
|
fn read(self) -> ConnectionReadResult<ConnectionState<NotRegistered>> {
|
||||||
|
Self::handle_handshake_result(self.mid_handshake.handshake())
|
||||||
|
}
|
||||||
|
|
||||||
|
fn handle_handshake_result(
|
||||||
|
handshake_result: WsHandshakeResult<TlsStream>,
|
||||||
|
) -> ConnectionReadResult<ConnectionState<NotRegistered>> {
|
||||||
|
match handshake_result {
|
||||||
|
Ok(web_socket) => {
|
||||||
|
let conn = ConnectionState::WsConnection(WsConnection {
|
||||||
|
web_socket,
|
||||||
|
phantom_data: PhantomData::default(),
|
||||||
});
|
});
|
||||||
|
|
||||||
(opt_new_self, stop_loop)
|
Ok(ConnectionReadStatus::Ok(conn))
|
||||||
|
}
|
||||||
|
Err(HandshakeError::Interrupted(mid_handshake)) => {
|
||||||
|
let conn = ConnectionState::WsHandshaking(WsHandshaking {
|
||||||
|
mid_handshake,
|
||||||
|
phantom_data: PhantomData::default(),
|
||||||
|
});
|
||||||
|
|
||||||
|
Ok(ConnectionReadStatus::WouldBlock(conn))
|
||||||
|
}
|
||||||
|
Err(HandshakeError::Failure(err)) => {
|
||||||
|
return Err(std::io::Error::new(ErrorKind::InvalidData, err))
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
#[inline]
|
fn register(
|
||||||
pub fn close(&mut self) {
|
mut self,
|
||||||
if let Either::Left(ref mut ews) = self.inner {
|
poll: &mut Poll,
|
||||||
if ews.ws.can_read() {
|
token: Token,
|
||||||
if let Err(err) = ews.ws.close(None) {
|
interest: Interest,
|
||||||
::log::info!("error closing ws: {}", err);
|
) -> WsHandshaking<Registered> {
|
||||||
}
|
let tcp_stream = &mut self.mid_handshake.get_mut().get_mut().sock;
|
||||||
|
|
||||||
// Required after ws.close()
|
poll.registry()
|
||||||
if let Err(err) = ews.ws.write_pending() {
|
.register(tcp_stream, token, interest)
|
||||||
::log::info!("error writing pending messages after closing ws: {}", err)
|
.unwrap();
|
||||||
}
|
|
||||||
}
|
WsHandshaking {
|
||||||
|
mid_handshake: self.mid_handshake,
|
||||||
|
phantom_data: PhantomData::default(),
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
pub fn deregister(&mut self, poll: &mut Poll) -> ::std::io::Result<()> {
|
fn close(mut self) {
|
||||||
use Either::{Left, Right};
|
::log::debug!("closing connection (WsHandshaking state)");
|
||||||
|
|
||||||
match self.inner {
|
let tcp_stream = &mut self.mid_handshake.get_mut().get_mut().sock;
|
||||||
Left(EstablishedWs { ref mut ws, .. }) => ws.get_mut().deregister(poll),
|
|
||||||
Right(HandshakeMachine::TcpStream(ref mut stream)) => {
|
let _ = tcp_stream.shutdown(Shutdown::Both);
|
||||||
poll.registry().deregister(stream)
|
|
||||||
}
|
|
||||||
Right(HandshakeMachine::TlsMidHandshake(ref mut handshake)) => {
|
|
||||||
poll.registry().deregister(handshake.get_mut())
|
|
||||||
}
|
|
||||||
Right(HandshakeMachine::TlsStream(ref mut stream)) => {
|
|
||||||
poll.registry().deregister(stream.get_mut())
|
|
||||||
}
|
|
||||||
Right(HandshakeMachine::WsMidHandshake(ref mut handshake)) => {
|
|
||||||
handshake.get_mut().get_mut().deregister(poll)
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
pub type ConnectionMap = HashMap<Token, Connection>;
|
impl WsHandshaking<Registered> {
|
||||||
|
fn deregister(mut self, poll: &mut Poll) -> WsHandshaking<NotRegistered> {
|
||||||
|
let tcp_stream = &mut self.mid_handshake.get_mut().get_mut().sock;
|
||||||
|
|
||||||
|
poll.registry().deregister(tcp_stream).unwrap();
|
||||||
|
|
||||||
|
WsHandshaking {
|
||||||
|
mid_handshake: self.mid_handshake,
|
||||||
|
phantom_data: PhantomData::default(),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn reregister(
|
||||||
|
&mut self,
|
||||||
|
poll: &mut Poll,
|
||||||
|
token: Token,
|
||||||
|
interest: Interest,
|
||||||
|
) -> std::io::Result<()> {
|
||||||
|
let tcp_stream = &mut self.mid_handshake.get_mut().get_mut().sock;
|
||||||
|
|
||||||
|
poll.registry().reregister(tcp_stream, token, interest)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
struct WsConnection<R: RegistryStatus> {
|
||||||
|
web_socket: tungstenite::WebSocket<TlsStream>,
|
||||||
|
phantom_data: PhantomData<R>,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl WsConnection<NotRegistered> {
|
||||||
|
fn read(mut self) -> ConnectionReadResult<ConnectionState<NotRegistered>> {
|
||||||
|
match self.web_socket.read_message() {
|
||||||
|
Ok(
|
||||||
|
message @ tungstenite::Message::Text(_) | message @ tungstenite::Message::Binary(_),
|
||||||
|
) => match InMessage::from_ws_message(message) {
|
||||||
|
Ok(message) => {
|
||||||
|
::log::debug!("received WebSocket message");
|
||||||
|
|
||||||
|
Ok(ConnectionReadStatus::Message(
|
||||||
|
ConnectionState::WsConnection(self),
|
||||||
|
message,
|
||||||
|
))
|
||||||
|
}
|
||||||
|
Err(err) => Err(std::io::Error::new(ErrorKind::InvalidData, err)),
|
||||||
|
},
|
||||||
|
Ok(message) => {
|
||||||
|
::log::info!("received unexpected WebSocket message: {}", message);
|
||||||
|
|
||||||
|
Err(std::io::Error::new(
|
||||||
|
ErrorKind::InvalidData,
|
||||||
|
"unexpected WebSocket message type",
|
||||||
|
))
|
||||||
|
}
|
||||||
|
Err(tungstenite::Error::Io(err)) if err.kind() == ErrorKind::WouldBlock => {
|
||||||
|
let conn = ConnectionState::WsConnection(self);
|
||||||
|
|
||||||
|
Ok(ConnectionReadStatus::WouldBlock(conn))
|
||||||
|
}
|
||||||
|
Err(tungstenite::Error::Io(err)) => Err(err),
|
||||||
|
Err(err) => Err(std::io::Error::new(ErrorKind::InvalidData, err)),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn register(
|
||||||
|
mut self,
|
||||||
|
poll: &mut Poll,
|
||||||
|
token: Token,
|
||||||
|
interest: Interest,
|
||||||
|
) -> WsConnection<Registered> {
|
||||||
|
poll.registry()
|
||||||
|
.register(self.web_socket.get_mut().get_mut(), token, interest)
|
||||||
|
.unwrap();
|
||||||
|
|
||||||
|
WsConnection {
|
||||||
|
web_socket: self.web_socket,
|
||||||
|
phantom_data: PhantomData::default(),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn close(mut self) {
|
||||||
|
::log::debug!("closing connection (WsConnection state)");
|
||||||
|
|
||||||
|
let _ = self.web_socket.close(None);
|
||||||
|
let _ = self.web_socket.write_pending();
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
impl WsConnection<Registered> {
|
||||||
|
fn deregister(mut self, poll: &mut Poll) -> WsConnection<NotRegistered> {
|
||||||
|
poll.registry()
|
||||||
|
.deregister(self.web_socket.get_mut().get_mut())
|
||||||
|
.unwrap();
|
||||||
|
|
||||||
|
WsConnection {
|
||||||
|
web_socket: self.web_socket,
|
||||||
|
phantom_data: PhantomData::default(),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn reregister(
|
||||||
|
&mut self,
|
||||||
|
poll: &mut Poll,
|
||||||
|
token: Token,
|
||||||
|
interest: Interest,
|
||||||
|
) -> std::io::Result<()> {
|
||||||
|
poll.registry()
|
||||||
|
.reregister(self.web_socket.get_mut().get_mut(), token, interest)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
|
||||||
|
|
@ -1,14 +1,13 @@
|
||||||
use std::io::ErrorKind;
|
use std::io::ErrorKind;
|
||||||
use std::time::Duration;
|
use std::sync::Arc;
|
||||||
use std::vec::Drain;
|
use std::time::{Duration, Instant};
|
||||||
|
|
||||||
|
use anyhow::Context;
|
||||||
use aquatic_common::access_list::AccessListQuery;
|
use aquatic_common::access_list::AccessListQuery;
|
||||||
use crossbeam_channel::Receiver;
|
|
||||||
use hashbrown::HashMap;
|
use hashbrown::HashMap;
|
||||||
use log::{debug, error, info};
|
|
||||||
use mio::net::TcpListener;
|
use mio::net::TcpListener;
|
||||||
use mio::{Events, Interest, Poll, Token};
|
use mio::{Events, Interest, Poll, Token};
|
||||||
use native_tls::TlsAcceptor;
|
use socket2::{Domain, Protocol, Socket, Type};
|
||||||
use tungstenite::protocol::WebSocketConfig;
|
use tungstenite::protocol::WebSocketConfig;
|
||||||
|
|
||||||
use aquatic_common::convert_ipv4_mapped_ipv6;
|
use aquatic_common::convert_ipv4_mapped_ipv6;
|
||||||
|
|
@ -17,13 +16,101 @@ use aquatic_ws_protocol::*;
|
||||||
use crate::common::*;
|
use crate::common::*;
|
||||||
use crate::config::Config;
|
use crate::config::Config;
|
||||||
|
|
||||||
|
pub mod connection;
|
||||||
|
|
||||||
use super::common::*;
|
use super::common::*;
|
||||||
|
|
||||||
pub mod connection;
|
use connection::{Connection, NotRegistered, Registered};
|
||||||
pub mod utils;
|
|
||||||
|
|
||||||
use connection::*;
|
struct ConnectionMap {
|
||||||
use utils::*;
|
token_counter: Token,
|
||||||
|
connections: HashMap<Token, Connection<Registered>>,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl Default for ConnectionMap {
|
||||||
|
fn default() -> Self {
|
||||||
|
Self {
|
||||||
|
token_counter: Token(2),
|
||||||
|
connections: Default::default(),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
impl ConnectionMap {
|
||||||
|
fn insert_and_register_new<F>(&mut self, poll: &mut Poll, connection_creator: F)
|
||||||
|
where
|
||||||
|
F: FnOnce(Token) -> Connection<NotRegistered>,
|
||||||
|
{
|
||||||
|
self.token_counter.0 = self.token_counter.0.wrapping_add(1);
|
||||||
|
|
||||||
|
// Don't assign LISTENER_TOKEN or CHANNEL_TOKEN
|
||||||
|
if self.token_counter.0 < 2 {
|
||||||
|
self.token_counter.0 = 2;
|
||||||
|
}
|
||||||
|
|
||||||
|
let token = self.token_counter;
|
||||||
|
|
||||||
|
// Remove, deregister and close any existing connection with this token.
|
||||||
|
// This shouldn't happen in practice.
|
||||||
|
if let Some(connection) = self.connections.remove(&token) {
|
||||||
|
::log::warn!(
|
||||||
|
"removing existing connection {} because of token reuse",
|
||||||
|
token.0
|
||||||
|
);
|
||||||
|
|
||||||
|
connection.deregister(poll).close();
|
||||||
|
}
|
||||||
|
|
||||||
|
let connection = connection_creator(token);
|
||||||
|
|
||||||
|
self.insert_and_register(poll, token, connection);
|
||||||
|
}
|
||||||
|
|
||||||
|
fn insert_and_register(
|
||||||
|
&mut self,
|
||||||
|
poll: &mut Poll,
|
||||||
|
key: Token,
|
||||||
|
conn: Connection<NotRegistered>,
|
||||||
|
) {
|
||||||
|
self.connections.insert(key, conn.register(poll, key));
|
||||||
|
}
|
||||||
|
|
||||||
|
fn remove_and_deregister(
|
||||||
|
&mut self,
|
||||||
|
poll: &mut Poll,
|
||||||
|
key: &Token,
|
||||||
|
) -> Option<Connection<NotRegistered>> {
|
||||||
|
if let Some(connection) = self.connections.remove(key) {
|
||||||
|
Some(connection.deregister(poll))
|
||||||
|
} else {
|
||||||
|
None
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn get_mut(&mut self, key: &Token) -> Option<&mut Connection<Registered>> {
|
||||||
|
self.connections.get_mut(key)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Close and remove inactive connections
|
||||||
|
fn clean(mut self, poll: &mut Poll) -> Self {
|
||||||
|
let now = Instant::now();
|
||||||
|
|
||||||
|
let mut retained_connections = HashMap::default();
|
||||||
|
|
||||||
|
for (token, connection) in self.connections.drain() {
|
||||||
|
if connection.valid_until.0 < now {
|
||||||
|
connection.deregister(poll).close();
|
||||||
|
} else {
|
||||||
|
retained_connections.insert(token, connection);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
ConnectionMap {
|
||||||
|
connections: retained_connections,
|
||||||
|
..self
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
pub fn run_socket_worker(
|
pub fn run_socket_worker(
|
||||||
config: Config,
|
config: Config,
|
||||||
|
|
@ -33,7 +120,7 @@ pub fn run_socket_worker(
|
||||||
poll: Poll,
|
poll: Poll,
|
||||||
in_message_sender: InMessageSender,
|
in_message_sender: InMessageSender,
|
||||||
out_message_receiver: OutMessageReceiver,
|
out_message_receiver: OutMessageReceiver,
|
||||||
opt_tls_acceptor: Option<TlsAcceptor>,
|
tls_config: Arc<rustls::ServerConfig>,
|
||||||
) {
|
) {
|
||||||
match create_listener(&config) {
|
match create_listener(&config) {
|
||||||
Ok(listener) => {
|
Ok(listener) => {
|
||||||
|
|
@ -47,7 +134,7 @@ pub fn run_socket_worker(
|
||||||
in_message_sender,
|
in_message_sender,
|
||||||
out_message_receiver,
|
out_message_receiver,
|
||||||
listener,
|
listener,
|
||||||
opt_tls_acceptor,
|
tls_config,
|
||||||
);
|
);
|
||||||
}
|
}
|
||||||
Err(err) => {
|
Err(err) => {
|
||||||
|
|
@ -57,7 +144,7 @@ pub fn run_socket_worker(
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
pub fn run_poll_loop(
|
fn run_poll_loop(
|
||||||
config: Config,
|
config: Config,
|
||||||
state: &State,
|
state: &State,
|
||||||
socket_worker_index: usize,
|
socket_worker_index: usize,
|
||||||
|
|
@ -65,13 +152,13 @@ pub fn run_poll_loop(
|
||||||
in_message_sender: InMessageSender,
|
in_message_sender: InMessageSender,
|
||||||
out_message_receiver: OutMessageReceiver,
|
out_message_receiver: OutMessageReceiver,
|
||||||
listener: ::std::net::TcpListener,
|
listener: ::std::net::TcpListener,
|
||||||
opt_tls_acceptor: Option<TlsAcceptor>,
|
tls_config: Arc<rustls::ServerConfig>,
|
||||||
) {
|
) {
|
||||||
let poll_timeout = Duration::from_micros(config.network.poll_timeout_microseconds);
|
let poll_timeout = Duration::from_micros(config.network.poll_timeout_microseconds);
|
||||||
let ws_config = WebSocketConfig {
|
let ws_config = WebSocketConfig {
|
||||||
max_message_size: Some(config.network.websocket_max_message_size),
|
max_message_size: Some(config.network.websocket_max_message_size),
|
||||||
max_frame_size: Some(config.network.websocket_max_frame_size),
|
max_frame_size: Some(config.network.websocket_max_frame_size),
|
||||||
max_send_queue: None,
|
max_send_queue: Some(2),
|
||||||
..Default::default()
|
..Default::default()
|
||||||
};
|
};
|
||||||
|
|
||||||
|
|
@ -82,10 +169,9 @@ pub fn run_poll_loop(
|
||||||
.register(&mut listener, LISTENER_TOKEN, Interest::READABLE)
|
.register(&mut listener, LISTENER_TOKEN, Interest::READABLE)
|
||||||
.unwrap();
|
.unwrap();
|
||||||
|
|
||||||
let mut connections: ConnectionMap = HashMap::new();
|
let mut connections = ConnectionMap::default();
|
||||||
let mut local_responses = Vec::new();
|
let mut local_responses = Vec::new();
|
||||||
|
|
||||||
let mut poll_token_counter = Token(0usize);
|
|
||||||
let mut iter_counter = 0usize;
|
let mut iter_counter = 0usize;
|
||||||
|
|
||||||
loop {
|
loop {
|
||||||
|
|
@ -97,41 +183,68 @@ pub fn run_poll_loop(
|
||||||
for event in events.iter() {
|
for event in events.iter() {
|
||||||
let token = event.token();
|
let token = event.token();
|
||||||
|
|
||||||
if token == LISTENER_TOKEN {
|
match token {
|
||||||
accept_new_streams(
|
LISTENER_TOKEN => {
|
||||||
ws_config,
|
accept_new_streams(
|
||||||
&mut listener,
|
&tls_config,
|
||||||
&mut poll,
|
ws_config,
|
||||||
&mut connections,
|
socket_worker_index,
|
||||||
valid_until,
|
&mut listener,
|
||||||
&mut poll_token_counter,
|
&mut poll,
|
||||||
);
|
&mut connections,
|
||||||
} else if token != CHANNEL_TOKEN {
|
valid_until,
|
||||||
run_handshakes_and_read_messages(
|
);
|
||||||
&config,
|
}
|
||||||
state,
|
CHANNEL_TOKEN => {
|
||||||
socket_worker_index,
|
write_or_queue_messages(
|
||||||
&mut local_responses,
|
&mut poll,
|
||||||
&in_message_sender,
|
out_message_receiver
|
||||||
&opt_tls_acceptor,
|
.try_iter()
|
||||||
&mut poll,
|
.take(out_message_receiver.len()),
|
||||||
&mut connections,
|
&mut connections,
|
||||||
token,
|
);
|
||||||
valid_until,
|
}
|
||||||
);
|
token => {
|
||||||
|
if event.is_writable() {
|
||||||
|
let mut remove_connection = false;
|
||||||
|
|
||||||
|
if let Some(connection) = connections.get_mut(&token) {
|
||||||
|
if let Err(err) = connection.write(&mut poll) {
|
||||||
|
::log::debug!("Connection::write error: {}", err);
|
||||||
|
|
||||||
|
remove_connection = true;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
if remove_connection {
|
||||||
|
if let Some(connection) =
|
||||||
|
connections.remove_and_deregister(&mut poll, &token)
|
||||||
|
{
|
||||||
|
connection.close();
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if event.is_readable() {
|
||||||
|
handle_stream_read_event(
|
||||||
|
&config,
|
||||||
|
state,
|
||||||
|
&mut local_responses,
|
||||||
|
&in_message_sender,
|
||||||
|
&mut poll,
|
||||||
|
&mut connections,
|
||||||
|
token,
|
||||||
|
valid_until,
|
||||||
|
);
|
||||||
|
}
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
send_out_messages(
|
write_or_queue_messages(&mut poll, local_responses.drain(..), &mut connections);
|
||||||
&mut poll,
|
|
||||||
local_responses.drain(..),
|
|
||||||
&out_message_receiver,
|
|
||||||
&mut connections,
|
|
||||||
);
|
|
||||||
}
|
}
|
||||||
|
|
||||||
// Remove inactive connections, but not every iteration
|
// Remove inactive connections, but not every iteration
|
||||||
if iter_counter % 128 == 0 {
|
if iter_counter % 128 == 0 {
|
||||||
remove_inactive_connections(&mut connections);
|
connections = connections.clean(&mut poll);
|
||||||
}
|
}
|
||||||
|
|
||||||
iter_counter = iter_counter.wrapping_add(1);
|
iter_counter = iter_counter.wrapping_add(1);
|
||||||
|
|
@ -139,194 +252,155 @@ pub fn run_poll_loop(
|
||||||
}
|
}
|
||||||
|
|
||||||
fn accept_new_streams(
|
fn accept_new_streams(
|
||||||
|
tls_config: &Arc<rustls::ServerConfig>,
|
||||||
ws_config: WebSocketConfig,
|
ws_config: WebSocketConfig,
|
||||||
|
socket_worker_index: usize,
|
||||||
listener: &mut TcpListener,
|
listener: &mut TcpListener,
|
||||||
poll: &mut Poll,
|
poll: &mut Poll,
|
||||||
connections: &mut ConnectionMap,
|
connections: &mut ConnectionMap,
|
||||||
valid_until: ValidUntil,
|
valid_until: ValidUntil,
|
||||||
poll_token_counter: &mut Token,
|
|
||||||
) {
|
) {
|
||||||
loop {
|
loop {
|
||||||
match listener.accept() {
|
match listener.accept() {
|
||||||
Ok((mut stream, _)) => {
|
Ok((stream, _)) => {
|
||||||
poll_token_counter.0 = poll_token_counter.0.wrapping_add(1);
|
let naive_peer_addr = if let Ok(peer_addr) = stream.peer_addr() {
|
||||||
|
peer_addr
|
||||||
|
} else {
|
||||||
|
continue;
|
||||||
|
};
|
||||||
|
|
||||||
if poll_token_counter.0 < 2 {
|
connections.insert_and_register_new(poll, move |token| {
|
||||||
poll_token_counter.0 = 2;
|
|
||||||
}
|
|
||||||
|
|
||||||
let token = *poll_token_counter;
|
|
||||||
|
|
||||||
remove_connection_if_exists(poll, connections, token);
|
|
||||||
|
|
||||||
poll.registry()
|
|
||||||
.register(&mut stream, token, Interest::READABLE)
|
|
||||||
.unwrap();
|
|
||||||
|
|
||||||
let connection = Connection::new(ws_config, valid_until, stream);
|
|
||||||
|
|
||||||
connections.insert(token, connection);
|
|
||||||
}
|
|
||||||
Err(err) => {
|
|
||||||
if err.kind() == ErrorKind::WouldBlock {
|
|
||||||
break;
|
|
||||||
}
|
|
||||||
|
|
||||||
info!("error while accepting streams: {}", err);
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
/// On the stream given by poll_token, get TLS (if requested) and tungstenite
|
|
||||||
/// up and running, then read messages and pass on through channel.
|
|
||||||
pub fn run_handshakes_and_read_messages(
|
|
||||||
config: &Config,
|
|
||||||
state: &State,
|
|
||||||
socket_worker_index: usize,
|
|
||||||
local_responses: &mut Vec<(ConnectionMeta, OutMessage)>,
|
|
||||||
in_message_sender: &InMessageSender,
|
|
||||||
opt_tls_acceptor: &Option<TlsAcceptor>, // If set, run TLS
|
|
||||||
poll: &mut Poll,
|
|
||||||
connections: &mut ConnectionMap,
|
|
||||||
poll_token: Token,
|
|
||||||
valid_until: ValidUntil,
|
|
||||||
) {
|
|
||||||
let access_list_mode = config.access_list.mode;
|
|
||||||
|
|
||||||
loop {
|
|
||||||
if let Some(established_ws) = connections
|
|
||||||
.get_mut(&poll_token)
|
|
||||||
.map(|c| {
|
|
||||||
// Ugly but works
|
|
||||||
c.valid_until = valid_until;
|
|
||||||
|
|
||||||
c
|
|
||||||
})
|
|
||||||
.and_then(Connection::get_established_ws)
|
|
||||||
{
|
|
||||||
use ::tungstenite::Error::Io;
|
|
||||||
|
|
||||||
match established_ws.ws.read_message() {
|
|
||||||
Ok(ws_message) => {
|
|
||||||
let naive_peer_addr = established_ws.peer_addr;
|
|
||||||
let converted_peer_ip = convert_ipv4_mapped_ipv6(naive_peer_addr.ip());
|
let converted_peer_ip = convert_ipv4_mapped_ipv6(naive_peer_addr.ip());
|
||||||
|
|
||||||
let meta = ConnectionMeta {
|
let meta = ConnectionMeta {
|
||||||
out_message_consumer_id: ConsumerId(socket_worker_index),
|
out_message_consumer_id: ConsumerId(socket_worker_index),
|
||||||
connection_id: ConnectionId(poll_token.0),
|
connection_id: ConnectionId(token.0),
|
||||||
naive_peer_addr,
|
naive_peer_addr,
|
||||||
converted_peer_ip,
|
converted_peer_ip,
|
||||||
pending_scrape_id: None, // FIXME
|
pending_scrape_id: None, // FIXME
|
||||||
};
|
};
|
||||||
|
|
||||||
debug!("read message");
|
Connection::new(tls_config.clone(), ws_config, stream, valid_until, meta)
|
||||||
|
});
|
||||||
match InMessage::from_ws_message(ws_message) {
|
|
||||||
Ok(InMessage::AnnounceRequest(ref request))
|
|
||||||
if !state
|
|
||||||
.access_list
|
|
||||||
.allows(access_list_mode, &request.info_hash.0) =>
|
|
||||||
{
|
|
||||||
let out_message = OutMessage::ErrorResponse(ErrorResponse {
|
|
||||||
failure_reason: "Info hash not allowed".into(),
|
|
||||||
action: Some(ErrorResponseAction::Announce),
|
|
||||||
info_hash: Some(request.info_hash),
|
|
||||||
});
|
|
||||||
|
|
||||||
local_responses.push((meta, out_message));
|
|
||||||
}
|
|
||||||
Ok(in_message) => {
|
|
||||||
if let Err(err) = in_message_sender.send((meta, in_message)) {
|
|
||||||
error!("InMessageSender: couldn't send message: {:?}", err);
|
|
||||||
}
|
|
||||||
}
|
|
||||||
Err(_) => {
|
|
||||||
// FIXME: maybe this condition just occurs when enough data hasn't been recevied?
|
|
||||||
/*
|
|
||||||
info!("error parsing message: {:?}", err);
|
|
||||||
let out_message = OutMessage::ErrorResponse(ErrorResponse {
|
|
||||||
failure_reason: "Error parsing message".into(),
|
|
||||||
action: None,
|
|
||||||
info_hash: None,
|
|
||||||
});
|
|
||||||
local_responses.push((meta, out_message));
|
|
||||||
*/
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
Err(Io(err)) if err.kind() == ErrorKind::WouldBlock => {
|
|
||||||
break;
|
|
||||||
}
|
|
||||||
Err(tungstenite::Error::ConnectionClosed) => {
|
|
||||||
remove_connection_if_exists(poll, connections, poll_token);
|
|
||||||
|
|
||||||
break;
|
|
||||||
}
|
|
||||||
Err(err) => {
|
|
||||||
info!("error reading messages: {}", err);
|
|
||||||
|
|
||||||
remove_connection_if_exists(poll, connections, poll_token);
|
|
||||||
|
|
||||||
break;
|
|
||||||
}
|
|
||||||
}
|
}
|
||||||
} else if let Some(connection) = connections.remove(&poll_token) {
|
Err(err) if err.kind() == ErrorKind::WouldBlock => {
|
||||||
let (opt_new_connection, stop_loop) =
|
|
||||||
connection.advance_handshakes(opt_tls_acceptor, valid_until);
|
|
||||||
|
|
||||||
if let Some(connection) = opt_new_connection {
|
|
||||||
connections.insert(poll_token, connection);
|
|
||||||
}
|
|
||||||
|
|
||||||
if stop_loop {
|
|
||||||
break;
|
break;
|
||||||
}
|
}
|
||||||
} else {
|
Err(err) => {
|
||||||
break;
|
::log::info!("error while accepting streams: {}", err);
|
||||||
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Read messages from channel, send to peers
|
fn handle_stream_read_event(
|
||||||
pub fn send_out_messages(
|
config: &Config,
|
||||||
|
state: &State,
|
||||||
|
local_responses: &mut Vec<(ConnectionMeta, OutMessage)>,
|
||||||
|
in_message_sender: &InMessageSender,
|
||||||
poll: &mut Poll,
|
poll: &mut Poll,
|
||||||
local_responses: Drain<(ConnectionMeta, OutMessage)>,
|
|
||||||
out_message_receiver: &Receiver<(ConnectionMeta, OutMessage)>,
|
|
||||||
connections: &mut ConnectionMap,
|
connections: &mut ConnectionMap,
|
||||||
|
token: Token,
|
||||||
|
valid_until: ValidUntil,
|
||||||
) {
|
) {
|
||||||
let len = out_message_receiver.len();
|
let access_list_mode = config.access_list.mode;
|
||||||
|
|
||||||
for (meta, out_message) in local_responses.chain(out_message_receiver.try_iter().take(len)) {
|
if let Some(mut connection) = connections.remove_and_deregister(poll, &token) {
|
||||||
let opt_established_ws = connections
|
let message_handler = &mut |meta, message| match message {
|
||||||
.get_mut(&Token(meta.connection_id.0))
|
InMessage::AnnounceRequest(ref request)
|
||||||
.and_then(Connection::get_established_ws);
|
if !state
|
||||||
|
.access_list
|
||||||
|
.allows(access_list_mode, &request.info_hash.0) =>
|
||||||
|
{
|
||||||
|
let out_message = OutMessage::ErrorResponse(ErrorResponse {
|
||||||
|
failure_reason: "Info hash not allowed".into(),
|
||||||
|
action: Some(ErrorResponseAction::Announce),
|
||||||
|
info_hash: Some(request.info_hash),
|
||||||
|
});
|
||||||
|
|
||||||
if let Some(established_ws) = opt_established_ws {
|
local_responses.push((meta, out_message));
|
||||||
if established_ws.peer_addr != meta.naive_peer_addr {
|
|
||||||
info!("socket worker error: peer socket addrs didn't match");
|
|
||||||
|
|
||||||
continue;
|
|
||||||
}
|
}
|
||||||
|
in_message => {
|
||||||
use ::tungstenite::Error::Io;
|
if let Err(err) = in_message_sender.send((meta, in_message)) {
|
||||||
|
::log::info!("InMessageSender: couldn't send message: {:?}", err);
|
||||||
let ws_message = out_message.to_ws_message();
|
|
||||||
|
|
||||||
match established_ws.ws.write_message(ws_message) {
|
|
||||||
Ok(()) => {
|
|
||||||
debug!("sent message");
|
|
||||||
}
|
|
||||||
Err(Io(err)) if err.kind() == ErrorKind::WouldBlock => {}
|
|
||||||
Err(tungstenite::Error::ConnectionClosed) => {
|
|
||||||
remove_connection_if_exists(poll, connections, Token(meta.connection_id.0));
|
|
||||||
}
|
|
||||||
Err(err) => {
|
|
||||||
info!("error writing ws message: {}", err);
|
|
||||||
|
|
||||||
remove_connection_if_exists(poll, connections, Token(meta.connection_id.0));
|
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
};
|
||||||
|
|
||||||
|
connection.valid_until = valid_until;
|
||||||
|
|
||||||
|
match connection.read(message_handler) {
|
||||||
|
Ok(connection) => {
|
||||||
|
connections.insert_and_register(poll, token, connection);
|
||||||
|
}
|
||||||
|
Err(_) => {}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
fn write_or_queue_messages<I>(poll: &mut Poll, responses: I, connections: &mut ConnectionMap)
|
||||||
|
where
|
||||||
|
I: Iterator<Item = (ConnectionMeta, OutMessage)>,
|
||||||
|
{
|
||||||
|
for (meta, out_message) in responses {
|
||||||
|
let token = Token(meta.connection_id.0);
|
||||||
|
|
||||||
|
let mut remove_connection = false;
|
||||||
|
|
||||||
|
if let Some(connection) = connections.get_mut(&token) {
|
||||||
|
if connection.get_meta().naive_peer_addr != meta.naive_peer_addr {
|
||||||
|
::log::warn!(
|
||||||
|
"socket worker error: connection socket addr {} didn't match channel {}. Token: {}.",
|
||||||
|
connection.get_meta().naive_peer_addr,
|
||||||
|
meta.naive_peer_addr,
|
||||||
|
token.0
|
||||||
|
);
|
||||||
|
|
||||||
|
remove_connection = true;
|
||||||
|
} else {
|
||||||
|
match connection.write_or_queue_message(poll, out_message) {
|
||||||
|
Ok(()) => {}
|
||||||
|
Err(err) => {
|
||||||
|
::log::debug!("Connection::write_or_queue_message error: {}", err);
|
||||||
|
|
||||||
|
remove_connection = true;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
if remove_connection {
|
||||||
|
connections.remove_and_deregister(poll, &token);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn create_listener(config: &Config) -> ::anyhow::Result<::std::net::TcpListener> {
|
||||||
|
let builder = if config.network.address.is_ipv4() {
|
||||||
|
Socket::new(Domain::IPV4, Type::STREAM, Some(Protocol::TCP))
|
||||||
|
} else {
|
||||||
|
Socket::new(Domain::IPV6, Type::STREAM, Some(Protocol::TCP))
|
||||||
|
}
|
||||||
|
.context("Couldn't create socket2::Socket")?;
|
||||||
|
|
||||||
|
if config.network.ipv6_only {
|
||||||
|
builder
|
||||||
|
.set_only_v6(true)
|
||||||
|
.context("Couldn't put socket in ipv6 only mode")?
|
||||||
|
}
|
||||||
|
|
||||||
|
builder
|
||||||
|
.set_nonblocking(true)
|
||||||
|
.context("Couldn't put socket in non-blocking mode")?;
|
||||||
|
builder
|
||||||
|
.set_reuse_port(true)
|
||||||
|
.context("Couldn't put socket in reuse_port mode")?;
|
||||||
|
builder
|
||||||
|
.bind(&config.network.address.into())
|
||||||
|
.with_context(|| format!("Couldn't bind socket to address {}", config.network.address))?;
|
||||||
|
builder
|
||||||
|
.listen(128)
|
||||||
|
.context("Couldn't listen for connections on socket")?;
|
||||||
|
|
||||||
|
Ok(builder.into())
|
||||||
|
}
|
||||||
|
|
|
||||||
|
|
@ -1,66 +0,0 @@
|
||||||
use std::time::Instant;
|
|
||||||
|
|
||||||
use anyhow::Context;
|
|
||||||
use mio::{Poll, Token};
|
|
||||||
use socket2::{Domain, Protocol, Socket, Type};
|
|
||||||
|
|
||||||
use crate::config::Config;
|
|
||||||
|
|
||||||
use super::connection::*;
|
|
||||||
|
|
||||||
pub fn create_listener(config: &Config) -> ::anyhow::Result<::std::net::TcpListener> {
|
|
||||||
let builder = if config.network.address.is_ipv4() {
|
|
||||||
Socket::new(Domain::IPV4, Type::STREAM, Some(Protocol::TCP))
|
|
||||||
} else {
|
|
||||||
Socket::new(Domain::IPV6, Type::STREAM, Some(Protocol::TCP))
|
|
||||||
}
|
|
||||||
.context("Couldn't create socket2::Socket")?;
|
|
||||||
|
|
||||||
if config.network.ipv6_only {
|
|
||||||
builder
|
|
||||||
.set_only_v6(true)
|
|
||||||
.context("Couldn't put socket in ipv6 only mode")?
|
|
||||||
}
|
|
||||||
|
|
||||||
builder
|
|
||||||
.set_nonblocking(true)
|
|
||||||
.context("Couldn't put socket in non-blocking mode")?;
|
|
||||||
builder
|
|
||||||
.set_reuse_port(true)
|
|
||||||
.context("Couldn't put socket in reuse_port mode")?;
|
|
||||||
builder
|
|
||||||
.bind(&config.network.address.into())
|
|
||||||
.with_context(|| format!("Couldn't bind socket to address {}", config.network.address))?;
|
|
||||||
builder
|
|
||||||
.listen(128)
|
|
||||||
.context("Couldn't listen for connections on socket")?;
|
|
||||||
|
|
||||||
Ok(builder.into())
|
|
||||||
}
|
|
||||||
|
|
||||||
pub fn remove_connection_if_exists(poll: &mut Poll, connections: &mut ConnectionMap, token: Token) {
|
|
||||||
if let Some(mut connection) = connections.remove(&token) {
|
|
||||||
connection.close();
|
|
||||||
|
|
||||||
if let Err(err) = connection.deregister(poll) {
|
|
||||||
::log::error!("couldn't deregister stream: {}", err);
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
// Close and remove inactive connections
|
|
||||||
pub fn remove_inactive_connections(connections: &mut ConnectionMap) {
|
|
||||||
let now = Instant::now();
|
|
||||||
|
|
||||||
connections.retain(|_, connection| {
|
|
||||||
if connection.valid_until.0 < now {
|
|
||||||
connection.close();
|
|
||||||
|
|
||||||
false
|
|
||||||
} else {
|
|
||||||
true
|
|
||||||
}
|
|
||||||
});
|
|
||||||
|
|
||||||
connections.shrink_to_fit();
|
|
||||||
}
|
|
||||||
|
|
@ -22,6 +22,7 @@ futures = "0.3"
|
||||||
futures-rustls = "0.22"
|
futures-rustls = "0.22"
|
||||||
glommio = { git = "https://github.com/DataDog/glommio.git", rev = "2efe2f2a08f54394a435b674e8e0125057cbff03" }
|
glommio = { git = "https://github.com/DataDog/glommio.git", rev = "2efe2f2a08f54394a435b674e8e0125057cbff03" }
|
||||||
hashbrown = { version = "0.11", features = ["serde"] }
|
hashbrown = { version = "0.11", features = ["serde"] }
|
||||||
|
log = "0.4"
|
||||||
mimalloc = { version = "0.1", default-features = false }
|
mimalloc = { version = "0.1", default-features = false }
|
||||||
rand = { version = "0.8", features = ["small_rng"] }
|
rand = { version = "0.8", features = ["small_rng"] }
|
||||||
rand_distr = "0.4"
|
rand_distr = "0.4"
|
||||||
|
|
|
||||||
|
|
@ -105,7 +105,7 @@ impl Connection {
|
||||||
*num_active_connections.borrow_mut() += 1;
|
*num_active_connections.borrow_mut() += 1;
|
||||||
|
|
||||||
if let Err(err) = connection.run_connection_loop().await {
|
if let Err(err) = connection.run_connection_loop().await {
|
||||||
eprintln!("connection error: {:?}", err);
|
eprintln!("connection error: {}", err);
|
||||||
}
|
}
|
||||||
|
|
||||||
*num_active_connections.borrow_mut() -= 1;
|
*num_active_connections.borrow_mut() -= 1;
|
||||||
|
|
@ -159,7 +159,26 @@ impl Connection {
|
||||||
}
|
}
|
||||||
|
|
||||||
async fn read_message(&mut self) -> anyhow::Result<()> {
|
async fn read_message(&mut self) -> anyhow::Result<()> {
|
||||||
match OutMessage::from_ws_message(self.stream.next().await.unwrap()?) {
|
let message = match self
|
||||||
|
.stream
|
||||||
|
.next()
|
||||||
|
.await
|
||||||
|
.ok_or_else(|| anyhow::anyhow!("stream finished"))??
|
||||||
|
{
|
||||||
|
message @ tungstenite::Message::Text(_) | message @ tungstenite::Message::Binary(_) => {
|
||||||
|
message
|
||||||
|
}
|
||||||
|
message => {
|
||||||
|
eprintln!(
|
||||||
|
"Received WebSocket message of unexpected type: {:?}",
|
||||||
|
message
|
||||||
|
);
|
||||||
|
|
||||||
|
return Ok(());
|
||||||
|
}
|
||||||
|
};
|
||||||
|
|
||||||
|
match OutMessage::from_ws_message(message) {
|
||||||
Ok(OutMessage::Offer(offer)) => {
|
Ok(OutMessage::Offer(offer)) => {
|
||||||
self.load_test_state
|
self.load_test_state
|
||||||
.statistics
|
.statistics
|
||||||
|
|
@ -205,7 +224,7 @@ impl Connection {
|
||||||
self.can_send = true;
|
self.can_send = true;
|
||||||
}
|
}
|
||||||
Err(err) => {
|
Err(err) => {
|
||||||
eprintln!("error deserializing offer: {:?}", err);
|
eprintln!("error deserializing message: {:?}", err);
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|
|
||||||
Loading…
Add table
Add a link
Reference in a new issue