Extend PASS and DROP opcode encoding to support incremental counter

This commit implements COUNT opcode by adding support for incremental
counters to the PASS and DROP opcodes. If counter number is present in the
instruction, the apf_interpreter will increment the relevant counter and
return PASS/DROP immediately, skipping over subsequent instructions.

* design doc: go/apf-v6-proposal

Bug: 293811969
Test: TH
Change-Id: I2430719782a384ad9a7147e42c447ec797e66c21
diff --git a/src/android/net/apf/ApfGenerator.java b/src/android/net/apf/ApfGenerator.java
index fe5628c..e63dddb 100644
--- a/src/android/net/apf/ApfGenerator.java
+++ b/src/android/net/apf/ApfGenerator.java
@@ -41,8 +41,15 @@
     }
     private enum Opcodes {
         LABEL(-1),
-        PASS(0),   // Unconditionally pass packet, requires R=0, LEN=0, e.g. "pass"
-        DROP(0),   // Unconditionally drop packet, requires R=1, LEN=0, e.g. "drop"
+        // Unconditionally pass (if R=0) or drop (if R=1) packet.
+        // An optional unsigned immediate value can be provided to encode the counter number.
+        // If the value is non-zero, the instruction increments the counter.
+        // The counter is located (-4 * counter number) bytes from the end of the data region.
+        // It is a U32 big-endian value and is always incremented by 1.
+        // This is more or less equivalent to: lddw R0, -N4; add R0,1; stdw R0, -N4; {pass,drop}
+        // e.g. "pass", "pass 1", "drop", "drop 1"
+        PASS(0),
+        DROP(0),
         LDB(1),    // Load 1 byte from immediate offset, e.g. "ldb R0, [5]"
         LDH(2),    // Load 2 bytes from immediate offset, e.g. "ldh R0, [5]"
         LDW(3),    // Load 4 bytes from immediate offset, e.g. "ldw R0, [5]"
@@ -932,6 +939,19 @@
     }
 
     /**
+     * Add an instruction to the end of the program to increment the counter value and
+     * immediately return PASS.
+     */
+    public ApfGenerator addCountAndPass(int counterNumber) throws IllegalInstructionException {
+        requireApfVersion(MIN_APF_VERSION_IN_DEV);
+        checkCounterNumber(counterNumber);
+        Instruction instruction = new Instruction(Opcodes.PASS, Register.R0);
+        instruction.addUnsignedImm(counterNumber);
+        addInstruction(instruction);
+        return this;
+    }
+
+    /**
      * Add an instruction to the end of the program to let the program immediately return DROP.
      */
     public ApfGenerator addDrop() throws IllegalInstructionException {
@@ -942,6 +962,19 @@
     }
 
     /**
+     * Add an instruction to the end of the program to increment the counter value and
+     * immediately return DROP.
+     */
+    public ApfGenerator addCountAndDrop(int counterNumber) throws IllegalInstructionException {
+        requireApfVersion(MIN_APF_VERSION_IN_DEV);
+        checkCounterNumber(counterNumber);
+        Instruction instruction = new Instruction(Opcodes.DROP, Register.R1);
+        instruction.addUnsignedImm(counterNumber);
+        addInstruction(instruction);
+        return this;
+    }
+
+    /**
      * Add an instruction to the end of the program to call the apf_allocate_buffer() function.
      *
      * @param register the register value contains the buffer size.
@@ -1128,6 +1161,13 @@
         }
     }
 
+    private void checkCounterNumber(int counterNumber) {
+        if (counterNumber < 1 || counterNumber > 1000) {
+            throw new IllegalArgumentException(
+                    "Counter number must be in range (0, 1000], counterNumber: " + counterNumber);
+        }
+    }
+
     /**
      * Add an instruction to the end of the program to load 32 bits from the data memory into
      * {@code register}. The source address is computed by adding the signed immediate
diff --git a/tests/unit/src/android/net/apf/ApfV5Test.kt b/tests/unit/src/android/net/apf/ApfV5Test.kt
index 447495c..1c21bb0 100644
--- a/tests/unit/src/android/net/apf/ApfV5Test.kt
+++ b/tests/unit/src/android/net/apf/ApfV5Test.kt
@@ -18,6 +18,7 @@
 import android.net.apf.ApfGenerator.IllegalInstructionException
 import androidx.test.filters.SmallTest
 import androidx.test.runner.AndroidJUnit4
+import java.lang.IllegalArgumentException
 import kotlin.test.assertContentEquals
 import kotlin.test.assertFailsWith
 import org.junit.Test
@@ -34,6 +35,15 @@
     fun testApfInstructionVersionCheck() {
         var gen = ApfGenerator(ApfGenerator.MIN_APF_VERSION)
         assertFailsWith<IllegalInstructionException> { gen.addDrop() }
+        assertFailsWith<IllegalInstructionException> { gen.addCountAndDrop(12) }
+        assertFailsWith<IllegalInstructionException> { gen.addCountAndPass(1000) }
+    }
+
+    @Test
+    fun testApfInstructionArgumentCheck() {
+        var gen = ApfGenerator(ApfGenerator.MIN_APF_VERSION_IN_DEV)
+        assertFailsWith<IllegalArgumentException> { gen.addCountAndPass(0) }
+        assertFailsWith<IllegalArgumentException> { gen.addCountAndDrop(0) }
     }
 
     @Test
@@ -42,13 +52,31 @@
         gen.addPass()
         var program = gen.generate()
         // encoding PASS opcode: opcode=0, imm_len=0, R=0
-        assertContentEquals(byteArrayOf(encodeInstruction(0, 0, 0)), program)
+        assertContentEquals(
+                byteArrayOf(encodeInstruction(opcode = 0, immLength = 0, register = 0)), program)
 
         gen = ApfGenerator(ApfGenerator.MIN_APF_VERSION_IN_DEV)
         gen.addDrop()
         program = gen.generate()
         // encoding DROP opcode: opcode=0, imm_len=0, R=1
-        assertContentEquals(byteArrayOf(encodeInstruction(0, 0, 1)), program)
+        assertContentEquals(
+                byteArrayOf(encodeInstruction(opcode = 0, immLength = 0, register = 1)), program)
+
+        gen = ApfGenerator(ApfGenerator.MIN_APF_VERSION_IN_DEV)
+        gen.addCountAndPass(129)
+        program = gen.generate()
+        // encoding COUNT(PASS) opcode: opcode=0, imm_len=size_of(imm), R=0, imm=counterNumber
+        assertContentEquals(
+                byteArrayOf(encodeInstruction(opcode = 0, immLength = 1, register = 0),
+                        0x81.toByte()), program)
+
+        gen = ApfGenerator(ApfGenerator.MIN_APF_VERSION_IN_DEV)
+        gen.addCountAndDrop(1000)
+        program = gen.generate()
+        // encoding COUNT(DROP) opcode: opcode=0, imm_len=size_of(imm), R=1, imm=counterNumber
+        assertContentEquals(
+                byteArrayOf(encodeInstruction(opcode = 0, immLength = 2, register = 1),
+                        0x03, 0xe8.toByte()), program)
 
         gen = ApfGenerator(ApfGenerator.MIN_APF_VERSION_IN_DEV)
         gen.addAlloc(ApfGenerator.Register.R0)