Test Failed
Pull Request — master (#271)
by
unknown
03:25
created

BinaryServer._await_closing_tasks()   A

Complexity

Conditions 3

Size

Total Lines 8
Code Lines 8

Duplication

Lines 0
Ratio 0 %

Importance

Changes 0
Metric Value
cc 3
eloc 8
nop 1
dl 0
loc 8
rs 10
c 0
b 0
f 0
1
"""
2
Socket server forwarding request to internal server
3
"""
4
import logging
5
import asyncio
6
from typing import Optional
7
8
from ..ua.ua_binary import header_from_binary
9
from ..common.utils import Buffer, NotEnoughData
10
from .uaprocessor import UaProcessor
11
from .internal_server import InternalServer
12
13
logger = logging.getLogger(__name__)
14
15
16
class OPCUAProtocol(asyncio.Protocol):
17
    """
18
    Instantiated for every connection.
19
    """
20
21
    def __init__(self, iserver: InternalServer, policies, clients, closing_tasks):
22
        self.peer_name = None
23
        self.transport = None
24
        self.processor = None
25
        self._buffer = b''
26
        self.iserver: InternalServer = iserver
27
        self.policies = policies
28
        self.clients = clients
29
        self.closing_tasks = closing_tasks
30
        self.messages = asyncio.Queue()
31
        self._task = None
32
33
    def __str__(self):
34
        return f'OPCUAProtocol({self.peer_name}, {self.processor.session})'
35
36
    __repr__ = __str__
37
38
    def connection_made(self, transport):
39
        self.peer_name = transport.get_extra_info('peername')
40
        logger.info('New connection from %s', self.peer_name)
41
        self.transport = transport
42
        self.processor = UaProcessor(self.iserver, self.transport)
43
        self.processor.set_policies(self.policies)
44
        self.iserver.asyncio_transports.append(transport)
45
        self.clients.append(self)
46
        self._task = self.iserver.loop.create_task(self._process_received_message_loop())
47
48
    def connection_lost(self, ex):
49
        logger.info('Lost connection from %s, %s', self.peer_name, ex)
50
        self.transport.close()
51
        self.iserver.asyncio_transports.remove(self.transport)
52
        closing_task = self.iserver.loop.create_task(self.processor.close())
53
        self.closing_tasks.append(closing_task)
54
        if self in self.clients:
55
            self.clients.remove(self)
56
        self.messages.put_nowait((None, None))
57
        self._task.cancel()
58
59
    def data_received(self, data):
60
        self._buffer += data
61
        # try to parse the incoming data
62
        while self._buffer:
63
            try:
64
                buf = Buffer(self._buffer)
65
                try:
66
                    header = header_from_binary(buf)
67
                except NotEnoughData:
68
                    logger.debug('Not enough data while parsing header from client, waiting for more')
69
                    return
70
                if len(buf) < header.body_size:
71
                    logger.debug('We did not receive enough data from client. Need %s got %s', header.body_size,
72
                                 len(buf))
73
                    return
74
                # we have a complete message
75
                self.messages.put_nowait((header, buf))
76
                self._buffer = self._buffer[(header.header_size + header.body_size):]
77
            except Exception:
78
                logger.exception('Exception raised while parsing message from client')
79
                return
80
81
    async def _process_received_message_loop(self):
82
        """
83
        Take message from the queue and try to process it.
84
        """
85
        while True:
86
            header, buf = await self.messages.get()
87
            if header is None and buf is None:
88
                # Connection was closed, end task
89
                break
90
            try:
91
                await self._process_one_msg(header, buf)
92
            except Exception:
93
                logger.exception('Exception raised while processing message from client')
94
95
    async def _process_one_msg(self, header, buf):
96
        logger.debug('_process_received_message %s %s', header.body_size, len(buf))
97
        ret = await self.processor.process(header, buf)
98
        if not ret:
99
            logger.info('processor returned False, we close connection from %s', self.peer_name)
100
            self.transport.close()
101
            return
102
103
104
class BinaryServer:
105
    def __init__(self, internal_server: InternalServer, hostname, port):
106
        self.logger = logging.getLogger(__name__)
107
        self.hostname = hostname
108
        self.port = port
109
        self.iserver: InternalServer = internal_server
110
        self._server: Optional[asyncio.AbstractServer] = None
111
        self._policies = []
112
        self.clients = []
113
        self.closing_tasks = []
114
115
    def set_policies(self, policies):
116
        self._policies = policies
117
118
    def _make_protocol(self):
119
        """Protocol Factory"""
120
        return OPCUAProtocol(
121
            iserver=self.iserver,
122
            policies=self._policies,
123
            clients=self.clients,
124
            closing_tasks=self.closing_tasks,
125
        )
126
127
    async def start(self):
128
        self._server = await self.iserver.loop.create_server(self._make_protocol, self.hostname, self.port)
129
        # get the port and the hostname from the created server socket
130
        # only relevant for dynamic port asignment (when self.port == 0)
131
        if self.port == 0 and len(self._server.sockets) == 1:
132
            # will work for AF_INET and AF_INET6 socket names
133
            # these are to only families supported by the create_server call
134
            sockname = self._server.sockets[0].getsockname()
135
            self.hostname = sockname[0]
136
            self.port = sockname[1]
137
        self.logger.info('Listening on %s:%s', self.hostname, self.port)
138
        await self._await_closing_tasks()
139
140
    async def stop(self):
141
        self.logger.info('Closing asyncio socket server')
142
        for transport in self.iserver.asyncio_transports:
143
            transport.close()
144
145
        # Wait for all transport closing tasks to complete
146
        results = await asyncio.gather(*self.closing_tasks, return_exceptions=True)
147
        for result in results:
148
            if isinstance(result, Exception):
149
                self.logger.error(f"An error ocurred while closing a transport: {result}")
150
        self.closing_tasks = []
151
152
        if self._server:
153
            self.iserver.loop.call_soon(self._server.close)
154
            await self._server.wait_closed()
155
156
    async def _await_closing_tasks(self):
157
        await asyncio.sleep(10)
158
        while self.closing_tasks:
159
            task = self.closing_tasks.pop()
160
            try:
161
                await task
162
            except Exception:
163
                logger.exception("Unexpected crash in BinaryServer._await_closing_tasks")