nbsocket.py 19.1 KB
Newer Older
Dmitry Volodin's avatar
Dmitry Volodin committed
1
2
# -*- coding: utf-8 -*-
##----------------------------------------------------------------------
Dmitry Volodin's avatar
Dmitry Volodin committed
3
## Clean and lightweight non-blocking socket I/O implementation
Dmitry Volodin's avatar
Dmitry Volodin committed
4
5
6
7
##----------------------------------------------------------------------
## Copyright (C) 2007-2009 The NOC Project
## See LICENSE for details
##----------------------------------------------------------------------
Dmitry Volodin's avatar
Dmitry Volodin committed
8
import socket,select,errno,time,logging,pty,os,signal,subprocess
Dmitry Volodin's avatar
Dmitry Volodin committed
9
10
from errno import EALREADY, EINPROGRESS, EWOULDBLOCK, ECONNRESET, \
     ENOTCONN, ESHUTDOWN, EINTR, EISCONN, ECONNREFUSED, EPIPE, errorcode
11
from noc.lib.debug import error_report
Dmitry Volodin's avatar
Dmitry Volodin committed
12

13
14
15
16
17
18
try:
    import ssl
    HAS_SSL=True
except ImportError:
    HAS_SSL=False

Dmitry Volodin's avatar
Dmitry Volodin committed
19
20
21
22
##
## Abstract non-blocking socket wrapper.
##
class Socket(object):
23
    TTL=None # maximum time to live in seconds
Dmitry Volodin's avatar
Dmitry Volodin committed
24
25
26
27
    def __init__(self,factory,sock):
        self.factory=factory
        self.socket=sock
        self.socket.setblocking(0)
28
        self.start_time=time.time()
29
        self.last_read=self.start_time+100
30
31
        self.factory.register_socket(self)
        self.name=None
Dmitry Volodin's avatar
Dmitry Volodin committed
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
    
    def can_read(self):
        return True
        
    def can_write(self):
        return True
        
    def handle_read(self):
        pass
        
    def handle_write(self):
        pass
        
    def on_close(self):
        pass
        
    def close(self):
        if self.socket:
            self.factory.unregister_socket(self)
            self.socket.close()
            self.socket=None
            self.on_close()
    
    def debug(self,msg):
Dmitry Volodin's avatar
Dmitry Volodin committed
56
        logging.debug("[%s(0x%x)] %s"%(self.__class__.__name__,id(self),msg))
57
58
59
60
    
    def set_name(self,name):
        self.name=name
        self.factory.register_socket(self,name)
61
62
63
64
65
    ##
    ## Update socket status to indicate socket still alive
    ##
    def update_status(self):
        self.last_read=time.time()
66
67
68
    # Stale sockets detection.
    # Called by SocketFactory.close_stale to determine should socket be closed forcefully
    def is_stale(self):
69
        return self.TTL and time.time()-self.last_read>=self.TTL
Dmitry Volodin's avatar
Dmitry Volodin committed
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
##
## Abstract Protocol Parser.
## Accepts data via feed method, polupates internal buffer (self.in_buffer).
## and calls parse_pdu.
## parse_pdu must parse self.in_buffer, retrieve all complete pdu,
## remove them from buffer and return a list of parsed PDU (or empty list if not found)
## All parsed pdu are returned via callback method
class Protocol(object):
    def __init__(self,callback):
        self.callback=callback
        self.in_buffer=""
        
    def feed(self,data):
        self.in_buffer+=data
        for pdu in self.parse_pdu():
            self.callback(pdu)
    
    def parse_pdu(self):
        return []
##
## Line protocol. PDUs are separated by "\n"
##
class LineProtocol(Protocol):
    def on_feed(self):
        pdus=self.in_buffer.split("\n")
        self.in_buffer=pdus.pop(-1)
        return pdus
##
## TCP Listener.
## Waits for connection and creates socket_class instance.
## socket_class should be subclass of AcceptedTCPSocket.
## Should not be used directly. Use SocketFactory.listen_tcp method instead
##
class ListenTCPSocket(Socket):
104
    def __init__(self,factory,address,port,socket_class,**kwargs):
Dmitry Volodin's avatar
Dmitry Volodin committed
105
106
107
108
109
110
        Socket.__init__(self,factory,socket.socket(socket.AF_INET,socket.SOCK_STREAM))
        self.socket.setsockopt(socket.SOL_SOCKET,socket.SO_REUSEADDR,
            self.socket.getsockopt(socket.SOL_SOCKET,socket.SO_REUSEADDR) | 1)
        self.socket.bind((address,port))
        self.socket.listen(5)
        self.socket_class=socket_class
111
112
        self.address=address
        self.port=port
113
        self.kwargs=kwargs
Dmitry Volodin's avatar
Dmitry Volodin committed
114
115
116
    
    def handle_read(self):
        s,addr=self.socket.accept()
117
        if self.socket_class.check_access(addr[0]):
118
            self.socket_class(self.factory,s,**self.kwargs)
Dmitry Volodin's avatar
Dmitry Volodin committed
119
120
121
        else:
            logging.error("Refusing connection from %s"%addr[0])
            s.close()
Dmitry Volodin's avatar
Dmitry Volodin committed
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
##
## UDP Listener
## Waits for UDP packets on specified ports
## and calls on_read for each packet
class ListenUDPSocket(Socket):
    def __init__(self,factory,address,port):
        Socket.__init__(self,factory,socket.socket(socket.AF_INET,socket.SOCK_DGRAM))
        self.socket.setsockopt(socket.SOL_SOCKET,socket.SO_REUSEADDR,
            self.socket.getsockopt(socket.SOL_SOCKET,socket.SO_REUSEADDR) | 1)
        self.socket.bind((address,port))
    
    def handle_read(self):
        msg,transport_address=self.socket.recvfrom(8192)
        if msg=="":
            return
        self.on_read(msg,transport_address[0],transport_address[1])
    
    def on_read(self,data,address,port):
        pass
141
142
143
    
    def can_write(self):
        return False
Dmitry Volodin's avatar
Dmitry Volodin committed
144
145
146
147
148
149
150
151
152
153
154
155
156

##
## Base class for TCP sockets. Should not be used directly.
## Use ConnectedTCPSocket and AcceptedTCPSocket instead
##
class TCPSocket(Socket):
    protocol_class=None
    def __init__(self,factory,socket):
        Socket.__init__(self,factory,socket)
        self.is_connected=False
        self.out_buffer=""
        if self.protocol_class:
            self.protocol=self.protocol_class(self.on_read)
157
158
159
160
161
162
163
        self.in_shutdown=False
        
    def close(self,flush=False):
        if flush and len(self.out_buffer)>0:
            self.in_shutdown=True
        else:
            super(TCPSocket,self).close()
Dmitry Volodin's avatar
Dmitry Volodin committed
164
165
166
167
        
    def handle_write(self):
        sent=self.socket.send(self.out_buffer)
        self.out_buffer=self.out_buffer[sent:]
168
169
        if self.in_shutdown and len(self.out_buffer)==0:
            self.close()
Dmitry Volodin's avatar
Dmitry Volodin committed
170
171
172
173
174
175
        
    def handle_connect(self):
        self.is_connected=True
        self.on_connect()
        
    def write(self,msg):
176
        self.debug("write(%s)"%repr(msg))
Dmitry Volodin's avatar
Dmitry Volodin committed
177
178
179
180
181
182
183
184
185
        self.out_buffer+=msg
        
    def on_read(self,data): pass
    
    def on_connect(self): pass

