Changing UpdateStreamPacket protobuf definition for chromoting

This code also changes the API for encoder and ClientConnection to eliminate
one less copy.

Review URL: http://codereview.chromium.org/2963003

git-svn-id: svn://svn.chromium.org/chrome/trunk/src@52561 0039d316-1c4b-4281-b951-d872f2087c98

18 files changed, 253 insertions, 140 deletions

diff --git a/remoting/base/protocol/chromotocol.proto b/remoting/base/protocol/chromotocol.proto
index e60716f..77bc924 100644
--- a/remoting/base/protocol/chromotocol.proto
+++ b/remoting/base/protocol/chromotocol.proto
@@ -51,9 +51,10 @@ enum PixelFormat {
   PixelFormatAscii = 9;
 }
 
-// A message with info about the update stream.
+// A message that denotes the beginning of an updating rectangle in an update
+// stream packet.
 // NEXT ID: 6
-message UpdateStreamPacketHeader {
+message UpdateStreamBeginRect {
   // X,Y coordinates (in screen pixels) for origin of this update.
   required int32 x = 1;
   required int32 y = 2;
@@ -69,12 +70,28 @@ message UpdateStreamPacketHeader {
   optional PixelFormat pixel_format = 6 [default=PixelFormatRgb24];
 }
 
-// A message to denote a partial update stream.
+// A message that contains partial data for updating an rectangle in an
+// update stream packet.
 // NEXT ID: 3
+message UpdateStreamRectData {
+  // The sequence number of the partial data for updating a rectangle.
+  optional int32 sequence_number = 1 [default=0];
+
+  // The partial data for updating a rectangle.
+  required bytes data = 2;
+}
+
+// A message that denotes the end of an updating rectangle.
+// NEXT ID: 1
+message UpdateStreamEndRect {
+}
+
+// A message to denote a partial update stream.
+// NEXT ID: 4
 message UpdateStreamPacketMessage {
-  // TODO(garykac): Make this required and fix unit tests.
-  optional UpdateStreamPacketHeader header = 2;
-  optional bytes data = 1;
+  optional UpdateStreamBeginRect begin_rect = 1;
+  optional UpdateStreamRectData rect_data = 2;
+  optional UpdateStreamEndRect end_rect = 3;
 }
 
 // Defines the message that is sent from the host to the client.
diff --git a/remoting/base/protocol_decoder_unittest.cc b/remoting/base/protocol_decoder_unittest.cc
index d1bd584..c2473a8 100644
--- a/remoting/base/protocol_decoder_unittest.cc
+++ b/remoting/base/protocol_decoder_unittest.cc
@@ -43,7 +43,10 @@ static void PrepareData(uint8** buffer, int* size) {
     AppendMessage(msg, &encoded_data);
     msg.Clear();
 
-    msg.mutable_update_stream_packet()->set_data(kTestData);
+    msg.mutable_update_stream_packet()->mutable_rect_data()->
+        set_sequence_number(0);
+    msg.mutable_update_stream_packet()->mutable_rect_data()->
+        set_data(kTestData);
     AppendMessage(msg, &encoded_data);
     msg.Clear();
 
@@ -101,7 +104,8 @@ TEST(ProtocolDecoderTest, BasicOperations) {
     } else if (type == 1) {
       // Partial update stream.
       EXPECT_TRUE(message_list[i]->has_update_stream_packet());
-      EXPECT_EQ(kTestData, message_list[i]->update_stream_packet().data());
+      EXPECT_EQ(kTestData,
+                message_list[i]->update_stream_packet().rect_data().data());
     } else if (type == 2) {
       // End update stream.
       EXPECT_TRUE(message_list[i]->has_end_update_stream());
diff --git a/remoting/client/decoder.h b/remoting/client/decoder.h
index ac2fe4d..4d083bd 100644
--- a/remoting/client/decoder.h
+++ b/remoting/client/decoder.h
@@ -57,9 +57,20 @@ class Decoder {
   // Give a HostMessage that contains the update stream packet that contains
   // the encoded data to the decoder.
   // The decoder will own |message| and is responsible for deleting it.
+  //
   // If the decoder has written something into |frame|,
   // |partial_decode_done_| is called with |frame| and updated regions.
   // Return true if the decoder can accept |message| and decode it.
+  //
+  // HostMessage returned by this method will contain a
+  // UpdateStreamPacketMessage.
+  // This message will contain either:
+  // 1. UpdateStreamBeginRect
+  // 2. UpdateStreamRectData
+  // 3. UpdateStreamEndRect
+  //
+  // See remoting/base/protocol/chromotocol.proto for more information about
+  // these messages.
   virtual bool PartialDecode(HostMessage* message) = 0;
 
   // Notify the decoder that we have received the last update stream packet.
@@ -68,6 +79,34 @@ class Decoder {
   // If the update stream is not received fully and this method is called the
   // decoder should also call |decode_done_| as soon as possible.
   virtual void EndDecode() = 0;
+
+ protected:
+  // Every decoder will have two internal states because there are three
+  // kinds of messages send to PartialDecode().
+  //
+  // Here's a state diagram:
+  //
+  //                UpdateStreamBeginRect       UpdateStreamRectData
+  //                    ..............              ............
+  //                   .              .            .            .
+  //                  .                v          .              .
+  // kWaitingForBeginRect         kWaitingForRectData            .
+  //                  ^                .          ^              .
+  //                   .              .            .            .
+  //                    ..............              ............
+  //                    UpdateStreaEndRect
+  enum State {
+    // In this state the decoder is waiting for UpdateStreamBeginRect.
+    // After receiving UpdateStreaBeginRect, the encoder will transit to
+    // to kWaitingForRectData state.
+    kWaitingForBeginRect,
+
+    // In this state the decoder is waiting for UpdateStreamRectData.
+    // The decode remains in this state if UpdateStreamRectData is received.
+    // The decoder will transit to kWaitingForBeginRect if UpdateStreamEndRect
+    // is received.
+    kWaitingForRectData,
+  };
 };
 
 }  // namespace remoting
diff --git a/remoting/client/decoder_verbatim.cc b/remoting/client/decoder_verbatim.cc
index e9b88ef..d27d265 100644
--- a/remoting/client/decoder_verbatim.cc
+++ b/remoting/client/decoder_verbatim.cc
@@ -9,7 +9,13 @@
 namespace remoting {
 
 DecoderVerbatim::DecoderVerbatim()
-    : updated_rects_(NULL),
+    : state_(kWaitingForBeginRect),
+      rect_x_(0),
+      rect_y_(0),
+      rect_width_(0),
+      rect_height_(0),
+      bytes_per_pixel_(0),
+      updated_rects_(NULL),
       reverse_rows_(true) {
 }
 
@@ -20,6 +26,7 @@ bool DecoderVerbatim::BeginDecode(scoped_refptr<media::VideoFrame> frame,
   DCHECK(!partial_decode_done_.get());
   DCHECK(!decode_done_.get());
   DCHECK(!updated_rects_);
+  DCHECK_EQ(kWaitingForBeginRect, state_);
 
   partial_decode_done_.reset(partial_decode_done);
   decode_done_.reset(decode_done);
@@ -33,54 +40,88 @@ bool DecoderVerbatim::BeginDecode(scoped_refptr<media::VideoFrame> frame,
 
 bool DecoderVerbatim::PartialDecode(HostMessage* message) {
   scoped_ptr<HostMessage> msg_deleter(message);
+  DCHECK(message->has_update_stream_packet());
+
+  bool ret = true;
+  if (message->update_stream_packet().has_begin_rect())
+    ret = HandleBeginRect(message);
+  if (ret && message->update_stream_packet().has_rect_data())
+    ret = HandleRectData(message);
+  if (ret && message->update_stream_packet().has_end_rect())
+    ret = HandleEndRect(message);
+  return ret;
+}
+
+void DecoderVerbatim::EndDecode() {
+  DCHECK_EQ(kWaitingForBeginRect, state_);
+  decode_done_->Run();
+
+  partial_decode_done_.reset();
+  decode_done_.reset();
+  frame_ = NULL;
+  updated_rects_ = NULL;
+}
+
+bool DecoderVerbatim::HandleBeginRect(HostMessage* message) {
+  DCHECK_EQ(kWaitingForBeginRect, state_);
+  state_ = kWaitingForRectData;
+
+  rect_width_ = message->update_stream_packet().begin_rect().width();
+  rect_height_ = message->update_stream_packet().begin_rect().height();
+  rect_x_ = message->update_stream_packet().begin_rect().x();
+  rect_y_ = message->update_stream_packet().begin_rect().y();
 
-  int width = message->update_stream_packet().header().width();
-  int height = message->update_stream_packet().header().height();
-  int x = message->update_stream_packet().header().x();
-  int y = message->update_stream_packet().header().y();
   PixelFormat pixel_format =
-      message->update_stream_packet().header().pixel_format();
+      message->update_stream_packet().begin_rect().pixel_format();
 
   if (static_cast<PixelFormat>(frame_->format()) != pixel_format) {
     NOTREACHED() << "Pixel format of message doesn't match the video frame. "
                     "Expected vs received = "
                  << frame_->format() << " vs " << pixel_format
                  << " Color space conversion required.";
+    return false;
   }
 
-  int bytes_per_pixel = GetBytesPerPixel(pixel_format);
+  bytes_per_pixel_ = GetBytesPerPixel(pixel_format);
+  return true;
+}
+
+bool DecoderVerbatim::HandleRectData(HostMessage* message) {
+  DCHECK_EQ(kWaitingForRectData, state_);
+  DCHECK_EQ(0,
+            message->update_stream_packet().rect_data().sequence_number());
+
   // Copy the data line by line.
-  const int src_stride = bytes_per_pixel * width;
-  const char* src = message->update_stream_packet().data().c_str();
+  const int src_stride = bytes_per_pixel_ * rect_width_;
+  const char* src =
+      message->update_stream_packet().rect_data().data().c_str();
   int src_stride_dir = src_stride;
   if (reverse_rows_) {
     // Copy rows from bottom to top to flip the image vertically.
-    src += (height - 1) * src_stride;
+    src += (rect_height_ - 1) * src_stride;
     // Change the direction of the stride to work bottom to top.
     src_stride_dir *= -1;
   }
   const int dest_stride = frame_->stride(media::VideoFrame::kRGBPlane);
   uint8* dest = frame_->data(media::VideoFrame::kRGBPlane) +
-      dest_stride * y + bytes_per_pixel * x;
-  for (int i = 0; i < height; ++i) {
+      dest_stride * rect_y_ + bytes_per_pixel_ * rect_x_;
+  for (int i = 0; i < rect_height_; ++i) {
     memcpy(dest, src, src_stride);
     dest += dest_stride;
     src += src_stride_dir;
   }
 
   updated_rects_->clear();
-  updated_rects_->push_back(gfx::Rect(x, y, width, height));
+  updated_rects_->push_back(gfx::Rect(rect_x_, rect_y_,
+                                      rect_width_, rect_height_));
   partial_decode_done_->Run();
   return true;
 }
 
-void DecoderVerbatim::EndDecode() {
-  decode_done_->Run();
-
-  partial_decode_done_.reset();
-  decode_done_.reset();
-  frame_ = NULL;
-  updated_rects_ = NULL;
+bool DecoderVerbatim::HandleEndRect(HostMessage* message) {
+  DCHECK_EQ(kWaitingForRectData, state_);
+  state_ = kWaitingForBeginRect;
+  return true;
 }
 
 }  // namespace remoting
diff --git a/remoting/client/decoder_verbatim.h b/remoting/client/decoder_verbatim.h
index 85b49c6..6efc732 100644
--- a/remoting/client/decoder_verbatim.h
+++ b/remoting/client/decoder_verbatim.h
@@ -22,6 +22,20 @@ class DecoderVerbatim : public Decoder {
   virtual void EndDecode();
 
  private:
+  bool HandleBeginRect(HostMessage* message);
+  bool HandleRectData(HostMessage* message);
+  bool HandleEndRect(HostMessage* message);
+
+  // The internal state of the decoder.
+  State state_;
+
+  // Keeps track of the updating rect.
+  int rect_x_;
+  int rect_y_;
+  int rect_width_;
+  int rect_height_;
+  int bytes_per_pixel_;
+
   // Tasks to call when decode is done.
   scoped_ptr<Task> partial_decode_done_;
   scoped_ptr<Task> decode_done_;
diff --git a/remoting/client/decoder_verbatim_unittest.cc b/remoting/client/decoder_verbatim_unittest.cc
index ca2c6f7..d541d22 100644
--- a/remoting/client/decoder_verbatim_unittest.cc
+++ b/remoting/client/decoder_verbatim_unittest.cc
@@ -20,13 +20,22 @@ TEST(DecoderVerbatimTest, SimpleDecode) {
   const size_t kHeight = 1;
   const char kData[] = "ABCDEFGHIJ";
   scoped_ptr<HostMessage> msg(new HostMessage());
-  msg->mutable_update_stream_packet()->mutable_header()->set_width(kWidth);
-  msg->mutable_update_stream_packet()->mutable_header()->set_height(kHeight);
-  msg->mutable_update_stream_packet()->mutable_header()->set_x(0);
-  msg->mutable_update_stream_packet()->mutable_header()->set_y(0);
-  msg->mutable_update_stream_packet()->mutable_header()->set_pixel_format(
-      PixelFormatAscii);
-  msg->mutable_update_stream_packet()->set_data(kData, sizeof(kData));
+
+  // Prepare the begin rect message.
+  UpdateStreamBeginRect* begin_rect =
+      msg->mutable_update_stream_packet()->mutable_begin_rect();
+  begin_rect->set_width(kWidth);
+  begin_rect->set_height(kHeight);
+  begin_rect->set_x(0);
+  begin_rect->set_y(0);
+  begin_rect->set_pixel_format(PixelFormatAscii);
+
+  // Prepare the rect data.
+  msg->mutable_update_stream_packet()->mutable_rect_data()->set_data(
+      kData, sizeof(kData));
+
+  // Prepare the end rect.
+  msg->mutable_update_stream_packet()->mutable_end_rect();
 
   scoped_refptr<media::VideoFrame> frame;
   media::VideoFrame::CreateFrame(media::VideoFrame::ASCII, kWidth, kHeight,
diff --git a/remoting/client/simple_client.cc b/remoting/client/simple_client.cc
index 98074bf..d1dd6e2 100644
--- a/remoting/client/simple_client.cc
+++ b/remoting/client/simple_client.cc
@@ -86,9 +86,13 @@ class SimpleHostEventCallback : public HostConnection::HostEventCallback {
 
   void HandleUpdateStreamPacketMessage(HostMessage* host_msg) {
     const UpdateStreamPacketMessage& msg = host_msg->update_stream_packet();
-    std::cout << "UpdateStreamPacket (" << msg.header().x()
-              << ", " << msg.header().y() << ") ["
-              << msg.header().width() << " x " << msg.header().height() << "]"
+    if (!msg.has_begin_rect())
+      return;
+
+    std::cout << "UpdateStreamPacket (" << msg.begin_rect().x()
+              << ", " << msg.begin_rect().y() << ") ["
+              << msg.begin_rect().width() << " x "
+              << msg.begin_rect().height() << "]"
               << std::endl;
   }
 
diff --git a/remoting/client/x11_view.cc b/remoting/client/x11_view.cc
index d380cd2..c5f9a88 100644
--- a/remoting/client/x11_view.cc
+++ b/remoting/client/x11_view.cc
@@ -20,6 +20,9 @@ X11View::X11View(Display* display, XID window, int width, int height)
       width_(width),
       height_(height),
       picture_(0) {
+  media::VideoFrame::CreateFrame(media::VideoFrame::RGB32, width_, height_,
+                                 base::TimeDelta(), base::TimeDelta(), &frame_);
+  DCHECK(frame_);
 }
 
 X11View::~X11View() {
@@ -30,7 +33,7 @@ void X11View::Paint() {
   all_update_rects_.clear();
 
   // If we have not initialized the render target then do it now.
-  if (!frame_)
+  if (!picture_)
     InitPaintTarget();
 
   // Upload the image to a pixmap. And then create a picture from the pixmap
@@ -112,11 +115,6 @@ void X11View::InitPaintTarget() {
 
   picture_ = XRenderCreatePicture(display_, window_, pictformat, 0, NULL);
   CHECK(picture_) << "Backing picture not created";
-
-  // Create the video frame to carry the decoded image.
-  media::VideoFrame::CreateFrame(media::VideoFrame::RGB32, width_, height_,
-                                 base::TimeDelta(), base::TimeDelta(), &frame_);
-  DCHECK(frame_);
 }
 
 void X11View::HandleBeginUpdateStream(HostMessage* msg) {
diff --git a/remoting/host/client_connection.cc b/remoting/host/client_connection.cc
index d69a326..6c8516a 100644
--- a/remoting/host/client_connection.cc
+++ b/remoting/host/client_connection.cc
@@ -35,6 +35,16 @@ ClientConnection::~ClientConnection() {
   // jingle channel.
 }
 
+// static
+scoped_refptr<media::DataBuffer>
+    ClientConnection::CreateWireFormatDataBuffer(
+        const HostMessage* msg) {
+  // TODO(hclam): Instead of serializing |msg| create an DataBuffer
+  // object that wraps around it.
+  scoped_ptr<const HostMessage> message_deleter(msg);
+  return SerializeAndFrameMessage(*msg);
+}
+
 void ClientConnection::SendInitClientMessage(int width, int height) {
   DCHECK_EQ(loop_, MessageLoop::current());
   DCHECK(!update_stream_size_);
@@ -68,7 +78,6 @@ void ClientConnection::SendBeginUpdateStreamMessage() {
 }
 
 void ClientConnection::SendUpdateStreamPacketMessage(
-    const UpdateStreamPacketHeader* header,
     scoped_refptr<DataBuffer> data) {
   DCHECK_EQ(loop_, MessageLoop::current());
 
@@ -76,16 +85,8 @@ void ClientConnection::SendUpdateStreamPacketMessage(
   if (!channel_)
     return;
 
-  HostMessage msg;
-  msg.mutable_update_stream_packet()->mutable_header()->CopyFrom(*header);
-  // TODO(hclam): This introduce one memory copy. Eliminate it.
-  msg.mutable_update_stream_packet()->set_data(
-      data->GetData(), data->GetDataSize());
-  DCHECK(msg.IsInitialized());
-
-  scoped_refptr<DataBuffer> encoded_data = SerializeAndFrameMessage(msg);
   update_stream_size_ += data->GetDataSize();
-  channel_->Write(encoded_data);
+  channel_->Write(data);
 }
 
 void ClientConnection::SendEndUpdateStreamMessage() {
diff --git a/remoting/host/client_connection.h b/remoting/host/client_connection.h
index 3e85cc9..c647eb6 100644
--- a/remoting/host/client_connection.h
+++ b/remoting/host/client_connection.h
@@ -62,6 +62,12 @@ class ClientConnection : public base::RefCountedThreadSafe<ClientConnection>,
 
   virtual ~ClientConnection();
 
+  // Creates a DataBuffer object that wraps around HostMessage. The DataBuffer
+  // object will be responsible for serializing and framing the message.
+  // DataBuffer will also own |msg| after this call.
+  static scoped_refptr<media::DataBuffer> CreateWireFormatDataBuffer(
+      const HostMessage* msg);
+
   virtual void set_jingle_channel(JingleChannel* channel) {
     channel_ = channel;
   }
@@ -75,10 +81,15 @@ class ClientConnection : public base::RefCountedThreadSafe<ClientConnection>,
   // Notifies the viewer the start of an update stream.
   virtual void SendBeginUpdateStreamMessage();
 
-  // Send encoded update stream data to the viewer. The viewer
-  // should not take ownership of the data.
+  // Send encoded update stream data to the viewer.
+  //
+  // |data| is the actual bytes in wire format. That means it is fully framed
+  // and serialized from a HostMessage. This is a special case only for
+  // UpdateStreamPacket to reduce the amount of memory copies.
+  //
+  // |data| should be created by calling to
+  // CreateWireFormatDataBuffer(HostMessage).
   virtual void SendUpdateStreamPacketMessage(
-      const UpdateStreamPacketHeader* header,
       scoped_refptr<media::DataBuffer> data);
 
   // Notifies the viewer the update stream has ended.
diff --git a/remoting/host/client_connection_unittest.cc b/remoting/host/client_connection_unittest.cc
index b7cb10c..9c3bcb3 100644
--- a/remoting/host/client_connection_unittest.cc
+++ b/remoting/host/client_connection_unittest.cc
@@ -55,14 +55,8 @@ TEST_F(ClientConnectionTest, SendUpdateStream) {
 
   // Then send the actual data.
   EXPECT_CALL(*channel_, Write(_));
-  scoped_ptr<UpdateStreamPacketHeader> header(new UpdateStreamPacketHeader);
-  header->set_x(0);
-  header->set_y(0);
-  header->set_width(640);
-  header->set_height(480);
-
   scoped_refptr<media::DataBuffer> data = new media::DataBuffer(10);
-  viewer_->SendUpdateStreamPacketMessage(header.get(), data);
+  viewer_->SendUpdateStreamPacketMessage(data);
 
   // Send the end of update message.
   EXPECT_CALL(*channel_, Write(_));
@@ -105,15 +99,8 @@ TEST_F(ClientConnectionTest, Close) {
   EXPECT_CALL(*channel_, Close());
   viewer_->Disconnect();
 
-  viewer_->SendBeginUpdateStreamMessage();
-  scoped_ptr<UpdateStreamPacketHeader> header(new UpdateStreamPacketHeader);
-  header->set_x(0);
-  header->set_y(0);
-  header->set_width(640);
-  header->set_height(480);
-
   scoped_refptr<media::DataBuffer> data = new media::DataBuffer(10);
-  viewer_->SendUpdateStreamPacketMessage(header.get(), data);
+  viewer_->SendUpdateStreamPacketMessage(data);
   viewer_->SendEndUpdateStreamMessage();
   viewer_->Disconnect();
 }
diff --git a/remoting/host/encoder.h b/remoting/host/encoder.h
index 09cbdde..ee43c94 100644
--- a/remoting/host/encoder.h
+++ b/remoting/host/encoder.h
@@ -17,6 +17,8 @@ namespace media {
 
 namespace remoting {
 
+class HostMessage;
+
 // A class to perform the task of encoding a continous stream of
 // images.
 // This class operates asynchronously to enable maximum throughput.
@@ -34,11 +36,11 @@ class Encoder {
   typedef int EncodingState;
 
   // DataAvailableCallback is called as blocks of data are made available
-  // from the encoder. The callback takes ownership of header and is
-  // responsible for deleting it.
-  typedef Callback3<const UpdateStreamPacketHeader*,
-                    const scoped_refptr<media::DataBuffer>&,
-                    EncodingState>::Type DataAvailableCallback;
+  // from the encoder. Data made available by the encoder is in the form
+  // of HostMessage to reduce the amount of memory copies.
+  // The callback takes ownership of the HostMessage and is responsible for
+  // deleting it.
+  typedef Callback2<HostMessage*, EncodingState>::Type DataAvailableCallback;
 
   virtual ~Encoder() {}
 
diff --git a/remoting/host/encoder_verbatim.cc b/remoting/host/encoder_verbatim.cc
index f1e2b46..fd9dadf 100644
--- a/remoting/host/encoder_verbatim.cc
+++ b/remoting/host/encoder_verbatim.cc
@@ -7,6 +7,7 @@
 #include "gfx/rect.h"
 #include "media/base/data_buffer.h"
 #include "remoting/base/protocol_util.h"
+#include "remoting/base/protocol/chromotocol.pb.h"
 
 namespace remoting {
 
@@ -17,13 +18,14 @@ void EncoderVerbatim::Encode(scoped_refptr<Capturer::CaptureData> capture_data,
                              DataAvailableCallback* data_available_callback) {
   int num_rects = capture_data->dirty_rects().size();
   for (int i = 0; i < num_rects; i++) {
-    scoped_refptr<DataBuffer> data;
     const gfx::Rect& dirty_rect = capture_data->dirty_rects()[i];
-    scoped_ptr<UpdateStreamPacketHeader> header(new UpdateStreamPacketHeader);
-    if (EncodeRect(dirty_rect,
-                   capture_data,
-                   header.get(),
-                   &data)) {
+    HostMessage* msg = new HostMessage();
+    UpdateStreamPacketMessage* packet = msg->mutable_update_stream_packet();
+
+    if (EncodeRect(dirty_rect, capture_data, packet)) {
+      // Prepare the end rect content.
+      packet->mutable_end_rect();
+
       EncodingState state = EncodingInProgress;
       if (i == 0) {
         state |= EncodingStarting;
@@ -31,9 +33,7 @@ void EncoderVerbatim::Encode(scoped_refptr<Capturer::CaptureData> capture_data,
       if (i == num_rects - 1) {
         state |= EncodingEnded;
       }
-      data_available_callback->Run(header.release(),
-                                   data,
-                                   state);
+      data_available_callback->Run(msg, state);
     }
   }
 
@@ -43,12 +43,18 @@ void EncoderVerbatim::Encode(scoped_refptr<Capturer::CaptureData> capture_data,
 bool EncoderVerbatim::EncodeRect(
     const gfx::Rect& dirty,
     const scoped_refptr<Capturer::CaptureData>& capture_data,
-    UpdateStreamPacketHeader* header,
-    scoped_refptr<DataBuffer>* output_data) {
-  int bytes_per_pixel = GetBytesPerPixel(capture_data->pixel_format());
-  int row_size = bytes_per_pixel * dirty.width();
+    UpdateStreamPacketMessage* packet) {
+  // Prepare the begin rect content.
+  packet->mutable_begin_rect()->set_x(dirty.x());
+  packet->mutable_begin_rect()->set_y(dirty.y());
+  packet->mutable_begin_rect()->set_width(dirty.width());
+  packet->mutable_begin_rect()->set_height(dirty.height());
+  packet->mutable_begin_rect()->set_encoding(EncodingNone);
+  packet->mutable_begin_rect()->set_pixel_format(capture_data->pixel_format());
 
   // Calculate the size of output.
+  int bytes_per_pixel = GetBytesPerPixel(capture_data->pixel_format());
+  int row_size = bytes_per_pixel * dirty.width();
   int output_size = 0;
   for (int i = 0; i < Capturer::DataPlanes::kPlaneCount; ++i) {
     // TODO(hclam): Handle YUV since the height would be different.
@@ -57,15 +63,10 @@ bool EncoderVerbatim::EncodeRect(
     output_size += row_size * dirty.height();
   }
 
-  header->set_x(dirty.x());
-  header->set_y(dirty.y());
-  header->set_width(dirty.width());
-  header->set_height(dirty.height());
-  header->set_encoding(EncodingNone);
-  header->set_pixel_format(capture_data->pixel_format());
-
-  *output_data = new DataBuffer(new uint8[output_size], output_size);
-  uint8* out = (*output_data)->GetWritableData();
+  // Resize the output data buffer.
+  packet->mutable_rect_data()->mutable_data()->resize(output_size);
+  uint8* out = reinterpret_cast<uint8*>(
+      &((*packet->mutable_rect_data()->mutable_data())[0]));
 
   for (int i = 0; i < Capturer::DataPlanes::kPlaneCount; ++i) {
     const uint8* in = capture_data->data_planes().data[i];
diff --git a/remoting/host/encoder_verbatim.h b/remoting/host/encoder_verbatim.h
index 745d544..175b954 100644
--- a/remoting/host/encoder_verbatim.h
+++ b/remoting/host/encoder_verbatim.h
@@ -9,6 +9,8 @@
 
 namespace remoting {
 
+class UpdateStreamPacket;
+
 // EncoderVerbatim implements Encoder and simply copies input to the output
 // buffer verbatim.
 class EncoderVerbatim : public Encoder {
@@ -21,12 +23,12 @@ class EncoderVerbatim : public Encoder {
                       DataAvailableCallback* data_available_callback);
 
  private:
-  // Encode a single dirty rect. Called by Encode().
+  // Encode a single dirty rect. Called by Encode(). Output is written
+  // to |msg|.
   // Returns false if there is an error.
   bool EncodeRect(const gfx::Rect& dirty,
                   const scoped_refptr<Capturer::CaptureData>& capture_data,
-                  UpdateStreamPacketHeader* header,
-                  scoped_refptr<media::DataBuffer>* output_data);
+                  UpdateStreamPacketMessage* msg);
 };
 
 }  // namespace remoting
diff --git a/remoting/host/mock_objects.h b/remoting/host/mock_objects.h
index 0585864..f4135f2 100644
--- a/remoting/host/mock_objects.h
+++ b/remoting/host/mock_objects.h
@@ -61,9 +61,8 @@ class MockClientConnection : public ClientConnection {
 
   MOCK_METHOD2(SendInitClientMessage, void(int width, int height));
   MOCK_METHOD0(SendBeginUpdateStreamMessage, void());
-  MOCK_METHOD2(SendUpdateStreamPacketMessage,
-               void(const UpdateStreamPacketHeader* header,
-                    scoped_refptr<media::DataBuffer> data));
+  MOCK_METHOD1(SendUpdateStreamPacketMessage,
+               void(scoped_refptr<media::DataBuffer> data));
   MOCK_METHOD0(SendEndUpdateStreamMessage, void());
   MOCK_METHOD0(GetPendingUpdateStreamMessages, int());
 
diff --git a/remoting/host/session_manager.cc b/remoting/host/session_manager.cc
index 8e1dd86..5be8b56d 100644
--- a/remoting/host/session_manager.cc
+++ b/remoting/host/session_manager.cc
@@ -322,25 +322,26 @@ void SessionManager::DoRateControl() {
   ScheduleNextRateControl();
 }
 
-void SessionManager::DoSendUpdate(const UpdateStreamPacketHeader* header,
-                                  const scoped_refptr<media::DataBuffer> data,
+void SessionManager::DoSendUpdate(HostMessage* message,
                                   Encoder::EncodingState state) {
-  // Take ownership of header.
-  scoped_ptr<const UpdateStreamPacketHeader> header_owner(header);
   DCHECK_EQ(network_loop_, MessageLoop::current());
 
+  // Create a data buffer in wire format from |message|.
+  scoped_refptr<media::DataBuffer> data =
+      ClientConnection::CreateWireFormatDataBuffer(message);
+
   for (ClientConnectionList::const_iterator i = clients_.begin();
-       i < clients_.end();
-       ++i) {
+       i < clients_.end(); ++i) {
+    // TODO(hclam): Merge BeginUpdateStreamMessage into |message|.
     if (state & Encoder::EncodingStarting) {
       (*i)->SendBeginUpdateStreamMessage();
     }
 
-    (*i)->SendUpdateStreamPacketMessage(header, data);
+    (*i)->SendUpdateStreamPacketMessage(data);
 
-    if (state & Encoder::EncodingEnded) {
+    // TODO(hclam): Merge EndUpdateStreamMessage into |message|.
+    if (state & Encoder::EncodingEnded)
       (*i)->SendEndUpdateStreamMessage();
-    }
   }
 }
 
@@ -390,9 +391,7 @@ void SessionManager::DoEncode(
 }
 
 void SessionManager::EncodeDataAvailableTask(
-    const UpdateStreamPacketHeader *header,
-    const scoped_refptr<media::DataBuffer>& data,
-    Encoder::EncodingState state) {
+    HostMessage* message, Encoder::EncodingState state) {
   DCHECK_EQ(encode_loop_, MessageLoop::current());
 
   // Before a new encode task starts, notify clients a new update
@@ -401,11 +400,7 @@ void SessionManager::EncodeDataAvailableTask(
   // task. The ownership will eventually pass to the ClientConnections.
   network_loop_->PostTask(
       FROM_HERE,
-      NewRunnableMethod(this,
-                        &SessionManager::DoSendUpdate,
-                        header,
-                        data,
-                        state));
+      NewRunnableMethod(this, &SessionManager::DoSendUpdate, message, state));
 
   if (state & Encoder::EncodingEnded) {
     capture_loop_->PostTask(
diff --git a/remoting/host/session_manager.h b/remoting/host/session_manager.h
index b56564f..c2a2149 100644
--- a/remoting/host/session_manager.h
+++ b/remoting/host/session_manager.h
@@ -131,8 +131,7 @@ class SessionManager : public base::RefCountedThreadSafe<SessionManager> {
   void DoRateControl();
 
   // DoSendUpdate takes ownership of header and is responsible for deleting it.
-  void DoSendUpdate(const UpdateStreamPacketHeader* header,
-                    const scoped_refptr<media::DataBuffer> data,
+  void DoSendUpdate(HostMessage* message,
                     Encoder::EncodingState state);
   void DoSendInit(scoped_refptr<ClientConnection> client,
                   int width, int height);
@@ -147,8 +146,7 @@ class SessionManager : public base::RefCountedThreadSafe<SessionManager> {
 
   // EncodeDataAvailableTask takes ownership of header and is responsible for
   // deleting it.
-  void EncodeDataAvailableTask(const UpdateStreamPacketHeader *header,
-                               const scoped_refptr<media::DataBuffer>& data,
+  void EncodeDataAvailableTask(HostMessage* message,
                                Encoder::EncodingState state);
 
   // Message loops used by this class.
diff --git a/remoting/host/session_manager_unittest.cc b/remoting/host/session_manager_unittest.cc
index 5a4e4eb..e6d8114 100644
--- a/remoting/host/session_manager_unittest.cc
+++ b/remoting/host/session_manager_unittest.cc
@@ -58,18 +58,11 @@ ACTION_P2(RunCallback, rects, data) {
   delete arg0;
 }
 
-ACTION_P3(FinishDecode, rects, buffer, header) {
-  gfx::Rect& rect = (*rects)[0];
+ACTION_P(FinishEncode, msg) {
   Encoder::EncodingState state = (Encoder::EncodingStarting |
                                   Encoder::EncodingInProgress |
                                   Encoder::EncodingEnded);
-  header->set_x(rect.x());
-  header->set_y(rect.y());
-  header->set_width(rect.width());
-  header->set_height(rect.height());
-  header->set_encoding(kEncoding);
-  header->set_pixel_format(kFormat);
-  arg2->Run(header, *buffer, state);
+  arg2->Run(msg, state);
 }
 
 ACTION_P(AssignCaptureData, data) {
@@ -111,15 +104,13 @@ TEST_F(SessionManagerTest, OneRecordCycle) {
       .WillOnce(RunCallback(update_rects, data));
 
   // Expect the encoder be called.
-  scoped_refptr<media::DataBuffer> buffer = new media::DataBuffer(0);
-  UpdateStreamPacketHeader *header = new UpdateStreamPacketHeader;
+  HostMessage* msg = new HostMessage();
   EXPECT_CALL(*encoder_, Encode(data, false, NotNull()))
-      .WillOnce(FinishDecode(&update_rects, &buffer, header));
+      .WillOnce(FinishEncode(msg));
 
   // Expect the client be notified.
   EXPECT_CALL(*client_, SendBeginUpdateStreamMessage());
-
-  EXPECT_CALL(*client_, SendUpdateStreamPacketMessage(header ,buffer));
+  EXPECT_CALL(*client_, SendUpdateStreamPacketMessage(_));
   EXPECT_CALL(*client_, SendEndUpdateStreamMessage());
   EXPECT_CALL(*client_, GetPendingUpdateStreamMessages())
       .Times(AtLeast(0))