Hostヘッダを任意に書き換えてHTTPサーバーに転送するプロキシ

首題の通り、HTTPのHostヘッダを任意に書き換えてHTTPサーバーに転送するプロキシ機能をpythonで書いてみたというもの。KubernetesのIngressに定義した仮想ホストにcurlで--resolve付けて接続先アプリを切り替えるのを、プロキシ側で行うようなものだ。

やってみようかなと思ったのは、RHEL8に付属するHAproxy 1.8が、backendへのcheckに期待したHostヘッダを付けてくれないという問題に対処するため。作ってみる前はもうちょっと、ちょちょいと出来るんじゃないのと思ったのだが。。
が、習作ではあるが作ったスクリプトは割と実用に耐える気もする。最大同時接続数は10で、接続のタイムアウトを60秒にしている、SIGINTとSIGTERMに反応してサービスは1秒くらいで終了する。
なお、作ってみた後に分かったが同じことをするもっと簡単な方法は他にもいくつかある。

実行は`python3 hcproxy.py`で。

```python: hcproxy.py
import socket
import threading
import time
import re
import signal

term_flag = False

def shstop(p):
  global term_flag
  if term_flag:
    return True
  if time.time() > p["timeout"]:
    return True
  if p["disconnected"]:
    return True
  return False

def readline(p, max_len):
  r = b""
  while True:
    if shstop(p):
      return b""
    try:
      b = p["to_client"].recv(1)
    except socket.timeout:
      continue
    except:
      raise
    if len(b) == 0:
      break
    r += b
    if b == b"\n":
      break
    if len(r) >= max_len:
      break
  return r

def client_to_server(p):
  p["to_client"].settimeout(1.0)
  while True:
    l = readline(p, 4096)
    if shstop(p):
      return
    if len(l) == 0:
      return
    if re.match(r'^Host\: .*', l.decode("utf-8"), re.I):
      l = b"Host: www.example.com:8081\r\n"
    p["to_server"].sendall(l)
    if l == b"\r\n":
      break
  while True:
    if shstop(p):
      return
    try:
      d = p["to_client"].recv(4096)
    except socket.timeout:
      continue
    except:
      raise
    if len(d) == 0:
      break
    p["to_server"].sendall(d)
  p["disconnected"] = True

def server_to_client(p):
  p["to_server"].settimeout(1.0)
  while True:
    if shstop(p):
      return
    try:
      d = p["to_server"].recv(4096)
    except socket.timeout:
      continue
    except:
      raise
    if len(d) == 0:
      break
    p["to_client"].sendall(d)
  p["disconnected"] = True

def signal_handler(sig, frame):
  global term_flag
  if sig == signal.SIGTERM or sig == signal.SIGINT:
    term_flag = True

def watchdog(pool, lock):
  global term_flag
  while not term_flag:
    with lock:
      for i, p in enumerate(pool):
        if not p["cts"].is_alive() and not p["stc"].is_alive():
          p["to_client"].close()
          p["to_server"].close()
          print("[" + p["cid"] + "] session closed")
          pool.pop(i)
    time.sleep(1)

if __name__ == "__main__":
  print("starting server")
  signal.signal(signal.SIGTERM, signal_handler)
  signal.signal(signal.SIGINT, signal_handler)
  pool = []
  lock = threading.Lock()
  cid = 0
  service = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
  service.bind(("localhost", 8080))
  service.listen(1)
  print("server started on localhost:8080")
  print("proxying to localhost:8081")
  print("with header 'Host: www.example.com:8081'")
  service.settimeout(1.0)
  wd = threading.Thread(target=watchdog, args=(pool, lock, ))
  wd.start()
  while True:
    if term_flag:
      print("terminating server")
      break
    try:
      to_client, addr = service.accept()
    except socket.timeout:
      continue
    except:
      raise
    if len(pool) > 10:
      print("too many sessions")
      to_client.close()
      continue
    print("[" + str(cid) + "] starting session")
    to_server = socket.create_connection(("localhost", 8081))
    p = {}
    p["to_client"] = to_client
    p["to_server"] = to_server
    p["timeout"] = time.time() + 60.0
    p["disconnected"] = False
    cts = threading.Thread(target=client_to_server, args=(p, ))
    cts.start()
    stc = threading.Thread(target=server_to_client, args=(p, ))
    stc.start()
    p["cts"] = cts
    p["stc"] = stc
    p["cid"] = str(cid)
    cid += 1
    with lock:
      pool.append(p)
```