Commit ef036b67 authored by Petr Špaček's avatar Petr Špaček

Deckard: temporarily move sendto/recvfrom methods to scenario

Further work requires importing pydnstest.scenario into testserver.
Socket functions were moved into scenario to avoid cyclic imports
between testserver and scenario modules.

We should be moving the socket functions in the other way,
from scenario into testserver. This will be subject of later
refactoring.
parent dc5fae77
...@@ -17,7 +17,6 @@ import random ...@@ -17,7 +17,6 @@ import random
import string import string
import time import time
from datetime import datetime from datetime import datetime
from pydnstest.testserver import recvfrom_msg, sendto_msg
# Global statistics # Global statistics
...@@ -84,6 +83,54 @@ def compare_sub(got, expected): ...@@ -84,6 +83,54 @@ def compare_sub(got, expected):
return True return True
def recvfrom_msg(stream, raw=False):
"""
Receive DNS message from TCP/UDP socket.
Returns:
if raw == False: (DNS message object, peer address)
if raw == True: (blob, peer address)
"""
if stream.type & socket.SOCK_DGRAM:
data, addr = stream.recvfrom(4096)
elif stream.type & socket.SOCK_STREAM:
data = stream.recv(2)
if len(data) == 0:
return None, None
msg_len = struct.unpack_from("!H", data)[0]
data = b""
received = 0
while received < msg_len:
next_chunk = stream.recv(4096)
if len(next_chunk) == 0:
return None, None
data += next_chunk
received += len(next_chunk)
addr = stream.getpeername()[0]
else:
raise NotImplementedError("[recvfrom_msg]: unknown socket type '%i'" % stream.type)
if not raw:
data = dns.message.from_wire(data, one_rr_per_rrset=True)
return data, addr
def sendto_msg(stream, message, addr=None):
""" Send DNS/UDP/TCP message. """
try:
if stream.type & socket.SOCK_DGRAM:
if addr is None:
stream.send(message)
else:
stream.sendto(message, addr)
elif stream.type & socket.SOCK_STREAM:
data = struct.pack("!H", len(message)) + message
stream.send(data)
else:
assert False, "[sendto_msg]: unknown socket type '%i'" % stream.type
except: # Failure to respond is OK, resolver should recover
pass
def replay_rrs(rrs, nqueries, destination, args=[]): def replay_rrs(rrs, nqueries, destination, args=[]):
""" Replay list of queries and report statistics. """ """ Replay list of queries and report statistics. """
navail, queries = len(rrs), [] navail, queries = len(rrs), []
......
...@@ -9,57 +9,8 @@ import time ...@@ -9,57 +9,8 @@ import time
import dns.message import dns.message
import dns.rdatatype import dns.rdatatype
import itertools import itertools
import struct
import binascii
def recvfrom_msg(stream, raw=False):
"""
Receive DNS message from TCP/UDP socket.
Returns:
if raw == False: (DNS message object, peer address)
if raw == True: (blob, peer address)
"""
if stream.type & socket.SOCK_DGRAM:
data, addr = stream.recvfrom(4096)
elif stream.type & socket.SOCK_STREAM:
data = stream.recv(2)
if len(data) == 0:
return None, None
msg_len = struct.unpack_from("!H", data)[0]
data = b""
received = 0
while received < msg_len:
next_chunk = stream.recv(4096)
if len(next_chunk) == 0:
return None, None
data += next_chunk
received += len(next_chunk)
addr = stream.getpeername()[0]
else:
raise NotImplementedError("[recvfrom_msg]: unknown socket type '%i'" % stream.type)
if not raw:
data = dns.message.from_wire(data, one_rr_per_rrset=True)
return data, addr
def sendto_msg(stream, message, addr=None):
""" Send DNS/UDP/TCP message. """
try:
if stream.type & socket.SOCK_DGRAM:
if addr is None:
stream.send(message)
else:
stream.sendto(message, addr)
elif stream.type & socket.SOCK_STREAM:
data = struct.pack("!H", len(message)) + message
stream.send(data)
else:
assert False, "[sendto_msg]: unknown socket type '%i'" % stream.type
except: # Failure to respond is OK, resolver should recover
pass
from pydnstest import scenario
def get_local_addr_str(family, iface): def get_local_addr_str(family, iface):
""" Returns pattern string for localhost address """ """ Returns pattern string for localhost address """
...@@ -179,7 +130,7 @@ class TestServer: ...@@ -179,7 +130,7 @@ class TestServer:
""" """
log = logging.getLogger('pydnstest.testserver.handle_query') log = logging.getLogger('pydnstest.testserver.handle_query')
server_addr = client.getsockname()[0] server_addr = client.getsockname()[0]
query, client_addr = recvfrom_msg(client) query, client_addr = scenario.recvfrom_msg(client)
if query is None: if query is None:
return False return False
log.debug('server %s received query from %s: %s', server_addr, client_addr, query) log.debug('server %s received query from %s: %s', server_addr, client_addr, query)
...@@ -204,7 +155,7 @@ class TestServer: ...@@ -204,7 +155,7 @@ class TestServer:
server_addr, server_addr,
'; '.join([str(rr) for rr in query.question])) '; '.join([str(rr) for rr in query.question]))
sendto_msg(client, data_to_wire, client_addr) scenario.sendto_msg(client, data_to_wire, client_addr)
return True return True
def query_io(self): def query_io(self):
......
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment