1 // Licensed to the Apache Software Foundation (ASF) under one 2 // or more contributor license agreements. See the NOTICE file 3 // distributed with this work for additional information 4 // regarding copyright ownership. The ASF licenses this file 5 // to you under the Apache License, Version 2.0 (the 6 // "License"); you may not use this file except in compliance 7 // with the License. You may obtain a copy of the License at 8 // 9 // http://www.apache.org/licenses/LICENSE-2.0 10 // 11 // Unless required by applicable law or agreed to in writing, 12 // software distributed under the License is distributed on an 13 // "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY 14 // KIND, either express or implied. See the License for the 15 // specific language governing permissions and limitations 16 // under the License. 17 18 use byteorder::{BigEndian, ReadBytesExt, WriteBytesExt}; 19 use std::cmp; 20 use std::io; 21 use std::io::{Read, Write}; 22 23 use super::{TReadTransport, TReadTransportFactory, TWriteTransport, TWriteTransportFactory}; 24 25 /// Default capacity of the read buffer in bytes. 26 const READ_CAPACITY: usize = 4096; 27 28 /// Default capacity of the write buffer in bytes. 29 const WRITE_CAPACITY: usize = 4096; 30 31 /// Transport that reads framed messages. 32 /// 33 /// A `TFramedReadTransport` maintains a fixed-size internal read buffer. 34 /// On a call to `TFramedReadTransport::read(...)` one full message - both 35 /// fixed-length header and bytes - is read from the wrapped channel and 36 /// buffered. Subsequent read calls are serviced from the internal buffer 37 /// until it is exhausted, at which point the next full message is read 38 /// from the wrapped channel. 39 /// 40 /// # Examples 41 /// 42 /// Create and use a `TFramedReadTransport`. 43 /// 44 /// ```no_run 45 /// use std::io::Read; 46 /// use thrift::transport::{TFramedReadTransport, TTcpChannel}; 47 /// 48 /// let mut c = TTcpChannel::new(); 49 /// c.open("localhost:9090").unwrap(); 50 /// 51 /// let mut t = TFramedReadTransport::new(c); 52 /// 53 /// t.read(&mut vec![0u8; 1]).unwrap(); 54 /// ``` 55 #[derive(Debug)] 56 pub struct TFramedReadTransport<C> 57 where 58 C: Read, 59 { 60 buf: Vec<u8>, 61 pos: usize, 62 cap: usize, 63 chan: C, 64 } 65 66 impl<C> TFramedReadTransport<C> 67 where 68 C: Read, 69 { 70 /// Create a `TFramedReadTransport` with a default-sized 71 /// internal read buffer that wraps the given `TIoChannel`. new(channel: C) -> TFramedReadTransport<C>72 pub fn new(channel: C) -> TFramedReadTransport<C> { 73 TFramedReadTransport::with_capacity(READ_CAPACITY, channel) 74 } 75 76 /// Create a `TFramedTransport` with an internal read buffer 77 /// of size `read_capacity` that wraps the given `TIoChannel`. with_capacity(read_capacity: usize, channel: C) -> TFramedReadTransport<C>78 pub fn with_capacity(read_capacity: usize, channel: C) -> TFramedReadTransport<C> { 79 TFramedReadTransport { 80 buf: vec![0; read_capacity], // FIXME: do I actually have to do this? 81 pos: 0, 82 cap: 0, 83 chan: channel, 84 } 85 } 86 } 87 88 impl<C> Read for TFramedReadTransport<C> 89 where 90 C: Read, 91 { read(&mut self, b: &mut [u8]) -> io::Result<usize>92 fn read(&mut self, b: &mut [u8]) -> io::Result<usize> { 93 if self.cap - self.pos == 0 { 94 let message_size = self.chan.read_i32::<BigEndian>()? as usize; 95 96 let buf_capacity = cmp::max(message_size, READ_CAPACITY); 97 self.buf.resize(buf_capacity, 0); 98 99 self.chan.read_exact(&mut self.buf[..message_size])?; 100 self.cap = message_size as usize; 101 self.pos = 0; 102 } 103 104 let nread = cmp::min(b.len(), self.cap - self.pos); 105 b[..nread].clone_from_slice(&self.buf[self.pos..self.pos + nread]); 106 self.pos += nread; 107 108 Ok(nread) 109 } 110 } 111 112 /// Factory for creating instances of `TFramedReadTransport`. 113 #[derive(Default)] 114 pub struct TFramedReadTransportFactory; 115 116 impl TFramedReadTransportFactory { new() -> TFramedReadTransportFactory117 pub fn new() -> TFramedReadTransportFactory { 118 TFramedReadTransportFactory {} 119 } 120 } 121 122 impl TReadTransportFactory for TFramedReadTransportFactory { 123 /// Create a `TFramedReadTransport`. create(&self, channel: Box<dyn Read + Send>) -> Box<dyn TReadTransport + Send>124 fn create(&self, channel: Box<dyn Read + Send>) -> Box<dyn TReadTransport + Send> { 125 Box::new(TFramedReadTransport::new(channel)) 126 } 127 } 128 129 /// Transport that writes framed messages. 130 /// 131 /// A `TFramedWriteTransport` maintains a fixed-size internal write buffer. All 132 /// writes are made to this buffer and are sent to the wrapped channel only 133 /// when `TFramedWriteTransport::flush()` is called. On a flush a fixed-length 134 /// header with a count of the buffered bytes is written, followed by the bytes 135 /// themselves. 136 /// 137 /// # Examples 138 /// 139 /// Create and use a `TFramedWriteTransport`. 140 /// 141 /// ```no_run 142 /// use std::io::Write; 143 /// use thrift::transport::{TFramedWriteTransport, TTcpChannel}; 144 /// 145 /// let mut c = TTcpChannel::new(); 146 /// c.open("localhost:9090").unwrap(); 147 /// 148 /// let mut t = TFramedWriteTransport::new(c); 149 /// 150 /// t.write(&[0x00]).unwrap(); 151 /// t.flush().unwrap(); 152 /// ``` 153 #[derive(Debug)] 154 pub struct TFramedWriteTransport<C> 155 where 156 C: Write, 157 { 158 buf: Vec<u8>, 159 channel: C, 160 } 161 162 impl<C> TFramedWriteTransport<C> 163 where 164 C: Write, 165 { 166 /// Create a `TFramedWriteTransport` with default-sized internal 167 /// write buffer that wraps the given `TIoChannel`. new(channel: C) -> TFramedWriteTransport<C>168 pub fn new(channel: C) -> TFramedWriteTransport<C> { 169 TFramedWriteTransport::with_capacity(WRITE_CAPACITY, channel) 170 } 171 172 /// Create a `TFramedWriteTransport` with an internal write buffer 173 /// of size `write_capacity` that wraps the given `TIoChannel`. with_capacity(write_capacity: usize, channel: C) -> TFramedWriteTransport<C>174 pub fn with_capacity(write_capacity: usize, channel: C) -> TFramedWriteTransport<C> { 175 TFramedWriteTransport { 176 buf: Vec::with_capacity(write_capacity), 177 channel, 178 } 179 } 180 } 181 182 impl<C> Write for TFramedWriteTransport<C> 183 where 184 C: Write, 185 { write(&mut self, b: &[u8]) -> io::Result<usize>186 fn write(&mut self, b: &[u8]) -> io::Result<usize> { 187 let current_capacity = self.buf.capacity(); 188 let available_space = current_capacity - self.buf.len(); 189 if b.len() > available_space { 190 let additional_space = cmp::max(b.len() - available_space, current_capacity); 191 self.buf.reserve(additional_space); 192 } 193 194 self.buf.extend_from_slice(b); 195 Ok(b.len()) 196 } 197 flush(&mut self) -> io::Result<()>198 fn flush(&mut self) -> io::Result<()> { 199 let message_size = self.buf.len(); 200 201 if let 0 = message_size { 202 return Ok(()); 203 } else { 204 self.channel.write_i32::<BigEndian>(message_size as i32)?; 205 } 206 207 // will spin if the underlying channel can't be written to 208 let mut byte_index = 0; 209 while byte_index < message_size { 210 let nwrite = self.channel.write(&self.buf[byte_index..message_size])?; 211 byte_index = cmp::min(byte_index + nwrite, message_size); 212 } 213 214 let buf_capacity = cmp::min(self.buf.capacity(), WRITE_CAPACITY); 215 self.buf.resize(buf_capacity, 0); 216 self.buf.clear(); 217 218 self.channel.flush() 219 } 220 } 221 222 /// Factory for creating instances of `TFramedWriteTransport`. 223 #[derive(Default)] 224 pub struct TFramedWriteTransportFactory; 225 226 impl TFramedWriteTransportFactory { new() -> TFramedWriteTransportFactory227 pub fn new() -> TFramedWriteTransportFactory { 228 TFramedWriteTransportFactory {} 229 } 230 } 231 232 impl TWriteTransportFactory for TFramedWriteTransportFactory { 233 /// Create a `TFramedWriteTransport`. create(&self, channel: Box<dyn Write + Send>) -> Box<dyn TWriteTransport + Send>234 fn create(&self, channel: Box<dyn Write + Send>) -> Box<dyn TWriteTransport + Send> { 235 Box::new(TFramedWriteTransport::new(channel)) 236 } 237 } 238 239 #[cfg(test)] 240 mod tests { 241 use super::*; 242 use crate::transport::mem::TBufferChannel; 243 244 // FIXME: test a forced reserve 245 246 #[test] must_read_message_smaller_than_initial_buffer_size()247 fn must_read_message_smaller_than_initial_buffer_size() { 248 let c = TBufferChannel::with_capacity(10, 10); 249 let mut t = TFramedReadTransport::with_capacity(8, c); 250 251 t.chan.set_readable_bytes(&[ 252 0x00, 0x00, 0x00, 0x04, /* message size */ 253 0x00, 0x01, 0x02, 0x03, /* message body */ 254 ]); 255 256 let mut buf = vec![0; 8]; 257 258 // we've read exactly 4 bytes 259 assert_eq!(t.read(&mut buf).unwrap(), 4); 260 assert_eq!(&buf[..4], &[0x00, 0x01, 0x02, 0x03]); 261 } 262 263 #[test] must_read_message_greater_than_initial_buffer_size()264 fn must_read_message_greater_than_initial_buffer_size() { 265 let c = TBufferChannel::with_capacity(10, 10); 266 let mut t = TFramedReadTransport::with_capacity(2, c); 267 268 t.chan.set_readable_bytes(&[ 269 0x00, 0x00, 0x00, 0x04, /* message size */ 270 0x00, 0x01, 0x02, 0x03, /* message body */ 271 ]); 272 273 let mut buf = vec![0; 8]; 274 275 // we've read exactly 4 bytes 276 assert_eq!(t.read(&mut buf).unwrap(), 4); 277 assert_eq!(&buf[..4], &[0x00, 0x01, 0x02, 0x03]); 278 } 279 280 #[test] must_read_multiple_messages_in_sequence_correctly()281 fn must_read_multiple_messages_in_sequence_correctly() { 282 let c = TBufferChannel::with_capacity(10, 10); 283 let mut t = TFramedReadTransport::with_capacity(2, c); 284 285 // 286 // 1st message 287 // 288 289 t.chan.set_readable_bytes(&[ 290 0x00, 0x00, 0x00, 0x04, /* message size */ 291 0x00, 0x01, 0x02, 0x03, /* message body */ 292 ]); 293 294 let mut buf = vec![0; 8]; 295 296 // we've read exactly 4 bytes 297 assert_eq!(t.read(&mut buf).unwrap(), 4); 298 assert_eq!(&buf, &[0x00, 0x01, 0x02, 0x03, 0x00, 0x00, 0x00, 0x00]); 299 300 // 301 // 2nd message 302 // 303 304 t.chan.set_readable_bytes(&[ 305 0x00, 0x00, 0x00, 0x01, /* message size */ 306 0x04, /* message body */ 307 ]); 308 309 let mut buf = vec![0; 8]; 310 311 // we've read exactly 1 byte 312 assert_eq!(t.read(&mut buf).unwrap(), 1); 313 assert_eq!(&buf, &[0x04, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00]); 314 } 315 316 #[test] must_write_message_smaller_than_buffer_size()317 fn must_write_message_smaller_than_buffer_size() { 318 let mem = TBufferChannel::with_capacity(0, 0); 319 let mut t = TFramedWriteTransport::with_capacity(20, mem); 320 321 let b = vec![0; 10]; 322 323 // should have written 10 bytes 324 assert_eq!(t.write(&b).unwrap(), 10); 325 } 326 327 #[test] must_return_zero_if_caller_calls_write_with_empty_buffer()328 fn must_return_zero_if_caller_calls_write_with_empty_buffer() { 329 let mem = TBufferChannel::with_capacity(0, 10); 330 let mut t = TFramedWriteTransport::with_capacity(10, mem); 331 332 let expected: [u8; 0] = []; 333 334 assert_eq!(t.write(&[]).unwrap(), 0); 335 assert_eq_transport_written_bytes!(t, expected); 336 } 337 338 #[test] must_write_to_inner_transport_on_flush()339 fn must_write_to_inner_transport_on_flush() { 340 let mem = TBufferChannel::with_capacity(10, 10); 341 let mut t = TFramedWriteTransport::new(mem); 342 343 let b: [u8; 5] = [0x00, 0x01, 0x02, 0x03, 0x04]; 344 assert_eq!(t.write(&b).unwrap(), 5); 345 assert_eq_transport_num_written_bytes!(t, 0); 346 347 assert!(t.flush().is_ok()); 348 349 let expected_bytes = [ 350 0x00, 0x00, 0x00, 0x05, /* message size */ 351 0x00, 0x01, 0x02, 0x03, 0x04, /* message body */ 352 ]; 353 354 assert_eq_transport_written_bytes!(t, expected_bytes); 355 } 356 357 #[test] must_write_message_greater_than_buffer_size_00()358 fn must_write_message_greater_than_buffer_size_00() { 359 let mem = TBufferChannel::with_capacity(0, 10); 360 361 // IMPORTANT: DO **NOT** CHANGE THE WRITE_CAPACITY OR THE NUMBER OF BYTES TO BE WRITTEN! 362 // these lengths were chosen to be just long enough 363 // that doubling the capacity is a **worse** choice than 364 // simply resizing the buffer to b.len() 365 366 let mut t = TFramedWriteTransport::with_capacity(1, mem); 367 let b = [0x00, 0x01, 0x02]; 368 369 // should have written 3 bytes 370 assert_eq!(t.write(&b).unwrap(), 3); 371 assert_eq_transport_num_written_bytes!(t, 0); 372 373 assert!(t.flush().is_ok()); 374 375 let expected_bytes = [ 376 0x00, 0x00, 0x00, 0x03, /* message size */ 377 0x00, 0x01, 0x02, /* message body */ 378 ]; 379 380 assert_eq_transport_written_bytes!(t, expected_bytes); 381 } 382 383 #[test] must_write_message_greater_than_buffer_size_01()384 fn must_write_message_greater_than_buffer_size_01() { 385 let mem = TBufferChannel::with_capacity(0, 10); 386 387 // IMPORTANT: DO **NOT** CHANGE THE WRITE_CAPACITY OR THE NUMBER OF BYTES TO BE WRITTEN! 388 // these lengths were chosen to be just long enough 389 // that doubling the capacity is a **better** choice than 390 // simply resizing the buffer to b.len() 391 392 let mut t = TFramedWriteTransport::with_capacity(2, mem); 393 let b = [0x00, 0x01, 0x02]; 394 395 // should have written 3 bytes 396 assert_eq!(t.write(&b).unwrap(), 3); 397 assert_eq_transport_num_written_bytes!(t, 0); 398 399 assert!(t.flush().is_ok()); 400 401 let expected_bytes = [ 402 0x00, 0x00, 0x00, 0x03, /* message size */ 403 0x00, 0x01, 0x02, /* message body */ 404 ]; 405 406 assert_eq_transport_written_bytes!(t, expected_bytes); 407 } 408 409 #[test] must_return_error_if_nothing_can_be_written_to_inner_transport_on_flush()410 fn must_return_error_if_nothing_can_be_written_to_inner_transport_on_flush() { 411 let mem = TBufferChannel::with_capacity(0, 0); 412 let mut t = TFramedWriteTransport::with_capacity(1, mem); 413 414 let b = vec![0; 10]; 415 416 // should have written 10 bytes 417 assert_eq!(t.write(&b).unwrap(), 10); 418 419 // let's flush 420 let r = t.flush(); 421 422 // this time we'll error out because the flush can't write to the underlying channel 423 assert!(r.is_err()); 424 } 425 426 #[test] must_write_successfully_after_flush()427 fn must_write_successfully_after_flush() { 428 // IMPORTANT: write capacity *MUST* be greater 429 // than message sizes used in this test + 4-byte frame header 430 let mem = TBufferChannel::with_capacity(0, 10); 431 let mut t = TFramedWriteTransport::with_capacity(5, mem); 432 433 // write and flush 434 let first_message: [u8; 5] = [0x00, 0x01, 0x02, 0x03, 0x04]; 435 assert_eq!(t.write(&first_message).unwrap(), 5); 436 assert!(t.flush().is_ok()); 437 438 let mut expected = Vec::new(); 439 expected.write_all(&[0x00, 0x00, 0x00, 0x05]).unwrap(); // message size 440 expected.extend_from_slice(&first_message); 441 442 // check the flushed bytes 443 assert_eq!(t.channel.write_bytes(), expected); 444 445 // reset our underlying transport 446 t.channel.empty_write_buffer(); 447 448 let second_message: [u8; 3] = [0x05, 0x06, 0x07]; 449 assert_eq!(t.write(&second_message).unwrap(), 3); 450 assert!(t.flush().is_ok()); 451 452 expected.clear(); 453 expected.write_all(&[0x00, 0x00, 0x00, 0x03]).unwrap(); // message size 454 expected.extend_from_slice(&second_message); 455 456 // check the flushed bytes 457 assert_eq!(t.channel.write_bytes(), expected); 458 } 459 } 460