##
## A socket wrapping accepted TCP connection.
## Following methods can be overrided for desired behavior:
186
## check_access(cls, address) - called before object creation
Dmitry Volodin's avatar
Dmitry Volodin committed
187
188
189
190
191
192
193
194
195
196
197
## on_connect(self)   - called when socket connected (just after creation)
## on_close(self)     - called when socket closed. Last method called
## on_read(self,data) - called when new portion of data available when protocol=None or when new PDU available
## Following methods used for operation:
## write(self,data)   - send data (Buffered, can be delayed or split into several send)
## close(self)        - close socket (Also implies on_close(self) event)
##
class AcceptedTCPSocket(TCPSocket):
    def __init__(self,factory,socket):
        TCPSocket.__init__(self,factory,socket)
        self.handle_connect()
198
199
200
201
    
    @classmethod
    def check_access(cls,address):
        return True
Dmitry Volodin's avatar
Dmitry Volodin committed
202
203
204
205
206
207
208
209
210
211
212
213
        
    def handle_read(self):
        try:
            data=self.socket.recv(8192)
        except socket.error,why:
            if why[0] in (ECONNRESET,ENOTCONN,ESHUTDOWN):
                self.close()
                return
            raise socket.error, why
        if data=="":
            self.close()
            return
214
        self.update_status()
Dmitry Volodin's avatar
Dmitry Volodin committed
215
216
217
218
219
220
221
222
        if self.protocol_class:
            self.protocol.feed(data)
        else:
            self.on_read(data)
        
    def can_write(self):
        return self.out_buffer
##
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
## SSL-enabled AcceptedTCPSocket
##
class AcceptedTCPSSLSocket(AcceptedTCPSocket):
    def __init__(self,factory,socket,cert):
        socket=ssl.wrap_socket(socket,server_side=True,do_handshake_on_connect=False,
            keyfile=cert,certfile=cert,
            ssl_version=ssl.PROTOCOL_TLSv1)
        self.ssl_handshake_passed=False
        TCPSocket.__init__(self,factory,socket)
    
    def handle_read(self):
        if self.ssl_handshake_passed:
            super(AcceptedTCPSSLSocket,self).handle_read()
        else: # Process SSL Handshake
            try:
                self.socket.do_handshake()
                self.ssl_handshake_passed=True
                self.debug("SSL Handshake passed: %s"%str(self.socket.cipher()))
                self.handle_connect() # handle_connect called after SSL negotiation
            except ssl.SSLError,err:
                if err.args[0] in [ssl.SSL_ERROR_WANT_READ,ssl.SSL_ERROR_WANT_WRITE]: # Incomplete handshake data
                    return
                logging.error("SSL Handshake failed: %s"%err[1])
                self.close()
##
Dmitry Volodin's avatar
Dmitry Volodin committed
248
249
250
251
252
253
254
255
256
257
258
## A socket wrapping connectiong TCP connection.
## Following methods can be overrided for desired behavior:
## on_connect(self)   - called when connection established and socket read to send data (first event)
## on_close(self)     - called when socket closed. Last method called
## on_read(self,data) - called when new portion of data available
## on_conn_refused(self) - called when connection refused (Also implies close(self). first event)
## Following methods used for operation:
## write(self,data)   - send data (Buffered, can be delayed or split into several send)
## close(self)        - close socket (Also implies on_close(self) event)
##
class ConnectedTCPSocket(TCPSocket):
259
    def __init__(self,factory,address,port,local_address=None):
Dmitry Volodin's avatar
Dmitry Volodin committed
260
        TCPSocket.__init__(self,factory,socket.socket(socket.AF_INET,socket.SOCK_STREAM))
261
262
        if local_address:
            self.socket.bind((local_address,0))
Dmitry Volodin's avatar
Dmitry Volodin committed
263
264
        e=self.socket.connect_ex((address,port))
        if e in (0, EISCONN):
265
            #self.handle_read()
Dmitry Volodin's avatar
Dmitry Volodin committed
266
267
268
269
            return
        elif e in (EINPROGRESS, EALREADY, EWOULDBLOCK):
            return
        else:
