bpf: constant map key tracking for array map lookuppr/bpf-tail-call-rebased2

Signed-off-by: Daniel Borkmann <daniel@iogearbox.net>

5 files changed, 58 insertions, 29 deletions

diff --git a/include/linux/bpf.h b/include/linux/bpf.h
index 62a369fb8d981..02253014067c7 100644
--- a/include/linux/bpf.h
+++ b/include/linux/bpf.h
@@ -72,9 +72,9 @@ struct bpf_map_ops {
 
 	/* Direct value access helpers. */
 	int (*map_direct_value_addr)(const struct bpf_map *map,
-				     u64 *imm, u32 off);
+				     u64 *imm, u32 idx, u32 off);
 	int (*map_direct_value_meta)(const struct bpf_map *map,
-				     u64 imm, u32 *off);
+				     u64 imm, u32 *idx, u32 *off);
 };
 
 struct bpf_map_memory {
diff --git a/kernel/bpf/arraymap.c b/kernel/bpf/arraymap.c
index b9ef993278c6a..5398f128aa54c 100644
--- a/kernel/bpf/arraymap.c
+++ b/kernel/bpf/arraymap.c
@@ -152,32 +152,32 @@ static void *array_map_lookup_elem(struct bpf_map *map, void *key)
 }
 
 static int array_map_direct_value_addr(const struct bpf_map *map, u64 *imm,
-				       u32 off)
+				       u32 idx, u32 off)
 {
 	struct bpf_array *array = container_of(map, struct bpf_array, map);
 
-	if (map->max_entries != 1)
-		return -ENOTSUPP;
-	if (off >= map->value_size)
+	if (idx >= map->max_entries || off >= map->value_size)
 		return -EINVAL;
 
-	*imm = (unsigned long)array->value;
+	*imm = (unsigned long)(array->value +
+			       array->elem_size * (idx & array->index_mask));
 	return 0;
 }
 
 static int array_map_direct_value_meta(const struct bpf_map *map, u64 imm,
-				       u32 *off)
+				       u32 *idx, u32 *off)
 {
 	struct bpf_array *array = container_of(map, struct bpf_array, map);
-	u64 base = (unsigned long)array->value;
-	u64 range = array->elem_size;
+	u64 rem, base = (unsigned long)array->value, slot = array->elem_size;
+	u64 range = slot * map->max_entries;
 
-	if (map->max_entries != 1)
-		return -ENOTSUPP;
 	if (imm < base || imm >= base + range)
 		return -ENOENT;
 
-	*off = imm - base;
+	base = imm - base;
+
+	*idx = div64_u64_rem(base, slot, &rem);
+	*off = rem;
 	return 0;
 }
 
