Add support for multiple IOHooks per StreamSocket

2 files changed, 85 insertions, 4 deletions

diff --git a/include/inspsocket.h b/include/inspsocket.h
index 0c5f3b3af..77d79bd95 100644
--- a/include/inspsocket.h
+++ b/include/inspsocket.h
@@ -240,6 +240,16 @@ class CoreExport StreamSocket : public EventHandler
 	 */
 	int ReadToRecvQ(std::string& rq);
 
+	/** Read data from a hook chain recursively, starting at 'hook'.
+	 * If 'hook' is NULL, the recvq is filled with data from SocketEngine::Recv(), otherwise it is filled with data from the
+	 * next hook in the chain.
+	 * @param hook Next IOHook in the chain, can be NULL
+	 * @param rq Receive queue to put incoming data into
+	 * @return < 0 on error or close, 0 if no new data is ready (but the socket is still connected), > 0 if data was read from
+	 the socket and put into the recvq
+	 */
+	int HookChainRead(IOHook* hook, std::string& rq);
+
  protected:
 	std::string recvq;
  public:
@@ -286,7 +296,7 @@ class CoreExport StreamSocket : public EventHandler
 	 */
 	bool GetNextLine(std::string& line, char delim = '\n');
 	/** Useful for implementing sendq exceeded */
-	size_t getSendQSize() const { return sendq.size(); }
+	size_t getSendQSize() const;
 
 	SendQueue& GetSendQ() { return sendq; }
 
@@ -376,5 +386,4 @@ class CoreExport BufferedSocket : public StreamSocket
 };
 
 inline IOHook* StreamSocket::GetIOHook() const { return iohook; }
-inline void StreamSocket::AddIOHook(IOHook* hook) { iohook = hook; }
 inline void StreamSocket::DelIOHook() { iohook = NULL; }
diff --git a/include/iohook.h b/include/iohook.h
index 576307963..f23648658 100644
--- a/include/iohook.h
+++ b/include/iohook.h
@@ -23,6 +23,8 @@ class StreamSocket;
 
 class IOHookProvider : public ServiceProvider
 {
+ 	const bool middlehook;
+
  public:
 	enum Type
 	{
@@ -32,8 +34,14 @@ class IOHookProvider : public ServiceProvider
 
 	const Type type;
 
-	IOHookProvider(Module* mod, const std::string& Name, Type hooktype = IOH_UNKNOWN)
-		: ServiceProvider(mod, Name, SERVICE_IOHOOK), type(hooktype) { }
+	IOHookProvider(Module* mod, const std::string& Name, Type hooktype = IOH_UNKNOWN, bool middle = false)
+		: ServiceProvider(mod, Name, SERVICE_IOHOOK), middlehook(middle), type(hooktype) { }
+
+	/** Check if the IOHook provided can appear in the non-last position of a hook chain.
+	 * That is the case if and only if the IOHook instances created are subclasses of IOHookMiddle.
+	 * @return True if the IOHooks provided are subclasses of IOHookMiddle
+	 */
+	bool IsMiddle() const { return middlehook; }
 
 	/** Called immediately after a connection is accepted. This is intended for raw socket
 	 * processing (e.g. modules which wrap the tcp connection within another library) and provides
@@ -87,3 +95,67 @@ class IOHook : public classbase
 	 */
 	virtual int OnStreamSocketRead(StreamSocket* sock, std::string& recvq) = 0;
 };
+
+class IOHookMiddle : public IOHook
+{
+	/** Data already processed by the IOHook waiting to go down the chain
+	 */
+	StreamSocket::SendQueue sendq;
+
+	/** Data waiting to go up the chain
+	 */
+	std::string precvq;
+
+	/** Next IOHook in the chain
+	 */
+	IOHook* nexthook;
+
+ protected:
+	/** Get all queued up data which has not yet been passed up the hook chain
+	 * @return RecvQ containing the data
+	 */
+	std::string& GetRecvQ() { return precvq; }
+
+	/** Get all queued up data which is ready to go down the hook chain
+	 * @return SendQueue containing all data waiting to go down the hook chain
+	 */
+	StreamSocket::SendQueue& GetSendQ() { return sendq; }
+
+ public:
+	/** Constructor
+	 * @param provider IOHookProvider that creates this object
+	 */
+	IOHookMiddle(IOHookProvider* provider)
+		: IOHook(provider)
+		, nexthook(NULL)
+	{
+	}
+
+	/** Get all queued up data which is ready to go down the hook chain
+	 * @return SendQueue containing all data waiting to go down the hook chain
+	 */
+	const StreamSocket::SendQueue& GetSendQ() const { return sendq; }
+
+	/** Get the next IOHook in the chain
+	 * @return Next hook in the chain or NULL if this is the last hook
+	 */
+	IOHook* GetNextHook() const { return nexthook; }
+
+	/** Set the next hook in the chain
+	 * @param hook Hook to set as the next hook in the chain
+	 */
+	void SetNextHook(IOHook* hook) { nexthook = hook; }
+
+	/** Check if a hook is capable of being the non-last hook in a hook chain and if so, cast it to an IOHookMiddle object.
+	 * @param hook IOHook to check
+	 * @return IOHookMiddle referring to the same hook or NULL
+	 */
+	static IOHookMiddle* ToMiddleHook(IOHook* hook)
+	{
+		if (hook->prov->IsMiddle())
+			return static_cast<IOHookMiddle*>(hook);
+		return NULL;
+	}
+
+	friend class StreamSocket;
+};