def test_decode_data_frame(self):
"""decode data frame"""
f = fstrm.FstrmCodec()
data_frame = f.encode(ctrl=fstrm.FSTRM_DATA_FRAME, payload=b"hello")
f.append(data=data_frame)
ft = -1
if f.process():
ft, ct, payload = f.decode()
self.assertEqual(ft, fstrm.FSTRM_DATA_FRAME)
def test_decode_control_stop(self):
"""decode control stop"""
f = fstrm.FstrmCodec()
data_bin = f.encode(ctrl=fstrm.FSTRM_CONTROL_STOP)
f.append(data=data_bin)
ft = -1
if f.process():
ft, ct, payload = f.decode()
self.assertEqual(ft, fstrm.FSTRM_CONTROL_STOP)
def test_control_finish(self):
"""decode control finish"""
f = fstrm.FstrmCodec()
data_bin = f.encode(ctrl=fstrm.FSTRM_CONTROL_FINISH)
f.append(data=data_bin)
ft = -1
if f.process():
ft, ct, payload = f.decode()
self.assertEqual(ft, fstrm.FSTRM_CONTROL_FINISH)
def test_control_ready(self):
"""decode control ready"""
f = fstrm.FstrmCodec()
data_bin = f.encode(ctrl=fstrm.FSTRM_CONTROL_READY,
ct=[b"protobuf:dnstap.Dnstap"])
f.append(data=data_bin)
ft = -1
if f.process():
ft, ct, payload = f.decode()
self.assertEqual(ft, fstrm.FSTRM_CONTROL_READY)
async def cb_onconnect(reader, writer, cfg, cfg_input, queues_list, stats,
geoip_reader, cache, start_shutdown):
"""callback when a connection is established"""
# get peer name
peername = writer.get_extra_info('peername')
if not len(peername):
peername = "(unix-socket)"
clogger.debug(f"Input handler: new connection from {peername}")
# access control list check
if "access-control-list" in cfg_input:
if len(writer.get_extra_info('peername')):
acls_network = []
for a in cfg_input["access-control-list"]:
acls_network.append(ipaddress.ip_network(a))
acl_allow = False
for acl in acls_network:
if ipaddress.ip_address(peername[0]) in acl:
acl_allow = True
if not acl_allow:
writer.close()
clogger.debug("Input handler: checking acl refused")
return
clogger.debug("Input handler: checking acl allowed")
# prepare frame streams decoder
content_type = b"protobuf:dnstap.Dnstap"
fstrm_handler = fstrm.FstrmCodec()
loop = asyncio.get_event_loop()
dnstap_protobuf = dnstap_pb.Dnstap()
try:
# syntax only works with python 3.8
# while data := await reader.read(fstrm_handler.pending_nb_bytes())
running = True
while running:
# read bytes
shutdown_wait_task = asyncio.create_task(start_shutdown.wait())
read_task = asyncio.create_task(
reader.read(fstrm_handler.pending_nb_bytes()))
done, pending = await asyncio.wait(
[shutdown_wait_task, read_task],
return_when=asyncio.FIRST_COMPLETED,
)
if shutdown_wait_task in done:
read_task.cancel()
return
else:
shutdown_wait_task.cancel()
data = await read_task
if not len(data):
running = False
break
# append data to the buffer
fstrm_handler.append(data=data)
# process the buffer, check if we have received a complete frame ?
if fstrm_handler.process():
# Ok, the frame is complete so let's decode it
ctrl, ct, payload = fstrm_handler.decode()
# handle the DATA frame
if ctrl == fstrm.FSTRM_DATA_FRAME:
await dnstap_decoder.cb_ondnstap(dnstap_protobuf, payload,
cfg, queues_list, stats,
geoip_reader, cache)
# handle the control frame READY
if ctrl == fstrm.FSTRM_CONTROL_READY:
clogger.debug(
f"Input handler: control ready received from {peername}"
)
if content_type not in ct:
raise Exception("content type error: %s" % ct)
# todo, checking content type
ctrl_accept = fstrm_handler.encode(
ctrl=fstrm.FSTRM_CONTROL_ACCEPT, ct=[content_type])
# respond with accept only if the content type is dnstap
writer.write(ctrl_accept)
await writer.drain()
clogger.debug(
f"Input handler: sending control accept to {peername}")
# handle the control frame READY
if ctrl == fstrm.FSTRM_CONTROL_START:
clogger.debug(
f"Input handler: control start received from {peername}"
)
# handle the control frame STOP
if ctrl == fstrm.FSTRM_CONTROL_STOP:
clogger.debug(
f"Input handler: control stop received from {peername}"
)
fstrm_handler.reset()
# send finish control
ctrl_finish = fstrm_handler.encode(
ctrl=fstrm.FSTRM_CONTROL_FINISH)
writer.write(ctrl_finish)
await writer.drain()
clogger.debug(
f"Input handler: sending control finish to {peername}")
except ConnectionError as e:
clogger.error(f'Input handler: {peername} - %s' % e)
except asyncio.CancelledError:
clogger.debug(f'Input handler: {peername} - closing connection.')
writer.close()
await writer.wait_closed()
except asyncio.IncompleteReadError:
clogger.debug(f'Input handler: {peername} - disconnected')
finally:
clogger.debug(f'Input handler: {peername} - closed')
async def tcp_client(cfg, cfg_input, queues_list, stats, geoip_reader, cache,
start_shutdown):
host, port = cfg_input["remote-address"], cfg_input["remote-port"]
clogger.debug("Input handler: connection to %s:%s" % (host, port))
reader, tcp_writer = await asyncio.open_connection(host, port)
clogger.debug("Input handler: connected")
content_type = b"protobuf:dnstap.Dnstap"
fstrm_handler = fstrm.FstrmCodec()
dnstap = dnstap_pb.Dnstap()
try:
# framestream - do handshake
ctrl_ready = fstrm_handler.encode(ctrl=fstrm.FSTRM_CONTROL_READY,
ct=[content_type])
tcp_writer.write(ctrl_ready)
shutdown_wait_task = asyncio.create_task(start_shutdown.wait())
while True:
read_task = asyncio.create_task(
reader.read(fstrm_handler.pending_nb_bytes()))
done, pending = await asyncio.wait(
[shutdown_wait_task, read_task],
return_when=asyncio.FIRST_COMPLETED,
)
if shutdown_wait_task in done:
read_task.cancel()
return
else:
shutdown_wait_task.cancel()
data = await read_task
if not len(data):
break
fstrm_handler.append(data=data)
# process the buffer, check if we have received a complete frame ?
if fstrm_handler.process():
# Ok, the frame is complete so let's decode it
ctrl, ct, payload = fstrm_handler.decode()
if ctrl == fstrm.FSTRM_CONTROL_ACCEPT:
ctrl_start = fstrm_handler.encode(
ctrl=fstrm.FSTRM_CONTROL_START)
tcp_writer.write(ctrl_start)
break
# waiting for incoming data
running = True
while running:
# read bytes
read_task = asyncio.create_task(
reader.read(fstrm_handler.pending_nb_bytes()))
done, pending = await asyncio.wait(
[shutdown_wait_task, read_task],
return_when=asyncio.FIRST_COMPLETED,
)
if shutdown_wait_task in done:
read_task.cancel()
return
else:
data = await read_task
if not len(data):
running = False
break
# append data to the buffer
fstrm_handler.append(data=data)
# process the buffer, check if we have received a complete frame ?
if fstrm_handler.process():
# Ok, the frame is complete so let's decode it
ctrl, ct, payload = fstrm_handler.decode()
# handle the DATA frame
if ctrl == fstrm.FSTRM_DATA_FRAME:
loop.create_task(
dnstap_decoder.cb_ondnstap(dnstap_protobuf, payload,
cfg, queues_list, stats,
geoip_reader, cache))
except ConnectionError as e:
clogger.error(f'Input handler: {peername} - %s' % e)
except asyncio.CancelledError:
clogger.debug(f'Input handler: {peername} - closing connection.')
writer.close()
except asyncio.IncompleteReadError:
clogger.debug(f'Input handler: {peername} - disconnected')
finally:
writer.close()
clogger.debug(f'Input handler: {peername} - closed')