Dmitry Volodin's avatar
Dmitry Volodin committed
270
            raise socket.error, (e, errorcode[e])
Dmitry Volodin's avatar
Dmitry Volodin committed
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
        
    def handle_read(self):
        if not self.is_connected:
            self.handle_connect()
            return
        try:
            data=self.socket.recv(8192)
        except socket.error,why:
            if why[0]==ECONNREFUSED:
                self.on_conn_refused()
                self.close()
                return
            if why[0] in (ECONNRESET,ENOTCONN,ESHUTDOWN):
                self.close()
                return
            raise socket.error,why
        if not data:
            self.close()
            return
290
        self.update_status()
Dmitry Volodin's avatar
Dmitry Volodin committed
291
292
293
294
295
296
297
298
299
300
301
302
303
        if self.protocol_class:
            self.protocol.feed(data)
        else:
            self.on_read(data)
    
    def can_write(self):
        return self.out_buffer or not self.is_connected
    
    def handle_write(self):
        if not self.is_connected:
            try:
                self.socket.send("")
            except socket.error,why:
304
                if why[0] in (EPIPE,ECONNREFUSED):
Dmitry Volodin's avatar
Dmitry Volodin committed
305
306
307
308
309
310
311
312
313
                    self.on_conn_refused()
                    self.close()
                    return
                raise socket.error,why
            self.handle_connect()
            return
        TCPSocket.handle_write(self)
    
    def on_conn_refused(self): pass
314
##
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
##
##
class ConnectedTCPSSLSocket(ConnectedTCPSocket):
    def __init__(self,factory,address,port,local_address=None):
        self.ssl_handshake_passed=False
        super(ConnectedTCPSSLSocket,self).__init__(factory,address,port,local_address)
        self.socket=ssl.wrap_socket(self.socket,server_side=False,do_handshake_on_connect=False,ssl_version=ssl.PROTOCOL_TLSv1)
    
    def handle_read(self):
        if self.ssl_handshake_passed:
            super(ConnectedTCPSSLSocket,self).handle_read()
        else: # Process SSL Handshake
            try:
                self.socket.do_handshake()
                self.ssl_handshake_passed=True
                self.debug("SSL Handshake passed: %s"%str(self.socket.cipher()))
            except ssl.SSLError,err:
                if err.args[0] in [ssl.SSL_ERROR_WANT_READ,ssl.SSL_ERROR_WANT_WRITE]: # Incomplete handshake data
                    return
                raise
##
Dmitry Volodin's avatar
Dmitry Volodin committed
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
## UDP send/receive socket
##
class UDPSocket(Socket):
    def __init__(self,factory):
        self.out_buffer=[]
        super(UDPSocket,self).__init__(factory,socket.socket(socket.AF_INET,socket.SOCK_DGRAM))
    
    def can_write(self):
        return self.out_buffer
    
    def handle_write(self):
        msg,addr=self.out_buffer.pop(0)
        self.socket.sendto(msg,addr)

    def handle_read(self):
        msg,transport_address=self.socket.recvfrom(8192)
        if msg=="":
            return
        self.on_read(msg,transport_address[0],transport_address[1])
    
    def on_read(self,data,address,port):
        pass

    def sendto(self,msg,addr):
        self.out_buffer+=[(msg,addr)]
    
##
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
## File wrapper to mimic behavior of socket
##
class FileWrapper(object):
    def __init__(self,fileno):
        self._fileno=fileno
        
    def fileno(self):
        return self._fileno
        
    def recv(self, *args):
        return os.read(self._fileno, *args)
        
    def send(self, *args):
        return os.write(self._fileno, *args)
        
    read = recv
    
    write = send
    
    def close(self):
        os.close(self._fileno)
        
    def setblocking(self,status):
        pass
##
## PTY Socket Emulation
Dmitry Volodin's avatar
Dmitry Volodin committed
389
## Events: on_read, on_close
390
391
392
##
class PTYSocket(Socket):
    def __init__(self,factory,argv):
