merchant_fault_proxy.py (12123B)
1 #!/usr/bin/env python3 2 """Merchant frontend that can reset selected requests before access log. 3 4 The proxy is deliberately small and HTTP/1.1-only. It accepts TCP or Unix 5 socket clients, optionally terminates TLS, keeps client and upstream 6 connections alive, and forwards one 7 request/response at a time. For the configured path, the first N requests are 8 reset after their headers have arrived but before a byte is sent upstream. 9 Consequently its FAULT log proves the transport attempt happened, while 10 neither its ACCESS log nor the merchant backend sees an HTTP request. 11 """ 12 13 import argparse 14 import os 15 import socket 16 import ssl 17 import struct 18 import threading 19 import time 20 from urllib.parse import urlsplit 21 22 23 def log(message): 24 print(message, flush=True) 25 26 27 def recv_until(sock, buffer, marker): 28 while marker not in buffer: 29 chunk = sock.recv(65536) 30 if not chunk: 31 return None, b"" 32 buffer += chunk 33 if len(buffer) > 1024 * 1024: 34 raise RuntimeError("HTTP header exceeded 1 MiB") 35 head, buffer = buffer.split(marker, 1) 36 return head + marker, buffer 37 38 39 def content_length(header): 40 for line in header.split(b"\r\n")[1:]: 41 name, sep, value = line.partition(b":") 42 if sep and name.strip().lower() == b"content-length": 43 return int(value.strip()) 44 return 0 45 46 47 def connection_closes(header): 48 for line in header.split(b"\r\n")[1:]: 49 name, sep, value = line.partition(b":") 50 if sep and name.strip().lower() == b"connection": 51 return b"close" in value.lower() 52 return False 53 54 55 def transfer_is_chunked(header): 56 for line in header.split(b"\r\n")[1:]: 57 name, sep, value = line.partition(b":") 58 if sep and name.strip().lower() == b"transfer-encoding": 59 return b"chunked" in value.lower() 60 return False 61 62 63 def relay_exact(source, destination, initial, length): 64 data = initial 65 if len(data) > length: 66 destination.sendall(data[:length]) 67 return data[length:] 68 if data: 69 destination.sendall(data) 70 length -= len(data) 71 while length: 72 chunk = source.recv(min(65536, length)) 73 if not chunk: 74 raise EOFError("connection closed in fixed-length HTTP body") 75 destination.sendall(chunk) 76 length -= len(chunk) 77 return b"" 78 79 80 def relay_chunked(source, destination, initial): 81 buffer = initial 82 while True: 83 line, buffer = recv_until(source, buffer, b"\r\n") 84 if line is None: 85 raise EOFError("connection closed in chunk-size line") 86 destination.sendall(line) 87 size_text = line[:-2].split(b";", 1)[0] 88 size = int(size_text, 16) 89 if size: 90 buffer = relay_exact(source, destination, buffer, size + 2) 91 continue 92 # The zero chunk is followed by zero or more trailer lines and one 93 # blank line. Reading a line at a time also handles the usual no- 94 # trailer form, which contains only that final CRLF. 95 while True: 96 trailer, buffer = recv_until(source, buffer, b"\r\n") 97 if trailer is None: 98 raise EOFError("connection closed in chunk trailers") 99 destination.sendall(trailer) 100 if trailer == b"\r\n": 101 return buffer 102 103 104 def reset_connection(sock): 105 try: 106 sock.setsockopt(socket.SOL_SOCKET, 107 socket.SO_LINGER, 108 struct.pack("ii", 1, 0)) 109 except OSError: 110 pass 111 try: 112 sock.close() 113 except OSError: 114 pass 115 116 117 class FaultBudget: 118 def __init__(self, path, count, control_file): 119 self.path = path 120 self.remaining = count 121 self.issued = 0 122 self.control_file = control_file 123 self.control_version = None 124 self.lock = threading.Lock() 125 126 def refresh(self): 127 if not self.control_file: 128 return 129 try: 130 stat = os.stat(self.control_file) 131 version = (stat.st_mtime_ns, stat.st_size) 132 if version == self.control_version: 133 return 134 with open(self.control_file, encoding="utf-8") as stream: 135 fields = stream.read().strip().split() 136 except FileNotFoundError: 137 return 138 if len(fields) != 2: 139 raise RuntimeError( 140 "fault control file must contain: PATH COUNT") 141 self.path = fields[0] 142 self.remaining = int(fields[1]) 143 self.issued = 0 144 self.control_version = version 145 log(f"ARM path={self.path} count={self.remaining}") 146 147 def consume(self, path): 148 with self.lock: 149 self.refresh() 150 if path != self.path or self.remaining == 0: 151 return None 152 self.remaining -= 1 153 self.issued += 1 154 return self.issued 155 156 157 class Proxy: 158 def __init__(self, args): 159 self.args = args 160 self.faults = FaultBudget(args.fault_path, 161 args.fault_count, 162 args.control_file) 163 self.context = None 164 if args.cert: 165 self.context = ssl.SSLContext(ssl.PROTOCOL_TLS_SERVER) 166 self.context.minimum_version = ssl.TLSVersion.TLSv1_2 167 self.context.set_alpn_protocols(["http/1.1"]) 168 self.context.load_cert_chain(args.cert, args.key) 169 self.connection_id = 0 170 self.connection_lock = threading.Lock() 171 172 def next_connection_id(self): 173 with self.connection_lock: 174 self.connection_id += 1 175 return self.connection_id 176 177 def handle(self, raw_client): 178 connection_id = self.next_connection_id() 179 stage = "TLS handshake" 180 if self.context: 181 try: 182 client = self.context.wrap_socket(raw_client, server_side=True) 183 except (ssl.SSLError, OSError): 184 raw_client.close() 185 return 186 else: 187 client = raw_client 188 try: 189 stage = "upstream connect" 190 upstream = socket.create_connection( 191 (self.args.upstream_host, self.args.upstream_port), timeout=2) 192 upstream.settimeout(self.args.io_timeout) 193 client.settimeout(self.args.io_timeout) 194 except OSError as exc: 195 log(f"UPSTREAM_CONNECT_ERROR connection={connection_id} error={exc}") 196 reset_connection(client) 197 return 198 199 client_buffer = b"" 200 upstream_buffer = b"" 201 prior_requests = 0 202 last_response = time.monotonic() 203 try: 204 while True: 205 stage = "request headers from client" 206 request, client_buffer = recv_until( 207 client, client_buffer, b"\r\n\r\n") 208 if request is None: 209 return 210 request_line = request.split(b"\r\n", 1)[0].decode( 211 "iso-8859-1", "replace") 212 parts = request_line.split(" ") 213 if len(parts) != 3: 214 raise RuntimeError(f"malformed request line {request_line!r}") 215 method, target, _ = parts 216 path = urlsplit(target).path 217 idle_ms = int((time.monotonic() - last_response) * 1000) 218 fault_index = self.faults.consume(path) 219 if fault_index is not None: 220 log("FAULT " 221 f"index={fault_index} connection={connection_id} " 222 f"path={path} reused={'yes' if prior_requests else 'no'} " 223 f"prior_requests={prior_requests} idle_ms={idle_ms}") 224 reset_connection(client) 225 upstream.close() 226 return 227 228 body_length = content_length(request) 229 stage = "request headers to upstream" 230 upstream.sendall(request) 231 stage = "request body to upstream" 232 client_buffer = relay_exact( 233 client, upstream, client_buffer, body_length) 234 prior_requests += 1 235 log("ACCESS " 236 f"connection={connection_id} request={prior_requests} " 237 f"method={method} path={path}") 238 239 stage = "response headers from upstream" 240 response, upstream_buffer = recv_until( 241 upstream, upstream_buffer, b"\r\n\r\n") 242 if response is None: 243 raise EOFError("upstream closed before response headers") 244 stage = "response headers to client" 245 client.sendall(response) 246 status_parts = response.split(b"\r\n", 1)[0].split(b" ") 247 no_body = (method == "HEAD" or 248 (len(status_parts) > 1 and 249 (status_parts[1][:1] == b"1" or 250 status_parts[1] in (b"204", b"304")))) 251 if transfer_is_chunked(response): 252 stage = "chunked response body to client" 253 upstream_buffer = relay_chunked( 254 upstream, client, upstream_buffer) 255 elif not no_body: 256 length = content_length(response) 257 if length: 258 stage = "fixed response body to client" 259 upstream_buffer = relay_exact( 260 upstream, client, upstream_buffer, length) 261 last_response = time.monotonic() 262 if connection_closes(request) or connection_closes(response): 263 return 264 except (EOFError, OSError, RuntimeError, ValueError) as exc: 265 log(f"CONNECTION_ERROR connection={connection_id} " 266 f"stage={stage} error={exc}") 267 finally: 268 try: 269 client.close() 270 except OSError: 271 pass 272 try: 273 upstream.close() 274 except OSError: 275 pass 276 277 def serve(self): 278 if self.args.listen_unix: 279 if os.path.exists(self.args.listen_unix): 280 raise RuntimeError( 281 f"Unix listener path already exists: {self.args.listen_unix}") 282 listener = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) 283 listener.bind(self.args.listen_unix) 284 listen_address = self.args.listen_unix 285 else: 286 listener = socket.socket(socket.AF_INET, socket.SOCK_STREAM) 287 listener.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) 288 listener.bind((self.args.listen_host, self.args.listen_port)) 289 listen_address = (f"{listener.getsockname()[0]}:" 290 f"{listener.getsockname()[1]}") 291 listener.listen(128) 292 log(f"READY {listen_address}") 293 while True: 294 client, _ = listener.accept() 295 thread = threading.Thread(target=self.handle, 296 args=(client,), 297 daemon=True) 298 thread.start() 299 300 301 def main(): 302 parser = argparse.ArgumentParser() 303 parser.add_argument("--listen-host", default="127.0.0.1") 304 listen = parser.add_mutually_exclusive_group(required=True) 305 listen.add_argument("--listen-port", type=int) 306 listen.add_argument("--listen-unix") 307 parser.add_argument("--upstream-host", default="127.0.0.1") 308 parser.add_argument("--upstream-port", type=int, required=True) 309 parser.add_argument("--fault-path") 310 parser.add_argument("--fault-count", type=int, default=2) 311 parser.add_argument("--control-file") 312 parser.add_argument("--io-timeout", type=float, default=300) 313 parser.add_argument("--cert") 314 parser.add_argument("--key") 315 args = parser.parse_args() 316 if bool(args.cert) != bool(args.key): 317 parser.error("--cert and --key must be specified together") 318 if not args.fault_path and not args.control_file: 319 parser.error("one of --fault-path or --control-file is required") 320 Proxy(args).serve() 321 322 323 if __name__ == "__main__": 324 main()