1 files changed, 53 insertions, 1 deletions

diff --git a/lib/server/SocketStream.h b/lib/server/SocketStream.h
index 2fb5e391..fd57af8f 100644
--- a/lib/server/SocketStream.h
+++ b/lib/server/SocketStream.h
@@ -10,6 +10,13 @@
 #ifndef SOCKETSTREAM__H
 #define SOCKETSTREAM__H
 
+#include <climits>
+
+#ifdef HAVE_SYS_POLL_H
+#	include <sys/poll.h>
+#endif
+
+#include "BoxTime.h"
 #include "IOStream.h"
 #include "Socket.h"
 
@@ -41,7 +48,16 @@ public:
 	void Attach(int socket);
 
 	virtual int Read(void *pBuffer, int NBytes, int Timeout = IOStream::TimeOutInfinite);
-	virtual void Write(const void *pBuffer, int NBytes);
+	virtual void Write(const void *pBuffer, int NBytes,
+		int Timeout = IOStream::TimeOutInfinite);
+
+	// Why not inherited from IOStream? Never mind, we want to enforce
+	// supplying a timeout for network operations anyway.
+	virtual void Write(const std::string& rBuffer, int Timeout)
+	{
+		IOStream::Write(rBuffer, Timeout);
+	}
+
 	virtual void Close();
 	virtual bool StreamDataLeft();
 	virtual bool StreamClosed();
@@ -53,6 +69,42 @@ public:
 protected:
 	void MarkAsReadClosed() {mReadClosed = true;}
 	void MarkAsWriteClosed() {mWriteClosed = true;}
+	void CheckForMissingTimeout(int Timeout);
+
+	// Converts a timeout in milliseconds (or IOStream::TimeOutInfinite)
+	// into one that can be passed to poll() (also in milliseconds), also
+	// compensating for time elapsed since the wait should have started,
+	// if known.
+	int PollTimeout(int timeout, box_time_t start_time)
+	{
+		if (timeout == IOStream::TimeOutInfinite)
+		{
+			return INFTIM;
+		}
+
+		if (start_time == 0)
+		{
+			return timeout; // no adjustment possible
+		}
+
+		box_time_t end_time = start_time + MilliSecondsToBoxTime(timeout);
+		box_time_t now = GetCurrentBoxTime();
+		box_time_t remaining = end_time - now;
+
+		if (remaining < 0)
+		{
+			return 0; // no delay
+		}
+		else if (BoxTimeToMilliSeconds(remaining) > INT_MAX)
+		{
+			return INT_MAX;
+		}
+		else
+		{
+			return (int) BoxTimeToMilliSeconds(remaining);
+		}
+	}
+	bool Poll(short Events, int Timeout);
 
 private:
 	tOSSocketHandle mSocketHandle;