lib: cgen: group same-key maps using bitmasks

src/libbpfilter/cgen/fixup.c

@@ -10,6 +10,7 @@
 
 #include <bpfilter/dump.h>
 #include <bpfilter/helper.h>
+#include <bpfilter/set.h>
 
 int bf_fixup_new(struct bf_fixup **fixup, enum bf_fixup_type type,
                  size_t insn_offset, const union bf_fixup_attr *attr)
@@ -74,7 +75,8 @@ void bf_fixup_dump(const struct bf_fixup *fixup, prefix_t *prefix)
         // No specific value to dump
         break;
     case BF_FIXUP_TYPE_SET_MAP_FD:
-        DUMP(prefix, "set_index: %lu", fixup->attr.set_index);
+        DUMP(prefix, "set: %s",
+             fixup->attr.set_ptr ? fixup->attr.set_ptr->name : "(null)");
         break;
     default:
         DUMP(prefix, "unsupported bf_fixup_type: %d", fixup->type);

src/libbpfilter/cgen/fixup.h

@@ -10,6 +10,8 @@
 #include <bpfilter/dump.h>
 #include <bpfilter/elfstub.h>
 
+struct bf_set;
+
 /**
  * Field to fixup in a @c bpf_insn structure.
  */
@@ -44,7 +46,9 @@ enum bf_fixup_type
 
 union bf_fixup_attr
 {
-    size_t set_index;
+    /** Set referenced by a `BF_FIXUP_TYPE_SET_MAP_FD` fixup. The fixup
+     * holds a non-owning pointer to a set in the generator's chain. */
+    const struct bf_set *set_ptr;
     enum bf_elfstub_id elfstub_id;
 };
 

src/libbpfilter/cgen/handle.h

@@ -51,8 +51,8 @@ struct bf_handle
     /** Log map. NULL if not created. */
     struct bf_map *lmap;
 
-    /** List of set maps. If a set is empty in the chain, NULL is inserted in
-     * this list instead of a `bf_map` to preserve sets indexes. */
+    /** List of set maps. Contains at most one map for each unique key
+     * format. */
     bf_list sets;
 };
 

src/libbpfilter/cgen/matcher/set.c

@@ -5,31 +5,70 @@
 
 #include "cgen/matcher/set.h"
 
+#include <limits.h>
+#include <stdint.h>
+
 #include <bpfilter/chain.h>
 #include <bpfilter/matcher.h>
 #include <bpfilter/set.h>
 
 #include "cgen/program.h"
 #include "cgen/stub.h"
 
