home · contact · privacy
Only send PONG to connection_id that sent the PING.
[plomrogue2-experiments] / new2 / plomrogue / io_tcp.py
1 import socketserver
2
3
4 # Avoid "Address already in use" errors.
5 socketserver.TCPServer.allow_reuse_address = True
6
7
8
9 from plomrogue.errors import BrokenSocketConnection
10 class PlomSocket:
11
12     def __init__(self, socket):
13         self.socket = socket
14
15     def send(self, message, silent_connection_break=False):
16         """Send via self.socket, encoded/delimited as way recv() expects.
17
18         In detail, all \ and $ in message are escaped with prefixed \,
19         and an unescaped $ is appended as a message delimiter. Then,
20         socket.send() is called as often as necessary to ensure
21         message is sent fully, as socket.send() due to buffering may
22         not send all of it right away.
23
24         Assuming socket is blocking, it's rather improbable that
25         socket.send() will be partial / return a positive value less
26         than the (byte) length of msg – but not entirely out of the
27         question. See: - <http://stackoverflow.com/q/19697218> -
28         <http://stackoverflow.com/q/2618736> -
29         <http://stackoverflow.com/q/8900474>
30
31         This also handles a socket.send() return value of 0, which
32         might be possible or not (?) for blocking sockets: -
33         <http://stackoverflow.com/q/34919846>
34
35         """
36         escaped_message = ''
37         for char in message:
38             if char in ('\\', '$'):
39                 escaped_message += '\\'
40             escaped_message += char
41         escaped_message += '$'
42         data = escaped_message.encode()
43         totalsent = 0
44         while totalsent < len(data):
45             socket_broken = False
46             try:
47                 sent = self.socket.send(data[totalsent:])
48                 socket_broken = sent == 0
49                 totalsent = totalsent + sent
50             except OSError as err:
51                 if err.errno == 9:  # "Bad file descriptor", when connection broken
52                     socket_broken = True
53                 else:
54                     raise err
55             if socket_broken and not silent_connection_break:
56                 raise BrokenSocketConnection
57
58     def recv(self):
59         """Get full send()-prepared message from self.socket.
60
61         In detail, socket.recv() is looped over for sequences of bytes
62         that can be decoded as a Unicode string delimited by an
63         unescaped $, with \ and $ escapable by \. If a sequence of
64         characters that ends in an unescaped $ cannot be decoded as
65         Unicode, None is returned as its representation. Stop once
66         socket.recv() returns nothing.
67
68         Under the hood, the TCP stack receives packets that construct
69         the input payload in an internal buffer; socket.recv(BUFSIZE)
70         pops up to BUFSIZE bytes from that buffer, without knowledge
71         either about the input's segmentation into packets, or whether
72         the input is segmented in any other meaningful way; that's why
73         we do our own message segmentation with $ as a delimiter.
74
75         """
76         esc = False
77         data = b''
78         msg = b''
79         while True:
80             try:
81                 data = self.socket.recv(1024)
82             except OSError as err:
83                 if err.errno == 9:  # "Bad file descriptor", when connection broken
84                     raise BrokenSocketConnection
85             if 0 == len(data):
86                 break
87             for c in data:
88                 if esc:
89                     msg += bytes([c])
90                     esc = False
91                 elif chr(c) == '\\':
92                     esc = True
93                 elif chr(c) == '$':
94                     try:
95                         yield msg.decode()
96                     except UnicodeDecodeError:
97                         yield None
98                     msg = b''
99                 else:
100                     msg += bytes([c])
101
102
103
104 class PlomSocketSSL(PlomSocket):
105
106     def __init__(self, *args, server_side=False, certfile=None, keyfile=None, **kwargs):
107         import ssl
108         super().__init__(*args, **kwargs)
109         if server_side:
110             self.socket = ssl.wrap_socket(self.socket, server_side=True,
111                                           certfile=certfile, keyfile=keyfile)
112         else:
113             self.socket = ssl.wrap_socket(self.socket)
114
115
116
117 class IO_Handler(socketserver.BaseRequestHandler):
118
119     def __init__(self, *args, **kwargs):
120         super().__init__(*args, **kwargs)
121
122     def handle(self):
123         """Move messages between network socket and game IO loop via queues.
124
125         On start (a new connection from client to server), sets up a
126         new queue, sends it via self.server.queue_out to the game IO
127         loop thread, and from then on receives messages to send back
128         from the game IO loop via that new queue.
129
130         At the same time, loops over socket's recv to get messages
131         from the outside into the game IO loop by way of
132         self.server.queue_out into the game IO. Ends connection once a
133         'QUIT' message is received from socket, and then also calls
134         for a kill of its own queue.
135
136         """
137
138         def send_queue_messages(plom_socket, queue_in, thread_alive):
139             """Send messages via socket from queue_in while thread_alive[0]."""
140             while thread_alive[0]:
141                 try:
142                     msg = queue_in.get(timeout=1)
143                 except queue.Empty:
144                     continue
145                 plom_socket.send(msg, True)
146
147         import uuid
148         import queue
149         import threading
150         if self.server.socket_class == PlomSocketSSL:
151             plom_socket = self.server.socket_class(self.request,
152                                                    server_side=True,
153                                                    certfile=self.server.certfile,
154                                                    keyfile=self.server.keyfile)
155         else:
156             plom_socket = self.server.socket_class(self.request)
157         print('CONNECTION FROM:', str(self.client_address))
158         connection_id = uuid.uuid4()
159         queue_in = queue.Queue()
160         self.server.clients[connection_id] = queue_in
161         thread_alive = [True]
162         t = threading.Thread(target=send_queue_messages,
163                              args=(plom_socket, queue_in, thread_alive))
164         t.start()
165         for message in plom_socket.recv():
166             if message is None:
167                 plom_socket.send('BAD MESSAGE', True)
168             elif 'QUIT' == message:
169                 plom_socket.send('BYE', True)
170                 break
171             else:
172                 self.server.queue_out.put((connection_id, message))
173         del self.server.clients[connection_id]
174         thread_alive[0] = False
175         print('CONNECTION CLOSED FROM:', str(self.client_address))
176         plom_socket.socket.close()
177
178
179
180 class PlomTCPServer(socketserver.ThreadingTCPServer):
181     """Bind together threaded IO handling server and message queue.
182
183     By default this only serves to localhost connections.  For remote
184     connections, consider using PlomTCPServerSSL for more security,
185     which defaults to serving all connections.
186
187     """
188
189     def __init__(self, queue, port, host='127.0.0.1', *args, **kwargs):
190         super().__init__((host, port), IO_Handler, *args, **kwargs)
191         self.socket_class = PlomSocket
192         self.queue_out = queue
193         self.daemon_threads = True  # Else, server's threads have daemon=False.
194         self.clients = {}
195
196
197
198 class PlomTCPServerSSL(PlomTCPServer):
199
200     def __init__(self, *args, certfile=None, keyfile=None, **kwargs):
201         super().__init__(*args, host='0.0.0.0', **kwargs)
202         self.certfile = certfile
203         self.keyfile = keyfile
204         self.socket_class = PlomSocketSSL