diff mbox series

[bpf-next,05/13] bpf: abstract away global subprog arg preparation logic from reg state setup

Message ID 20231204233931.49758-6-andrii@kernel.org (mailing list archive)
State Superseded
Delegated to: BPF
Headers show
Series Enhance BPF global subprogs with argument tags | expand

Checks

Context Check Description
bpf/vmtest-bpf-next-PR success PR summary
bpf/vmtest-bpf-next-VM_Test-0 success Logs for Lint
bpf/vmtest-bpf-next-VM_Test-8 success Logs for aarch64-gcc / veristat
bpf/vmtest-bpf-next-VM_Test-2 success Logs for Validate matrix.py
bpf/vmtest-bpf-next-VM_Test-1 success Logs for ShellCheck
bpf/vmtest-bpf-next-VM_Test-3 success Logs for aarch64-gcc / build / build for aarch64 with gcc
bpf/vmtest-bpf-next-VM_Test-7 success Logs for aarch64-gcc / test (test_verifier, false, 360) / test_verifier on aarch64 with gcc
bpf/vmtest-bpf-next-VM_Test-5 success Logs for aarch64-gcc / test (test_progs, false, 360) / test_progs on aarch64 with gcc
bpf/vmtest-bpf-next-VM_Test-6 success Logs for aarch64-gcc / test (test_progs_no_alu32, false, 360) / test_progs_no_alu32 on aarch64 with gcc
bpf/vmtest-bpf-next-VM_Test-4 success Logs for aarch64-gcc / test (test_maps, false, 360) / test_maps on aarch64 with gcc
bpf/vmtest-bpf-next-VM_Test-9 success Logs for s390x-gcc / build / build for s390x with gcc
bpf/vmtest-bpf-next-VM_Test-14 success Logs for s390x-gcc / veristat
bpf/vmtest-bpf-next-VM_Test-15 success Logs for set-matrix
bpf/vmtest-bpf-next-VM_Test-16 success Logs for x86_64-gcc / build / build for x86_64 with gcc
bpf/vmtest-bpf-next-VM_Test-17 success Logs for x86_64-gcc / test (test_maps, false, 360) / test_maps on x86_64 with gcc
bpf/vmtest-bpf-next-VM_Test-18 success Logs for x86_64-gcc / test (test_progs, false, 360) / test_progs on x86_64 with gcc
bpf/vmtest-bpf-next-VM_Test-19 success Logs for x86_64-gcc / test (test_progs_no_alu32, false, 360) / test_progs_no_alu32 on x86_64 with gcc
bpf/vmtest-bpf-next-VM_Test-20 success Logs for x86_64-gcc / test (test_progs_no_alu32_parallel, true, 30) / test_progs_no_alu32_parallel on x86_64 with gcc
bpf/vmtest-bpf-next-VM_Test-21 success Logs for x86_64-gcc / test (test_progs_parallel, true, 30) / test_progs_parallel on x86_64 with gcc
bpf/vmtest-bpf-next-VM_Test-22 success Logs for x86_64-gcc / test (test_verifier, false, 360) / test_verifier on x86_64 with gcc
bpf/vmtest-bpf-next-VM_Test-23 success Logs for x86_64-gcc / veristat / veristat on x86_64 with gcc
bpf/vmtest-bpf-next-VM_Test-24 success Logs for x86_64-llvm-16 / build / build for x86_64 with llvm-16
bpf/vmtest-bpf-next-VM_Test-25 success Logs for x86_64-llvm-16 / test (test_maps, false, 360) / test_maps on x86_64 with llvm-16
bpf/vmtest-bpf-next-VM_Test-28 success Logs for x86_64-llvm-16 / test (test_verifier, false, 360) / test_verifier on x86_64 with llvm-16
bpf/vmtest-bpf-next-VM_Test-29 success Logs for x86_64-llvm-16 / veristat
bpf/vmtest-bpf-next-VM_Test-26 success Logs for x86_64-llvm-16 / test (test_progs, false, 360) / test_progs on x86_64 with llvm-16
bpf/vmtest-bpf-next-VM_Test-27 success Logs for x86_64-llvm-16 / test (test_progs_no_alu32, false, 360) / test_progs_no_alu32 on x86_64 with llvm-16
bpf/vmtest-bpf-next-VM_Test-13 success Logs for s390x-gcc / test (test_verifier, false, 360) / test_verifier on s390x with gcc
bpf/vmtest-bpf-next-VM_Test-11 success Logs for s390x-gcc / test (test_progs, false, 360) / test_progs on s390x with gcc
bpf/vmtest-bpf-next-VM_Test-10 success Logs for s390x-gcc / test (test_maps, false, 360) / test_maps on s390x with gcc
bpf/vmtest-bpf-next-VM_Test-12 success Logs for s390x-gcc / test (test_progs_no_alu32, false, 360) / test_progs_no_alu32 on s390x with gcc
netdev/series_format success Posting correctly formatted
netdev/tree_selection success Clearly marked for bpf-next, async
netdev/ynl success SINGLE THREAD; Generated files up to date; no warnings/errors; no diff in generated;
netdev/fixes_present success Fixes tag not required for -next series
netdev/header_inline success No static functions without inline keyword in header files
netdev/build_32bit success Errors and warnings before: 2659 this patch: 2659
netdev/cc_maintainers warning 8 maintainers not CCed: haoluo@google.com jolsa@kernel.org kpsingh@kernel.org john.fastabend@gmail.com martin.lau@linux.dev song@kernel.org sdf@google.com yonghong.song@linux.dev
netdev/build_clang success Errors and warnings before: 1276 this patch: 1276
netdev/verify_signedoff success Signed-off-by tag matches author and committer
netdev/deprecated_api success None detected
netdev/check_selftest success No net selftest shell script
netdev/verify_fixes success No Fixes tag
netdev/build_allmodconfig_warn success Errors and warnings before: 2736 this patch: 2736
netdev/checkpatch warning WARNING: line length of 87 exceeds 80 columns WARNING: line length of 94 exceeds 80 columns
netdev/build_clang_rust success No Rust files in patch. Skipping build
netdev/kdoc success Errors and warnings before: 0 this patch: 0
netdev/source_inline success Was 0 now: 0