+/**
+ * @brief Emit bitmask check instructions after a map lookup.
+ *
+ * Set map values are bitmasks, sized to fit one bit per set in the
+ * group. After map_lookup_elem returns a non-NULL pointer in BPF_REG_0,
+ * this function emits instructions to:
+ *   1. Load the byte containing this set's bit from the pointer.
+ *   2. AND it with the set's bit.
+ *   3. Jump to the next rule if the bit is not set.
+ */
+static int _bf_emit_bitmask_check(struct bf_program *program, size_t bit_index)
+{
+    assert(program);
+
+    /* r0 points to the map value. Load the byte containing our bit. */
+    EMIT(program, BPF_LDX_MEM(BPF_B, BPF_REG_0, BPF_REG_0,
+                              (int16_t)(bit_index / CHAR_BIT)));
+
+    /* AND with the set's bit within that byte. */
+    EMIT(program,
+         BPF_ALU32_IMM(BPF_AND, BPF_REG_0, 1U << (bit_index % CHAR_BIT)));
+
+    /* Jump to next rule if the result is 0 (bit not set). */
+    EMIT_FIXUP_JMP_NEXT_RULE(program, BPF_JMP_IMM(BPF_JEQ, BPF_REG_0, 0, 0));
+
+    return 0;
+}
+
 static int _bf_matcher_generate_set_trie(struct bf_program *program,
                                          const struct bf_matcher *matcher)
 {
     assert(program);
     assert(matcher);
 
+    uint32_t set_index = *(uint32_t *)bf_matcher_payload(matcher);
     const struct bf_set *set =
         bf_chain_get_set_for_matcher(program->runtime.chain, matcher);
-    enum bf_matcher_type type = set->key[0];
-    const struct bf_matcher_meta *meta = bf_matcher_get_meta(type);
+    enum bf_matcher_type type;
+    const struct bf_matcher_meta *meta;
+    size_t bit_index;
     int r;
 
     if (!set) {
-        return bf_err_r(-ENOENT, "set #%u not found in %s",
-                        *(uint32_t *)bf_matcher_payload(matcher),
+        return bf_err_r(-ENOENT, "set #%u not found in %s", set_index,
                         program->runtime.chain->name);
     }
 
+    type = set->key[0];
+    meta = bf_matcher_get_meta(type);
+
+    r = bf_program_set_bit_index(program, set, &bit_index);
+    if (r)
+        return bf_err_r(r, "set '%s' not assigned to any group", set->name);
+
     r = bf_stub_rule_check_protocol(program, meta);
     if (r)
         return bf_err_r(r, "failed to check for protocol");
@@ -67,15 +106,15 @@ static int _bf_matcher_generate_set_trie(struct bf_program *program,
                         bf_matcher_type_to_str(type), type);
     }
 
-    EMIT_LOAD_SET_FD_FIXUP(program, BPF_REG_1,
-                           *(uint32_t *)bf_matcher_payload(matcher));
+    EMIT_LOAD_SET_FD_FIXUP(program, BPF_REG_1, set);
     EMIT(program, BPF_MOV64_REG(BPF_REG_2, BPF_REG_10));
     EMIT(program, BPF_ALU64_IMM(BPF_ADD, BPF_REG_2, BF_PROG_SCR_OFF(4)));
     EMIT(program, BPF_EMIT_CALL(BPF_FUNC_map_lookup_elem));
 
-    // Jump to the next rule if map_lookup_elem returned 0
+    // Jump to the next rule if map_lookup_elem returned NULL
     EMIT_FIXUP_JMP_NEXT_RULE(program, BPF_JMP_IMM(BPF_JEQ, BPF_REG_0, 0, 0));
-    return 0;
+
+    return _bf_emit_bitmask_check(program, bit_index);
 }
 
 int bf_matcher_generate_set(struct bf_program *program,
@@ -86,18 +125,23 @@ int bf_matcher_generate_set(struct bf_program *program,
 
     const struct bf_set *set =
         bf_chain_get_set_for_matcher(program->runtime.chain, matcher);
+    uint32_t set_index = *(uint32_t *)bf_matcher_payload(matcher);
+    size_t bit_index;
     size_t offset = 0;
     int r;
 
     if (!set) {
-        return bf_err_r(-ENOENT, "set #%u not found in %s",
-                        *(uint32_t *)bf_matcher_payload(matcher),
+        return bf_err_r(-ENOENT, "set #%u not found in %s", set_index,
                         program->runtime.chain->name);
     }
 
     if (set->use_trie)
         return _bf_matcher_generate_set_trie(program, matcher);
 
+    r = bf_program_set_bit_index(program, set, &bit_index);
+    if (r)
+        return bf_err_r(r, "set '%s' not assigned to any group", set->name);
+
     // Ensure the packet uses the required protocols
     for (size_t i = 0; i < set->n_comps; ++i) {
         enum bf_matcher_type type = set->key[i];
@@ -136,14 +180,13 @@ int bf_matcher_generate_set(struct bf_program *program,
         offset += meta->hdr_payload_size;
     }
 
-    EMIT_LOAD_SET_FD_FIXUP(program, BPF_REG_1,
-                           *(uint32_t *)bf_matcher_payload(matcher));
+    EMIT_LOAD_SET_FD_FIXUP(program, BPF_REG_1, set);
     EMIT(program, BPF_MOV64_REG(BPF_REG_2, BPF_REG_10));
     EMIT(program, BPF_ALU64_IMM(BPF_ADD, BPF_REG_2, BF_PROG_SCR_OFF(0)));
     EMIT(program, BPF_EMIT_CALL(BPF_FUNC_map_lookup_elem));
 
-    // Jump to the next rule if map_lookup_elem returned 0
+    // Jump to the next rule if map_lookup_elem returned NULL
     EMIT_FIXUP_JMP_NEXT_RULE(program, BPF_JMP_IMM(BPF_JEQ, BPF_REG_0, 0, 0));
 
-    return 0;
+    return _bf_emit_bitmask_check(program, bit_index);
 }

src/libbpfilter/cgen/prog/map.c

@@ -211,7 +211,8 @@ int bf_map_new(struct bf_map **map, const char *name, enum bf_map_type type,
 }
 
 int bf_map_new_from_set(struct bf_map **map, const char *name,
-                        const struct bf_set *set)
+                        const struct bf_set *set, size_t set_size,
+                        size_t value_size)
 {
     assert(map);
     assert(name);
@@ -220,7 +221,7 @@ int bf_map_new_from_set(struct bf_map **map, const char *name,
     return _bf_map_new(map, name, BF_MAP_TYPE_SET,
                        set->use_trie ? BF_BPF_MAP_TYPE_LPM_TRIE :
                                        BF_BPF_MAP_TYPE_HASH,
-                       set->elem_size, 1, bf_list_size(&set->elems));
+                       set->elem_size, value_size, set_size);
 }
 
 int bf_map_new_from_pack(struct bf_map **map, int dir_fd, bf_rpack_node_t node)

src/libbpfilter/cgen/prog/map.h

@@ -71,10 +71,13 @@ int bf_map_new(struct bf_map **map, const char *name, enum bf_map_type type,
  *        also as filename when pinning the map to the system. Can't be NULL or
  *        empty.
  * @param set Set to create the map from. Can't be NULL.
+ * @param set_size Number of elements to reserve room for in the map.
+ * @param value_size Size (in bytes) of the value field.
  * @return 0 on success, or a negative error value on error.
  */
 int bf_map_new_from_set(struct bf_map **map, const char *name,
-                        const struct bf_set *set);
+                        const struct bf_set *set, size_t set_size,
+                        size_t value_size);
 
 /**
  * @brief Allocate and initialize a new map from serialized data.