forked from mirrors/gecko-dev
114 lines
3.4 KiB
Python
114 lines
3.4 KiB
Python
# vim: set ts=4 et sw=4 tw=80
|
|
# This Source Code Form is subject to the terms of the Mozilla Public
|
|
# License, v. 2.0. If a copy of the MPL was not distributed with this
|
|
# file, You can obtain one at http://mozilla.org/MPL/2.0/.
|
|
|
|
from __future__ import absolute_import, print_function
|
|
|
|
from twisted.internet import protocol, reactor
|
|
from twisted.internet.task import LoopingCall
|
|
import txws
|
|
import psutil
|
|
|
|
import argparse
|
|
import six
|
|
import sys
|
|
import os
|
|
|
|
# maps a command issued via websocket to running an executable with args
|
|
commands = {
|
|
"iceserver": [sys.executable, "-u", os.path.join("iceserver", "iceserver.py")]
|
|
}
|
|
|
|
|
|
class ProcessSide(protocol.ProcessProtocol):
|
|
"""Handles the spawned process (I/O, process termination)"""
|
|
|
|
def __init__(self, socketSide):
|
|
self.socketSide = socketSide
|
|
|
|
def outReceived(self, data):
|
|
if self.socketSide:
|
|
lines = data.splitlines()
|
|
for line in lines:
|
|
self.socketSide.transport.write(line)
|
|
|
|
def errReceived(self, data):
|
|
self.outReceived(data)
|
|
|
|
def processEnded(self, reason):
|
|
if self.socketSide:
|
|
self.outReceived(str(reason))
|
|
self.socketSide.processGone()
|
|
|
|
def socketGone(self):
|
|
self.socketSide = None
|
|
self.transport.loseConnection()
|
|
self.transport.signalProcess("KILL")
|
|
|
|
|
|
class SocketSide(protocol.Protocol):
|
|
"""
|
|
Handles the websocket (I/O, closed connection), and spawning the process
|
|
"""
|
|
|
|
def __init__(self):
|
|
self.processSide = None
|
|
|
|
def dataReceived(self, data):
|
|
if not self.processSide:
|
|
self.processSide = ProcessSide(self)
|
|
# We deliberately crash if |data| isn't on the "menu",
|
|
# or there is some problem spawning.
|
|
data = six.ensure_str(data)
|
|
reactor.spawnProcess(
|
|
self.processSide, commands[data][0], commands[data], env=os.environ
|
|
)
|
|
|
|
def connectionLost(self, reason):
|
|
if self.processSide:
|
|
self.processSide.socketGone()
|
|
|
|
def processGone(self):
|
|
self.processSide = None
|
|
self.transport.loseConnection()
|
|
|
|
|
|
class ProcessSocketBridgeFactory(protocol.Factory):
|
|
"""Builds sockets that can launch/bridge to a process"""
|
|
|
|
def buildProtocol(self, addr):
|
|
return SocketSide()
|
|
|
|
|
|
# Parent process could have already exited, so this is slightly racy. Only
|
|
# alternative is to set up a pipe between parent and child, but that requires
|
|
# special cooperation from the parent.
|
|
parent_process = psutil.Process(os.getpid()).parent()
|
|
|
|
|
|
def check_parent():
|
|
""" Checks if parent process is still alive, and exits if not """
|
|
if not parent_process.is_running():
|
|
print("websocket/process bridge exiting because parent process is gone")
|
|
reactor.stop()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
parser = argparse.ArgumentParser(description="Starts websocket/process bridge.")
|
|
parser.add_argument(
|
|
"--port",
|
|
type=str,
|
|
dest="port",
|
|
default="8191",
|
|
help="Port for websocket/process bridge. Default 8191.",
|
|
)
|
|
args = parser.parse_args()
|
|
|
|
parent_checker = LoopingCall(check_parent)
|
|
parent_checker.start(1)
|
|
|
|
bridgeFactory = ProcessSocketBridgeFactory()
|
|
reactor.listenTCP(int(args.port), txws.WebSocketFactory(bridgeFactory))
|
|
print("websocket/process bridge listening on port %s" % args.port)
|
|
reactor.run()
|