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>

diff --git a/config-example.xml b/config-example.xml
index 43cd05c612f9157411dd5a445579441c5b243bd8..d1cd43eea6d4b3af7dce43e2d43dbc18a09116db 100644 (file)
--- 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 5e5f625c735226abe50b89297bb8e5f2c02f8641..bd513d0a1dbe7db30b0533a2d67166c769dc25a8 100644 (file)
--- 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 a91610180401613bf7cf31a1b8ef11a1c0c9540f..86fd3d1c6a7d0beeff0562abb332acf6c22c7c28 100644 (file)
--- 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 (file)
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)
+
+
+