transports/stratum_tcp.py

   1 import json
   2 import Queue as queue
   3 import socket
   4 import threading
   5 import time
   6 import traceback, sys
   7
   8 from processor import Session, Dispatcher
   9 from utils import print_log
  10
  11
  12 class TcpSession(Session):
  13
  14     def __init__(self, dispatcher, connection, address, use_ssl, ssl_certfile, ssl_keyfile):
  15         Session.__init__(self, dispatcher)
  16         self.use_ssl = use_ssl
  17         if use_ssl:
  18             import ssl
  19             self._connection = ssl.wrap_socket(
  20                 connection,
  21                 server_side=True,
  22                 certfile=ssl_certfile,
  23                 keyfile=ssl_keyfile,
  24                 ssl_version=ssl.PROTOCOL_SSLv23,
  25                 do_handshake_on_connect=False)
  26         else:
  27             self._connection = connection
  28
  29         self.address = address[0] + ":%d"%address[1]
  30         self.name = "TCP " if not use_ssl else "SSL "
  31         self.timeout = 1000
  32         self.response_queue = queue.Queue()
  33         self.dispatcher.add_session(self)
  34
  35     def do_handshake(self):
  36         if self.use_ssl:
  37             self._connection.do_handshake()
  38
  39     def connection(self):
  40         if self.stopped():
  41             raise Exception("Session was stopped")
  42         else:
  43             return self._connection
  44
  45     def shutdown(self):
  46         try:
  47             self._connection.shutdown(socket.SHUT_RDWR)
  48         except:
  49             # print_log("problem shutting down", self.address)
  50             # traceback.print_exc(file=sys.stdout)
  51             pass
  52
  53         self._connection.close()
  54
  55     def send_response(self, response):
  56         self.response_queue.put(response)
  57
  58
  59 class TcpClientResponder(threading.Thread):
  60
  61     def __init__(self, session):
  62         self.session = session
  63         threading.Thread.__init__(self)
  64
  65     def run(self):
  66         while not self.session.stopped():
  67             try:
  68                 response = self.session.response_queue.get(timeout=10)
  69             except queue.Empty:
  70                 continue
  71             data = json.dumps(response) + "\n"
  72             try:
  73                 while data:
  74                     l = self.session.connection().send(data)
  75                     data = data[l:]
  76             except:
  77                 self.session.stop()
  78
  79
  80
  81 class TcpClientRequestor(threading.Thread):
  82
  83     def __init__(self, dispatcher, session):
  84         self.shared = dispatcher.shared
  85         self.dispatcher = dispatcher
  86         self.message = ""
  87         self.session = session
  88         threading.Thread.__init__(self)
  89
  90     def run(self):
  91         try:
  92             self.session.do_handshake()
  93         except:
  94             self.session.stop()
  95             return
  96
  97         while not self.shared.stopped():
  98
  99             data = self.receive()
 100             if not data:
 101                 self.session.stop()
 102                 break
 103
 104             self.message += data
 105             self.session.time = time.time()
 106
 107             while self.parse():
 108                 pass
 109
 110
 111     def receive(self):
 112         try:
 113             return self.session.connection().recv(2048)
 114         except:
 115             return ''
 116
 117     def parse(self):
 118         raw_buffer = self.message.find('\n')
 119         if raw_buffer == -1:
 120             return False
 121
 122         raw_command = self.message[0:raw_buffer].strip()
 123         self.message = self.message[raw_buffer + 1:]
 124         if raw_command == 'quit':
 125             self.session.stop()
 126             return False
 127
 128         try:
 129             command = json.loads(raw_command)
 130         except:
 131             self.dispatcher.push_response({"error": "bad JSON", "request": raw_command})
 132             return True
 133
 134         try:
 135             # Try to load vital fields, and return an error if
 136             # unsuccessful.
 137             message_id = command['id']
 138             method = command['method']
 139         except KeyError:
 140             # Return an error JSON in response.
 141             self.dispatcher.push_response({"error": "syntax error", "request": raw_command})
 142         else:
 143             self.dispatcher.push_request(self.session, command)
 144             # sleep a bit to prevent a single session from DOSing the queue
 145             time.sleep(0.01)
 146
 147         return True
 148
 149
 150 class TcpServer(threading.Thread):
 151
 152     def __init__(self, dispatcher, host, port, use_ssl, ssl_certfile, ssl_keyfile):
 153         self.shared = dispatcher.shared
 154         self.dispatcher = dispatcher.request_dispatcher
 155         threading.Thread.__init__(self)
 156         self.daemon = True
 157         self.host = host
 158         self.port = port
 159         self.lock = threading.Lock()
 160         self.use_ssl = use_ssl
 161         self.ssl_keyfile = ssl_keyfile
 162         self.ssl_certfile = ssl_certfile
 163
 164     def run(self):
 165         print_log( ("SSL" if self.use_ssl else "TCP") + " server started on port %d"%self.port)
 166         sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
 167         sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
 168         sock.bind((self.host, self.port))
 169         sock.listen(5)
 170
 171         while not self.shared.stopped():
 172
 173             #if self.use_ssl: print_log("SSL: socket listening")
 174             try:
 175                 connection, address = sock.accept()
 176             except:
 177                 traceback.print_exc(file=sys.stdout)
 178                 time.sleep(0.1)
 179                 continue
 180
 181             #if self.use_ssl: print_log("SSL: new session", address)
 182             try:
 183                 session = TcpSession(self.dispatcher, connection, address, use_ssl=self.use_ssl, ssl_certfile=self.ssl_certfile, ssl_keyfile=self.ssl_keyfile)
 184             except BaseException, e:
 185                 error = str(e)
 186                 print_log("cannot start TCP session", error, address)
 187                 connection.close()
 188                 time.sleep(0.1)
 189                 continue
 190
 191             client_req = TcpClientRequestor(self.dispatcher, session)
 192             client_req.start()
 193             responder = TcpClientResponder(session)
 194             responder.start()