393
        self.out_buffer=""
394
395
396
397
        self.pid,fd=pty.fork()
        if self.pid==0:
            os.execv(argv[0],argv)
        else:
398
            super(PTYSocket,self).__init__(factory,FileWrapper(fd))
399
400
            
    def handle_read(self):
401
402
403
404
405
        try:
            data=self.socket.read(8192)
        except OSError:
            self.close()
            return
406
        self.update_status()
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
        if data:
            self.on_read(data)
        else:
            self.close()

    def can_write(self):
        return self.out_buffer
        
    def handle_write(self):
        sent=self.socket.send(self.out_buffer)
        self.out_buffer=self.out_buffer[sent:]

    def handle_connect(self):
        self.is_connected=True
        self.on_connect()

    def write(self,msg):
Dmitry Volodin's avatar
Dmitry Volodin committed
424
        self.debug("write(%s)"%repr(msg))
425
426
427
428
429
430
431
432
433
        self.out_buffer+=msg
    
    def close(self):
        Socket.close(self)
        try:
            pid,status=os.waitpid(self.pid,os.WNOHANG)
        except os.error:
            return
        if pid:
Dmitry Volodin's avatar
Dmitry Volodin committed
434
            self.debug("Child pid=%d is already terminated. Zombie released"%pid)
435
        else:
Dmitry Volodin's avatar
Dmitry Volodin committed
436
            self.debug("Child pid=%d is not terminated. Killing"%self.pid)
437
438
439
            try:
                os.kill(self.pid,signal.SIGKILL)
            except:
Dmitry Volodin's avatar
Dmitry Volodin committed
440
                self.debug("Child pid=%d was killed from another place"%self.pid)
441
442
443
    
    def on_read(self,data):
        pass
Dmitry Volodin's avatar
Dmitry Volodin committed
444
445
446
447
448
449
450
451
452
453
##
## PopenSocket
## Events: on_read, on_close
##
class PopenSocket(Socket):
    def __init__(self,factory,argv):
        self.debug("Launching %s"%argv)
        self.p=subprocess.Popen(argv,stdout=subprocess.PIPE)
        super(PopenSocket,self).__init__(factory,FileWrapper(self.p.stdout.fileno()))
        self.update_status()
Dmitry Volodin's avatar
Dmitry Volodin committed
454

Dmitry Volodin's avatar
Dmitry Volodin committed
455
456
457
458
459
460
461
462
463
464
465
466
467
468
    def handle_read(self):
        try:
            data=self.socket.read(8192)
        except OSError:
            self.close()
            return
        self.update_status()
        if data:
            self.on_read(data)
        else:
            self.close()

    def close(self):
        Socket.close(self)
Dmitry Volodin's avatar
Dmitry Volodin committed
469
470
471
472
473
474
475
476
##
## Socket Factory.
## Methods:
## listen_tcp  - create new TCP listener
## connect_tcp - create new TCP connection
## run         - main event loop
##
class SocketFactory(object):
477
    def __init__(self,tick_callback=None):
Dmitry Volodin's avatar
Dmitry Volodin committed
478
        self.sockets={}
479
480
        self.socket_name={}
        self.name_socket={}
481
        self.tick_callback=tick_callback
Dmitry Volodin's avatar
Dmitry Volodin committed
482
    
483
484
    def register_socket(self,socket,name=None):
        logging.debug("register_socket(%s,%s)"%(socket,name))
Dmitry Volodin's avatar
Dmitry Volodin committed
485
        self.sockets[socket.socket.fileno()]=socket
486
487
488
489
490
491
492
493
494
        if socket in self.socket_name:
            # Socket was registred
            old_name=self.socket_name[socket]
            del self.socket_name[socket]
            if old_name:
                del self.name_socket[old_name]
        self.socket_name[socket]=name
        if name:
            self.name_socket[name]=socket
