unsigned long arg)
 {
        struct task_struct *t;
-       unsigned long size, unused;
+       unsigned long flags, unused;
        int mode, i;
        struct kcov_remote_arg *remote_arg;
        struct kcov_remote *remote;
-       unsigned long flags;
 
        switch (cmd) {
-       case KCOV_INIT_TRACE:
-               /*
-                * Enable kcov in trace mode and setup buffer size.
-                * Must happen before anything else.
-                */
-               if (kcov->mode != KCOV_MODE_DISABLED)
-                       return -EBUSY;
-               /*
-                * Size must be at least 2 to hold current position and one PC.
-                * Later we allocate size * sizeof(unsigned long) memory,
-                * that must not overflow.
-                */
-               size = arg;
-               if (size < 2 || size > INT_MAX / sizeof(unsigned long))
-                       return -EINVAL;
-               kcov->size = size;
-               kcov->mode = KCOV_MODE_INIT;
-               return 0;
        case KCOV_ENABLE:
                /*
                 * Enable coverage for the current task.
        struct kcov_remote_arg *remote_arg = NULL;
        unsigned int remote_num_handles;
        unsigned long remote_arg_size;
-       unsigned long flags;
+       unsigned long size, flags;
 
-       if (cmd == KCOV_REMOTE_ENABLE) {
+       kcov = filep->private_data;
+       switch (cmd) {
+       case KCOV_INIT_TRACE:
+               /*
+                * Enable kcov in trace mode and setup buffer size.
+                * Must happen before anything else.
+                *
+                * First check the size argument - it must be at least 2
+                * to hold the current position and one PC. Later we allocate
+                * size * sizeof(unsigned long) memory, that must not overflow.
+                */
+               size = arg;
+               if (size < 2 || size > INT_MAX / sizeof(unsigned long))
+                       return -EINVAL;
+               spin_lock_irqsave(&kcov->lock, flags);
+               if (kcov->mode != KCOV_MODE_DISABLED) {
+                       spin_unlock_irqrestore(&kcov->lock, flags);
+                       return -EBUSY;
+               }
+               kcov->size = size;
+               kcov->mode = KCOV_MODE_INIT;
+               spin_unlock_irqrestore(&kcov->lock, flags);
+               return 0;
+       case KCOV_REMOTE_ENABLE:
                if (get_user(remote_num_handles, (unsigned __user *)(arg +
                                offsetof(struct kcov_remote_arg, num_handles))))
                        return -EFAULT;
                        return -EINVAL;
                }
                arg = (unsigned long)remote_arg;
+               fallthrough;
+       default:
+               /*
+                * All other commands can be normally executed under a spin lock, so we
+                * obtain and release it here in order to simplify kcov_ioctl_locked().
+                */
+               spin_lock_irqsave(&kcov->lock, flags);
+               res = kcov_ioctl_locked(kcov, cmd, arg);
+               spin_unlock_irqrestore(&kcov->lock, flags);
+               kfree(remote_arg);
+               return res;
        }
-
-       kcov = filep->private_data;
-       spin_lock_irqsave(&kcov->lock, flags);
-       res = kcov_ioctl_locked(kcov, cmd, arg);
-       spin_unlock_irqrestore(&kcov->lock, flags);
-
-       kfree(remote_arg);
-
-       return res;
 }
 
 static const struct file_operations kcov_fops = {