improve logging and support configurable mime types

This commit is contained in:
Kay Faraday 2021-09-08 03:07:44 +00:00
parent 44f956c1a3
commit 50c22cd5b9
2 changed files with 29 additions and 13 deletions

40
app.py
View File

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

View File

@ -13,6 +13,8 @@ upstream = 'http://localhost:3001'
# these can be regexes # these can be regexes
pattern = 'foo' pattern = 'foo'
repl = 'bar' 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 # just pass through site2.example unmodified
[hosts."site2.example"] [hosts."site2.example"]