Dmitry Volodin's avatar
Dmitry Volodin committed
495
496
497
498
        
    def unregister_socket(self,socket):
        logging.debug("unregister_socket(%s)"%socket)
        del self.sockets[socket.socket.fileno()]
499
500
501
502
        old_name=self.socket_name[socket]
        del self.socket_name[socket]
        if old_name:
            del self.name_socket[old_name]
Dmitry Volodin's avatar
Dmitry Volodin committed
503
        
504
    def listen_tcp(self,address,port,socket_class,**kwargs):
Dmitry Volodin's avatar
Dmitry Volodin committed
505
506
        if not issubclass(socket_class,AcceptedTCPSocket):
            raise "socket_class should be a AcceptedTCPSocket subclass"
507
        l=ListenTCPSocket(self,address,port,socket_class,**kwargs)
508
        l.set_name("listen-tcp-%s:%d"%(address,port))
509
        return l
Dmitry Volodin's avatar
Dmitry Volodin committed
510
511
512
513
514
515
    
    def connect_tcp(self,address,port,socket_class):
        if not issubclass(socket_class,ConnectedTCPSocket):
            raise "socket_class should be a ConnectedTCPSocket subclass"
        return socket_class(self,address,port)
        
516
517
518
519
520
    def get_socket_by_name(self,name):
        return self.name_socket[name]
        
    def get_name_by_socket(self,socket):
        return self.socket_name[socket]
521
522
523
    
    def __len__(self):
        return len(self.sockets)
524
        
525
526
527
528
529
    def close_stale(self):
        for s in [s for s in self.sockets.values() if s.is_stale()]:
            logging.debug("Closing stale socket %s"%s)
            s.close()
        
Dmitry Volodin's avatar
Dmitry Volodin committed
530
531
532
533
534
    def loop(self,timeout=1):
        if self.sockets:
            r=[f for f,s in self.sockets.items() if s.can_read()]
            w=[f for f,s in self.sockets.items() if s.can_write()]
            if r or w:
535
536
537
538
539
540
                try:
                    r,w,x=select.select(r,w,[],timeout)
                except select.error,why:
                    if why[0]==EINTR:
                        return
                    raise
Dmitry Volodin's avatar
Dmitry Volodin committed
541
542
543
544
545
546
547
548
                if r or w:
                    # Write events processed before read events
                    # to catch connection refused causes
                    for f in w:
                        if f in self.sockets:
                            try:
                                self.sockets[f].handle_write()
                            except:
549
                                error_report()
550
551
552
553
                                try:
                                    self.sockets[f].close()
                                except:
                                    pass
Dmitry Volodin's avatar
Dmitry Volodin committed
554
555
556
557
558
                    for f in r:
                        if f in self.sockets:
                            try:
                                self.sockets[f].handle_read()
                            except:
559
                                error_report()
560
561
562
563
                                try:
                                    self.sockets[f].close()
                                except:
                                    pass
Dmitry Volodin's avatar
Dmitry Volodin committed
564
565
566
567
568
            else:
                time.sleep(timeout)
        else:
            time.sleep(timeout)
    
569
570
    def run(self,run_forever=False):
        if run_forever:
571
            cond=lambda:True
572
        else:
573
            cond=lambda:len(self.sockets)>0
574
575
576
            # Wait for any socket
            while not cond():
                time.sleep(1)
577
578
579
580
581
582
583
584
        last_tick=time.time()
        last_stale=time.time()
        while cond():
            self.loop(1)
            t=time.time()
            if self.tick_callback and t-last_tick>=1:
                self.tick_callback()
                last_tick=t
585
            if t-last_stale>3:
586
587
                self.close_stale()
                last_stale=t
588
589
590
591
592
593
    
    ##
    ## Return amount of active sockets which are descendants from sclass
    ##
    def count_subclass_sockets(self,sclass):
        return len([s for s in self.sockets.values() if issubclass(s.__class__,sclass)])