#!/usr/bin/env python3
"""Apply the combined [ebusy_force + 13114dc-adapted] changes to a COPY of the
pristine v6.6.18 net/tls/tls_sw.c, to verify that the 13114dc fix closes the
failed-backlog-decryption UAF bypass of the 32b55c5 (free_sgout) fix.

The 13114dc upstream patch was authored against a slightly different base, so we
adapt its intent to v6.6.18's tls_do_decryption structure:
  - add async_done to tls_decrypt_arg.inargs
  - on the -EBUSY backlog path, set async_done=true / async=false and return the
    wait result directly (the async callback tls_decrypt_done already ran and
    freed the pages + kfree()'d the aead_req/dctx block)
  - in tls_decrypt_sg, on error after tls_do_decryption, if async_done skip
    exit_free_pages/exit_free (goto exit_free_skb) to avoid the double-free
  - after the async-hold check, if async_done return 0 (data ready immediately)
The ebusy_force change (force -EINPROGRESS -> -EBUSY for async decrypts) is also
included so the backlog path is reachable on demand.
"""
import sys, shutil

SRC = "/data/pruva/project-cache/6cba72b6-b104-4ebf-be7b-01f748e87ccf/linux-6.6.18/net/tls/tls_sw.c.orig"
DST = "/data/pruva/project-cache/6cba72b6-b104-4ebf-be7b-01f748e87ccf/linux-6.6.18/net/tls/tls_sw.c"

with open(SRC) as f:
    s = f.read()

def replace(old, new, label):
    global s
    if old not in s:
        print(f"FAIL: pattern not found for {label!r}", file=sys.stderr)
        sys.exit(1)
    if s.count(old) != 1:
        print(f"FAIL: pattern not unique for {label!r} (count={s.count(old)})", file=sys.stderr)
        sys.exit(1)
    s = s.replace(old, new, 1)
    print(f"OK: {label}")

# 1) add async_done to inargs
replace(
    "\tbool async;\n\tu8 tail;",
    "\tbool async;\n\tbool async_done;\n\tu8 tail;",
    "inargs async_done",
)

# 2) tls_do_decryption: ebusy_force + 13114dc -EBUSY handling
replace(
    "\tret = crypto_aead_decrypt(aead_req);\n"
    "\tif (ret == -EBUSY) {\n"
    "\t\tret = tls_decrypt_async_wait(ctx);\n"
    "\t\tret = ret ?: -EINPROGRESS;\n"
    "\t}\n"
    "\tif (ret == -EINPROGRESS) {\n"
    "\t\tif (darg->async)\n"
    "\t\t\treturn 0;\n\n"
    "\t\tret = crypto_wait_req(ret, &ctx->async_wait);\n"
    "\t}\n"
    "\tdarg->async = false;\n",
    "\tret = crypto_aead_decrypt(aead_req);\n"
    "\t/* [VARIANT] simulate crypto backlog congestion for async decrypts */\n"
    "\tif (darg->async && ret == -EINPROGRESS)\n"
    "\t\tret = -EBUSY;\n"
    "\tif (ret == -EBUSY) {\n"
    "\t\tret = tls_decrypt_async_wait(ctx);\n"
    "\t\t/* [FIX 13114dc] async callback already ran and freed the memory */\n"
    "\t\tdarg->async_done = true;\n"
    "\t\tdarg->async = false;\n"
    "\t\treturn ret;\n"
    "\t}\n"
    "\tif (ret == -EINPROGRESS) {\n"
    "\t\tif (darg->async)\n"
    "\t\t\treturn 0;\n\n"
    "\t\tret = crypto_wait_req(ret, &ctx->async_wait);\n"
    "\t}\n"
    "\tdarg->async = false;\n",
    "tls_do_decryption ebusy+13114dc",
)

# 3) tls_decrypt_sg error path: skip double-free when async_done
replace(
    "\tif (err)\n"
    "\t\tgoto exit_free_pages;\n",
    "\tif (err) {\n"
    "\t\t/* [FIX 13114dc] callback already freed pages/mem; skip double-free */\n"
    "\t\tif (darg->async_done)\n"
    "\t\t\tgoto exit_free_skb;\n"
    "\t\tgoto exit_free_pages;\n"
    "\t}\n",
    "tls_decrypt_sg error async_done skip",
)

# 4) tls_decrypt_sg after async-hold: return 0 when async_done
replace(
    "\t\treturn err;\n"
    "\t}\n\n"
    "\tif (prot->tail_size)\n"
    "\t\tdarg->tail = dctx->tail;\n",
    "\t\treturn err;\n"
    "\t}\n\n"
    "\t/* [FIX 13114dc] backlog already completed; data ready for immediate copy */\n"
    "\tif (unlikely(darg->async_done))\n"
    "\t\treturn 0;\n\n"
    "\tif (prot->tail_size)\n"
    "\t\tdarg->tail = dctx->tail;\n",
    "tls_decrypt_sg async_done return0",
)

with open(DST, "w") as f:
    f.write(s)
print("WROTE", DST)