├── README.md ├── Cargo.toml ├── Makefile ├── .gitignore ├── integration_test.sh ├── LICENSE └── src ├── multiplexer_tests.rs └── main.rs /README.md: -------------------------------------------------------------------------------- 1 | # max_tcp 2 | MaxTCP proxies a TCP connection with multiple TCP connections. 3 | -------------------------------------------------------------------------------- /Cargo.toml: -------------------------------------------------------------------------------- 1 | [package] 2 | name = "max_tcp" 3 | version = "0.1.0" 4 | authors = ["imos "] 5 | 6 | [dependencies] 7 | mio = "0.6" 8 | getopts = "0.2" 9 | env_logger = "0.5.2" 10 | log = "0.4.0" 11 | rand = "0.4" 12 | lazy_static = "1.0" 13 | -------------------------------------------------------------------------------- /Makefile: -------------------------------------------------------------------------------- 1 | usage: 2 | @echo 'Usage: make (build|test)' 3 | .PHONY: usage 4 | 5 | build: 6 | cargo build --release 7 | .PHONY: build 8 | 9 | test: cargo-test integration-test 10 | .PHONY: test 11 | 12 | cargo-test: 13 | cargo test --release 14 | .PHONY: cargo-test 15 | 16 | integration-test: 17 | bash ./integration_test.sh 18 | .PHONY: integration-test 19 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | # Generated by Cargo 2 | # will have compiled files and executables 3 | /target/ 4 | 5 | # Remove Cargo.lock from gitignore if creating an executable, leave it for libraries 6 | # More information here http://doc.crates.io/guide.html#cargotoml-vs-cargolock 7 | Cargo.lock 8 | 9 | # These are backup files generated by rustfmt 10 | **/*.rs.bk 11 | /target/ 12 | **/*.rs.bk 13 | 14 | # Remove Mac invisible files. 15 | .DS_Store 16 | **/.DS_Store 17 | -------------------------------------------------------------------------------- /integration_test.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | set -e -u 3 | 4 | server() { 5 | ./target/release/max_tcp --server 216.239.32.21:80 127.0.0.1:8888 6 | } 7 | 8 | client() { 9 | ./target/release/max_tcp 127.0.0.1:8888 127.0.0.1:8889 10 | } 11 | 12 | cargo build --release 13 | RUST_BACKTRACE=1 RUST_LOG=info ./target/release/max_tcp --server 216.239.32.21:80 127.0.0.1:8888 & 14 | SERVER="$!" 15 | RUST_BACKTRACE=1 RUST_LOG=info ./target/release/max_tcp 127.0.0.1:8888 127.0.0.1:8889 & 16 | CLIENT="$!" 17 | ACTUAL="$(curl -H 'Host: imoz.jp' http://127.0.0.1:8889/img/profile.png | openssl md5)" 18 | EXPECTED='eec15e3d1e47bfd46c19dc7a32a8d4c0' 19 | sleep 0.3 20 | kill -TERM "${SERVER}" 21 | kill -TERM "${CLIENT}" 22 | wait 23 | if [ "${ACTUAL}" == "${EXPECTED}" ]; then 24 | echo 'Passed.' 25 | else 26 | echo "Actual: ${RESULT}" 27 | echo "Expected: eec15e3d1e47bfd46c19dc7a32a8d4c0" 28 | exit 1 29 | fi 30 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2018 Kentaro IMAJO 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. 22 | -------------------------------------------------------------------------------- /src/multiplexer_tests.rs: -------------------------------------------------------------------------------- 1 | #![cfg(test)] 2 | 3 | use super::*; 4 | use rand::Rng; 5 | 6 | #[test] 7 | fn test_init() { 8 | let _ = Multiplexer::new("", 0); 9 | let _ = Multiplexer::new("", 1); 10 | let _ = Multiplexer::new("", 100); 11 | } 12 | 13 | trait ToString { 14 | fn to_string(&self) -> String; 15 | } 16 | 17 | impl ToString for VecDeque { 18 | fn to_string(&self) -> String { 19 | String::from_utf8(self.iter().map(|byte| *byte).collect()).unwrap() 20 | } 21 | } 22 | 23 | trait VecTestUtil { 24 | fn push_back_str(&mut self, data: &str); 25 | } 26 | 27 | impl VecTestUtil for VecDeque { 28 | fn push_back_str(&mut self, data: &str) { 29 | data.bytes().for_each(|byte| self.push_back(byte)); 30 | } 31 | } 32 | 33 | impl VecTestUtil for Vec { 34 | fn push_back_str(&mut self, data: &str) { 35 | data.bytes().for_each(|byte| self.push(byte)); 36 | } 37 | } 38 | 39 | #[derive(Copy, Clone, PartialEq, PartialOrd)] 40 | struct Clock(usize); 41 | 42 | struct FakeClock(std::sync::Mutex); 43 | 44 | impl FakeClock { 45 | fn new() -> FakeClock { 46 | FakeClock(std::sync::Mutex::new(Clock(0))) 47 | } 48 | 49 | fn set(&self, clock: Clock) { 50 | *self.0.lock().unwrap() = clock; 51 | } 52 | 53 | fn get(&self) -> Clock { 54 | *self.0.lock().unwrap() 55 | } 56 | } 57 | 58 | struct FakeStream<'a> { 59 | fake_clock: &'a FakeClock, 60 | is_dead: bool, 61 | reads: VecDeque<(Clock, std::io::Result>)>, 62 | writes: VecDeque, 63 | } 64 | 65 | impl<'a> FakeStream<'a> { 66 | pub fn new(fake_clock: &'a FakeClock) -> FakeStream<'a> { 67 | FakeStream { 68 | fake_clock: fake_clock, 69 | is_dead: false, 70 | reads: VecDeque::new(), 71 | writes: VecDeque::new(), 72 | } 73 | } 74 | 75 | pub fn push_str_to_read(&mut self, clock: Clock, data: &str) { 76 | self.reads 77 | .push_back((clock, Ok(data.bytes().collect::>()))); 78 | } 79 | 80 | pub fn mark_as_dead(&mut self) { 81 | self.is_dead = true; 82 | } 83 | } 84 | 85 | impl<'a> std::io::Read for FakeStream<'a> { 86 | fn read(&mut self, buf: &mut [u8]) -> std::io::Result { 87 | if self.reads.len() == 0 && self.is_dead { 88 | return Ok(0); 89 | } else if self.reads.len() == 0 || self.fake_clock.get() < self.reads[0].0 { 90 | return Err(std::io::Error::new( 91 | std::io::ErrorKind::WouldBlock, 92 | "Waiting for data.", 93 | )); 94 | } 95 | let (read_clock, read_buffer) = match self.reads.pop_front() { 96 | Some(read) => read, 97 | None => return Ok(0), 98 | }; 99 | let mut read_buffer = match read_buffer { 100 | Ok(read_buffer) => read_buffer, 101 | Err(error) => return Err(error), 102 | }; 103 | let mut length = 0; 104 | while length < buf.len() { 105 | buf[length] = match read_buffer.pop_front() { 106 | Some(byte) => byte, 107 | None => break, 108 | }; 109 | length += 1; 110 | } 111 | if !read_buffer.is_empty() { 112 | self.reads.push_front((read_clock, Ok(read_buffer))); 113 | } 114 | Ok(length) 115 | } 116 | } 117 | 118 | impl<'a> std::io::Write for FakeStream<'a> { 119 | fn write(&mut self, buf: &[u8]) -> std::io::Result { 120 | if self.is_dead { 121 | return Err(std::io::Error::new( 122 | std::io::ErrorKind::BrokenPipe, 123 | "Connection is dead.", 124 | )); 125 | } 126 | for byte in buf { 127 | self.writes.push_back(*byte); 128 | } 129 | Ok(buf.len()) 130 | } 131 | 132 | fn flush(&mut self) -> std::io::Result<()> { 133 | Ok(()) 134 | } 135 | } 136 | 137 | impl<'a> mio::Evented for FakeStream<'a> { 138 | fn register( 139 | &self, 140 | _: &mio::Poll, 141 | _: mio::Token, 142 | _: mio::Ready, 143 | _: mio::PollOpt, 144 | ) -> std::io::Result<()> { 145 | Ok(()) 146 | } 147 | 148 | fn reregister( 149 | &self, 150 | _: &mio::Poll, 151 | _: mio::Token, 152 | _: mio::Ready, 153 | _: mio::PollOpt, 154 | ) -> std::io::Result<()> { 155 | Ok(()) 156 | } 157 | 158 | fn deregister(&self, _: &mio::Poll) -> std::io::Result<()> { 159 | Ok(()) 160 | } 161 | } 162 | 163 | #[test] 164 | fn test_master_read() { 165 | let mut multiplexer = Multiplexer::new("", 0); 166 | let fake_clock = FakeClock::new(); 167 | let mut fake_master = FakeStream::new(&fake_clock); 168 | 169 | fake_master.push_str_to_read(Clock(0), "TEST1"); 170 | fake_master.push_str_to_read(Clock(1), "TEST2"); 171 | 172 | multiplexer.is_dirty = false; 173 | multiplexer.read_master(&mut fake_master).unwrap(); 174 | assert_eq!(multiplexer.master_read_buffer.to_string(), "TEST1"); 175 | assert!(multiplexer.is_dirty, "read should mark it as dirty."); 176 | 177 | fake_clock.set(Clock(1)); 178 | multiplexer.is_dirty = false; 179 | multiplexer.read_master(&mut fake_master).unwrap(); 180 | assert_eq!(multiplexer.master_read_buffer.to_string(), "TEST1TEST2"); 181 | assert!(multiplexer.is_dirty, "read should mark it as dirty."); 182 | 183 | fake_clock.set(Clock(2)); 184 | multiplexer.is_dirty = false; 185 | multiplexer.read_master(&mut fake_master).unwrap(); 186 | assert!(!multiplexer.is_dirty, "No read should happen."); 187 | 188 | assert!(multiplexer.master_connectivity); 189 | } 190 | 191 | #[test] 192 | fn test_master_read_should_not_recieve_excess_data() { 193 | let mut multiplexer = Multiplexer::new("", 0); 194 | let fake_clock = FakeClock::new(); 195 | let mut fake_master = FakeStream::new(&fake_clock); 196 | 197 | for _ in 0..Multiplexer::MASTER_BUFFER_SIZE / 4 { 198 | fake_master.push_str_to_read(Clock(0), "0123456789abcdef0123456789abcdef"); 199 | } 200 | 201 | multiplexer.read_master(&mut fake_master).unwrap(); 202 | assert!( 203 | multiplexer.master_read_buffer.len() >= multiplexer.master_read_buffer.capacity() - 4, 204 | "master_read_buffer must be full (actual: {}, expected: _ >= {}).", 205 | multiplexer.master_read_buffer.len(), 206 | multiplexer.master_read_buffer.capacity() - 4 207 | ); 208 | 209 | multiplexer.master_read_buffer.clear(); 210 | multiplexer.read_master(&mut fake_master).unwrap(); 211 | assert!( 212 | multiplexer.master_read_buffer.len() >= multiplexer.master_read_buffer.capacity() - 4, 213 | "master_read_buffer must be full (actual: {}, expected: _ >= {}).", 214 | multiplexer.master_read_buffer.len(), 215 | multiplexer.master_read_buffer.capacity() - 4 216 | ); 217 | } 218 | 219 | #[test] 220 | fn test_master_read_can_detect_eof() { 221 | let mut multiplexer = Multiplexer::new("", 0); 222 | let fake_clock = FakeClock::new(); 223 | let mut fake_master = FakeStream::new(&fake_clock); 224 | 225 | fake_master.push_str_to_read(Clock(0), "TEST1"); 226 | fake_master.mark_as_dead(); 227 | 228 | multiplexer.is_dirty = false; 229 | assert!(multiplexer.master_connectivity); 230 | multiplexer.read_master(&mut fake_master).unwrap(); 231 | assert_eq!(multiplexer.master_read_buffer.to_string(), "TEST1"); 232 | assert!(multiplexer.is_dirty, "read should mark it as dirty."); 233 | assert!(!multiplexer.master_connectivity); 234 | } 235 | 236 | #[test] 237 | fn test_master_write() { 238 | let mut multiplexer = Multiplexer::new("", 0); 239 | let fake_clock = FakeClock::new(); 240 | let mut fake_master = FakeStream::new(&fake_clock); 241 | 242 | multiplexer.master_write_buffers[0].0 = true; 243 | multiplexer.master_write_buffers[0].1.push_back_str("TEST1"); 244 | multiplexer.master_write_buffers[2].0 = true; 245 | multiplexer.master_write_buffers[2].1.push_back_str("TEST3"); 246 | 247 | multiplexer.is_dirty = false; 248 | multiplexer.write_master(&mut fake_master).unwrap(); 249 | assert_eq!(fake_master.writes.to_string(), "TEST1"); 250 | assert!(multiplexer.is_dirty, "write should mark it as dirty."); 251 | assert_eq!(multiplexer.gather_index, 1); 252 | 253 | multiplexer.master_write_buffers[0].0 = true; 254 | multiplexer.master_write_buffers[0].1.push_back_str("TEST2"); 255 | 256 | multiplexer.is_dirty = false; 257 | multiplexer.write_master(&mut fake_master).unwrap(); 258 | assert_eq!(fake_master.writes.to_string(), "TEST1TEST2TEST3"); 259 | assert!(multiplexer.is_dirty, "write should mark it as dirty."); 260 | assert_eq!(multiplexer.gather_index, 3); 261 | 262 | multiplexer.is_dirty = false; 263 | multiplexer.write_master(&mut fake_master).unwrap(); 264 | assert!(!multiplexer.is_dirty, "No write should happen."); 265 | } 266 | 267 | #[test] 268 | fn test_slave_read() { 269 | let fake_clock = FakeClock::new(); 270 | let mut fake_slaves = Vec::new(); 271 | for _ in 0..3 { 272 | fake_slaves.push(FakeStream::new(&fake_clock)); 273 | } 274 | let mut multiplexer = Multiplexer::new("", fake_slaves.len()); 275 | 276 | fake_slaves[0].push_str_to_read(Clock(1), "TEST1"); 277 | fake_slaves[0].mark_as_dead(); 278 | fake_slaves[1].push_str_to_read(Clock(0), "TEST2"); 279 | fake_slaves[2].push_str_to_read(Clock(1), "TEST3"); 280 | fake_slaves[2].push_str_to_read(Clock(2), "TEST4"); 281 | fake_slaves[2].mark_as_dead(); 282 | 283 | multiplexer.is_dirty = false; 284 | multiplexer.read_slaves(&mut fake_slaves).unwrap(); 285 | assert_eq!(multiplexer.slave_read_buffers[0].to_string(), ""); 286 | assert_eq!(multiplexer.slave_read_buffers[1].to_string(), "TEST2"); 287 | assert_eq!(multiplexer.slave_read_buffers[2].to_string(), ""); 288 | assert_eq!(multiplexer.num_of_alive_slaves, 3); 289 | assert!(multiplexer.is_dirty, "read should mark it as dirty."); 290 | 291 | fake_clock.set(Clock(1)); 292 | multiplexer.is_dirty = false; 293 | multiplexer.read_slaves(&mut fake_slaves).unwrap(); 294 | assert_eq!(multiplexer.slave_read_buffers[0].to_string(), "TEST1"); 295 | assert_eq!(multiplexer.slave_read_buffers[1].to_string(), "TEST2"); 296 | assert_eq!(multiplexer.slave_read_buffers[2].to_string(), "TEST3"); 297 | // Slave #0 should close. 298 | assert_eq!(multiplexer.num_of_alive_slaves, 2); 299 | assert!(multiplexer.is_dirty, "read should mark it as dirty."); 300 | 301 | fake_clock.set(Clock(2)); 302 | multiplexer.is_dirty = false; 303 | multiplexer.read_slaves(&mut fake_slaves).unwrap(); 304 | assert_eq!(multiplexer.slave_read_buffers[0].to_string(), "TEST1"); 305 | assert_eq!(multiplexer.slave_read_buffers[1].to_string(), "TEST2"); 306 | assert_eq!(multiplexer.slave_read_buffers[2].to_string(), "TEST3TEST4"); 307 | // Slave #2 should close. 308 | assert_eq!(multiplexer.num_of_alive_slaves, 1); 309 | assert!(multiplexer.is_dirty, "read should mark it as dirty."); 310 | 311 | fake_clock.set(Clock(3)); 312 | multiplexer.is_dirty = false; 313 | multiplexer.read_slaves(&mut fake_slaves).unwrap(); 314 | assert_eq!(multiplexer.slave_read_buffers[0].to_string(), "TEST1"); 315 | assert_eq!(multiplexer.slave_read_buffers[1].to_string(), "TEST2"); 316 | assert_eq!(multiplexer.slave_read_buffers[2].to_string(), "TEST3TEST4"); 317 | assert_eq!(multiplexer.num_of_alive_slaves, 1); 318 | assert!(!multiplexer.is_dirty, "No read should happen."); 319 | } 320 | 321 | #[test] 322 | fn test_slave_write() { 323 | let fake_clock = FakeClock::new(); 324 | let mut fake_slaves = Vec::new(); 325 | for _ in 0..3 { 326 | fake_slaves.push(FakeStream::new(&fake_clock)); 327 | } 328 | let mut multiplexer = Multiplexer::new("", fake_slaves.len()); 329 | 330 | multiplexer.slave_write_buffers[1].push_back_str("TEST1"); 331 | 332 | multiplexer.is_dirty = false; 333 | multiplexer.write_slaves(&mut fake_slaves).unwrap(); 334 | assert_eq!(fake_slaves[0].writes.to_string(), ""); 335 | assert_eq!(fake_slaves[1].writes.to_string(), "TEST1"); 336 | assert_eq!(fake_slaves[2].writes.to_string(), ""); 337 | assert!(multiplexer.is_dirty, "write should mark it as dirty."); 338 | 339 | multiplexer.slave_write_buffers[0].push_back_str("TEST2"); 340 | multiplexer.slave_write_buffers[1].push_back_str("TEST3"); 341 | multiplexer.slave_write_buffers[2].push_back_str("TEST4"); 342 | 343 | multiplexer.is_dirty = false; 344 | multiplexer.write_slaves(&mut fake_slaves).unwrap(); 345 | assert_eq!(fake_slaves[0].writes.to_string(), "TEST2"); 346 | assert_eq!(fake_slaves[1].writes.to_string(), "TEST1TEST3"); 347 | assert_eq!(fake_slaves[2].writes.to_string(), "TEST4"); 348 | assert!(multiplexer.is_dirty, "write should mark it as dirty."); 349 | 350 | multiplexer.is_dirty = false; 351 | multiplexer.write_slaves(&mut fake_slaves).unwrap(); 352 | assert_eq!(fake_slaves[0].writes.to_string(), "TEST2"); 353 | assert_eq!(fake_slaves[1].writes.to_string(), "TEST1TEST3"); 354 | assert_eq!(fake_slaves[2].writes.to_string(), "TEST4"); 355 | assert!(!multiplexer.is_dirty, "No write should happen."); 356 | } 357 | 358 | #[test] 359 | fn test_scatter() { 360 | let mut multiplexer = Multiplexer::new("", 2); 361 | 362 | { 363 | multiplexer.master_read_buffer.push_back_str("TESTTEST1"); 364 | 365 | multiplexer.is_dirty = false; 366 | multiplexer.scatter(); 367 | assert_eq!(multiplexer.slave_write_buffers[0].to_string(), ""); 368 | assert_eq!( 369 | multiplexer.slave_write_buffers[1].to_string(), 370 | "\x00\x00\x00\x00\x09\x00TESTTEST1" 371 | ); 372 | assert!(multiplexer.is_dirty, "scatter should mark it as dirty."); 373 | } 374 | { 375 | multiplexer.master_read_buffer.push_back_str("TEST2"); 376 | 377 | multiplexer.is_dirty = false; 378 | multiplexer.scatter(); 379 | assert_eq!( 380 | multiplexer.slave_write_buffers[0].to_string(), 381 | "\x01\x00\x00\x00\x05\x00TEST2" 382 | ); 383 | assert_eq!( 384 | multiplexer.slave_write_buffers[1].to_string(), 385 | "\x00\x00\x00\x00\x09\x00TESTTEST1" 386 | ); 387 | assert!(multiplexer.is_dirty, "scatter should mark it as dirty."); 388 | } 389 | { 390 | multiplexer.master_read_buffer.push_back_str("TEST3"); 391 | 392 | multiplexer.is_dirty = false; 393 | multiplexer.scatter(); 394 | assert_eq!( 395 | multiplexer.slave_write_buffers[0].to_string(), 396 | "\x01\x00\x00\x00\x05\x00TEST2\x02\x00\x00\x00\x05\x00TEST3" 397 | ); 398 | assert_eq!( 399 | multiplexer.slave_write_buffers[1].to_string(), 400 | "\x00\x00\x00\x00\x09\x00TESTTEST1" 401 | ); 402 | assert!(multiplexer.is_dirty, "scatter should mark it as dirty."); 403 | } 404 | { 405 | multiplexer.is_dirty = false; 406 | multiplexer.scatter(); 407 | assert!(!multiplexer.is_dirty, "No scatter should happen."); 408 | } 409 | } 410 | 411 | #[test] 412 | fn test_gather() { 413 | let mut multiplexer = Multiplexer::new("", 2); 414 | let fake_clock = FakeClock::new(); 415 | let mut fake_master = FakeStream::new(&fake_clock); 416 | 417 | multiplexer.slave_read_buffers[0].push_back_str("\x00\x00\x00\x00\x06\x00TEST1X"); 418 | multiplexer.slave_read_buffers[0].push_back_str("\x03\x00\x00\x00\x09\x00TEST4XXXX"); 419 | // Incomplete chunk cannot be accepted. 420 | multiplexer.slave_read_buffers[0].push_back_str("\x04\x00\x00\x00\x0A\x00TEST5XXXX"); 421 | 422 | // Receive the first chunk only. 423 | { 424 | multiplexer.is_dirty = false; 425 | multiplexer.gather(); 426 | assert!( 427 | multiplexer.is_dirty, 428 | "gather should mark the multiplexer as dirty." 429 | ); 430 | multiplexer.write_master(&mut fake_master).unwrap(); 431 | assert_eq!(fake_master.writes.to_string(), "TEST1X"); 432 | } 433 | 434 | multiplexer.slave_read_buffers[1].push_back_str("\x01\x00\x00\x00\x07\x00TEST2XX"); 435 | multiplexer.slave_read_buffers[1].push_back_str("\x02\x00\x00\x00\x08\x00TEST3XXX"); 436 | 437 | // Receive chunks 3 more chunks. 438 | { 439 | multiplexer.is_dirty = false; 440 | multiplexer.gather(); 441 | assert!( 442 | multiplexer.is_dirty, 443 | "gather should mark the multiplexer as dirty." 444 | ); 445 | multiplexer.write_master(&mut fake_master).unwrap(); 446 | assert_eq!( 447 | fake_master.writes.to_string(), 448 | "TEST1XTEST2XXTEST3XXXTEST4XXXX" 449 | ); 450 | } 451 | 452 | // Cannot receive more complete chunks because the next chunk is incomplete. 453 | { 454 | multiplexer.is_dirty = false; 455 | multiplexer.gather(); 456 | assert!(!multiplexer.is_dirty, "gather should not happen."); 457 | multiplexer.write_master(&mut fake_master).unwrap(); 458 | assert_eq!( 459 | fake_master.writes.to_string(), 460 | "TEST1XTEST2XXTEST3XXXTEST4XXXX" 461 | ); 462 | } 463 | 464 | // Complete the chunk, then it can be received. 465 | multiplexer.slave_read_buffers[0].push_back_str("X"); 466 | { 467 | multiplexer.is_dirty = false; 468 | multiplexer.gather(); 469 | assert!( 470 | multiplexer.is_dirty, 471 | "gather should mark the multiplexer as dirty." 472 | ); 473 | multiplexer.write_master(&mut fake_master).unwrap(); 474 | assert_eq!( 475 | fake_master.writes.to_string(), 476 | "TEST1XTEST2XXTEST3XXXTEST4XXXXTEST5XXXXX" 477 | ); 478 | } 479 | } 480 | 481 | #[test] 482 | fn test_integration() { 483 | use rand::SeedableRng; 484 | 485 | let fake_clock = FakeClock::new(); 486 | let mut fake_slaves = Vec::new(); 487 | for _ in 0..3 { 488 | fake_slaves.push(FakeStream::new(&fake_clock)); 489 | } 490 | let mut fake_master = FakeStream::new(&fake_clock); 491 | let mut multiplexer = Multiplexer::new("", fake_slaves.len()); 492 | let mut send_rng = rand::XorShiftRng::from_seed([1; 4]); 493 | let mut recv_rng = rand::XorShiftRng::from_seed([1; 4]); 494 | let mut parameter_rng = rand::XorShiftRng::from_seed([2; 4]); 495 | let mut num_of_sent_bytes = 0; 496 | 497 | for senario in 0..3 { 498 | let (master_speed, slave_speed) = match senario { 499 | 0 => (100 * fake_slaves.len(), 100), 500 | 1 => (100 * fake_slaves.len(), 10), 501 | 2 => (0, 1000), 502 | _ => unreachable!(), 503 | }; 504 | for _ in 0..1000 { 505 | let mut data = VecDeque::new(); 506 | if master_speed > 0 { 507 | for _ in 0..parameter_rng.gen_range(0, master_speed) { 508 | data.push_back(send_rng.gen()); 509 | num_of_sent_bytes += 1; 510 | } 511 | } 512 | fake_master.reads.push_back((fake_clock.get(), Ok(data))); 513 | 514 | multiplexer 515 | .multiplex_while_dirty(&mut fake_master, &mut fake_slaves) 516 | .unwrap(); 517 | 518 | for fake_slave in fake_slaves.iter_mut() { 519 | let length = std::cmp::min( 520 | fake_slave.writes.len(), 521 | parameter_rng.gen_range(0, slave_speed), 522 | ); 523 | if length > 0 { 524 | fake_slave.reads.push_back(( 525 | fake_clock.get(), 526 | Ok(fake_slave.writes.drain(..length).collect()), 527 | )); 528 | } 529 | } 530 | } 531 | println!( 532 | "Send: {} bytes, Recv: {} bytes", 533 | num_of_sent_bytes, 534 | fake_master.writes.len() 535 | ); 536 | } 537 | 538 | for (index, byte) in fake_master.writes.iter().enumerate() { 539 | assert_eq!( 540 | *byte, 541 | recv_rng.gen::(), 542 | "Received unexpected data in #{}.", 543 | index 544 | ); 545 | } 546 | assert_eq!(num_of_sent_bytes, fake_master.writes.len()); 547 | } 548 | -------------------------------------------------------------------------------- /src/main.rs: -------------------------------------------------------------------------------- 1 | extern crate env_logger; 2 | extern crate getopts; 3 | #[macro_use] 4 | extern crate log; 5 | extern crate mio; 6 | extern crate rand; 7 | 8 | use rand::Rng; 9 | use std::collections::BinaryHeap; 10 | use std::collections::VecDeque; 11 | use std::io::Read; 12 | use std::io::Write; 13 | 14 | mod multiplexer_tests; 15 | 16 | trait VecDequeUtil { 17 | fn push_back_usize(&mut self, value: usize, length: usize); 18 | fn peek_usize(&self, position: usize, length: usize) -> usize; 19 | } 20 | 21 | impl VecDequeUtil for VecDeque { 22 | fn push_back_usize(&mut self, value: usize, length: usize) { 23 | assert!( 24 | length * 8 < usize::max_value().count_ones() as usize, 25 | "length is too big for usize (actual: {}, expected: _ < {}).", 26 | length * 8, 27 | usize::max_value().count_ones() 28 | ); 29 | let mut value = value; 30 | for _ in 0..length { 31 | self.push_back(value as u8); 32 | value >>= 8; 33 | } 34 | } 35 | 36 | fn peek_usize(&self, start: usize, length: usize) -> usize { 37 | assert!( 38 | length * 8 < usize::max_value().count_ones() as usize, 39 | "length is too big for usize (actual: {}, expected: _ < {}).", 40 | length * 8, 41 | usize::max_value().count_ones() 42 | ); 43 | assert!( 44 | start + length <= self.len(), 45 | "Out of bounds in peek_usize (actual: {}, expected: _ <= {}).", 46 | start + length, 47 | self.len() 48 | ); 49 | 50 | let mut result = 0; 51 | let mut base = 1; 52 | for position in start..(start + length) { 53 | result += self[position] as usize * base; 54 | base <<= 8; 55 | } 56 | result 57 | } 58 | } 59 | 60 | trait ErrorUtil { 61 | fn prepend_err(self, message: &str) -> Self; 62 | } 63 | 64 | impl ErrorUtil for std::io::Result { 65 | fn prepend_err(self, message: &str) -> Self { 66 | self.map_err(|error| { 67 | std::io::Error::new(error.kind(), format!("{}\n- {}", message, error)) 68 | }) 69 | } 70 | } 71 | 72 | #[derive(Clone, Debug)] 73 | struct MultiplexerStats { 74 | time: std::time::Instant, 75 | flush: usize, 76 | master_read: usize, 77 | master_read_packets: usize, 78 | master_write: usize, 79 | master_write_packets: usize, 80 | } 81 | 82 | impl MultiplexerStats { 83 | fn new(time: std::time::Instant) -> Self { 84 | Self { 85 | time: time, 86 | flush: 0, 87 | master_read: 0, 88 | master_read_packets: 0, 89 | master_write: 0, 90 | master_write_packets: 0, 91 | } 92 | } 93 | } 94 | 95 | pub struct Multiplexer { 96 | /// Name of the multiplexer connection. Used for debugging. 97 | name: String, 98 | /// Read buffers, each of which is bound to a slave connection. 99 | pub(crate) slave_read_buffers: Vec>, 100 | /// Write buffers, each of which is bound to a slave connection. 101 | pub(crate) slave_write_buffers: Vec>, 102 | /// Priority queue of (capacity, index). 103 | pub(crate) slave_write_capacity_queue: BinaryHeap<(usize, usize)>, 104 | /// Connectivity for each slave. 105 | pub(crate) slave_connectivities: Vec, 106 | /// A read buffer for the master connection. 107 | pub(crate) master_read_buffer: VecDeque, 108 | /// Write buffers, each of which represents a pair of existence and a chunk. 109 | pub(crate) master_write_buffers: VecDeque<(bool, Vec)>, 110 | /// Master's connectivity. 111 | pub(crate) master_connectivity: bool, 112 | /// # of alive connections. 113 | pub(crate) num_of_alive_slaves: usize, 114 | /// Chunk index to read from slaves. 115 | pub(crate) gather_index: usize, 116 | /// Chunk index to write to slaves. 117 | pub(crate) scatter_index: usize, 118 | /// Is the state changed. 119 | pub(crate) is_dirty: bool, 120 | /// Temporary buffer. 121 | temporary_buffer: [u8; 1024], 122 | /// Stats. 123 | stats: MultiplexerStats, 124 | /// Stats which are output in the last time. 125 | last_stats: MultiplexerStats, 126 | } 127 | 128 | impl Multiplexer { 129 | const SLAVE_BUFFER_SIZE: usize = 16 * 1024; 130 | const MASTER_BUFFER_SIZE: usize = 128 * 1024; 131 | const PACKET_SIZE: usize = 1024; 132 | 133 | const INDEX_HEADER_SIZE: usize = 4; 134 | const LENGTH_HEADER_SIZE: usize = 2; 135 | const HEADER_SIZE: usize = Self::INDEX_HEADER_SIZE + Self::LENGTH_HEADER_SIZE; 136 | 137 | pub fn new(name: &str, num_of_slaves: usize) -> Multiplexer { 138 | info!( 139 | "Multiplexer {}: It will allocate {} bytes for buffering.", 140 | name, 141 | Self::SLAVE_BUFFER_SIZE * num_of_slaves * 2 + Self::MASTER_BUFFER_SIZE * 2 142 | ); 143 | 144 | // Assert constants. 145 | assert!( 146 | Self::PACKET_SIZE.count_ones() == 1, 147 | "PACKET_SIZE must be a power of 2, but {}.", 148 | Self::PACKET_SIZE 149 | ); 150 | assert!( 151 | Self::MASTER_BUFFER_SIZE % Self::PACKET_SIZE == 0, 152 | "MASTER_BUFFER_SIZE must be divisible by PACKET_SIZE, but the remainder is {}.", 153 | Self::MASTER_BUFFER_SIZE % Self::PACKET_SIZE 154 | ); 155 | 156 | // Prepare a skeleton of Multiplexer. 157 | let current_time = std::time::Instant::now(); 158 | let mut multiplexer = Multiplexer { 159 | name: name.into(), 160 | slave_read_buffers: Vec::with_capacity(num_of_slaves), 161 | slave_write_buffers: Vec::with_capacity(num_of_slaves), 162 | slave_write_capacity_queue: BinaryHeap::with_capacity(num_of_slaves), 163 | slave_connectivities: Vec::with_capacity(num_of_slaves), 164 | master_read_buffer: VecDeque::with_capacity(Self::MASTER_BUFFER_SIZE), 165 | master_write_buffers: VecDeque::with_capacity( 166 | Self::MASTER_BUFFER_SIZE / Self::PACKET_SIZE, 167 | ), 168 | master_connectivity: true, 169 | num_of_alive_slaves: num_of_slaves, 170 | gather_index: 0, 171 | scatter_index: 0, 172 | is_dirty: true, 173 | temporary_buffer: [0; Self::PACKET_SIZE], 174 | stats: MultiplexerStats::new(current_time), 175 | last_stats: MultiplexerStats::new(current_time), 176 | }; 177 | 178 | // Populate 2-dimensional vectors. 179 | for _ in 0..num_of_slaves { 180 | multiplexer 181 | .slave_read_buffers 182 | .push(VecDeque::with_capacity(Self::SLAVE_BUFFER_SIZE)); 183 | multiplexer 184 | .slave_write_buffers 185 | .push(VecDeque::with_capacity(Self::SLAVE_BUFFER_SIZE)); 186 | multiplexer.slave_connectivities.push(true); 187 | } 188 | while multiplexer.master_write_buffers.len() < multiplexer.master_write_buffers.capacity() { 189 | multiplexer 190 | .master_write_buffers 191 | .push_back((false, Vec::with_capacity(Self::PACKET_SIZE))); 192 | } 193 | 194 | multiplexer 195 | } 196 | 197 | pub fn read_master( 198 | &mut self, 199 | master: &mut Reader, 200 | ) -> std::io::Result<()> { 201 | let master_read_buffer = &mut self.master_read_buffer; 202 | while master_read_buffer.len() < master_read_buffer.capacity() { 203 | let length_limit = std::cmp::min( 204 | master_read_buffer.capacity() - master_read_buffer.len(), 205 | self.temporary_buffer.len(), 206 | ); 207 | let length = match master.read(&mut self.temporary_buffer[0..length_limit]) { 208 | Ok(0) => { 209 | if self.master_connectivity { 210 | info!("Multiplexer {}: Connection to master is closed.", self.name); 211 | self.master_connectivity = false; 212 | } 213 | break; 214 | } 215 | Ok(length) => length, 216 | Err(ref error) if error.kind() == std::io::ErrorKind::WouldBlock => break, 217 | Err(error) => { 218 | return Err(error).prepend_err(&format!( 219 | "Failed to try reading {} bytes from a master.", 220 | length_limit 221 | )) 222 | } 223 | }; 224 | self.is_dirty = true; 225 | self.stats.master_read += length; 226 | self.stats.master_read_packets += 1; 227 | self.temporary_buffer[..length].iter().for_each(|byte| { 228 | master_read_buffer.push_back(*byte); 229 | }); 230 | assert!( 231 | master_read_buffer.capacity() <= Self::MASTER_BUFFER_SIZE * 2, 232 | "Multiplexer {}: master_read_buffer's capacity is too big \ 233 | (actual: {}, expected: _ <= {}).", 234 | self.name, 235 | master_read_buffer.capacity(), 236 | Self::MASTER_BUFFER_SIZE * 2 237 | ); 238 | } 239 | 240 | Ok(()) 241 | } 242 | 243 | pub fn write_master( 244 | &mut self, 245 | master: &mut Writer, 246 | ) -> std::io::Result<()> { 247 | loop { 248 | // If consuming the current buffer completes, advance the gather index. 249 | if { 250 | let (existence, ref mut master_write_buffer) = self.master_write_buffers[0]; 251 | 252 | // Finish if the next buffer is unavailable. 253 | if !existence { 254 | break; 255 | } 256 | 257 | // Write data and remove corresponding ones from the buffer. 258 | let length = match master.write(master_write_buffer.as_slice()) { 259 | Ok(length) => length, 260 | Err(ref error) if error.kind() == std::io::ErrorKind::WouldBlock => break, 261 | Err(error) => { 262 | return Err(error).prepend_err(&format!( 263 | "Failed to write {} bytes to a master socket.", 264 | master_write_buffer.len() 265 | )) 266 | } 267 | }; 268 | self.is_dirty = true; 269 | self.stats.master_write += length; 270 | self.stats.master_write_packets += 1; 271 | master_write_buffer.drain(..length); 272 | master_write_buffer.is_empty() 273 | } { 274 | // Advance master_write_buffers along with gather_index increment. 275 | let master_write_buffer = match self.master_write_buffers.pop_front() { 276 | Some((_, master_write_buffer)) => master_write_buffer, 277 | _ => unreachable!(), 278 | }; 279 | assert!( 280 | master_write_buffer.capacity() == Self::PACKET_SIZE, 281 | "Multiplexer {}: Unexpected capacity for a master write buffer \ 282 | (actual: {}, expected: {}).", 283 | self.name, 284 | master_write_buffer.capacity(), 285 | Self::PACKET_SIZE 286 | ); 287 | self.master_write_buffers 288 | .push_back((false, master_write_buffer)); 289 | self.gather_index += 1; 290 | } 291 | } 292 | 293 | Ok(()) 294 | } 295 | 296 | pub fn read_slaves( 297 | &mut self, 298 | slaves: &mut Vec, 299 | ) -> std::io::Result<()> { 300 | assert!( 301 | self.slave_read_buffers.len() == slaves.len(), 302 | "Multiplexer {}: Inconsistent number of slaves (expected: {}, actual: {}).", 303 | self.name, 304 | self.slave_read_buffers.len(), 305 | slaves.len() 306 | ); 307 | 308 | // For each slave, receive bytes from the slave and store them into 309 | // a buffer as much as possible. 310 | for (index, slave_read_buffer) in self.slave_read_buffers.iter_mut().enumerate() { 311 | while slave_read_buffer.len() < slave_read_buffer.capacity() { 312 | let length = std::cmp::min( 313 | self.temporary_buffer.len(), 314 | slave_read_buffer.capacity() - slave_read_buffer.len(), 315 | ); 316 | let length = match slaves[index].read(&mut self.temporary_buffer[..length]) { 317 | Ok(0) => { 318 | if self.slave_connectivities[index] { 319 | info!( 320 | "Multiplexer {}: Connection to slave #{} is dead.", 321 | self.name, 322 | index 323 | ); 324 | self.slave_connectivities[index] = false; 325 | self.num_of_alive_slaves -= 1; 326 | } 327 | break; 328 | } 329 | Ok(length) => length, 330 | Err(ref error) if error.kind() == std::io::ErrorKind::WouldBlock => break, 331 | Err(error) => { 332 | return Err(error).prepend_err(&format!( 333 | "Failed to read {} bytes from slave #{}.", 334 | length, 335 | index 336 | )); 337 | } 338 | }; 339 | self.is_dirty = true; 340 | self.temporary_buffer[..length].iter().for_each(|byte| { 341 | slave_read_buffer.push_back(*byte); 342 | }); 343 | assert!( 344 | slave_read_buffer.capacity() < Self::SLAVE_BUFFER_SIZE * 2, 345 | "Multiplexer {}: slave_read_buffer's capacity is too big \ 346 | (actual: {}, expected: _ < {}).", 347 | self.name, 348 | slave_read_buffer.capacity(), 349 | Self::SLAVE_BUFFER_SIZE * 2 350 | ); 351 | } 352 | } 353 | 354 | Ok(()) 355 | } 356 | 357 | pub fn write_slaves( 358 | &mut self, 359 | slaves: &mut Vec, 360 | ) -> std::io::Result<()> { 361 | assert!( 362 | self.slave_write_buffers.len() == slaves.len(), 363 | "Multiplexer {}: Inconsistent number of slaves (expected: {}, actual: {}).", 364 | self.name, 365 | self.slave_write_buffers.len(), 366 | slaves.len() 367 | ); 368 | 369 | // For each slave, write the contents of a buffer into the slave as much as possible. 370 | for (index, slave_write_buffer) in self.slave_write_buffers.iter_mut().enumerate() { 371 | while !slave_write_buffer.is_empty() { 372 | assert!( 373 | !slave_write_buffer.as_slices().0.is_empty(), 374 | "VecDequeue.as_slices().0 is not expected to be empty." 375 | ); 376 | let length = match slaves[index].write(slave_write_buffer.as_slices().0) { 377 | Ok(length) => length, 378 | Err(ref error) if error.kind() == std::io::ErrorKind::WouldBlock => break, 379 | Err(error) => { 380 | return Err(error).prepend_err(&format!( 381 | "Failed to write {} bytes to slave #{}.", 382 | slave_write_buffer.as_slices().0.len(), 383 | index 384 | )); 385 | } 386 | }; 387 | self.is_dirty = true; 388 | slave_write_buffer.drain(..length); 389 | } 390 | } 391 | 392 | Ok(()) 393 | } 394 | 395 | pub fn scatter(&mut self) { 396 | // Eearly return in case that master_read_buffer is empty. 397 | if self.master_read_buffer.is_empty() { 398 | return; 399 | } 400 | 401 | // Build slave write capacity queue. 402 | self.slave_write_capacity_queue.clear(); 403 | for (index, slave_write_buffer) in self.slave_write_buffers.iter_mut().enumerate() { 404 | self.slave_write_capacity_queue.push(( 405 | slave_write_buffer.capacity() - slave_write_buffer.len(), 406 | index, 407 | )); 408 | } 409 | 410 | // Scatter master read buffer to slave write buffers. 411 | while let Some((capacity, index)) = self.slave_write_capacity_queue.pop() { 412 | // Stop scattering if master read buffer is empty or no capacity in slave write buffers. 413 | if self.master_read_buffer.is_empty() || capacity < Self::PACKET_SIZE { 414 | break; 415 | } 416 | 417 | // Try later if master read buffer is not enough to fill a packet. 418 | if self.is_dirty 419 | && self.master_read_buffer.len() < Self::PACKET_SIZE - Self::HEADER_SIZE 420 | { 421 | break; 422 | } 423 | 424 | // Borrow the slave buffer. 425 | let slave_write_buffer = &mut self.slave_write_buffers[index]; 426 | assert!( 427 | capacity == slave_write_buffer.capacity() - slave_write_buffer.len(), 428 | "Multiplexer {}: Unexpected capacity (expected: {}, actual: {}).", 429 | self.name, 430 | capacity, 431 | slave_write_buffer.capacity() - slave_write_buffer.len() 432 | ); 433 | 434 | // Determine the length of a chunk to move. 435 | let length = std::cmp::min( 436 | std::cmp::min(Self::PACKET_SIZE, capacity) - Self::HEADER_SIZE, 437 | self.master_read_buffer.len(), 438 | ); 439 | assert!( 440 | 0 < length && length <= Self::PACKET_SIZE - Self::HEADER_SIZE, 441 | "Multiplexer {}: Unexpected length to scatter (actual: {}, expected: 0 < _ <= {}).", 442 | self.name, 443 | length, 444 | Self::PACKET_SIZE - Self::HEADER_SIZE 445 | ); 446 | 447 | self.is_dirty = true; 448 | if self.scatter_index == 0 { 449 | info!( 450 | "Multiplexer {}: Master's first packet ({} bytes) is arrived.", 451 | self.name, 452 | length 453 | ); 454 | } 455 | 456 | // Write the header. 457 | slave_write_buffer.push_back_usize(self.scatter_index, Self::INDEX_HEADER_SIZE); 458 | self.scatter_index += 1; 459 | slave_write_buffer.push_back_usize(length, Self::LENGTH_HEADER_SIZE); 460 | 461 | // Write the contents. 462 | self.master_read_buffer.drain(..length).for_each(|byte| { 463 | slave_write_buffer.push_back(byte); 464 | }); 465 | 466 | // Return the slave into the queue again. 467 | self.slave_write_capacity_queue.push(( 468 | slave_write_buffer.capacity() - slave_write_buffer.len(), 469 | index, 470 | )); 471 | assert!( 472 | slave_write_buffer.capacity() < Self::SLAVE_BUFFER_SIZE * 2, 473 | "Multiplexer {}: slave_write_buffer's capacity is too big \ 474 | (actual: {}, expected: _ < {}).", 475 | self.name, 476 | slave_write_buffer.capacity(), 477 | Self::SLAVE_BUFFER_SIZE * 2 478 | ); 479 | } 480 | } 481 | 482 | pub fn gather(&mut self) { 483 | // For each slave, move chunks which are targetted in the master write buffer. 484 | for slave_read_buffer in self.slave_read_buffers.iter_mut() { 485 | while slave_read_buffer.len() >= Self::HEADER_SIZE { 486 | let length = 487 | slave_read_buffer.peek_usize(Self::INDEX_HEADER_SIZE, Self::LENGTH_HEADER_SIZE); 488 | // Skip if the first chunk is incomplete. 489 | if slave_read_buffer.len() < length + Self::HEADER_SIZE { 490 | break; 491 | } 492 | 493 | let index = slave_read_buffer 494 | .peek_usize(0, Self::INDEX_HEADER_SIZE) 495 | .wrapping_sub(self.gather_index) 496 | & ((1 << (Self::INDEX_HEADER_SIZE * 8)) - 1); 497 | // Skip if the chunk is not targetted in the master write buffer. 498 | if self.master_write_buffers.len() <= index { 499 | break; 500 | } 501 | 502 | // Point to a buffer to write to. 503 | let master_write_buffer = &mut self.master_write_buffers[index]; 504 | if master_write_buffer.0 { 505 | error!( 506 | "Multiplexer {}: Write buffer #{} duplicates.", 507 | self.name, 508 | self.gather_index + index 509 | ); 510 | } 511 | 512 | self.is_dirty = true; 513 | 514 | // Populate the write buffer. 515 | master_write_buffer.0 = true; 516 | slave_read_buffer 517 | .drain(..Self::HEADER_SIZE + length) 518 | .skip(Self::HEADER_SIZE) 519 | .for_each(|byte| { 520 | master_write_buffer.1.push(byte); 521 | }); 522 | } 523 | } 524 | } 525 | 526 | pub(crate) fn multiplex_once< 527 | Stream: std::marker::Sized + std::io::Read + std::io::Write + mio::Evented, 528 | >( 529 | &mut self, 530 | master: &mut Stream, 531 | slaves: &mut Vec, 532 | ) -> std::io::Result { 533 | self.is_dirty = false; 534 | self.write_master(master)?; 535 | self.write_slaves(slaves)?; 536 | self.read_master(master)?; 537 | self.read_slaves(slaves)?; 538 | self.gather(); 539 | self.scatter(); 540 | Ok(self.is_dirty) 541 | } 542 | 543 | pub(crate) fn multiplex_while_dirty< 544 | Stream: std::marker::Sized + std::io::Read + std::io::Write + mio::Evented, 545 | >( 546 | &mut self, 547 | master: &mut Stream, 548 | slaves: &mut Vec, 549 | ) -> std::io::Result<()> { 550 | while self.multiplex_once(master, slaves)? {} 551 | self.stats.flush += 1; 552 | Ok(()) 553 | } 554 | 555 | fn display_stats(&mut self) { 556 | self.stats.time = std::time::Instant::now(); 557 | let interval = self.stats.time - self.last_stats.time; 558 | if interval < std::time::Duration::from_millis(100) { 559 | return; 560 | } 561 | let interval = interval.as_secs() as f64 + interval.subsec_nanos() as f64 * 1.0e-9; 562 | let bytes_to_mbps = 8.0 / interval / 1024.0 / 1024.0; 563 | info!( 564 | "Multiplexer {}: Stats ({:.0} ms): \ 565 | flush = {}; \ 566 | master_read = {} bytes ({:.2} Mbps), {} packets; \ 567 | master_write = {} bytes ({:.2} Mbps), {} packets.", 568 | self.name, 569 | interval * 1000.0, 570 | self.stats.flush, 571 | self.stats.master_read, 572 | (self.stats.master_read - self.last_stats.master_read) as f64 * bytes_to_mbps, 573 | self.stats.master_read_packets, 574 | self.stats.master_write, 575 | (self.stats.master_write - self.last_stats.master_write) as f64 * bytes_to_mbps, 576 | self.stats.master_write_packets, 577 | ); 578 | self.last_stats = self.stats.clone(); 579 | } 580 | 581 | pub fn multiplex( 582 | &mut self, 583 | master: &mut Stream, 584 | slaves: &mut Vec, 585 | ) -> std::io::Result<()> { 586 | info!( 587 | "Multiplexer {}: Started a multiplexer with {} slaves.", 588 | self.name, 589 | slaves.len() 590 | ); 591 | let poll = mio::Poll::new().prepend_err("Failed to Poll::new().")?; 592 | poll.register( 593 | master, 594 | mio::Token(0), 595 | mio::Ready::readable() | mio::Ready::writable(), 596 | mio::PollOpt::edge(), 597 | ).prepend_err("Failed to register master to poll.")?; 598 | for slave in slaves.iter() { 599 | poll.register( 600 | slave, 601 | mio::Token(0), 602 | mio::Ready::readable() | mio::Ready::writable(), 603 | mio::PollOpt::edge(), 604 | ).prepend_err("Failed to register a slave to poll.")?; 605 | } 606 | while self.master_connectivity && self.num_of_alive_slaves > 0 { 607 | if !self.is_dirty { 608 | // Wait for a new event. 609 | let polling_start_time = std::time::Instant::now(); 610 | poll.poll( 611 | &mut mio::Events::with_capacity(1), 612 | Some(std::time::Duration::from_secs(5)), 613 | ).prepend_err("Failed to poll.")?; 614 | 615 | // Wiat for a short while to prevent packets from being fragmented. 616 | let polling_time = polling_start_time.elapsed(); 617 | if polling_time < std::time::Duration::from_millis(5) { 618 | std::thread::sleep(std::time::Duration::from_millis(5) - polling_time); 619 | } 620 | 621 | if self.stats.time.elapsed() > std::time::Duration::from_secs(10) { 622 | self.display_stats(); 623 | } 624 | } 625 | 626 | self.multiplex_while_dirty(master, slaves)?; 627 | } 628 | info!("Multiplexer {}: Multiplexer is closed.", self.name); 629 | self.display_stats(); 630 | 631 | Ok(()) 632 | } 633 | } 634 | 635 | struct MaxTcp { 636 | config: MaxTcpConfig, 637 | } 638 | 639 | #[derive(Clone)] 640 | struct MaxTcpConfig { 641 | remote_address: std::net::SocketAddr, 642 | bind_address: std::net::SocketAddr, 643 | server_mode: bool, 644 | socket_buffer_size: usize, 645 | num_of_slaves: usize, 646 | } 647 | 648 | impl MaxTcpConfig { 649 | fn configure_stream(&self, stream: &mio::net::TcpStream) -> std::io::Result<()> { 650 | stream 651 | .set_keepalive(Some(std::time::Duration::from_secs(10))) 652 | .prepend_err("Failed to set_keepalive.")?; 653 | stream 654 | .set_nodelay(true) 655 | .prepend_err("Failed to set_nodelay.")?; 656 | stream 657 | .set_send_buffer_size(self.socket_buffer_size) 658 | .prepend_err("Failed to set_send_buffer_size.")?; 659 | stream 660 | .set_recv_buffer_size(self.socket_buffer_size) 661 | .prepend_err("Failed to set_recv_buffer_size.")?; 662 | Ok(()) 663 | } 664 | } 665 | 666 | struct ServerSocket { 667 | stream: mio::net::TcpStream, 668 | remote_address: std::net::SocketAddr, 669 | buffer: VecDeque, 670 | } 671 | 672 | impl MaxTcp { 673 | const HEADER_SIZE: usize = 16; 674 | 675 | fn with_args(args: &Vec) -> std::result::Result { 676 | let mut opts = getopts::Options::new(); 677 | opts.optflag("h", "help", "Print this help."); 678 | opts.optflag("s", "server", "Enable server mode."); 679 | opts.optopt( 680 | "l", 681 | "listen", 682 | "Address to listen. (default: 0.0.0.0:2200)", 683 | "ADDRESS:PORT", 684 | ); 685 | opts.optopt( 686 | "c", 687 | "connections", 688 | "# of TCP connections per connection.", 689 | "NUMBER", 690 | ); 691 | opts.optopt( 692 | "b", 693 | "socket_buffer_size", 694 | "Buffer size for socket buffers.", 695 | "SIZE", 696 | ); 697 | 698 | let matches = match opts.parse(&args[1..]) { 699 | Ok(m) => m, 700 | Err(f) => return Err(format!("Invalid argument: {}", f.to_string())), 701 | }; 702 | if matches.opt_present("h") { 703 | print!( 704 | "{}", 705 | opts.usage(&format!( 706 | "Usage: {} [options] REMOTE_ADDRESS:PORT (BIND_ADDRESS:PORT)", 707 | args[0] 708 | )) 709 | ); 710 | std::process::exit(0); 711 | } 712 | 713 | let (remote_address, bind_address): (&str, &str) = match matches.free.len() { 714 | 1 => (&matches.free[0], "0.0.0.0:2200"), 715 | 2 => (&matches.free[0], &matches.free[1]), 716 | 0 => return Err("Missing arguments.".into()), 717 | num_of_arguments => { 718 | return Err(format!( 719 | "One or two arguments must be given, but {} arguments are found.", 720 | num_of_arguments 721 | )) 722 | } 723 | }; 724 | let remote_address: std::net::SocketAddr = match remote_address.parse() { 725 | Ok(address) => address, 726 | Err(error) => { 727 | return Err(format!( 728 | "Failed to parse a remote address `{}`: {}", 729 | remote_address, 730 | error 731 | )) 732 | } 733 | }; 734 | let bind_address: std::net::SocketAddr = match bind_address.parse() { 735 | Ok(address) => address, 736 | Err(error) => { 737 | return Err(format!( 738 | "Failed to parse a remote address `{}`: {}", 739 | remote_address, 740 | error 741 | )) 742 | } 743 | }; 744 | 745 | Ok(Self { 746 | config: MaxTcpConfig { 747 | remote_address: remote_address, 748 | bind_address: bind_address, 749 | server_mode: matches.opt_present("s"), 750 | socket_buffer_size: matches 751 | .opt_str("socket_buffer_size") 752 | .unwrap_or("4096".into()) 753 | .parse() 754 | .map_err(|error| { 755 | format!("Failed to parse --socket_buffer_size: {}", error) 756 | })?, 757 | num_of_slaves: matches 758 | .opt_str("connections") 759 | .unwrap_or("16".into()) 760 | .parse() 761 | .map_err(|error| format!("Failed to parse --connections: {}", error))?, 762 | }, 763 | }) 764 | } 765 | 766 | fn server_thread( 767 | config: &MaxTcpConfig, 768 | multiplexer_name: &str, 769 | slaves: Vec, 770 | slave_address: std::net::SocketAddr, 771 | ) -> std::io::Result<()> { 772 | let master_address = config.remote_address; 773 | info!("Starting connection from: {}", slave_address); 774 | let mut master = mio::net::TcpStream::connect(&master_address) 775 | .prepend_err(&format!("Failed to connect to {}.", &master_address))?; 776 | { 777 | let poll = mio::Poll::new().prepend_err("Failed to Poll::new().")?; 778 | let master = master.try_clone().prepend_err("Failed to clone master.")?; 779 | poll.register( 780 | &master, 781 | mio::Token(0), 782 | mio::Ready::writable(), 783 | mio::PollOpt::edge(), 784 | ).prepend_err("Failed to register a master socket to connect.")?; 785 | poll.poll(&mut mio::Events::with_capacity(1), None) 786 | .prepend_err("Failed to poll a master socket.")?; 787 | } 788 | let mut slaves = slaves; 789 | let mut multiplexer = Multiplexer::new(multiplexer_name, slaves.len()); 790 | multiplexer.multiplex(&mut master, &mut slaves) 791 | } 792 | 793 | fn server(config: &MaxTcpConfig, listener: mio::net::TcpListener) -> std::io::Result<()> { 794 | info!("Starting server targeting to {}.", config.remote_address); 795 | let mut next_token = mio::Token(1); 796 | let mut pending_sockets = std::collections::HashMap::::new(); 797 | let mut ready_sockets = std::collections::HashMap::<[u8; 16], Vec>::new(); 798 | loop { 799 | let poll = mio::Poll::new().prepend_err("Failed to Poll::new().")?; 800 | let mut polling_sockets = Vec::new(); 801 | let listener = listener 802 | .try_clone() 803 | .prepend_err("Failed to clone a master socket.")?; 804 | poll.register( 805 | &listener, 806 | mio::Token(0), 807 | mio::Ready::readable(), 808 | mio::PollOpt::edge(), 809 | ).prepend_err("Failed to register to poll a master socket.")?; 810 | for (token, socket) in &pending_sockets { 811 | let socket = socket 812 | .stream 813 | .try_clone() 814 | .prepend_err("Failed to clone a slave socket.")?; 815 | poll.register( 816 | &socket, 817 | *token, 818 | mio::Ready::readable(), 819 | mio::PollOpt::edge(), 820 | ).prepend_err("Failed to register to poll a slave candidate.")?; 821 | polling_sockets.push(socket); 822 | } 823 | let mut events = mio::Events::with_capacity(1024); 824 | poll.poll(&mut events, None) 825 | .prepend_err("Failed to poll a lister and pending sockets.")?; 826 | for event in &events { 827 | match event.token() { 828 | mio::Token(0) => { 829 | let (stream, remote_address) = 830 | listener.accept().prepend_err("Failed to accept.")?; 831 | config 832 | .configure_stream(&stream) 833 | .prepend_err("Failed to configure a slave socket.")?; 834 | pending_sockets.insert( 835 | next_token, 836 | ServerSocket { 837 | stream: stream, 838 | remote_address: remote_address, 839 | buffer: VecDeque::new(), 840 | }, 841 | ); 842 | next_token = mio::Token(next_token.0 + 1); 843 | } 844 | token => { 845 | { 846 | let socket = match pending_sockets.get_mut(&token) { 847 | Some(socket) => socket, 848 | None => unreachable!(), 849 | }; 850 | let mut buffer = [0; Self::HEADER_SIZE]; 851 | let length = socket 852 | .stream 853 | .read(&mut buffer[0..(Self::HEADER_SIZE - socket.buffer.len())]) 854 | .prepend_err("Failed to read a header from a pending socket.")?; 855 | buffer[..length] 856 | .iter() 857 | .for_each(|byte| socket.buffer.push_back(*byte)); 858 | assert!( 859 | socket.buffer.len() <= Self::HEADER_SIZE, 860 | "Read too many data (actual: {}, expected: _ <= {}).", 861 | socket.buffer.len(), 862 | Self::HEADER_SIZE 863 | ); 864 | if socket.buffer.len() < Self::HEADER_SIZE { 865 | continue; 866 | } 867 | } 868 | let key = if let Some(socket) = pending_sockets.remove(&token) { 869 | let mut key = [0; Self::HEADER_SIZE]; 870 | for position in 0..Self::HEADER_SIZE { 871 | key[position] = socket.buffer[position]; 872 | } 873 | let value = ready_sockets.entry(key).or_insert(Vec::new()); 874 | value.push(socket); 875 | if value.len() != key[0] as usize { 876 | continue; 877 | } 878 | key 879 | } else { 880 | unreachable!(); 881 | }; 882 | if let Some(sockets) = ready_sockets.remove(&key) { 883 | let config = config.clone(); 884 | let mut slaves = Vec::new(); 885 | let socket_address = sockets[0].remote_address; 886 | for socket in sockets { 887 | slaves.push(socket.stream); 888 | } 889 | let multiplexer_name = 890 | format!("#{:<02x}{:<02x}{:<02x}[S]", key[1], key[2], key[3]); 891 | let _ = std::thread::spawn(move || { 892 | match Self::server_thread( 893 | &config, 894 | &multiplexer_name, 895 | slaves, 896 | socket_address, 897 | ) { 898 | Ok(_) => info!("Connection closed successfully."), 899 | Err(error) => info!("Connection closed: {}", error), 900 | } 901 | }); 902 | } 903 | } 904 | }; 905 | } 906 | } 907 | } 908 | 909 | fn client_thread( 910 | config: &MaxTcpConfig, 911 | stream: mio::net::TcpStream, 912 | socket_addr: std::net::SocketAddr, 913 | ) -> std::io::Result<()> { 914 | info!("Starting a client connection from: {}", socket_addr); 915 | let mut stream = stream; 916 | config 917 | .configure_stream(&stream) 918 | .prepend_err("Failed to configure a master socket.")?; 919 | let mut sockets = Vec::new(); 920 | let num_of_slaves = config.num_of_slaves; 921 | let mut header = [0; Self::HEADER_SIZE]; 922 | rand::thread_rng().fill_bytes(&mut header); 923 | header[0] = num_of_slaves as u8; 924 | let multiplexer_name = 925 | format!("#{:<02x}{:<02x}{:<02x}[C]", header[1], header[2], header[3]); 926 | 927 | for slave_id in 0..num_of_slaves { 928 | let socket = 929 | mio::net::TcpStream::connect(&config.remote_address).prepend_err(&format!( 930 | "Failed to connect slave #{} of {} slaves.", 931 | slave_id, 932 | num_of_slaves 933 | ))?; 934 | config 935 | .configure_stream(&socket) 936 | .prepend_err("Failed to configure a slave socket.")?; 937 | sockets.push(socket); 938 | } 939 | for (slave_id, socket) in sockets.iter_mut().enumerate() { 940 | let poll = mio::Poll::new().prepend_err("Failed to Poll::new().")?; 941 | let polling_socket = socket 942 | .try_clone() 943 | .prepend_err("Failed to clone a slave socket.")?; 944 | poll.register( 945 | &polling_socket, 946 | mio::Token(0), 947 | mio::Ready::writable(), 948 | mio::PollOpt::edge(), 949 | ).prepend_err("Failed to register a slave socket to write.")?; 950 | poll.poll(&mut mio::Events::with_capacity(1), None) 951 | .prepend_err("Failed to poll to write.")?; 952 | socket.write(&header).prepend_err(&format!( 953 | "Failed to write header to slave #{} of {} slaves.", 954 | slave_id, 955 | num_of_slaves 956 | ))?; 957 | } 958 | let mut multiplexer = Multiplexer::new(&multiplexer_name, sockets.len()); 959 | multiplexer 960 | .multiplex(&mut stream, &mut sockets) 961 | .prepend_err("Multiplexer for client failed.") 962 | } 963 | 964 | fn client(config: &MaxTcpConfig, listener: mio::net::TcpListener) -> std::io::Result<()> { 965 | let poll = mio::Poll::new().prepend_err("Failed to Poll::new().")?; 966 | poll.register( 967 | &listener, 968 | mio::Token(0), 969 | mio::Ready::readable(), 970 | mio::PollOpt::edge(), 971 | ).prepend_err("Failed to register listener to poll.")?; 972 | loop { 973 | poll.poll(&mut mio::Events::with_capacity(1), None) 974 | .prepend_err("Failed to poll to accept.")?; 975 | let connection = listener 976 | .accept() 977 | .prepend_err("Failed to accept a new connection.")?; 978 | info!("Accepting a new connection from {}...", connection.1); 979 | let config = config.clone(); 980 | std::thread::spawn(move || { 981 | match Self::client_thread(&config, connection.0, connection.1) { 982 | Ok(_) => info!("Connection closed successfully."), 983 | Err(error) => info!("Connection closed by error:\n- {}", error), 984 | } 985 | }); 986 | } 987 | } 988 | 989 | fn run(&self) -> std::io::Result<()> { 990 | info!("Listening address `{}`", &self.config.bind_address); 991 | let listener = match mio::net::TcpListener::bind(&self.config.bind_address) { 992 | Ok(listener) => listener, 993 | Err(error) => { 994 | return Err(std::io::Error::new( 995 | error.kind(), 996 | format!( 997 | "Failed to bind address `{}`: {}", 998 | &self.config.bind_address, 999 | error 1000 | ), 1001 | )); 1002 | } 1003 | }; 1004 | if self.config.server_mode { 1005 | MaxTcp::server(&self.config, listener).prepend_err("MaxTcp::server failed.") 1006 | } else { 1007 | MaxTcp::client(&self.config, listener).prepend_err("MaxTcp::client failed.") 1008 | } 1009 | } 1010 | } 1011 | 1012 | fn main() { 1013 | env_logger::init(); 1014 | 1015 | let error = match MaxTcp::with_args(&std::env::args().collect()) { 1016 | Ok(max_tcp) => match max_tcp.run() { 1017 | Err(error) => format!("{}", error), 1018 | Ok(_) => std::process::exit(0), 1019 | }, 1020 | Err(error) => format!("{}", error), 1021 | }; 1022 | eprintln!("{}", error); 1023 | std::process::exit(1); 1024 | } 1025 | --------------------------------------------------------------------------------