count-non-zero

Problem

This problem comes from pwn.college: Assembly Crash Course ⤴.

In this level, you will be working with control flow manipulation. This involves using instructions to both indirectly and directly control the special register rip, the instruction pointer. You will use instructions such as jmp, call, cmp, and their alternatives to implement the requested behavior.

We will be testing your code multiple times in this level with dynamic values! This means we will be running your code in a variety of random ways to verify that the logic is robust enough to survive normal use.

In previous levels, you discovered the for-loop to iterate for a number of times, both dynamically and statically known, but what happens when you want to iterate until you meet a condition?

A second loop structure exists called the while-loop to fill this demand. In the while-loop, you iterate until a condition is met.

As an example, say we had a location in memory with adjacent numbers and we wanted to get the average of all the numbers until we find one bigger or equal to 0xff:

average = 0
i = 0
while x[i] < 0xff:
  average += x[i]
  i += 1
average /= i

Using the above knowledge, please perform the following:

Count the consecutive non-zero bytes in a contiguous region of memory, where:

• rdi = memory address of the 1st byte
• rax = number of consecutive non-zero bytes

Additionally, if rdi = 0, then set rax = 0 (we will check)!

An example test-case, let:

• rdi = 0x1000
• [0x1000] = 0x41
• [0x1001] = 0x42
• [0x1002] = 0x43
• [0x1003] = 0x00

Then: rax = 3 should be set.

Solution

I decided that rax will serve as both the counter and byte offset. By doing this, we can pull a couple of tricks.

The first trick is that we immediately set rax to zero. If rdi is also zero, we’re done! Else, rax is just a zeroed out counter.

The second trick was already mentioned. rax is a counter, and we can use that as an offset for rdi. Thus everytime we find a consecutive non-zero byte, the counter goes up which also allows use to conveniently check the next byte.

.intel_syntax noprefix
.global _start

_start:
  # Set rax = 0. If rdi = 0, then we're good!
  xor rax, rax
  cmp rdi, 0
  je end

check:
  # rax serves as both counter and byte offset
  cmp qword ptr [rdi + rax], 0
  je end
  
  add rax, 1
  jmp check

end: