From 7a20a14a0647fb2856898b6994cd35917c10791f Mon Sep 17 00:00:00 2001
From: Mark Slee <mcslee@apache.org>
Date: Thu, 1 Feb 2007 21:39:50 +0000
Subject: [PATCH] Fix bufferedroutertransport overflow bugs

Reviewed By: aditya


git-svn-id: https://svn.apache.org/repos/asf/incubator/thrift/trunk@664964 13f79535-47bb-0310-9956-ffa450edef68
---
 .../transport/TBufferedRouterTransport.cpp    | 34 ++++++++-----------
 .../src/transport/TBufferedRouterTransport.h  |  5 ++-
 2 files changed, 16 insertions(+), 23 deletions(-)

diff --git a/lib/cpp/src/transport/TBufferedRouterTransport.cpp b/lib/cpp/src/transport/TBufferedRouterTransport.cpp
index ad6a28ff..60ab594a 100644
--- a/lib/cpp/src/transport/TBufferedRouterTransport.cpp
+++ b/lib/cpp/src/transport/TBufferedRouterTransport.cpp
@@ -47,32 +47,26 @@ void TBufferedRouterTransport::write(const uint8_t* buf, uint32_t len) {
     return;
   }
 
-  if (len + wLen_ >= wBufSize_) {
-    uint32_t copy = wBufSize_ - wLen_;
-    memcpy(wBuf_ + wLen_, buf, copy);
-    trans_->write(wBuf_+wPos_, wBufSize_-wPos_);
-    wLen_ += copy;
-    wPos_ = wLen_;
-
-    uint32_t left = len-copy;
-    if (left > 0) {
-      // double the size of the write buffer
-      wBuf_ = (uint8_t *)realloc(wBuf_, sizeof(uint8_t) * wBufSize_ * 2);
-      memcpy(wBuf_ + wLen_, buf+copy, left);
-      wLen_ += left;
-      wBufSize_*=2;
+  // Make the buffer as big as it needs to be
+  if ((len + wLen_) >= wBufSize_) {
+    uint32_t newBufSize = wBufSize_*2;
+    while ((len + wLen_) >= newBufSize) {
+      newBufSize *= 2;
     }
-  } else {
-    memcpy(wBuf_+wLen_, buf, len);
-    wLen_ += len;
+    wBuf_ = (uint8_t *)realloc(wBuf_, sizeof(uint8_t) * newBufSize);
+    wBufSize_ = newBufSize;
   }
+
+  // Copy into the buffer
+  memcpy(wBuf_ + wLen_, buf, len);
+  wLen_ += len;
 }
 
 void TBufferedRouterTransport::flush()  {
   // Write out any data waiting in the write buffer
-  if (wLen_-wPos_ > 0) {
-    trans_->write(wBuf_+wPos_, wLen_-wPos_);
-    wPos_ = wLen_;
+  if (wLen_ > 0) {
+    trans_->write(wBuf_, wLen_);
+    wLen_ = 0;
   }
 
   // Flush the underlying transport
diff --git a/lib/cpp/src/transport/TBufferedRouterTransport.h b/lib/cpp/src/transport/TBufferedRouterTransport.h
index add31072..0b4577c4 100644
--- a/lib/cpp/src/transport/TBufferedRouterTransport.h
+++ b/lib/cpp/src/transport/TBufferedRouterTransport.h
@@ -25,7 +25,7 @@ class TBufferedRouterTransport : public TTransport {
     trans_(trans),
     rtrans_(rtrans),
     rBufSize_(512), rPos_(0), rLen_(0),
-    wBufSize_(512), wPos_(0), wLen_(0) {
+    wBufSize_(512), wLen_(0) {
 
     rBuf_ = (uint8_t*) malloc(sizeof(uint8_t) * rBufSize_);
     wBuf_ = (uint8_t*) malloc(sizeof(uint8_t) * wBufSize_);
@@ -35,7 +35,7 @@ class TBufferedRouterTransport : public TTransport {
     trans_(trans),
     rtrans_(rtrans),
     rBufSize_(512), rPos_(0), rLen_(0),
-    wBufSize_(sz), wPos_(0), wLen_(0) {
+    wBufSize_(sz), wLen_(0) {
 
     rBuf_ = (uint8_t*) malloc(sizeof(uint8_t) * rBufSize_);
     wBuf_ = (uint8_t*) malloc(sizeof(uint8_t) * wBufSize_);
@@ -98,7 +98,6 @@ class TBufferedRouterTransport : public TTransport {
 
   uint8_t* wBuf_;
   uint32_t wBufSize_;
-  uint32_t wPos_;
   uint32_t wLen_;
 };
 
-- 
2.17.1