diff --git a/lib/tfw/networking/server/controller_responder.py b/lib/tfw/networking/server/controller_responder.py index 32567d7..a0924e7 100644 --- a/lib/tfw/networking/server/controller_responder.py +++ b/lib/tfw/networking/server/controller_responder.py @@ -5,21 +5,28 @@ from tfw.networking.serialization import deserialize_all, serialize_all class ControllerResponder: def __init__(self, fsm): self.fsm = fsm + self.token = None self.controller_connector = ControllerConnector() self.controller_connector.register_callback(self.handle_controller_request) self.controller_request_handlers = { 'solution_check': self.handle_solution_check_request, - 'test': self.handle_test_request + 'test': self.handle_test_request, + 'token': self.handle_token_request } def handle_controller_request(self, stream, msg_parts): key, data = deserialize_all(*msg_parts) response = self.controller_request_handlers[key](data) - stream.send_multipart(serialize_all(key, response)) + stream.send_multipart(serialize_all(self.token, response)) def handle_test_request(self, data): return 'OK' + def handle_token_request(self, data): + if self.token is None: + self.token = data + return {'token': self.token} + def handle_solution_check_request(self, data): return { 'solved': self.fsm.is_solved(), diff --git a/src/controller/app.py b/src/controller/app.py index 67b2faa..a6ae9b6 100644 --- a/src/controller/app.py +++ b/src/controller/app.py @@ -1,3 +1,4 @@ +import secrets from tornado.ioloop import IOLoop from tornado.web import Application @@ -7,13 +8,24 @@ from handlers import SolutionCheckHandler, TestHandler from tfw.networking.async_solvable_connector import AsyncSolvableConnector from tfw.config.logs import logging + log = logging.getLogger(__name__) + +async def setup_token(solvable_connector, token): + solvable_connector.send('token', token) + key, data = await solvable_connector.recv() + log.debug('Token setup is done.') + if __name__ == '__main__': - solvable_connector = AsyncSolvableConnector() + kwargs = { + 'solvable_connector': AsyncSolvableConnector(), + 'token': secrets.token_hex(32) + } + IOLoop.instance().add_callback(setup_token, **kwargs) routes = [ - (r'/{secret}/?'.format(secret=taoenv.SECRET), SolutionCheckHandler, {'solvable_connector': solvable_connector}), - (r'/{secret}/test/?'.format(secret=taoenv.SECRET), TestHandler, {'solvable_connector': solvable_connector}) + (r'/{secret}/?'.format(secret=taoenv.SECRET), SolutionCheckHandler, kwargs), + (r'/{secret}/test/?'.format(secret=taoenv.SECRET), TestHandler, kwargs) ] app = Application( routes diff --git a/src/controller/handlers/solution_check_handler.py b/src/controller/handlers/solution_check_handler.py index cfe2c1c..e84dc80 100644 --- a/src/controller/handlers/solution_check_handler.py +++ b/src/controller/handlers/solution_check_handler.py @@ -1,15 +1,19 @@ -from tornado.web import RequestHandler +import secrets +from tornado.web import RequestHandler, HTTPError from tfw.config.logs import logging log = logging.getLogger(__name__) class SolutionCheckHandler(RequestHandler): - def initialize(self, solvable_connector): + def initialize(self, solvable_connector, token): self.solvable_connector = solvable_connector + self.token = token async def get(self): log.debug('Sending request to solvable') self.solvable_connector.send('solution_check', {}) - resp_key, resp_data = await self.solvable_connector.recv() + resp_token, resp_data = await self.solvable_connector.recv() + if not secrets.compare_digest(self.token, resp_token): + raise HTTPError(500, 'Solvable didn\'t provide initial token.') log.debug('Received answer from solvable') self.write(resp_data) diff --git a/src/controller/handlers/test_handler.py b/src/controller/handlers/test_handler.py index 564bcdc..85c675c 100644 --- a/src/controller/handlers/test_handler.py +++ b/src/controller/handlers/test_handler.py @@ -1,11 +1,15 @@ -from tornado.web import RequestHandler +import secrets +from tornado.web import RequestHandler, HTTPError class TestHandler(RequestHandler): - def initialize(self, solvable_connector): + def initialize(self, solvable_connector, token): self.solvable_connector = solvable_connector + self.token = token async def get(self): self.solvable_connector.send('test', {}) - resp_key, resp_data = await self.solvable_connector.recv() + resp_token, resp_data = await self.solvable_connector.recv() + if not secrets.compare_digest(self.token, resp_token): + raise HTTPError(500, 'Solvable didn\'t provide initial token.') self.write(resp_data)