Commit Message

Andrii Nakryiko Dec. 4, 2023, 11:39 p.m. UTC 
btf_prepare_func_args() is used to understand expectations and
restrictions on global subprog arguments. But current implementation is
hard to extend, as it intermixes BTF-based func prototype parsing and
interpretation logic with setting up register state at subprog entry.

Worse still, those registers are not completely set up inside
btf_prepare_func_args(), requiring some more logic later in
do_check_common(). Like calling mark_reg_unknown() and similar
initialization operations.

This intermixing of BTF interpretation and register state setup is
problematic. First, it causes duplication of BTF parsing logic for global
subprog verification (to set up initial state of global subprog) and
global subprog call sites analysis (when we need to check that whatever
is being passed into global subprog matches expectations), performed in
btf_check_subprog_call().

Given we want to extend global func argument with tags later, this
duplication is problematic. So refactor btf_prepare_func_args() to do
only BTF-based func proto and args parsing, returning high-level
argument "expectations" only, with no regard to specifics of register
state. I.e., if it's a context argument, instead of setting register
state to PTR_TO_CTX, we return ARG_PTR_TO_CTX enum for that argument as
"an argument specification" for further processing inside
do_check_common(). Similarly for SCALAR arguments, PTR_TO_MEM, etc.

This allows to reuse btf_prepare_func_args() in following patches at
global subprog call site analysis time. It also keeps register setup
code consistently in one place, do_check_common().

Besides all this, we cache this argument specs information inside
env->subprog_info, eliminating the need to redo these potentially
expensive BTF traversals, especially if BPF program's BTF is big and/or
there are lots of global subprog calls.

