Make direct stream API a bit safer

Provide some safety checks when directly accessing the underlying
pointer of streams.

7 files changed, 42 insertions, 38 deletions

diff --git a/common/rdr/HexInStream.cxx b/common/rdr/HexInStream.cxx
index 3800961d..322432c0 100644
--- a/common/rdr/HexInStream.cxx
+++ b/common/rdr/HexInStream.cxx
@@ -76,9 +76,8 @@ bool HexInStream::fillBuffer(size_t maxSize, bool wait) {
   if (!in_stream.check(2, wait))
     return false;
 
-  const U8* iptr = in_stream.getptr();
-  const U8* eptr = in_stream.getend();
-  size_t length = min((eptr - iptr)/2, maxSize);
+  size_t length = min(in_stream.avail()/2, maxSize);
+  const U8* iptr = in_stream.getptr(length*2);
 
   U8* optr = (U8*) end;
   for (size_t i=0; i<length; i++) {
@@ -88,7 +87,7 @@ bool HexInStream::fillBuffer(size_t maxSize, bool wait) {
     optr[i] = v;
   }
 
-  in_stream.setptr(iptr + length*2);
+  in_stream.setptr(length*2);
   end += length;
 
   return true;
diff --git a/common/rdr/HexOutStream.cxx b/common/rdr/HexOutStream.cxx
index 88153c54..798f6bdf 100644
--- a/common/rdr/HexOutStream.cxx
+++ b/common/rdr/HexOutStream.cxx
@@ -66,17 +66,15 @@ void
 HexOutStream::writeBuffer() {
   U8* pos = start;
   while (pos != ptr) {
-    out_stream.check(2);
-    U8* optr = out_stream.getptr();
-    U8* oend = out_stream.getend();
-    size_t length = min(ptr-pos, (oend-optr)/2);
+    U8* optr = out_stream.getptr(2);
+    size_t length = min(ptr-pos, out_stream.avail()/2);
 
     for (size_t i=0; i<length; i++) {
       optr[i*2] = intToHex((pos[i] >> 4) & 0xf);
       optr[i*2+1] = intToHex(pos[i] & 0xf);
     }
 
-    out_stream.setptr(optr + length*2);
+    out_stream.setptr(length*2);
     pos += length;
   }
   offset += ptr - start;
diff --git a/common/rdr/InStream.h b/common/rdr/InStream.h
index ce3a344d..46af74c9 100644
--- a/common/rdr/InStream.h
+++ b/common/rdr/InStream.h
@@ -25,6 +25,7 @@
 #define __RDR_INSTREAM_H__
 
 #include <rdr/types.h>
+#include <rdr/Exception.h>
 #include <string.h> // for memcpy
 
 namespace rdr {
@@ -119,13 +120,14 @@ namespace rdr {
 
     virtual size_t pos() = 0;
 
-    // getptr(), getend() and setptr() are "dirty" methods which allow you to
-    // manipulate the buffer directly.  This is useful for a stream which is a
-    // wrapper around an underlying stream.
+    // getptr() and setptr() are "dirty" methods which allow you direct access
+    // to the buffer. This is useful for a stream which is a wrapper around an
+    // some other stream API.
 
-    inline const U8* getptr() const { return ptr; }
-    inline const U8* getend() const { return end; }
-    inline void setptr(const U8* p) { ptr = p; }
+    inline const U8* getptr(size_t length) { check(length); return ptr; }
+    inline void setptr(size_t length) { if (length > avail())
+                                          throw Exception("Input stream overflow");
+                                        skip(length); }
 
   private:
 
diff --git a/common/rdr/OutStream.h b/common/rdr/OutStream.h
index 63c43169..0d5f113e 100644
--- a/common/rdr/OutStream.h
+++ b/common/rdr/OutStream.h
@@ -25,6 +25,7 @@
 #define __RDR_OUTSTREAM_H__
 
 #include <rdr/types.h>
+#include <rdr/Exception.h>
 #include <rdr/InStream.h>
 #include <string.h> // for memcpy
 
@@ -132,13 +133,16 @@ namespace rdr {
 
     virtual void cork(bool enable) { corked = enable; flush(); }
 
-    // getptr(), getend() and setptr() are "dirty" methods which allow you to
-    // manipulate the buffer directly.  This is useful for a stream which is a
-    // wrapper around an underlying stream.
+    // getptr() and setptr() are "dirty" methods which allow you direct access
+    // to the buffer. This is useful for a stream which is a wrapper around an
+    // some other stream API. Note that setptr() should not called with a value
+    // larger than the bytes actually written as doing so can result in
+    // security issues. Use pad() in such cases instead.
 
-    inline U8* getptr() { return ptr; }
-    inline U8* getend() { return end; }
-    inline void setptr(U8* p) { ptr = p; }
+    inline U8* getptr(size_t length) { check(length); return ptr; }
+    inline void setptr(size_t length) { if (length > avail())
+                                          throw Exception("Output stream overflow");
+                                        ptr += length; }
 
   private:
 
diff --git a/common/rdr/ZlibInStream.cxx b/common/rdr/ZlibInStream.cxx
index b9e772d5..26977228 100644
--- a/common/rdr/ZlibInStream.cxx
+++ b/common/rdr/ZlibInStream.cxx
@@ -95,18 +95,19 @@ bool ZlibInStream::fillBuffer(size_t maxSize, bool wait)
 
   size_t n = underlying->check(1, wait);
   if (n == 0) return false;
-  zs->next_in = (U8*)underlying->getptr();
-  zs->avail_in = underlying->avail();
-  if (zs->avail_in > bytesIn)
-    zs->avail_in = bytesIn;
+  size_t length = underlying->avail();
+  if (length > bytesIn)
+    length = bytesIn;
+  zs->next_in = (U8*)underlying->getptr(length);
+  zs->avail_in = length;
 
   int rc = inflate(zs, Z_SYNC_FLUSH);
   if (rc < 0) {
     throw Exception("ZlibInStream: inflate failed");
   }
 
-  bytesIn -= zs->next_in - underlying->getptr();
+  bytesIn -= length - zs->avail_in;
   end = zs->next_out;
-  underlying->setptr(zs->next_in);
+  underlying->setptr(length - zs->avail_in);
   return true;
 }
diff --git a/common/rdr/ZlibOutStream.cxx b/common/rdr/ZlibOutStream.cxx
index 0eb89222..ba6f1785 100644
--- a/common/rdr/ZlibOutStream.cxx
+++ b/common/rdr/ZlibOutStream.cxx
@@ -144,9 +144,9 @@ void ZlibOutStream::deflate(int flush)
     return;
 
   do {
-    underlying->check(1);
-    zs->next_out = underlying->getptr();
-    zs->avail_out = underlying->avail();
+    size_t chunk;
+    zs->next_out = underlying->getptr(1);
+    zs->avail_out = chunk = underlying->avail();
 
 #ifdef ZLIBOUT_DEBUG
     vlog.debug("calling deflate, avail_in %d, avail_out %d",
@@ -167,7 +167,7 @@ void ZlibOutStream::deflate(int flush)
                zs->next_out-underlying->getptr());
 #endif
 
-    underlying->setptr(zs->next_out);
+    underlying->setptr(chunk - zs->avail_out);
   } while (zs->avail_out == 0);
 }
 
diff --git a/common/rfb/JpegCompressor.cxx b/common/rfb/JpegCompressor.cxx
index 5cc3855d..64c6d3f0 100644
--- a/common/rfb/JpegCompressor.cxx
+++ b/common/rfb/JpegCompressor.cxx
@@ -75,6 +75,7 @@ JpegOutputMessage(j_common_ptr cinfo)
 struct JPEG_DEST_MGR {
   struct jpeg_destination_mgr pub;
   JpegCompressor *instance;
+  size_t chunkSize;
 };
 
 static void
@@ -84,8 +85,8 @@ JpegInitDestination(j_compress_ptr cinfo)
   JpegCompressor *jc = dest->instance;
 
   jc->clear();
-  dest->pub.next_output_byte = jc->getptr();
-  dest->pub.free_in_buffer = jc->avail();
+  dest->pub.next_output_byte = jc->getptr(jc->length());
+  dest->pub.free_in_buffer = dest->chunkSize = jc->avail();
 }
 
 static boolean
@@ -94,10 +95,9 @@ JpegEmptyOutputBuffer(j_compress_ptr cinfo)
   JPEG_DEST_MGR *dest = (JPEG_DEST_MGR *)cinfo->dest;
   JpegCompressor *jc = dest->instance;
 
-  jc->setptr(jc->getend());
-  jc->check(jc->length());
-  dest->pub.next_output_byte = jc->getptr();
-  dest->pub.free_in_buffer = jc->avail();
+  jc->setptr(jc->avail());
+  dest->pub.next_output_byte = jc->getptr(jc->length());
+  dest->pub.free_in_buffer = dest->chunkSize = jc->avail();
 
   return TRUE;
 }
@@ -108,7 +108,7 @@ JpegTermDestination(j_compress_ptr cinfo)
   JPEG_DEST_MGR *dest = (JPEG_DEST_MGR *)cinfo->dest;
   JpegCompressor *jc = dest->instance;
 
-  jc->setptr(dest->pub.next_output_byte);
+  jc->setptr(dest->chunkSize - dest->pub.free_in_buffer);
 }
 
 JpegCompressor::JpegCompressor(int bufferLen) : MemOutStream(bufferLen)