serve.py (6143B)
1 #!/usr/bin/env python 2 3 # This Source Code Form is subject to the terms of the Mozilla Public 4 # License, v. 2.0. If a copy of the MPL was not distributed with this 5 # file, You can obtain one at http://mozilla.org/MPL/2.0/. 6 7 """Spawns necessary HTTP servers for testing Marionette in child 8 processes. 9 10 """ 11 12 import argparse 13 import multiprocessing 14 import os 15 import sys 16 from collections import defaultdict 17 18 from . import httpd 19 20 __all__ = [ 21 "default_doc_root", 22 "iter_proc", 23 "iter_url", 24 "registered_servers", 25 "servers", 26 "start", 27 "where_is", 28 ] 29 here = os.path.abspath(os.path.dirname(__file__)) 30 31 32 class BlockingChannel: 33 def __init__(self, channel): 34 self.chan = channel 35 self.lock = multiprocessing.Lock() 36 37 def call(self, func, args=()): 38 self.send((func, args)) 39 return self.recv() 40 41 def send(self, *args): 42 try: 43 self.lock.acquire() 44 self.chan.send(args) 45 finally: 46 self.lock.release() 47 48 def recv(self): 49 try: 50 self.lock.acquire() 51 payload = self.chan.recv() 52 if isinstance(payload, tuple) and len(payload) == 1: 53 return payload[0] 54 return payload 55 except KeyboardInterrupt: 56 return ("stop", ()) 57 finally: 58 self.lock.release() 59 60 61 class ServerProxy(multiprocessing.Process, BlockingChannel): 62 def __init__(self, channel, init_func, *init_args, **init_kwargs): 63 multiprocessing.Process.__init__(self) 64 BlockingChannel.__init__(self, channel) 65 self.init_func = init_func 66 self.init_args = init_args 67 self.init_kwargs = init_kwargs 68 69 def run(self): 70 try: 71 server = self.init_func(*self.init_args, **self.init_kwargs) 72 server.start() 73 self.send(("ok", ())) 74 75 while True: 76 # ["func", ("arg", ...)] 77 # ["prop", ()] 78 sattr, fargs = self.recv() 79 attr = getattr(server, sattr) 80 81 # apply fargs to attr if it is a function 82 if callable(attr): 83 rv = attr(*fargs) 84 85 # otherwise attr is a property 86 else: 87 rv = attr 88 89 self.send(rv) 90 91 if sattr == "stop": 92 return 93 94 except Exception as e: 95 self.send(("stop", e)) 96 97 except KeyboardInterrupt: 98 server.stop() 99 100 101 class ServerProc(BlockingChannel): 102 def __init__(self, init_func): 103 self._init_func = init_func 104 self.proc = None 105 106 parent_chan, self.child_chan = multiprocessing.Pipe() 107 BlockingChannel.__init__(self, parent_chan) 108 109 def start(self, doc_root, ssl_config, **kwargs): 110 self.proc = ServerProxy( 111 self.child_chan, self._init_func, doc_root, ssl_config, **kwargs 112 ) 113 self.proc.daemon = True 114 self.proc.start() 115 116 res, exc = self.recv() 117 if res == "stop": 118 raise exc 119 120 def get_url(self, url): 121 return self.call("get_url", (url,)) 122 123 @property 124 def doc_root(self): 125 return self.call("doc_root", ()) 126 127 def stop(self): 128 self.call("stop") 129 if not self.is_alive: 130 return 131 self.proc.join() 132 133 def kill(self): 134 if not self.is_alive: 135 return 136 self.proc.terminate() 137 self.proc.join(0) 138 139 @property 140 def is_alive(self): 141 if self.proc is not None: 142 return self.proc.is_alive() 143 return False 144 145 146 def http_server(doc_root, ssl_config, host="127.0.0.1", **kwargs): 147 return httpd.FixtureServer(doc_root, url=f"http://{host}:0/", **kwargs) 148 149 150 def https_server(doc_root, ssl_config, host="127.0.0.1", **kwargs): 151 return httpd.FixtureServer( 152 doc_root, 153 url=f"https://{host}:0/", 154 ssl_key=ssl_config["key_path"], 155 ssl_cert=ssl_config["cert_path"], 156 **kwargs, 157 ) 158 159 160 def start_servers(doc_root, ssl_config, **kwargs): 161 servers = defaultdict() 162 for schema, builder_fn in registered_servers: 163 proc = ServerProc(builder_fn) 164 proc.start(doc_root, ssl_config, **kwargs) 165 servers[schema] = (proc.get_url("/"), proc) 166 return servers 167 168 169 def start(doc_root=None, **kwargs): 170 """Start all relevant test servers. 171 172 If no `doc_root` is given the default 173 testing/marionette/harness/marionette_harness/www directory will be used. 174 175 Additional keyword arguments can be given which will be passed on 176 to the individual ``FixtureServer``'s in httpd.py. 177 178 """ 179 doc_root = doc_root or default_doc_root 180 ssl_config = { 181 "cert_path": httpd.default_ssl_cert, 182 "key_path": httpd.default_ssl_key, 183 } 184 185 global servers 186 servers = start_servers(doc_root, ssl_config, **kwargs) 187 return servers 188 189 190 def where_is(uri, on="http"): 191 """Returns the full URL, including scheme, hostname, and port, for 192 a fixture resource from the server associated with the ``on`` key. 193 It will by default look for the resource in the "http" server. 194 195 """ 196 return servers.get(on)[1].get_url(uri) 197 198 199 def iter_proc(servers): 200 for _, (_, proc) in servers.items(): 201 yield proc 202 203 204 def iter_url(servers): 205 for _, (url, _) in servers.items(): 206 yield url 207 208 209 default_doc_root = os.path.join(os.path.dirname(here), "www") 210 registered_servers = [("http", http_server), ("https", https_server)] 211 servers = defaultdict() 212 213 214 def main(args): 215 global servers 216 217 parser = argparse.ArgumentParser() 218 parser.add_argument( 219 "-r", dest="doc_root", help="Path to document root. Overrides default." 220 ) 221 args = parser.parse_args() 222 223 servers = start(args.doc_root) 224 for url in iter_url(servers): 225 print(f"{sys.argv[0]}: listening on {url}", file=sys.stderr) 226 227 try: 228 while any(proc.is_alive for proc in iter_proc(servers)): 229 for proc in iter_proc(servers): 230 proc.proc.join(1) 231 except KeyboardInterrupt: 232 for proc in iter_proc(servers): 233 proc.kill() 234 235 236 if __name__ == "__main__": 237 main(sys.argv[1:])