diff --git a/tonic/examples/server.rs b/tonic/examples/server.rs index 6b28ced..af5b228 100644 --- a/tonic/examples/server.rs +++ b/tonic/examples/server.rs @@ -1,12 +1,16 @@ #![feature(async_await, type_alias_impl_trait)] -use std::future::Future; -use tonic::{server::{UnaryService, Grpc}, Status, Request, Response, body}; -use tower_service::Service; -use tower_h2::{RecvBody, Server}; -use tokio::net::TcpListener; use futures_util::future; +use std::future::Future; use std::task::{Context, Poll}; +use tokio::net::TcpListener; +use tonic::{ + body, + server::{Grpc, UnaryService}, + Request, Response, Status, +}; +use tower_h2::{RecvBody, Server}; +use tower_service::Service; #[derive(Clone, PartialEq, prost::Message)] pub struct HelloRequest { @@ -56,7 +60,7 @@ async fn main() -> Result<(), Box> { } } - Ok(()) + Ok(()) } #[derive(Debug)] diff --git a/tonic/src/codec.rs b/tonic/src/codec.rs index 80c9605..1064dd1 100644 --- a/tonic/src/codec.rs +++ b/tonic/src/codec.rs @@ -35,9 +35,24 @@ where loop { match source.try_next().await { Ok(Some(item)) => { + + + buf.reserve(5); + unsafe { + buf.advance_mut(5); + } encoder.encode(item, &mut buf).map_err(drop).unwrap(); - let len = buf.len(); - yield Ok(buf.split_to(len).freeze().into_buf()); + + // now that we know length, we can write the header + let len = buf.len() - 5; + assert!(len <= ::std::u32::MAX as usize); + { + let mut cursor = ::std::io::Cursor::new(&mut buf[..5]); + cursor.put_u8(0); // byte must be 0, reserve doesn't auto-zero + cursor.put_u32_be(len as u32); + } + + yield Ok(buf.split_to(len + 5).freeze().into_buf()); }, Ok(None) => break, Err(status) => yield Err(status), @@ -105,6 +120,7 @@ where let mut buf = (&buf1[..]).into_buf(); if let State::ReadHeader = state { + println!("reading header"); if buf.remaining() < 5 { return Ok(None); } @@ -135,10 +151,14 @@ where } if let State::ReadBody { len, .. } = state { + println!("reading body"); if buf.remaining() < *len { return Ok(None); } + // advance past the header + buf1.advance(5); + match decoder.decode(buf1) { Ok(Some(msg)) => { *state = State::ReadHeader;