diff --git a/tftp/bootstrap.py b/tftp/bootstrap.py index e34f436..57e061e 100644 --- a/tftp/bootstrap.py +++ b/tftp/bootstrap.py @@ -50,10 +50,13 @@ class TFTPBootstrap(DatagramProtocol): @ivar backend: L{IReader} or L{IWriter} provider, that is used for this transfer @type backend: L{IReader} or L{IWriter} provider + @ivar max_blocksize: the maximum block size allowed for this transfer. Used to + cap the negotiated block size if a higher value is requested by the peer. + @type max_blocksize: C{int} """ supported_options = (b'blksize', b'timeout', b'tsize') - def __init__(self, remote, backend, options=None, _clock=None): + def __init__(self, remote, backend, options=None, max_blocksize=MAX_BLOCK_SIZE, _clock=None): if options is None: self.options = OrderedDict() else: @@ -62,6 +65,7 @@ def __init__(self, remote, backend, options=None, _clock=None): self.remote = remote self.timeout_watchdog = succeed(None) self.backend = backend + self.max_blocksize = max_blocksize if _clock is not None: self._clock = _clock else: @@ -91,7 +95,7 @@ def processOptions(self, options): def option_blksize(self, val): """Process the block size option. Valid range is between 8 and 65464, - inclusive. If the value is more, than L{MAX_BLOCK_SIZE}, L{MAX_BLOCK_SIZE} + inclusive. If the value is more, than `self.max_blocksize`, `self.max_blocksize` is returned instead. @param val: value of the option @@ -107,7 +111,7 @@ def option_blksize(self, val): return None if int_blksize < 8 or int_blksize > 65464: return None - int_blksize = min((int_blksize, MAX_BLOCK_SIZE)) + int_blksize = min((int_blksize, self.max_blocksize)) return intToBytes(int_blksize) def option_timeout(self, val): @@ -219,8 +223,8 @@ class LocalOriginWriteSession(TFTPBootstrap): a read from a remote server """ - def __init__(self, remote, writer, options=None, _clock=None): - TFTPBootstrap.__init__(self, remote, writer, options, _clock) + def __init__(self, remote, writer, options=None, max_blocksize=MAX_BLOCK_SIZE, _clock=None): + TFTPBootstrap.__init__(self, remote, writer, options, max_blocksize, _clock=_clock) self.session = WriteSession(writer, self._clock) def startProtocol(self): @@ -263,8 +267,8 @@ class RemoteOriginWriteSession(TFTPBootstrap): """ timeout = (1, 3, 7) - def __init__(self, remote, writer, options=None, _clock=None): - TFTPBootstrap.__init__(self, remote, writer, options, _clock) + def __init__(self, remote, writer, options=None, max_blocksize=MAX_BLOCK_SIZE, _clock=None): + TFTPBootstrap.__init__(self, remote, writer, options, max_blocksize, _clock=_clock) self.session = WriteSession(writer, self._clock) def startProtocol(self): @@ -299,8 +303,8 @@ class LocalOriginReadSession(TFTPBootstrap): a write to a remote server. """ - def __init__(self, remote, reader, options=None, _clock=None): - TFTPBootstrap.__init__(self, remote, reader, options, _clock) + def __init__(self, remote, reader, options=None, max_blocksize=MAX_BLOCK_SIZE, _clock=None): + TFTPBootstrap.__init__(self, remote, reader, options, max_blocksize, _clock=_clock) self.session = ReadSession(reader, self._clock) def startProtocol(self): @@ -344,8 +348,8 @@ class RemoteOriginReadSession(TFTPBootstrap): """ timeout = (1, 3, 7) - def __init__(self, remote, reader, options=None, _clock=None): - TFTPBootstrap.__init__(self, remote, reader, options, _clock) + def __init__(self, remote, reader, options=None, max_blocksize=MAX_BLOCK_SIZE, _clock=None): + TFTPBootstrap.__init__(self, remote, reader, options, max_blocksize, _clock=_clock) self.session = ReadSession(reader, self._clock) def option_tsize(self, val): diff --git a/tftp/protocol.py b/tftp/protocol.py index e03d574..fa54487 100644 --- a/tftp/protocol.py +++ b/tftp/protocol.py @@ -14,6 +14,8 @@ from twisted.python import log from twisted.python.context import call +from tftp.session import MAX_BLOCK_SIZE + class TFTP(DatagramProtocol): """TFTP dispatch protocol. Handles read requests (RRQ) and write requests (WRQ) @@ -23,9 +25,13 @@ class TFTP(DatagramProtocol): local resources @type backend: L{IBackend} provider + @ivar max_blocksize: the maximum block size allowed for this transfer. Used to + cap the negotiated block size if a higher value is requested by the peer. + @type max_blocksize: C{int} """ - def __init__(self, backend, _clock=None): + def __init__(self, backend, max_blocksize=MAX_BLOCK_SIZE, _clock=None): self.backend = backend + self.max_blocksize = max_blocksize if _clock is None: self._clock = reactor else: @@ -90,13 +96,13 @@ def _startSession(self, datagram, addr, mode): if mode == b'netascii': fs_interface = NetasciiReceiverProxy(fs_interface) session = RemoteOriginWriteSession(addr, fs_interface, - datagram.options, _clock=self._clock) + datagram.options, self.max_blocksize, _clock=self._clock) reactor.listenUDP(0, session) returnValue(session) elif datagram.opcode == OP_RRQ: if mode == b'netascii': fs_interface = NetasciiSenderProxy(fs_interface) session = RemoteOriginReadSession(addr, fs_interface, - datagram.options, _clock=self._clock) + datagram.options, self.max_blocksize, _clock=self._clock) reactor.listenUDP(0, session) returnValue(session) diff --git a/tftp/test/test_bootstrap.py b/tftp/test/test_bootstrap.py index 739a925..190f9ec 100644 --- a/tftp/test/test_bootstrap.py +++ b/tftp/test/test_bootstrap.py @@ -157,6 +157,17 @@ def test_multiple_options(self): actual_options[b'timeout'] = b'123' self.assertEqual(list(opts.items()), list(actual_options.items())) +class TestMaxBlocksizeProcessing(unittest.TestCase): + + def setUp(self): + self.proto = TFTPBootstrap(('127.0.0.1', 1111), None, max_blocksize=1400) + + def test_blksize(self): + self.s = MockSession() + opts = self.proto.processOptions(OrderedDict({b'blksize':b'65464'})) + self.proto.applyOptions(self.s, opts) + self.assertEqual(self.s.block_size, 1400) + self.assertEqual(opts, OrderedDict({b'blksize':b'1400'})) class BootstrapLocalOriginWrite(unittest.TestCase): diff --git a/tftp/test/test_protocol.py b/tftp/test/test_protocol.py index 8367ec9..917af16 100644 --- a/tftp/test/test_protocol.py +++ b/tftp/test/test_protocol.py @@ -237,7 +237,7 @@ def setUp(self): with self.temp_dir.child(b'nonempty').open('w') as fd: fd.write(b'Something uninteresting') self.backend = FilesystemAsyncBackend(self.temp_dir, self.clock) - self.tftp = TFTP(self.backend, self.clock) + self.tftp = TFTP(self.backend, _clock=self.clock) def test_get_reader_defers(self): rrq_datagram = RRQDatagram(b'nonempty', b'NetASCiI', {}) diff --git a/twisted/plugins/tftp_plugin.py b/twisted/plugins/tftp_plugin.py index 8ff7d3f..0fcc371 100644 --- a/twisted/plugins/tftp_plugin.py +++ b/twisted/plugins/tftp_plugin.py @@ -22,7 +22,8 @@ class TFTPOptions(usage.Options): ] optParameters = [ ['port', 'p', 1069, 'Port number to listen on.', int], - ['root-directory', 'd', None, 'Root directory for this server.', to_path] + ['root-directory', 'd', None, 'Root directory for this server.', to_path], + ['max_blocksize', 'b', 8192, 'Max blocksize for the FTFP datagram.', int] ] def postOptions(self): @@ -40,6 +41,6 @@ def makeService(self, options): backend = FilesystemSynchronousBackend(options["root-directory"], can_read=options['enable-reading'], can_write=options['enable-writing']) - return internet.UDPServer(options['port'], TFTP(backend)) + return internet.UDPServer(options['port'], TFTP(backend, options['max_blocksize'])) serviceMaker = TFTPServiceCreator()