From 8b5d6869bed4f592d633fc7891160df5d316a61c Mon Sep 17 00:00:00 2001
From: jamesbunton <jamesbunton@55fbd22a-6204-0410-b2f0-b6c764c7e90a>
Date: Sun, 5 Feb 2006 06:54:54 +0000
Subject: [PATCH] Partial rate-limiting. Need to try and throttle down
 receiving of data too.

git-svn-id: http://delx.cjb.net/svn/pymsnt/trunk@118 55fbd22a-6204-0410-b2f0-b6c764c7e90a

committer: jamesbunton <jamesbunton@55fbd22a-6204-0410-b2f0-b6c764c7e90a>
---
 config-example.xml   |  2 ++
 src/config.py        |  1 +
 src/ft.py            | 28 ++++++++++++++++++++++------
 src/tlib/throttle.py | 35 +++++++++++++++++++++++++++++++++++
 4 files changed, 60 insertions(+), 6 deletions(-)
 create mode 100644 src/tlib/throttle.py

diff --git a/config-example.xml b/config-example.xml
index 43cd05c..d1cd43e 100644
--- a/config-example.xml
+++ b/config-example.xml
@@ -53,6 +53,8 @@ Do not include the jid of the transport -->
 <!-- File transfer settings -->
 <!-- The maximum size of a file transfer (in bytes). For unlimited, comment out, or set to 0 -->
 <ftSizeLimit>524288</ftSizeLimit>
+<!-- The maximum rate for file transfer (in bytes). For unlimited, comment out, or set to 0 -->
+<ftRateLimit>2048</ftRateLimit>
 <!-- Please give the port to listen for Jabber socks5 transfers on. -->
 <ftJabberPort>8010</ftJabberPort>
 <!-- Please give the port to listen for HTTP GETs here (Used in OOB file transfers). -->
diff --git a/src/config.py b/src/config.py
index 5e5f625..bd513d0 100644
--- a/src/config.py
+++ b/src/config.py
@@ -24,6 +24,7 @@ ftJabberPort = ""
 ftOOBPort = ""
 ftOOBRoot = "http://" + host + "/"
 ftSizeLimit = "0"
+ftRateLimit = "0"
 
 admins = []
 
diff --git a/src/ft.py b/src/ft.py
index a916101..86fd3d1 100644
--- a/src/ft.py
+++ b/src/ft.py
@@ -1,6 +1,7 @@
 # Copyright 2005 James Bunton <james@delx.cjb.net>
 # Licensed for distribution under the GPL version 2, check COPYING for details
 
+from tlib.throttle import Throttler
 from tlib.xmlw import Element
 from twisted.internet import protocol
 
@@ -14,6 +15,17 @@ import random
 import sys
 
 
+def doRateLimit(setConsumer, consumer):
+	try:
+		rateLimit = int(config.ftRateLimit)
+	except ValueError:
+		rateLimit = 0
+	if rateLimit > 0:
+		throttler = Throttler(consumer, rateLimit)
+		setConsumer(throttler)
+	else:
+		setConsumer(consumer)
+
 def checkSizeOk(size):
 	try:
 		size = int(size)
@@ -44,11 +56,15 @@ class FTSend:
 		session.legacycon.sendFile(to, self)
 	
 	def accept(self, legacyFileSend):
-		self.startTransfer(legacyFileSend)
+		doRateLimit(self.startTransfer, legacyFileSend)
+		self.cleanup()
 	
 	def reject(self):
-		del self.startTransfer
 		self.cancelTransfer()
+		self.cleanup()
+	
+	def cleanup(self):
+		del self.startTransfer, self.cancelTransfer
 
 
 try:
@@ -380,17 +396,17 @@ class Proxy65(protocol.Factory):
 			assert address not in self.activeConns
 			self.activeConns[address] = None
 			
-			if not isinstance(olist[0], (JEP65ConnectionReceive, JEP65ConnectionSend)):
+			if not isinstance(olist[0], JEP65ConnectionReceive):
 				legacyftp = olist[0]
 				connection = olist[1]
-			elif not isinstance(olist[1], (JEP65ConnectionReceive, JEP65ConnectionSend)):
+			elif not isinstance(olist[1], JEP65ConnectionReceive):
 				legacyftp = olist[1]
 				connection = olist[0]
 			else:
 				LogEvent(WARN, '', "No JEP65Connection")
 				return
 
-			legacyftp.accept(connection.transport)
+			doRateLimit(legacyftp.accept, connection.transport)
 		else:
 			LogEvent(WARN, '', "No pending connection.")
 	
@@ -426,7 +442,7 @@ from debug import LogEvent, INFO, WARN, ERROR
 class OOBReceiveConnector:
 	def __init__(self, ftReceive, ftHttpPush):
 		self.ftReceive, self.ftHttpPush = ftReceive, ftHttpPush
-		self.ftReceive.legacyftp.accept(self)
+		doRateLimit(self.ftReceive.legacyftp.accept, self)
 	
 	def write(self, data):
 		self.ftHttpPush.write(data)
diff --git a/src/tlib/throttle.py b/src/tlib/throttle.py
new file mode 100644
index 0000000..65df425
--- /dev/null
+++ b/src/tlib/throttle.py
@@ -0,0 +1,35 @@
+from twisted.internet import task
+
+
+
+class Throttler:
+	def __init__(self, consumer, speed):
+		self.consumer = consumer
+
+		self.buffer = ""
+		self.speed = speed # Bytes/second
+		self.canClose = False
+
+		self.loopTask = task.LoopingCall(self.loopFunc)
+		self.loopTask.start(1.0)
+	
+	def write(self, data):
+		if not self.consumer:
+			raise ValueError, "I/O operation on closed 'file'"
+		self.buffer += data
+	
+	def close(self):
+		self.canClose = True
+	
+	def loopFunc(self):
+		if self.canClose and len(self.buffer) == 0:
+			self.consumer.close()
+			self.consumer = None
+			self.loopTask.stop()
+
+		else:
+			chunk, self.buffer = self.buffer[:self.speed], self.buffer[self.speed:]
+			self.consumer.write(chunk)
+
+
+
-- 
2.39.2