| OLD | NEW |
| 1 # Copyright 2014 The LUCI Authors. All rights reserved. | 1 # Copyright 2014 The LUCI Authors. All rights reserved. |
| 2 # Use of this source code is governed under the Apache License, Version 2.0 | 2 # Use of this source code is governed under the Apache License, Version 2.0 |
| 3 # that can be found in the LICENSE file. | 3 # that can be found in the LICENSE file. |
| 4 | 4 |
| 5 import BaseHTTPServer | |
| 6 import base64 | 5 import base64 |
| 7 import hashlib | 6 import hashlib |
| 8 import json | 7 import json |
| 9 import logging | 8 import logging |
| 10 import re | 9 import re |
| 11 import sys | |
| 12 import threading | |
| 13 import urllib2 | |
| 14 import urlparse | |
| 15 import zlib | 10 import zlib |
| 16 | 11 |
| 12 import httpserver_mock |
| 13 |
| 17 ALGO = hashlib.sha1 | 14 ALGO = hashlib.sha1 |
| 18 | 15 |
| 19 | 16 |
| 20 def hash_content(content): | 17 def hash_content(content): |
| 21 return ALGO(content).hexdigest() | 18 return ALGO(content).hexdigest() |
| 22 | 19 |
| 23 | 20 |
| 24 class FakeSigner(object): | 21 class FakeSigner(object): |
| 25 | 22 |
| 26 @classmethod | 23 @classmethod |
| 27 def generate(cls, message, embedded): | 24 def generate(cls, message, embedded): |
| 28 return '%s_<<<%s>>>' % (repr(message), json.dumps(embedded)) | 25 return '%s_<<<%s>>>' % (repr(message), json.dumps(embedded)) |
| 29 | 26 |
| 30 @classmethod | 27 @classmethod |
| 31 def validate(cls, ticket, message): | 28 def validate(cls, ticket, message): |
| 32 a = re.match(r'^' + repr(message) + r'_<<<(.*)>>>$', ticket, re.DOTALL) | 29 a = re.match(r'^' + repr(message) + r'_<<<(.*)>>>$', ticket, re.DOTALL) |
| 33 if not a: | 30 if not a: |
| 34 raise ValueError('Message %s cannot validate ticket %s' % ( | 31 raise ValueError('Message %s cannot validate ticket %s' % ( |
| 35 repr(message), ticket)) | 32 repr(message), ticket)) |
| 36 return json.loads(a.groups()[0]) | 33 return json.loads(a.groups()[0]) |
| 37 | 34 |
| 38 | 35 |
| 39 class MockHandler(BaseHTTPServer.BaseHTTPRequestHandler): | 36 class IsolateServerHandler(httpserver_mock.MockHandler): |
| 40 def _json(self, data): | |
| 41 """Sends a JSON response.""" | |
| 42 self.send_response(200) | |
| 43 self.send_header('Content-type', 'application/json') | |
| 44 self.end_headers() | |
| 45 json.dump(data, self.wfile) | |
| 46 | |
| 47 def _octet_stream(self, data): | |
| 48 """Sends a binary response.""" | |
| 49 self.send_response(200) | |
| 50 self.send_header('Content-type', 'application/octet-stream') | |
| 51 self.end_headers() | |
| 52 self.wfile.write(data) | |
| 53 | |
| 54 def _read_body(self): | |
| 55 """Reads the request body.""" | |
| 56 return self.rfile.read(int(self.headers['Content-Length'])) | |
| 57 | |
| 58 def _drop_body(self): | |
| 59 """Reads the request body.""" | |
| 60 size = int(self.headers['Content-Length']) | |
| 61 while size: | |
| 62 chunk = min(4096, size) | |
| 63 self.rfile.read(chunk) | |
| 64 size -= chunk | |
| 65 | |
| 66 def log_message(self, fmt, *args): | |
| 67 logging.info( | |
| 68 '%s - - [%s] %s', self.address_string(), self.log_date_time_string(), | |
| 69 fmt % args) | |
| 70 | |
| 71 | |
| 72 class IsolateServerHandler(MockHandler): | |
| 73 """An extremely minimal implementation of the isolate server API v1.0.""" | 37 """An extremely minimal implementation of the isolate server API v1.0.""" |
| 74 | 38 |
| 75 def _should_push_to_gs(self, isolated, size): | 39 def _should_push_to_gs(self, isolated, size): |
| 76 max_memcache = 500 * 1024 | 40 max_memcache = 500 * 1024 |
| 77 min_direct_gs = 501 | 41 min_direct_gs = 501 |
| 78 if isolated and size <= max_memcache: | 42 if isolated and size <= max_memcache: |
| 79 return False | 43 return False |
| 80 return size >= min_direct_gs | 44 return size >= min_direct_gs |
| 81 | 45 |
| 82 def _generate_signed_url(self, digest, namespace='default'): | 46 def _generate_signed_url(self, digest, namespace='default'): |
| (...skipping 92 matching lines...) Expand 10 before | Expand all | Expand 10 after Loading... |
| 175 else: | 139 else: |
| 176 body = self._read_body() | 140 body = self._read_body() |
| 177 if self.path.startswith('/FAKE_GCS/'): | 141 if self.path.startswith('/FAKE_GCS/'): |
| 178 namespace, h = self.path[len('/FAKE_GCS/'):].split('/', 1) | 142 namespace, h = self.path[len('/FAKE_GCS/'):].split('/', 1) |
| 179 self.server.contents.setdefault(namespace, {})[h] = body | 143 self.server.contents.setdefault(namespace, {})[h] = body |
| 180 self._octet_stream('') | 144 self._octet_stream('') |
| 181 else: | 145 else: |
| 182 raise NotImplementedError(self.path) | 146 raise NotImplementedError(self.path) |
| 183 | 147 |
| 184 | 148 |
| 185 class MockServer(object): | 149 class MockIsolateServer(httpserver_mock.MockServer): |
| 186 _HANDLER_CLS = None | |
| 187 | |
| 188 def __init__(self): | |
| 189 self._closed = False | |
| 190 self._server = BaseHTTPServer.HTTPServer( | |
| 191 ('127.0.0.1', 0), self._HANDLER_CLS) | |
| 192 self._server.url = self.url = 'http://localhost:%d' % ( | |
| 193 self._server.server_port) | |
| 194 self._thread = threading.Thread(target=self._run, name='httpd') | |
| 195 self._thread.daemon = True | |
| 196 self._thread.start() | |
| 197 logging.info('%s', self.url) | |
| 198 | |
| 199 def close(self): | |
| 200 self.close_start() | |
| 201 self.close_end() | |
| 202 | |
| 203 def close_start(self): | |
| 204 assert not self._closed | |
| 205 self._closed = True | |
| 206 urllib2.urlopen(self.url + '/on/quit') | |
| 207 | |
| 208 def close_end(self): | |
| 209 assert self._closed | |
| 210 self._thread.join() | |
| 211 | |
| 212 def _run(self): | |
| 213 while not self._closed: | |
| 214 self._server.handle_request() | |
| 215 | |
| 216 | |
| 217 class MockIsolateServer(MockServer): | |
| 218 _HANDLER_CLS = IsolateServerHandler | 150 _HANDLER_CLS = IsolateServerHandler |
| 219 | 151 |
| 220 def __init__(self): | 152 def __init__(self): |
| 221 super(MockIsolateServer, self).__init__() | 153 super(MockIsolateServer, self).__init__() |
| 222 self._server.contents = {} | 154 self._server.contents = {} |
| 223 self._server.discard_content = False | 155 self._server.discard_content = False |
| 224 | 156 |
| 225 def discard_content(self): | 157 def discard_content(self): |
| 226 """Stops saving content in memory. Used to test large files.""" | 158 """Stops saving content in memory. Used to test large files.""" |
| 227 self._server.discard_content = True | 159 self._server.discard_content = True |
| (...skipping 10 matching lines...) Expand all Loading... |
| 238 zlib.compress(content)) | 170 zlib.compress(content)) |
| 239 return h | 171 return h |
| 240 | 172 |
| 241 def add_content(self, namespace, content): | 173 def add_content(self, namespace, content): |
| 242 assert not self._server.discard_content | 174 assert not self._server.discard_content |
| 243 h = hash_content(content) | 175 h = hash_content(content) |
| 244 logging.info('add_content(%s, %s)', namespace, h) | 176 logging.info('add_content(%s, %s)', namespace, h) |
| 245 self._server.contents.setdefault(namespace, {})[h] = base64.b64encode( | 177 self._server.contents.setdefault(namespace, {})[h] = base64.b64encode( |
| 246 content) | 178 content) |
| 247 return h | 179 return h |
| OLD | NEW |