4 files changed, 35 insertions, 3 deletions

diff --git a/src/connection.cpp b/src/connection.cpp
index 5b79dfe..0e7f111 100644
--- a/src/connection.cpp
+++ b/src/connection.cpp
@@ -9,6 +9,7 @@
 #include <netinet/in.h>
 #include <netdb.h>
 #include <arpa/inet.h>
+#include "exceptions.h"
 
 Connection::Connection()
 {
@@ -244,10 +245,9 @@ bool Connection::readInput( int nSec, int nUSec, int *pnSecBack, int *pnUSecBack
         struct timeval tv;
         int retval;
         
-        /* Watch stdin (fd 0) to see when it has input. */
         FD_ZERO(&rfds);
         FD_SET(nSocket, &rfds);
-        /* Wait up to five seconds. */
+
         tv.tv_sec = nSec;
         tv.tv_usec = nUSec;
 
@@ -256,7 +256,7 @@ bool Connection::readInput( int nSec, int nUSec, int *pnSecBack, int *pnUSecBack
         if( pnSecBack ) (*pnSecBack) = tv.tv_sec;
         if( pnUSecBack ) (*pnUSecBack) = tv.tv_usec;
 
-        if (retval == -1)
+        if( retval == -1 )
         {
                 // Oh my god!!! some kind of horrible problem!!!!
                 return false;
@@ -272,6 +272,24 @@ bool Connection::readInput( int nSec, int nUSec, int *pnSecBack, int *pnUSecBack
         }
 }
 
+void Connection::waitForInput( int nBytesIn, int nSec, int nUSec )
+{
+        int rlen = getInputAmnt();
+
+        if( rlen >= nBytesIn )
+                return;
+
+        while( rlen < nBytesIn )
+        {
+                if( nSec == 0 && nUSec == 0 )
+                {
+                        throw ConnectionException("Socket Timeout");
+                }
+                readInput( nSec, nUSec, &nSec, &nUSec );
+                rlen = getInputAmnt();
+        }
+}
+
 bool Connection::clearOutput()
 {
         return xOutputBuf.clearData();
diff --git a/src/connection.h b/src/connection.h
index 29fcc25..5e86ff4 100644
--- a/src/connection.h
+++ b/src/connection.h
@@ -283,6 +283,18 @@ public:
          */
         bool readInput( int nSec, int nUSec, int *pnSecBack=NULL, int *pnUSecBack=NULL );
 
+        /**
+         * Waits until at least nBytesIn are read into the input buffer and ready
+         * to be used.  Wait at most nSec seconds plus nUSec micro seconds.
+         * If the timeout is exceeded, this function throws an exception.  If this
+         * function returns normally, you are guranteed to have at least nBytesIn
+         * bytes in your input buffer.
+         *@param nBytesIn Number of bytes to read.
+         *@param nSec The max seconds to wait.
+         *@param sUSec The max microseconds to wait.
+         */
+        void waitForInput( int nBytesIn, int nSec, int nUSec );
+
         /** Writes all data that is pending to the socket.
           *@returns True if all data was written succesfully, false otherwise.
           */
diff --git a/src/exceptions.cpp b/src/exceptions.cpp
index 8d30296..9696af8 100644
--- a/src/exceptions.cpp
+++ b/src/exceptions.cpp
@@ -3,3 +3,4 @@
 
 subExceptionDef( XmlException )
 subExceptionDef( FileException )
+subExceptionDef( ConnectionException )
diff --git a/src/exceptions.h b/src/exceptions.h
index 3bf2e0b..27aec3c 100644
--- a/src/exceptions.h
+++ b/src/exceptions.h
@@ -6,5 +6,6 @@
 
 subExceptionDecl( XmlException )
 subExceptionDecl( FileException )
+subExceptionDecl( ConnectionException )
 
 #endif