diff --git a/kernel/bpf/syscall.c b/kernel/bpf/syscall.c
index 81057e5922d54..7dd04d8c91e87 100644
--- a/kernel/bpf/syscall.c
+++ b/kernel/bpf/syscall.c
@@ -2294,13 +2294,13 @@ static int bpf_map_get_fd_by_id(const union bpf_attr *attr)
 }
 
 static const struct bpf_map *bpf_map_from_imm(const struct bpf_prog *prog,
-					      unsigned long addr, u32 *off,
-					      u32 *type)
+					      unsigned long addr, u32 *idx,
+					      u32 *off, u32 *type)
 {
 	const struct bpf_map *map;
 	int i;
 
-	for (i = 0, *off = 0; i < prog->aux->used_map_cnt; i++) {
+	for (i = 0, *off = *idx = 0; i < prog->aux->used_map_cnt; i++) {
 		map = prog->aux->used_maps[i];
 		if (map == (void *)addr) {
 			*type = BPF_PSEUDO_MAP_FD;
@@ -2308,7 +2308,7 @@ static const struct bpf_map *bpf_map_from_imm(const struct bpf_prog *prog,
 		}
 		if (!map->ops->map_direct_value_meta)
 			continue;
-		if (!map->ops->map_direct_value_meta(map, addr, off)) {
+		if (!map->ops->map_direct_value_meta(map, addr, idx, off)) {
 			*type = BPF_PSEUDO_MAP_VALUE;
 			return map;
 		}
@@ -2321,7 +2321,7 @@ static struct bpf_insn *bpf_insn_prepare_dump(const struct bpf_prog *prog)
 {
 	const struct bpf_map *map;
 	struct bpf_insn *insns;
-	u32 off, type;
+	u32 off, idx, type;
 	u64 imm;
 	int i;
 
@@ -2349,11 +2349,13 @@ static struct bpf_insn *bpf_insn_prepare_dump(const struct bpf_prog *prog)
 			continue;
 
 		imm = ((u64)insns[i + 1].imm << 32) | (u32)insns[i].imm;
-		map = bpf_map_from_imm(prog, imm, &off, &type);
+		map = bpf_map_from_imm(prog, imm, &idx, &off, &type);
 		if (map) {
 			insns[i].src_reg = type;
 			insns[i].imm = map->id;
+			insns[i].off = (u16)idx;
 			insns[i + 1].imm = off;
+			insns[i + 1].off = (u16)(idx >> 16);
 			continue;
 		}
 	}
diff --git a/kernel/bpf/verifier.c b/kernel/bpf/verifier.c
index 2463528e73e66..4f29b687bd9df 100644
--- a/kernel/bpf/verifier.c
+++ b/kernel/bpf/verifier.c
@@ -2810,9 +2810,10 @@ static int bpf_map_direct_read(struct bpf_map *map, int off, int size, u64 *val)
 	u64 addr;
 	int err;
 
-	err = map->ops->map_direct_value_addr(map, &addr, off);
+	err = map->ops->map_direct_value_addr(map, &addr, 0, off);
 	if (err)
 		return err;
+
 	ptr = (void *)(long)addr + off;
 
 	switch (size) {
@@ -4124,16 +4125,21 @@ record_func_key(struct bpf_verifier_env *env, struct bpf_call_arg_meta *meta,
 	struct bpf_map *map = meta->map_ptr;
 	u64 val;
 
-	if (func_id != BPF_FUNC_tail_call)
+	if (func_id != BPF_FUNC_tail_call &&
+	    func_id != BPF_FUNC_map_lookup_elem)
 		return 0;
+
 	if (!map) {
 		verbose(env, "kernel subsystem misconfigured verifier\n");
 		return -EINVAL;
 	}
-	if (map->map_type != BPF_MAP_TYPE_PROG_ARRAY)
+
+	if (map->map_type != BPF_MAP_TYPE_PROG_ARRAY &&
+	    map->map_type != BPF_MAP_TYPE_ARRAY)
 		return 0;
 
-	reg = &regs[BPF_REG_3];
+	reg = func_id == BPF_FUNC_tail_call ?
+	      &regs[BPF_REG_3] : &regs[BPF_REG_2];
 	if (!register_is_const(reg) || !tnum_in(range, reg->var_off)) {
 		bpf_map_key_store(aux, BPF_MAP_KEY_POISON);
 		return 0;
@@ -8208,7 +8214,7 @@ static int replace_map_fd_with_map_ptr(struct bpf_verifier_env *env)
 					return -EINVAL;
 				}
 
-				err = map->ops->map_direct_value_addr(map, &addr, off);
+				err = map->ops->map_direct_value_addr(map, &addr, 0, off);
 				if (err) {
 					verbose(env, "invalid access to map value pointer, value_size=%u off=%u\n",
 						map->value_size, off);
@@ -9121,7 +9127,8 @@ static int fixup_bpf_calls(struct bpf_verifier_env *env)
 	struct bpf_insn insn_buf[16];
 	struct bpf_prog *new_prog;
 	struct bpf_map *map_ptr;
-	int i, cnt, delta = 0;
+	int ret, i, cnt, delta = 0;
+	u32 map_key;
 
 	for (i = 0; i < insn_cnt; i++, insn++) {
 		if (insn->code == (BPF_ALU64 | BPF_MOD | BPF_X) ||
@@ -9270,8 +9277,6 @@ static int fixup_bpf_calls(struct bpf_verifier_env *env)
 			    !bpf_map_ptr_poisoned(aux) &&
 			    !bpf_map_ptr_unpriv(aux)) {
 				struct bpf_jit_poke_descriptor desc;
-				u32 map_key;
-				int ret;
 
 				map_key = bpf_map_key_immediate(aux);
 				map_ptr = BPF_MAP_PTR(aux->map_state);
@@ -9354,7 +9359,26 @@ static int fixup_bpf_calls(struct bpf_verifier_env *env)
 			ops = map_ptr->ops;
 			if (insn->imm == BPF_FUNC_map_lookup_elem &&
 			    ops->map_gen_lookup) {
-				cnt = ops->map_gen_lookup(map_ptr, insn_buf);
+				u64 addr = 0;
+
+				if (!bpf_map_key_poisoned(aux) &&
+				    ops->map_direct_value_addr) {
+					map_key = bpf_map_key_immediate(aux);
+					ret = ops->map_direct_value_addr(map_ptr,
+									 &addr,
+									 map_key, 0);
+					if (WARN_ON_ONCE(ret < 0 && addr))
+						addr = 0;
+				}
+				if (addr) {
+					struct bpf_insn tmp[] = {
+						BPF_LD_IMM64(BPF_REG_0, addr),
+					};
+					memcpy(insn_buf, tmp, sizeof(tmp));
+					cnt = ARRAY_SIZE(tmp);
+				} else {
+					cnt = ops->map_gen_lookup(map_ptr, insn_buf);
+				}
 				if (cnt == 0 || cnt >= ARRAY_SIZE(insn_buf)) {
 					verbose(env, "bpf verifier is misconfigured\n");
 					return -EINVAL;
diff --git a/tools/bpf/bpftool/xlated_dumper.c b/tools/bpf/bpftool/xlated_dumper.c
index 494d7ae3614de..a9be79f6dd6e8 100644
--- a/tools/bpf/bpftool/xlated_dumper.c
+++ b/tools/bpf/bpftool/xlated_dumper.c
@@ -195,7 +195,10 @@ static const char *print_imm(void *private_data,
 			 "map[id:%u]", insn->imm);
 	else if (insn->src_reg == BPF_PSEUDO_MAP_VALUE)
 		snprintf(dd->scratch_buff, sizeof(dd->scratch_buff),
-			 "map[id:%u][0]+%u", insn->imm, (insn + 1)->imm);
+			 "map[id:%u][%u]+%u", insn->imm,
+			 ((__u32)(__u16)insn[0].off) |
+			 ((__u32)(__u16)insn[1].off) << 16,
+			 (insn + 1)->imm);
 	else
 		snprintf(dd->scratch_buff, sizeof(dd->scratch_buff),
 			 "0x%llx", (unsigned long long)full_imm);