testing/tools/websocketprocessbridge/websocketprocessbridge.py

   1 # vim: set ts=4 et sw=4 tw=80
   2 # This Source Code Form is subject to the terms of the Mozilla Public
   3 # License, v. 2.0. If a copy of the MPL was not distributed with this
   4 # file, You can obtain one at http://mozilla.org/MPL/2.0/.
   5
   6 from twisted.internet import protocol, reactor
   7 from twisted.internet.task import LoopingCall
   8 from autobahn.twisted.websocket import WebSocketServerProtocol, WebSocketServerFactory
   9
  10 import psutil
  11
  12 import argparse
  13 import six
  14 import sys
  15 import os
  16
  17 # maps a command issued via websocket to running an executable with args
  18 commands = {
  19     "iceserver": [sys.executable, "-u", os.path.join("iceserver", "iceserver.py")]
  20 }
  21
  22
  23 class ProcessSide(protocol.ProcessProtocol):
  24     """Handles the spawned process (I/O, process termination)"""
  25
  26     def __init__(self, socketSide):
  27         self.socketSide = socketSide
  28
  29     def outReceived(self, data):
  30         data = six.ensure_str(data)
  31         if self.socketSide:
  32             lines = data.splitlines()
  33             for line in lines:
  34                 self.socketSide.sendMessage(line.encode("utf8"), False)
  35
  36     def errReceived(self, data):
  37         self.outReceived(data)
  38
  39     def processEnded(self, reason):
  40         if self.socketSide:
  41             self.outReceived(reason.getTraceback())
  42             self.socketSide.processGone()
  43
  44     def socketGone(self):
  45         self.socketSide = None
  46         self.transport.loseConnection()
  47         self.transport.signalProcess("KILL")
  48
  49
  50 class SocketSide(WebSocketServerProtocol):
  51     """
  52     Handles the websocket (I/O, closed connection), and spawning the process
  53     """
  54
  55     def __init__(self):
  56         super(SocketSide, self).__init__()
  57         self.processSide = None
  58
  59     def onConnect(self, request):
  60         return None
  61
  62     def onOpen(self):
  63         return None
  64
  65     def onMessage(self, payload, isBinary):
  66         # We only expect a single message, which tells us what kind of process
  67         # we're supposed to launch. ProcessSide pipes output to us for sending
  68         # back to the websocket client.
  69         if not self.processSide:
  70             self.processSide = ProcessSide(self)
  71             # We deliberately crash if |data| isn't on the "menu",
  72             # or there is some problem spawning.
  73             data = six.ensure_str(payload)
  74             try:
  75                 reactor.spawnProcess(
  76                     self.processSide, commands[data][0], commands[data], env=os.environ
  77                 )
  78             except BaseException as e:
  79                 print(e.str())
  80                 self.sendMessage(e.str())
  81                 self.processGone()
  82
  83     def onClose(self, wasClean, code, reason):
  84         if self.processSide:
  85             self.processSide.socketGone()
  86
  87     def processGone(self):
  88         self.processSide = None
  89         self.transport.loseConnection()
  90
  91
  92 # Parent process could have already exited, so this is slightly racy. Only
  93 # alternative is to set up a pipe between parent and child, but that requires
  94 # special cooperation from the parent.
  95 parent_process = psutil.Process(os.getpid()).parent()
  96
  97
  98 def check_parent():
  99     """Checks if parent process is still alive, and exits if not"""
 100     if not parent_process.is_running():
 101         print("websocket/process bridge exiting because parent process is gone")
 102         reactor.stop()
 103
 104
 105 if __name__ == "__main__":
 106     parser = argparse.ArgumentParser(description="Starts websocket/process bridge.")
 107     parser.add_argument(
 108         "--port",
 109         type=str,
 110         dest="port",
 111         default="8191",
 112         help="Port for websocket/process bridge. Default 8191.",
 113     )
 114     args = parser.parse_args()
 115
 116     parent_checker = LoopingCall(check_parent)
 117     parent_checker.start(1)
 118
 119     bridgeFactory = WebSocketServerFactory()
 120     bridgeFactory.protocol = SocketSide
 121     reactor.listenTCP(int(args.port), bridgeFactory)
 122     print("websocket/process bridge listening on port %s" % args.port)
 123     reactor.run()