}
 
        if (class == BPF_STX) {
-               if (reg->type != SCALAR_VALUE)
+               /* BPF_STX (including atomic variants) has multiple source
+                * operands, one of which is a ptr. Check whether the caller is
+                * asking about it.
+                */
+               if (t == SRC_OP && reg->type != SCALAR_VALUE)
                        return true;
                return BPF_SIZE(code) == BPF_DW;
        }
        return true;
 }
 
-/* Return TRUE if INSN doesn't have explicit value define. */
-static bool insn_no_def(struct bpf_insn *insn)
+/* Return the regno defined by the insn, or -1. */
+static int insn_def_regno(const struct bpf_insn *insn)
 {
-       u8 class = BPF_CLASS(insn->code);
-
-       return (class == BPF_JMP || class == BPF_JMP32 ||
-               class == BPF_STX || class == BPF_ST);
+       switch (BPF_CLASS(insn->code)) {
+       case BPF_JMP:
+       case BPF_JMP32:
+       case BPF_ST:
+               return -1;
+       case BPF_STX:
+               if (BPF_MODE(insn->code) == BPF_ATOMIC &&
+                   (insn->imm & BPF_FETCH)) {
+                       if (insn->imm == BPF_CMPXCHG)
+                               return BPF_REG_0;
+                       else
+                               return insn->src_reg;
+               } else {
+                       return -1;
+               }
+       default:
+               return insn->dst_reg;
+       }
 }
 
 /* Return TRUE if INSN has defined any 32-bit value explicitly. */
 static bool insn_has_def32(struct bpf_verifier_env *env, struct bpf_insn *insn)
 {
-       if (insn_no_def(insn))
+       int dst_reg = insn_def_regno(insn);
+
+       if (dst_reg == -1)
                return false;
 
-       return !is_reg64(env, insn, insn->dst_reg, NULL, DST_OP);
+       return !is_reg64(env, insn, dst_reg, NULL, DST_OP);
 }
 
 static void mark_insn_zext(struct bpf_verifier_env *env,
        for (i = 0; i < len; i++) {
                int adj_idx = i + delta;
                struct bpf_insn insn;
-               u8 load_reg;
+               int load_reg;
 
                insn = insns[adj_idx];
+               load_reg = insn_def_regno(&insn);
                if (!aux[adj_idx].zext_dst) {
                        u8 code, class;
                        u32 imm_rnd;
 
                        code = insn.code;
                        class = BPF_CLASS(code);
-                       if (insn_no_def(&insn))
+                       if (load_reg == -1)
                                continue;
 
                        /* NOTE: arg "reg" (the fourth one) is only used for
-                        *       BPF_STX which has been ruled out in above
-                        *       check, it is safe to pass NULL here.
+                        *       BPF_STX + SRC_OP, so it is safe to pass NULL
+                        *       here.
                         */
-                       if (is_reg64(env, &insn, insn.dst_reg, NULL, DST_OP)) {
+                       if (is_reg64(env, &insn, load_reg, NULL, DST_OP)) {
                                if (class == BPF_LD &&
                                    BPF_MODE(code) == BPF_IMM)
                                        i++;
                        imm_rnd = get_random_int();
                        rnd_hi32_patch[0] = insn;
                        rnd_hi32_patch[1].imm = imm_rnd;
-                       rnd_hi32_patch[3].dst_reg = insn.dst_reg;
+                       rnd_hi32_patch[3].dst_reg = load_reg;
                        patch = rnd_hi32_patch;
                        patch_len = 4;
                        goto apply_patch_buffer;
                if (!bpf_jit_needs_zext())
                        continue;
 
-               /* zext_dst means that we want to zero-extend whatever register
-                * the insn defines, which is dst_reg most of the time, with
-                * the notable exception of BPF_STX + BPF_ATOMIC + BPF_FETCH.
-                */
-               if (BPF_CLASS(insn.code) == BPF_STX &&
-                   BPF_MODE(insn.code) == BPF_ATOMIC) {
-                       /* BPF_STX + BPF_ATOMIC insns without BPF_FETCH do not
-                        * define any registers, therefore zext_dst cannot be
-                        * set.
-                        */
-                       if (WARN_ON(!(insn.imm & BPF_FETCH)))
-                               return -EINVAL;
-                       load_reg = insn.imm == BPF_CMPXCHG ? BPF_REG_0
-                                                          : insn.src_reg;
-               } else {
-                       load_reg = insn.dst_reg;
+               if (WARN_ON(load_reg == -1)) {
+                       verbose(env, "verifier bug. zext_dst is set, but no reg is defined\n");
+                       return -EFAULT;
                }
 
                zext_patch[0] = insn;