Signed-off-by: Andrii Nakryiko <andrii@kernel.org>
---
 include/linux/bpf.h          |  4 ++--
 include/linux/bpf_verifier.h | 16 ++++++++++++++++
 kernel/bpf/btf.c             | 31 ++++++++++++++++---------------
 kernel/bpf/verifier.c        | 36 +++++++++++++++++++++---------------
 4 files changed, 55 insertions(+), 32 deletions(-)

Comments

Eduard Zingerman Dec. 5, 2023, 11:21 p.m. UTC | #1 
On Mon, 2023-12-04 at 15:39 -0800, Andrii Nakryiko wrote:

Acked-by: Eduard Zingerman <eddyz87@gmail.com>

[...]
> diff --git a/include/linux/bpf.h b/include/linux/bpf.h
> index 379ac0a28405..c3a5d0fe3cdf 100644
> --- a/include/linux/bpf.h
> +++ b/include/linux/bpf.h
> @@ -704,6 +704,7 @@ enum bpf_arg_type {
>  
>  	ARG_PTR_TO_CTX,		/* pointer to context */
>  	ARG_ANYTHING,		/* any (initialized) argument is ok */
> +	ARG_SCALAR = ARG_ANYTHING, /* scalar value */

Nit: I agree that use of ARG_ANYTHING to denote scalars is confusing,
     but having two names for the same thing seems even more confusing.

[...]

> diff --git a/kernel/bpf/btf.c b/kernel/bpf/btf.c
> index d56433bf8aba..33a62df9c5a8 100644
> --- a/kernel/bpf/btf.c
> +++ b/kernel/bpf/btf.c
> @@ -6955,9 +6955,9 @@ int btf_check_subprog_call(struct bpf_verifier_env *env, int subprog,
>   * 0 - Successfully converted BTF into bpf_reg_state
>   * (either PTR_TO_CTX or SCALAR_VALUE).
>   */
> -int btf_prepare_func_args(struct bpf_verifier_env *env, int subprog,
> -			  struct bpf_reg_state *regs, u32 *arg_cnt)
> +int btf_prepare_func_args(struct bpf_verifier_env *env, int subprog)

Could you please also update the comment above this function?
It currently says: "Convert BTF of a function into bpf_reg_state if possible".

[...]

> diff --git a/kernel/bpf/verifier.c b/kernel/bpf/verifier.c
> index ee707736ce6b..16d5550eda4d 100644
> --- a/kernel/bpf/verifier.c
> +++ b/kernel/bpf/verifier.c
[...]
> @@ -19860,33 +19855,44 @@ static int do_check_common(struct bpf_verifier_env *env, int subprog)
[...]
>  		for (i = BPF_REG_1; i <= BPF_REG_5; i++) {
> -			if (regs[i].type == PTR_TO_CTX)
> +			arg = &sub->args[i - BPF_REG_1];
> +			reg = &regs[i];
> +
> +			if (arg->arg_type == ARG_PTR_TO_CTX) {
> +				reg->type = PTR_TO_CTX;
>  				mark_reg_known_zero(env, regs, i);
> -			else if (regs[i].type == SCALAR_VALUE)
> +			} else if (arg->arg_type == ARG_SCALAR) {
> +				reg->type = SCALAR_VALUE;
>  				mark_reg_unknown(env, regs, i);
> -			else if (base_type(regs[i].type) == PTR_TO_MEM) {
> -				const u32 mem_size = regs[i].mem_size;
> -
> +			} else if (base_type(arg->arg_type) == ARG_PTR_TO_MEM) {
> +				reg->type = PTR_TO_MEM;
> +				if (arg->arg_type & PTR_MAYBE_NULL)
> +					reg->type |= PTR_MAYBE_NULL;
>  				mark_reg_known_zero(env, regs, i);
> -				regs[i].mem_size = mem_size;
> -				regs[i].id = ++env->id_gen;
> +				reg->mem_size = arg->mem_size;
> +				reg->id = ++env->id_gen;
>  			}

Nit: maybe add an else branch here and report an error if unexpected
     argument type is returned by btf_prepare_func_args()?
Andrii Nakryiko Dec. 6, 2023, 5:59 p.m. UTC | #2 
On Tue, Dec 5, 2023 at 3:21 PM Eduard Zingerman <eddyz87@gmail.com> wrote:
>
> On Mon, 2023-12-04 at 15:39 -0800, Andrii Nakryiko wrote:
>
> Acked-by: Eduard Zingerman <eddyz87@gmail.com>
>
> [...]
> > diff --git a/include/linux/bpf.h b/include/linux/bpf.h
> > index 379ac0a28405..c3a5d0fe3cdf 100644
> > --- a/include/linux/bpf.h
> > +++ b/include/linux/bpf.h
> > @@ -704,6 +704,7 @@ enum bpf_arg_type {
> >
> >       ARG_PTR_TO_CTX,         /* pointer to context */
> >       ARG_ANYTHING,           /* any (initialized) argument is ok */
> > +     ARG_SCALAR = ARG_ANYTHING, /* scalar value */
>
> Nit: I agree that use of ARG_ANYTHING to denote scalars is confusing,
>      but having two names for the same thing seems even more confusing.
>

fair enough, I was undecided on this, I'll revert and use ARG_ANYTHING for now

> [...]
>
> > diff --git a/kernel/bpf/btf.c b/kernel/bpf/btf.c
> > index d56433bf8aba..33a62df9c5a8 100644
> > --- a/kernel/bpf/btf.c
> > +++ b/kernel/bpf/btf.c
> > @@ -6955,9 +6955,9 @@ int btf_check_subprog_call(struct bpf_verifier_env *env, int subprog,
> >   * 0 - Successfully converted BTF into bpf_reg_state
> >   * (either PTR_TO_CTX or SCALAR_VALUE).
> >   */
> > -int btf_prepare_func_args(struct bpf_verifier_env *env, int subprog,
> > -                       struct bpf_reg_state *regs, u32 *arg_cnt)
> > +int btf_prepare_func_args(struct bpf_verifier_env *env, int subprog)
>
> Could you please also update the comment above this function?
> It currently says: "Convert BTF of a function into bpf_reg_state if possible".
>

absolutely, will do

> [...]
>
> > diff --git a/kernel/bpf/verifier.c b/kernel/bpf/verifier.c
> > index ee707736ce6b..16d5550eda4d 100644
> > --- a/kernel/bpf/verifier.c
> > +++ b/kernel/bpf/verifier.c
> [...]
> > @@ -19860,33 +19855,44 @@ static int do_check_common(struct bpf_verifier_env *env, int subprog)
> [...]
> >               for (i = BPF_REG_1; i <= BPF_REG_5; i++) {
> > -                     if (regs[i].type == PTR_TO_CTX)
> > +                     arg = &sub->args[i - BPF_REG_1];
> > +                     reg = &regs[i];
> > +
> > +                     if (arg->arg_type == ARG_PTR_TO_CTX) {
> > +                             reg->type = PTR_TO_CTX;
> >                               mark_reg_known_zero(env, regs, i);
> > -                     else if (regs[i].type == SCALAR_VALUE)
> > +                     } else if (arg->arg_type == ARG_SCALAR) {
> > +                             reg->type = SCALAR_VALUE;
> >                               mark_reg_unknown(env, regs, i);
> > -                     else if (base_type(regs[i].type) == PTR_TO_MEM) {
> > -                             const u32 mem_size = regs[i].mem_size;
> > -
> > +                     } else if (base_type(arg->arg_type) == ARG_PTR_TO_MEM) {
> > +                             reg->type = PTR_TO_MEM;
> > +                             if (arg->arg_type & PTR_MAYBE_NULL)
> > +                                     reg->type |= PTR_MAYBE_NULL;
> >                               mark_reg_known_zero(env, regs, i);
> > -                             regs[i].mem_size = mem_size;
> > -                             regs[i].id = ++env->id_gen;
> > +                             reg->mem_size = arg->mem_size;
> > +                             reg->id = ++env->id_gen;
> >                       }
>
> Nit: maybe add an else branch here and report an error if unexpected
>      argument type is returned by btf_prepare_func_args()?

true, not sure why I didn't do it here, adding...

>
>
diff mbox series

Patch

diff --git a/include/linux/bpf.h b/include/linux/bpf.h
index 379ac0a28405..c3a5d0fe3cdf 100644
--- a/include/linux/bpf.h
+++ b/include/linux/bpf.h
@@ -704,6 +704,7 @@  enum bpf_arg_type {
 
 	ARG_PTR_TO_CTX,		/* pointer to context */
 	ARG_ANYTHING,		/* any (initialized) argument is ok */
+	ARG_SCALAR = ARG_ANYTHING, /* scalar value */
 	ARG_PTR_TO_SPIN_LOCK,	/* pointer to bpf_spin_lock */
 	ARG_PTR_TO_SOCK_COMMON,	/* pointer to sock_common */
 	ARG_PTR_TO_INT,		/* pointer to int */
@@ -2430,8 +2431,7 @@  int btf_check_subprog_arg_match(struct bpf_verifier_env *env, int subprog,
 				struct bpf_reg_state *regs);
 int btf_check_subprog_call(struct bpf_verifier_env *env, int subprog,
 			   struct bpf_reg_state *regs);
-int btf_prepare_func_args(struct bpf_verifier_env *env, int subprog,
-			  struct bpf_reg_state *reg, u32 *nargs);
+int btf_prepare_func_args(struct bpf_verifier_env *env, int subprog);
 int btf_check_type_match(struct bpf_verifier_log *log, const struct bpf_prog *prog,
 			 struct btf *btf, const struct btf_type *t);
 const char *btf_find_decl_tag_value(const struct btf *btf, const struct btf_type *pt,
diff --git a/include/linux/bpf_verifier.h b/include/linux/bpf_verifier.h
index 78d3f93b3802..23d054c9d1c8 100644
--- a/include/linux/bpf_verifier.h
+++ b/include/linux/bpf_verifier.h
@@ -574,6 +574,13 @@  static inline bool bpf_verifier_log_needed(const struct bpf_verifier_log *log)
 
 #define BPF_MAX_SUBPROGS 256
 
+struct bpf_subprog_arg_info {
+	enum bpf_arg_type arg_type;
+	union {
+		u32 mem_size;
+	};
+};
+
 struct bpf_subprog_info {
 	/* 'start' has to be the first field otherwise find_subprog() won't work */
 	u32 start; /* insn idx of function entry point */
@@ -585,6 +592,10 @@  struct bpf_subprog_info {
 	bool is_cb: 1;
 	bool is_async_cb: 1;
 	bool is_exception_cb: 1;
+	bool args_cached: 1;
+
+	u8 arg_cnt;
+	struct bpf_subprog_arg_info args[MAX_BPF_FUNC_REG_ARGS];
 };
 
 struct bpf_verifier_env;
@@ -690,6 +701,11 @@  struct bpf_verifier_env {
 	char tmp_str_buf[TMP_STR_BUF_LEN];
 };
 
+static inline struct bpf_subprog_info *subprog_info(struct bpf_verifier_env *env, int subprog)
+{
+	return &env->subprog_info[subprog];
+}
+
 __printf(2, 0) void bpf_verifier_vlog(struct bpf_verifier_log *log,
 				      const char *fmt, va_list args);
 __printf(2, 3) void bpf_verifier_log_write(struct bpf_verifier_env *env,
diff --git a/kernel/bpf/btf.c b/kernel/bpf/btf.c
index d56433bf8aba..33a62df9c5a8 100644
--- a/kernel/bpf/btf.c
+++ b/kernel/bpf/btf.c
@@ -6955,9 +6955,9 @@  int btf_check_subprog_call(struct bpf_verifier_env *env, int subprog,
  * 0 - Successfully converted BTF into bpf_reg_state
  * (either PTR_TO_CTX or SCALAR_VALUE).
  */
-int btf_prepare_func_args(struct bpf_verifier_env *env, int subprog,
-			  struct bpf_reg_state *regs, u32 *arg_cnt)
+int btf_prepare_func_args(struct bpf_verifier_env *env, int subprog)
 {
+	struct bpf_subprog_info *sub = subprog_info(env, subprog);
 	struct bpf_verifier_log *log = &env->log;
 	struct bpf_prog *prog = env->prog;
 	enum bpf_prog_type prog_type = prog->type;
@@ -6967,6 +6967,9 @@  int btf_prepare_func_args(struct bpf_verifier_env *env, int subprog,
 	u32 i, nargs, btf_id;
 	const char *tname;
 
+	if (sub->args_cached)
+		return 0;
+
 	if (!prog->aux->func_info ||
 	    prog->aux->func_info_aux[subprog].linkage != BTF_FUNC_GLOBAL) {
 		bpf_log(log, "Verifier bug\n");
@@ -6990,10 +6993,6 @@  int btf_prepare_func_args(struct bpf_verifier_env *env, int subprog,
 	}
 	tname = btf_name_by_offset(btf, t->name_off);
 
-	if (log->level & BPF_LOG_LEVEL)
-		bpf_log(log, "Validating %s() func#%d...\n",
-			tname, subprog);
-
 	if (prog->aux->func_info_aux[subprog].unreliable) {
 		bpf_log(log, "Verifier bug in function %s()\n", tname);
 		return -EFAULT;
@@ -7013,7 +7012,6 @@  int btf_prepare_func_args(struct bpf_verifier_env *env, int subprog,
 			tname, nargs, MAX_BPF_FUNC_REG_ARGS);
 		return -EINVAL;
 	}
-	*arg_cnt = nargs;
 	/* check that function returns int, exception cb also requires this */
 	t = btf_type_by_id(btf, t->type);
 	while (btf_type_is_modifier(t))
@@ -7028,24 +7026,24 @@  int btf_prepare_func_args(struct bpf_verifier_env *env, int subprog,
 	 * Only PTR_TO_CTX and SCALAR are supported atm.
 	 */
 	for (i = 0; i < nargs; i++) {
-		struct bpf_reg_state *reg = &regs[i + 1];
-
 		t = btf_type_by_id(btf, args[i].type);
 		while (btf_type_is_modifier(t))
 			t = btf_type_by_id(btf, t->type);
 		if (btf_type_is_int(t) || btf_is_any_enum(t)) {
-			reg->type = SCALAR_VALUE;
+			sub->args[i].arg_type = ARG_SCALAR;
 			continue;
 		}
 		if (btf_type_is_ptr(t)) {
+			u32 mem_size;
+
 			if (btf_get_prog_ctx_type(log, btf, t, prog_type, i)) {
-				reg->type = PTR_TO_CTX;
+				sub->args[i].arg_type = ARG_PTR_TO_CTX;
 				continue;
 			}
 
 			t = btf_type_skip_modifiers(btf, t->type, NULL);
 
-			ref_t = btf_resolve_size(btf, t, &reg->mem_size);
+			ref_t = btf_resolve_size(btf, t, &mem_size);
 			if (IS_ERR(ref_t)) {
 				bpf_log(log,
 				    "arg#%d reference type('%s %s') size cannot be determined: %ld\n",
@@ -7054,15 +7052,18 @@  int btf_prepare_func_args(struct bpf_verifier_env *env, int subprog,
 				return -EINVAL;
 			}
 
-			reg->type = PTR_TO_MEM | PTR_MAYBE_NULL;
-			reg->id = ++env->id_gen;
-
+			sub->args[i].arg_type = ARG_PTR_TO_MEM_OR_NULL;
+			sub->args[i].mem_size = mem_size;
 			continue;
 		}
 		bpf_log(log, "Arg#%d type %s in %s() is not supported yet.\n",
 			i, btf_type_str(t), tname);
 		return -EINVAL;
 	}
+
+	sub->arg_cnt = nargs;
+	sub->args_cached = true;
+
 	return 0;
 }
 
diff --git a/kernel/bpf/verifier.c b/kernel/bpf/verifier.c
index ee707736ce6b..16d5550eda4d 100644
--- a/kernel/bpf/verifier.c
+++ b/kernel/bpf/verifier.c
@@ -442,11 +442,6 @@  static struct bpf_func_info_aux *subprog_aux(const struct bpf_verifier_env *env,
 	return &env->prog->aux->func_info_aux[subprog];
 }
 
-static struct bpf_subprog_info *subprog_info(struct bpf_verifier_env *env, int subprog)
-{
-	return &env->subprog_info[subprog];
-}
-
 static void mark_subprog_exc_cb(struct bpf_verifier_env *env, int subprog)
 {
 	struct bpf_subprog_info *info = subprog_info(env, subprog);
@@ -19860,33 +19855,44 @@  static int do_check_common(struct bpf_verifier_env *env, int subprog)
 
 	regs = state->frame[state->curframe]->regs;
 	if (subprog || env->prog->type == BPF_PROG_TYPE_EXT) {
-		u32 nargs;
+		struct bpf_subprog_info *sub = subprog_info(env, subprog);
+		const char *sub_name = subprog_name(env, subprog);
+		struct bpf_subprog_arg_info *arg;
+		struct bpf_reg_state *reg;
 
-		ret = btf_prepare_func_args(env, subprog, regs, &nargs);
+		verbose(env, "Validating %s() func#%d...\n", sub_name, subprog);
+		ret = btf_prepare_func_args(env, subprog);
 		if (ret)
 			goto out;
+
 		if (subprog_is_exc_cb(env, subprog)) {
 			state->frame[0]->in_exception_callback_fn = true;
 			/* We have already ensured that the callback returns an integer, just
 			 * like all global subprogs. We need to determine it only has a single
 			 * scalar argument.
 			 */
-			if (nargs != 1 || regs[BPF_REG_1].type != SCALAR_VALUE) {
+			if (sub->arg_cnt != 1 || sub->args[0].arg_type != ARG_SCALAR) {
 				verbose(env, "exception cb only supports single integer argument\n");
 				return -EINVAL;
 			}
 		}
 		for (i = BPF_REG_1; i <= BPF_REG_5; i++) {
-			if (regs[i].type == PTR_TO_CTX)
+			arg = &sub->args[i - BPF_REG_1];
+			reg = &regs[i];
+
+			if (arg->arg_type == ARG_PTR_TO_CTX) {
+				reg->type = PTR_TO_CTX;
 				mark_reg_known_zero(env, regs, i);
-			else if (regs[i].type == SCALAR_VALUE)
+			} else if (arg->arg_type == ARG_SCALAR) {
+				reg->type = SCALAR_VALUE;
 				mark_reg_unknown(env, regs, i);
-			else if (base_type(regs[i].type) == PTR_TO_MEM) {
-				const u32 mem_size = regs[i].mem_size;
-
+			} else if (base_type(arg->arg_type) == ARG_PTR_TO_MEM) {
+				reg->type = PTR_TO_MEM;
+				if (arg->arg_type & PTR_MAYBE_NULL)
+					reg->type |= PTR_MAYBE_NULL;
 				mark_reg_known_zero(env, regs, i);
-				regs[i].mem_size = mem_size;
-				regs[i].id = ++env->id_gen;
+				reg->mem_size = arg->mem_size;
+				reg->id = ++env->id_gen;
 			}
 		}
 	} else {