memcpy(buf, req->iv, AES_BLOCK_SIZE);
 
        err = skcipher_walk_aead_encrypt(&walk, req, false);
-       if (unlikely(err))
-               return err;
 
        kernel_neon_begin();
 
        if (req->assoclen)
                ccm_calculate_auth_mac(req, mac);
 
-       do {
+       while (walk.nbytes) {
                u32 tail = walk.nbytes % AES_BLOCK_SIZE;
+               bool final = walk.nbytes == walk.total;
 
-               if (walk.nbytes == walk.total)
+               if (final)
                        tail = 0;
 
                ce_aes_ccm_encrypt(walk.dst.virt.addr, walk.src.virt.addr,
                                   walk.nbytes - tail, ctx->key_enc,
                                   num_rounds(ctx), mac, walk.iv);
 
-               if (walk.nbytes == walk.total)
-                       ce_aes_ccm_final(mac, buf, ctx->key_enc, num_rounds(ctx));
+               if (!final)
+                       kernel_neon_end();
+               err = skcipher_walk_done(&walk, tail);
+               if (!final)
+                       kernel_neon_begin();
+       }
 
-               kernel_neon_end();
+       ce_aes_ccm_final(mac, buf, ctx->key_enc, num_rounds(ctx));
 
-               if (walk.nbytes) {
-                       err = skcipher_walk_done(&walk, tail);
-                       if (unlikely(err))
-                               return err;
-                       if (unlikely(walk.nbytes))
-                               kernel_neon_begin();
-               }
-       } while (walk.nbytes);
+       kernel_neon_end();
 
        /* copy authtag to end of dst */
        scatterwalk_map_and_copy(mac, req->dst, req->assoclen + req->cryptlen,
                                 crypto_aead_authsize(aead), 1);
 
-       return 0;
+       return err;
 }
 
 static int ccm_decrypt(struct aead_request *req)
        memcpy(buf, req->iv, AES_BLOCK_SIZE);
 
        err = skcipher_walk_aead_decrypt(&walk, req, false);
-       if (unlikely(err))
-               return err;
 
        kernel_neon_begin();
 
        if (req->assoclen)
                ccm_calculate_auth_mac(req, mac);
 
-       do {
+       while (walk.nbytes) {
                u32 tail = walk.nbytes % AES_BLOCK_SIZE;
+               bool final = walk.nbytes == walk.total;
 
-               if (walk.nbytes == walk.total)
+               if (final)
                        tail = 0;
 
                ce_aes_ccm_decrypt(walk.dst.virt.addr, walk.src.virt.addr,
                                   walk.nbytes - tail, ctx->key_enc,
                                   num_rounds(ctx), mac, walk.iv);
 
-               if (walk.nbytes == walk.total)
-                       ce_aes_ccm_final(mac, buf, ctx->key_enc, num_rounds(ctx));
+               if (!final)
+                       kernel_neon_end();
+               err = skcipher_walk_done(&walk, tail);
+               if (!final)
+                       kernel_neon_begin();
+       }
 
-               kernel_neon_end();
+       ce_aes_ccm_final(mac, buf, ctx->key_enc, num_rounds(ctx));
 
-               if (walk.nbytes) {
-                       err = skcipher_walk_done(&walk, tail);
-                       if (unlikely(err))
-                               return err;
-                       if (unlikely(walk.nbytes))
-                               kernel_neon_begin();
-               }
-       } while (walk.nbytes);
+       kernel_neon_end();
+
+       if (unlikely(err))
+               return err;
 
        /* compare calculated auth tag with the stored one */
        scatterwalk_map_and_copy(buf, req->src,