improve logging and support configurable mime types

app.py

@@ -20,11 +20,12 @@ with open(sys.argv[1] if len(sys.argv) > 1 else 'config.toml') as f:
 	config = toml.load(f)
 
 for host in config['hosts'].values():
-	if not host.get('replace', True): continue
+	if not host.setdefault('replace', True): continue
 	host['pattern_decoded'] = host['pattern']
 	host['pattern'] = re.compile(host['pattern'].encode())
 	host['repl_decoded'] = host['repl']
 	host['repl'] = host['repl'].encode()
+	host['mime_types'] = frozenset(host.get('mime_types', ()))
 
 dprint = build_dprint(config.get('debug'))
 http = ContextVar('http')
@@ -41,25 +42,38 @@ async def handler(request):
 		# proxy redirects as-is
 		allow_redirects=False,
 	) as upstream_resp:
-		print('proceeding for', request.host, request.path, '→', hconfig['upstream'])
+		print(request.method, request.host, request.path, '→', hconfig['upstream'])
 		headers = upstream_resp.headers.copy()
 		# we're not using gzip here so don't confuse our client
 		with contextlib.suppress(KeyError): del headers['Content-Encoding']
 		resp = web.StreamResponse(status=upstream_resp.status, headers=headers)
 		await resp.prepare(request)
 
-		if upstream_resp.content_type == 'text/html' and hconfig.get('replace', True):
-			# iter_lines when
-			print('replacing', repr(hconfig['pattern_decoded']), 'with', repr(hconfig['repl_decoded']))
-			while (line := await upstream_resp.content.readline()):
-				await resp.write(hconfig['pattern'].sub(hconfig['repl'], line))
-		else:
-			print('not replacing')
-			async for chunk in upstream_resp.content.iter_chunked(io.DEFAULT_BUFFER_SIZE):
-				await resp.write(chunk)
+		if not hconfig['replace']:
+			print('Not replacing for this host')
+			return await proxy_passthrough(upstream_resp, resp)
+		if hconfig['mime_types'] and upstream_resp.content_type not in hconfig['mime_types']:
+			print('Not configured to replace for MIME type', upstream_resp.content_type)
+			return await proxy_passthrough(upstream_resp, resp)
 
-		await resp.write_eof()
-		return resp
+		return await proxy_replace(hconfig, upstream_resp, resp)
+
+async def proxy_replace(hconfig, upstream_resp, resp):
+	# iter_lines when
+	while (line := await upstream_resp.content.readline()):
+		await resp.write(hconfig['pattern'].sub(hconfig['repl'], line))
+
+	return await finalize_resp(resp)
+
+async def proxy_passthrough(upstream_resp, resp):
+	async for chunk in upstream_resp.content.iter_chunked(io.DEFAULT_BUFFER_SIZE):
+		await resp.write(chunk)
+
+	return await finalize_resp(resp)
+
+async def finalize_resp(resp):
+	await resp.write_eof()
+	return resp
 
 async def amain():
 	# It's not our job to hand out 413s, that's the job of upstream and downstream

config.example.toml

@@ -13,6 +13,8 @@ upstream = 'http://localhost:3001'
 # these can be regexes
 pattern = 'foo'
 repl = 'bar'
+# which mime types to replace for. defaults to all mime types.
+mime_types = ['text/html', 'application/json', 'application/activity+json', 'application/ld+json']
 
 # just pass through site2.example unmodified
 [hosts."site2.example"]