Ticket #6839: edns-protocols-6839.diff

File edns-protocols-6839.diff, 10.2 KB (added by Richard Wall, 2 years ago)

Trying to reacquaint myself with the remaining EDNS tasks

  • new file twisted/names/edns.py

    diff --git twisted/names/edns.py twisted/names/edns.py
    new file mode 100644
    index 0000000..9dd4085
    - +  
     1# -*- test-case-name: twisted.names.test.test_dns -*-
     2# Copyright (c) Twisted Matrix Laboratories.
     3# See LICENSE for details.
     4
     5"""
     6EDNS protocol implementation.
     7"""
     8import struct
     9from twisted.internet.defer import succeed, Deferred
     10from twisted.internet.interfaces import IProtocol
     11from twisted.internet.endpoints import TCP4ClientEndpoint, connectProtocol
     12from twisted.internet.protocol import Protocol
     13from twisted.python.components import proxyForInterface
     14from twisted.names.dns import _EDNSMessage, Query
     15
     16import attr
     17from zope.interface import implementer, Interface
     18
     19
     20class EDNSMessage(_EDNSMessage):
     21    """
     22    """
     23
     24
     25class IEDNSMessage(Interface):
     26    """
     27    """
     28
     29
     30class IEDNSMessageDecoder(Interface):
     31    """
     32    """
     33    def decode(message_bytes):
     34        """
     35        :param bytes message_bytes: The bytes of an encoded message.
     36        :rtype: IEDNSMessage
     37        """
     38
     39
     40class IEDNSMessageEncoder(Interface):
     41    """
     42    """
     43    def encode(message):
     44        """
     45        :param bytes message: An IEDNSMessage.
     46        :rtype: bytes of the encoded message.
     47        """
     48
     49
     50class IEDNSTransmitter(Interface):
     51    """
     52    """
     53    def transmit(message):
     54        """
     55        :param IEDNSMessage message: A message to transmit.
     56        :returns: A Deferred firing with the IEDNSMessage response for ``message``.
     57        """
     58
     59    def for_destination(reactor, destination):
     60        """
     61        """
     62
     63
     64class IEDNSResponseHandler(Interface):
     65    """
     66    """
     67    def expect(message):
     68        """
     69        :param IEDNSMessage message: A message expected.
     70        """
     71
     72    def receive(message):
     73        """
     74        :param IEDNSMessage message: A message received.
     75        """
     76
     77
     78class IEDNSResponder(Interface):
     79    """
     80    """
     81    def respond(message):
     82        """
     83        :param IEDNSMessage message: The received message.
     84        :returns: A Deferred that fires with IEDNSMessage.
     85        """
     86
     87
     88class IEDNSClient(Interface):
     89    """
     90    """
     91    def query():
     92        """
     93
     94        :returns: A Deferred that fires with IEDNSMessage
     95        """
     96
     97
     98@implementer(IProtocol)
     99class EDNSStreamProtocol(Protocol, object):
     100    """
     101    """
     102    buffer = b''
     103    length = None
     104
     105    def __init__(self, receiver):
     106        self.receiver = receiver
     107
     108    def dataReceived(self, data):
     109        self.buffer += data
     110
     111        while self.buffer:
     112            if self.length is None and len(self.buffer) >= 2:
     113                self.length = struct.unpack('!H', self.buffer[:2])[0]
     114                self.buffer = self.buffer[2:]
     115
     116            if len(self.buffer) >= self.length:
     117                chunk = self.buffer[:self.length]
     118                m = EDNSMessage()
     119                m.fromStr(chunk)
     120                self.receiver(m)
     121                self.buffer = self.buffer[self.length:]
     122                self.length = None
     123            else:
     124                break
     125
     126
     127@implementer(IEDNSTransmitter)
     128@attr.s
     129class EDNSDatagramTransmitter(object):
     130    reactor = attr.ib()
     131    receiver = attr.ib()
     132    destination = attr.ib()
     133
     134    def transmit(self, message):
     135        """
     136        """
     137        proto = EDNSDatagramProtocol(receiver=self._receiver.receive)
     138        self.reactor.listenUDP(0, proto)
     139        proto.transport.write(message.toStr(), self.destination)
     140        d = self.receiver.expect(message)
     141        return d
     142
     143    @classmethod
     144    def for_destination(cls, reactor, destination):
     145        """
     146        """
     147        receiver = EDNSResponseHandler()
     148        return succeed(cls(reactor, destination, receiver))
     149
     150
     151@implementer(IEDNSTransmitter)
     152@attr.s
     153class EDNSStreamTransmitter(object):
     154    reactor = attr.ib()
     155    receiver = attr.ib()
     156    protocol = attr.ib()
     157
     158    def transmit(self, message):
     159        """
     160        """
     161        d = self.receiver.expect(message)
     162        s = message.toStr()
     163        self.protocol.transport.write(struct.pack('!H', len(s)) + s)
     164        return d
     165
     166    @classmethod
     167    def for_destination(cls, reactor, destination):
     168        """
     169        """
     170        receiver = EDNSResponseHandler()
     171        client_endpoint = TCP4ClientEndpoint(
     172            reactor=reactor,
     173            host=destination.host,
     174            port=destination.port,
     175            timeout=1,
     176        )
     177        d = connectProtocol(client_endpoint, EDNSStreamProtocol(receiver=receiver.receive))
     178        d.addCallback(
     179            lambda protocol: cls(
     180                reactor=reactor,
     181                receiver=receiver,
     182                protocol=protocol
     183            )
     184        )
     185        return d
     186
     187
     188
     189@implementer(IEDNSResponseHandler)
     190class EDNSResponseHandler(object):
     191    def __init__(self):
     192        self._expected = {}
     193
     194    def expect(self, message):
     195        d = Deferred()
     196        self._expected[message.id] = d
     197        return d
     198
     199    def receive(self, message):
     200        try:
     201            d = self._expected.pop(message.id)
     202        except KeyError as e:
     203            raise
     204            # log.err("unexpected message", message)
     205        else:
     206            d.callback(message)
     207
     208
     209def query():
     210    return EDNSMessage(
     211        id=1234,
     212        recDes=True,
     213        queries=[
     214            Query(name=b"www.example.com")
     215        ]
     216    )
     217
     218
     219@implementer(IEDNSClient)
     220@attr.s
     221class EDNSClient(object):
     222    """
     223    """
     224    transmitter = attr.ib()
     225
     226    def query(self, **kwargs):
     227        """
     228        """
     229        message = query(**kwargs)
     230        d = self.transmitter.transmit(message)
     231
     232        def examine_response(message):
     233            import pdb; pdb.set_trace()
     234            if message.errors:
     235                return failure(message)
     236            else:
     237                return message
     238        d.addCallback(examine_response)
     239
     240        return d
     241
     242
     243def dig():
     244    pass
     245
     246
     247class EDNSServerFactory(object):
     248    """
     249    """
  • new file twisted/names/test/test_edns.py

    diff --git twisted/names/test/test_edns.py twisted/names/test/test_edns.py
    new file mode 100644
    index 0000000..efd156f
    - +  
     1# test-case-name: twisted.names.test.test_dns
     2# Copyright (c) Twisted Matrix Laboratories.
     3# See LICENSE for details.
     4
     5"""
     6Tests for twisted.names.edns.
     7"""
     8
     9from twisted.internet import reactor
     10from twisted.internet.address import IPv4Address
     11from twisted.internet.endpoints import (
     12    TCP4ServerEndpoint, TCP4ClientEndpoint, connectProtocol
     13)
     14from twisted.internet.protocol import ServerFactory
     15from twisted.trial.unittest import TestCase, SynchronousTestCase
     16
     17from zope.interface.verify import verifyObject
     18
     19from twisted.names.dns import (
     20    _EDNSMessage, DNSProtocol
     21)
     22from twisted.names.edns import (
     23    EDNSStreamProtocol, EDNSClient, EDNSServerFactory,
     24    IEDNSClient, EDNSDatagramTransmitter, EDNSStreamTransmitter,
     25)
     26
     27
     28def server_for_test(test, reactor, server_factory):
     29    server_starting = TCP4ServerEndpoint(
     30        reactor=reactor,
     31        port=0,
     32        interface=u"127.0.0.1"
     33    ).listen(
     34        protocolFactory=server_factory
     35    )
     36
     37    def started(listening_port):
     38        test.addCleanup(listening_port.stopListening)
     39        return listening_port.getHost()
     40    return server_starting.addCallback(started)
     41
     42
     43def client_for_test(test, reactor, client_protocol, server_address):
     44    def connect(server_address):
     45        client_endpoint = TCP4ClientEndpoint(
     46            reactor=reactor,
     47            host=server_address.host,
     48            port=server_address.port,
     49            timeout=1,
     50        )
     51        return connectProtocol(client_endpoint, client_protocol())
     52    connecting = connect(server_address)
     53
     54    def connected(client_protocol):
     55        test.addCleanup(client_protocol.transport.loseConnection)
     56        return client_protocol
     57    return connecting.addCallback(connected)
     58
     59
     60class EDNSRoundTripTests(TestCase):
     61    """
     62    Tests client server interaction.
     63    """
     64    def test_roundtrip(self):
     65        """
     66        """
     67        class Controller(object):
     68            def connectionMade(self, protocol):
     69                pass
     70
     71            def connectionLost(self, protocol):
     72                pass
     73
     74        starting = server_for_test(
     75            test=self,
     76            reactor=reactor,
     77            server_factory=ServerFactory.forProtocol(
     78                lambda: EDNSStreamProtocol(
     79                    receiver=lambda message: None
     80                )
     81            )
     82        )
     83
     84        def started(server_address):
     85            return client_for_test(
     86                test=self,
     87                reactor=reactor,
     88                client_protocol=lambda: DNSProtocol(
     89                    controller=Controller(),
     90                    reactor=reactor,
     91                ),
     92                server_address=server_address,
     93            )
     94        connecting = starting.addCallback(started)
     95
     96        def query(dns_protocol):
     97            dns_protocol.transport.write(b'x')
     98        return connecting.addCallback(query)
     99
     100    test_roundtrip.timeout = 2
     101
     102
     103class IEDNSClientTestsMixin(object):
     104    """
     105    """
     106    def test_interface(self):
     107        """
     108        The ``client`` provides ``IEDNSClient``.
     109        """
     110        self.assertTrue(
     111            verifyObject(IEDNSClient, self.client)
     112        )
     113    test_interface.timeout = 2
     114
     115    def test_happy(self):
     116        """
     117        ``query`` returns a deferred that fires with the ``IEDNSMessage``
     118        """
     119        d = self.client.query()
     120        return d
     121    test_happy.timeout = 2
     122
     123
     124def make_iedns_client_tests(edns_client):
     125    class Tests(IEDNSClientTestsMixin, TestCase):
     126        def setUp(self):
     127            d = edns_client()
     128            d.addCallback(lambda client: setattr(self, 'client', client))
     129            return d
     130    return Tests
     131
     132
     133class EDNSClientInterfaceTests(
     134        make_iedns_client_tests(
     135            edns_client=lambda: EDNSStreamTransmitter.for_destination(
     136                    reactor=reactor,
     137                    destination=IPv4Address('TCP', '8.8.8.8', 53)
     138                ).addCallback(lambda transmitter: EDNSClient(
     139                    transmitter=transmitter
     140                )
     141            )
     142        )
     143):
     144    """
     145    """