/*
 * Copyright 2002-2020, Axel Dörfler, axeld@pinc-software.de.
 * Copyright 2012, Andreas Henriksson, sausageboy@gmail.com
 * This file may be used under the terms of the MIT License.
 */


//! File system error checking


#include "CheckVisitor.h"

#include "BlockAllocator.h"
#include "BPlusTree.h"
#include "Inode.h"
#include "Volume.h"


struct check_index {
        check_index()
                :
                inode(NULL)
        {
        }

        char                            name[B_FILE_NAME_LENGTH];
        block_run                       run;
        Inode*                          inode;
};


CheckVisitor::CheckVisitor(Volume* volume)
        :
        FileSystemVisitor(volume),
        fCheckBitmap(NULL)
{
}


CheckVisitor::~CheckVisitor()
{
        free(fCheckBitmap);
}


status_t
CheckVisitor::StartBitmapPass()
{
        if (!_ControlValid())
                return B_BAD_VALUE;

        // Lock the volume's journal and block allocator
        GetVolume()->GetJournal(0)->Lock(NULL, true);
        recursive_lock_lock(&GetVolume()->Allocator().Lock());

        size_t size = _BitmapSize();
        fCheckBitmap = (uint32*)malloc(size);
        if (fCheckBitmap == NULL) {
                recursive_lock_unlock(&GetVolume()->Allocator().Lock());
                GetVolume()->GetJournal(0)->Unlock(NULL, true);
                return B_NO_MEMORY;
        }

        memset(&Control().stats, 0, sizeof(check_control::stats));

        // initialize bitmap
        memset(fCheckBitmap, 0, size);
        for (off_t block = GetVolume()->ToBlock(GetVolume()->Log())
                + GetVolume()->Log().Length(); block-- > 0;) {
                _SetCheckBitmapAt(block);
        }

        Control().pass = BFS_CHECK_PASS_BITMAP;
        Control().stats.block_size = GetVolume()->BlockSize();

        // TODO: check reserved area in bitmap!

        Start(VISIT_REGULAR | VISIT_INDICES | VISIT_REMOVED
                | VISIT_ATTRIBUTE_DIRECTORIES);

        return B_OK;
}


status_t
CheckVisitor::WriteBackCheckBitmap()
{
        if (GetVolume()->IsReadOnly())
                return B_OK;

        // calculate the number of used blocks in the check bitmap
        size_t size = _BitmapSize();
        off_t usedBlocks = 0LL;

        // TODO: update the allocation groups used blocks info
        for (uint32 i = size >> 2; i-- > 0;) {
                uint32 compare = 1;
                // Count the number of bits set
                for (int16 j = 0; j < 32; j++, compare <<= 1) {
                        if ((compare & fCheckBitmap[i]) != 0)
                                usedBlocks++;
                }
        }

        Control().stats.freed = GetVolume()->UsedBlocks() - usedBlocks
                + Control().stats.missing;
        if (Control().stats.freed < 0)
                Control().stats.freed = 0;

        // Should we fix errors? Were there any errors we can fix?
        if ((Control().flags & BFS_FIX_BITMAP_ERRORS) != 0
                && (Control().stats.freed != 0 || Control().stats.missing != 0)) {
                // If so, write the check bitmap back over the original one,
                // and use transactions here to play safe - we even use several
                // transactions, so that we don't blow the maximum log size
                // on large disks, since we don't need to make this atomic.
#if 0
                // prints the blocks that differ
                off_t block = 0;
                for (int32 i = 0; i < fNumGroups; i++) {
                        AllocationBlock cached(fVolume);
                        for (uint32 j = 0; j < fGroups[i].NumBlocks(); j++) {
                                cached.SetTo(fGroups[i], j);
                                for (uint32 k = 0; k < cached.NumBlockBits(); k++) {
                                        if (cached.IsUsed(k) != _CheckBitmapIsUsedAt(block)) {
                                                dprintf("differ block %lld (should be %d)\n", block,
                                                        _CheckBitmapIsUsedAt(block));
                                        }
                                        block++;
                                }
                        }
                }
#endif

                GetVolume()->SuperBlock().used_blocks
                        = HOST_ENDIAN_TO_BFS_INT64(usedBlocks);

                size_t blockSize = GetVolume()->BlockSize();
                off_t numBitmapBlocks = GetVolume()->NumBitmapBlocks();

                for (uint32 i = 0; i < numBitmapBlocks; i += 512) {
                        Transaction transaction(GetVolume(), 1 + i);

                        uint32 blocksToWrite = 512;
                        if (blocksToWrite + i > numBitmapBlocks)
                                blocksToWrite = numBitmapBlocks - i;

                        status_t status = transaction.WriteBlocks(1 + i,
                                (uint8*)fCheckBitmap + i * blockSize, blocksToWrite);
                        if (status < B_OK) {
                                FATAL(("error writing bitmap: %s\n", strerror(status)));
                                return status;
                        }
                        transaction.Done();
                }
        }

        return B_OK;
}


status_t
CheckVisitor::StartIndexPass()
{
        // if we don't have indices to rebuild, this pass is done
        if (indices.IsEmpty())
                return B_ENTRY_NOT_FOUND;

        Control().pass = BFS_CHECK_PASS_INDEX;

        status_t status = _PrepareIndices();
        if (status != B_OK) {
                Control().status = status;
                return status;
        }

        Start(VISIT_REGULAR);
        return Next();
}


status_t
CheckVisitor::StopChecking()
{
        if (GetVolume()->IsReadOnly()) {
                // We can't fix errors on this volume
                Control().flags = 0;
        }

        if (Control().status != B_ENTRY_NOT_FOUND)
                FATAL(("CheckVisitor didn't run through\n"));

        _FreeIndices();

        recursive_lock_unlock(&GetVolume()->Allocator().Lock());
        GetVolume()->GetJournal(0)->Unlock(NULL, true);
        return B_OK;
}


status_t
CheckVisitor::VisitDirectoryEntry(Inode* inode, Inode* parent,
        const char* treeName)
{
        Control().inode = inode->ID();
        Control().mode = inode->Mode();

        if (Pass() != BFS_CHECK_PASS_BITMAP)
                return B_OK;

        // check if the inode's name is the same as in the b+tree
        if (inode->IsRegularNode()) {
                RecursiveLocker locker(inode->SmallDataLock());
                NodeGetter node(GetVolume());
                status_t status = node.SetTo(inode);
                if (status != B_OK) {
                        Control().errors |= BFS_COULD_NOT_OPEN;
                        Control().status = status;
                        return B_OK;
                }

                const char* localName = inode->Name(node.Node());
                if (localName == NULL || strcmp(localName, treeName)) {
                        Control().errors |= BFS_NAMES_DONT_MATCH;
                        FATAL(("Names differ: tree \"%s\", inode \"%s\"\n", treeName,
                                localName));

                        if ((Control().flags & BFS_FIX_NAME_MISMATCHES) != 0) {
                                // Rename the inode
                                Transaction transaction(GetVolume(), inode->BlockNumber());

                                // Note, this may need extra blocks, but the inode will
                                // only be checked afterwards, so that it won't be lost
                                status_t status = inode->SetName(transaction, treeName);
                                if (status == B_OK)
                                        status = inode->WriteBack(transaction);
                                if (status == B_OK)
                                        status = transaction.Done();
                                if (status != B_OK) {
                                        Control().status = status;
                                        return B_OK;
                                }
                        }
                }
        }

        // Check for the correct mode of the node (if the mode of the
        // file don't fit to its parent, there is a serious problem)
        if (((parent->Mode() & S_ATTR_DIR) != 0
                        && !inode->IsAttribute())
                || ((parent->Mode() & S_INDEX_DIR) != 0
                        && !inode->IsIndex())
                || (is_directory(parent->Mode())
                        && !inode->IsRegularNode())) {
                FATAL(("inode at %" B_PRIdOFF " is of wrong type: %o (parent "
                        "%o at %" B_PRIdOFF ")!\n", inode->BlockNumber(),
                        inode->Mode(), parent->Mode(), parent->BlockNumber()));

                // if we are allowed to fix errors, we should remove the file
                if ((Control().flags & BFS_REMOVE_WRONG_TYPES) != 0
                        && (Control().flags & BFS_FIX_BITMAP_ERRORS) != 0) {
                        Control().status = _RemoveInvalidNode(parent, NULL, inode,
                                treeName);
                } else
                        Control().status = B_ERROR;

                Control().errors |= BFS_WRONG_TYPE;
        }
        return B_OK;
}


status_t
CheckVisitor::VisitInode(Inode* inode, const char* treeName)
{
        Control().inode = inode->ID();
        Control().mode = inode->Mode();
                // (we might have set these in VisitDirectoryEntry already)

        // set name
        if (treeName == NULL) {
                if (inode->GetName(Control().name) < B_OK) {
                        if (inode->IsContainer())
                                strcpy(Control().name, "(dir has no name)");
                        else
                                strcpy(Control().name, "(node has no name)");
                }
        } else
                strcpy(Control().name, treeName);

        status_t status = B_OK;

        switch (Pass()) {
                case BFS_CHECK_PASS_BITMAP:
                {
                        status = _CheckInodeBlocks(inode, NULL);
                        if (status != B_OK)
                                return status;

                        // Check the B+tree as well
                        if (inode->IsContainer()) {
                                bool repairErrors = (Control().flags & BFS_FIX_BPLUSTREES) != 0;
                                bool errorsFound = false;

                                status = inode->Tree()->Validate(repairErrors, errorsFound);

                                if (errorsFound) {
                                        Control().errors |= BFS_INVALID_BPLUSTREE;
                                        if (inode->IsIndex() && treeName != NULL && repairErrors) {
                                                // We completely rebuild corrupt indices
                                                check_index* index = new(std::nothrow) check_index;
                                                if (index == NULL)
                                                        return B_NO_MEMORY;

                                                strlcpy(index->name, treeName, sizeof(index->name));
                                                index->run = inode->BlockRun();
                                                Indices().Push(index);
                                        }
                                }
                        }

                        break;
                }

                case BFS_CHECK_PASS_INDEX:
                        status = _AddInodeToIndex(inode);
                        break;
        }

        Control().status = status;
        return B_OK;
}


status_t
CheckVisitor::OpenInodeFailed(status_t reason, ino_t id, Inode* parent,
                char* treeName, TreeIterator* iterator)
{
        FATAL(("Could not open inode at %" B_PRIdOFF ": %s\n", id,
                strerror(reason)));

        if (treeName != NULL)
                strlcpy(Control().name, treeName, B_FILE_NAME_LENGTH);
        else
                strcpy(Control().name, "(node has no name)");

        Control().inode = id;
        Control().errors = BFS_COULD_NOT_OPEN;

        // TODO: check other error codes; B_IO_ERROR might be a temporary
        // issue, so it should be guarded by a force mode
        if (reason == B_BAD_VALUE || reason == B_BAD_DATA || reason == B_IO_ERROR) {
                // Remove inode from the tree if we can
                if (parent != NULL && iterator != NULL
                        && (Control().flags & BFS_REMOVE_INVALID) != 0) {
                        Control().status = _RemoveInvalidNode(parent, iterator->Tree(),
                                NULL, treeName);
                } else
                        Control().status = B_ERROR;
        } else {
                Control().status = B_OK;
        }

        return B_OK;
}


status_t
CheckVisitor::OpenBPlusTreeFailed(Inode* inode)
{
        FATAL(("Could not open b+tree from inode at %" B_PRIdOFF "\n",
                inode->ID()));
        return B_OK;
}


status_t
CheckVisitor::TreeIterationFailed(status_t reason, Inode* parent)
{
        // Iterating over the B+tree failed - we let the checkfs run
        // fail completely, as we would delete all files we cannot
        // access.
        // TODO: maybe have a force parameter that actually does that.
        // TODO: we also need to be able to repair broken B+trees!
        return reason;
}


status_t
CheckVisitor::_RemoveInvalidNode(Inode* parent, BPlusTree* tree,
        Inode* inode, const char* name)
{
        // It's safe to start a transaction, because Inode::Remove()
        // won't touch the block bitmap (which we hold the lock for)
        // if we set the INODE_DONT_FREE_SPACE flag - since we fix
        // the bitmap anyway.
        Transaction transaction(GetVolume(), parent->BlockNumber());
        status_t status;

        if (inode != NULL) {
                inode->Node().flags |= HOST_ENDIAN_TO_BFS_INT32(INODE_DONT_FREE_SPACE);

                status = parent->Remove(transaction, name, NULL, false, true);
        } else {
                parent->WriteLockInTransaction(transaction);

                // does the file even exist?
                off_t id;
                status = tree->Find((uint8*)name, (uint16)strlen(name), &id);
                if (status == B_OK)
                        status = tree->Remove(transaction, name, id);
        }

        if (status == B_OK) {
                entry_cache_remove(GetVolume()->ID(), parent->ID(), name);
                transaction.Done();
        }

        return status;
}


bool
CheckVisitor::_ControlValid()
{
        if (Control().magic != BFS_IOCTL_CHECK_MAGIC) {
                FATAL(("invalid check_control!\n"));
                return false;
        }

        return true;
}


bool
CheckVisitor::_CheckBitmapIsUsedAt(off_t block) const
{
        size_t size = _BitmapSize();
        uint32 index = block / 32;      // 32bit resolution
        if (index > size / 4)
                return false;

        return BFS_ENDIAN_TO_HOST_INT32(fCheckBitmap[index])
                & (1UL << (block & 0x1f));
}


void
CheckVisitor::_SetCheckBitmapAt(off_t block)
{
        size_t size = _BitmapSize();
        uint32 index = block / 32;      // 32bit resolution
        if (index > size / 4)
                return;

        fCheckBitmap[index] |= HOST_ENDIAN_TO_BFS_INT32(1UL << (block & 0x1f));
}


size_t
CheckVisitor::_BitmapSize() const
{
        return GetVolume()->BlockSize() * GetVolume()->NumBitmapBlocks();
}


status_t
CheckVisitor::_CheckInodeBlocks(Inode* inode, const char* name)
{
        status_t status = _CheckAllocated(inode->BlockRun(), "inode");
        if (status != B_OK)
                return status;

        if (inode->IsSymLink() && (inode->Flags() & INODE_LONG_SYMLINK) == 0) {
                // symlinks may not have a valid data stream
                if (strlen(inode->Node().short_symlink) >= SHORT_SYMLINK_NAME_LENGTH)
                        return B_BAD_DATA;

                return B_OK;
        }

        data_stream* data = &inode->Node().data;

        // check the direct range

        if (data->max_direct_range) {
                for (int32 i = 0; i < NUM_DIRECT_BLOCKS; i++) {
                        if (data->direct[i].IsZero())
                                break;

                        status = _CheckAllocated(data->direct[i], "direct");
                        if (status < B_OK)
                                return status;

                        Control().stats.direct_block_runs++;
                        Control().stats.blocks_in_direct
                                += data->direct[i].Length();
                }
        }

        CachedBlock cached(GetVolume());

        // check the indirect range

        if (data->max_indirect_range) {
                status = _CheckAllocated(data->indirect, "indirect");
                if (status != B_OK)
                        return status;

                off_t block = GetVolume()->ToBlock(data->indirect);

                for (int32 i = 0; i < data->indirect.Length(); i++) {
                        status = cached.SetTo(block + i);
                        if (status != B_OK)
                                RETURN_ERROR(status);

                        block_run* runs = (block_run*)cached.Block();
                        int32 runsPerBlock = GetVolume()->BlockSize() / sizeof(block_run);
                        int32 index = 0;
                        for (; index < runsPerBlock; index++) {
                                if (runs[index].IsZero())
                                        break;

                                status = _CheckAllocated(runs[index], "indirect->run");
                                if (status < B_OK)
                                        return status;

                                Control().stats.indirect_block_runs++;
                                Control().stats.blocks_in_indirect
                                        += runs[index].Length();
                        }
                        Control().stats.indirect_array_blocks++;

                        if (index < runsPerBlock)
                                break;
                }
        }

        // check the double indirect range

        if (data->max_double_indirect_range) {
                status = _CheckAllocated(data->double_indirect, "double indirect");
                if (status != B_OK)
                        return status;

                int32 runsPerBlock = runs_per_block(GetVolume()->BlockSize());
                int32 runsPerArray = runsPerBlock * data->double_indirect.Length();

                CachedBlock cachedDirect(GetVolume());

                for (int32 indirectIndex = 0; indirectIndex < runsPerArray;
                                indirectIndex++) {
                        // get the indirect array block
                        status = cached.SetTo(GetVolume()->ToBlock(data->double_indirect)
                                        + indirectIndex / runsPerBlock);
                        if (status != B_OK)
                                return status;

                        block_run* array = (block_run*)cached.Block();
                        block_run indirect = array[indirectIndex % runsPerBlock];
                        // are we finished yet?
                        if (indirect.IsZero())
                                return B_OK;

                        status = _CheckAllocated(indirect, "double indirect->runs");
                        if (status != B_OK)
                                return status;

                        int32 maxIndex
                                = ((uint32)indirect.Length() << GetVolume()->BlockShift())
                                        / sizeof(block_run);

                        for (int32 index = 0; index < maxIndex; ) {
                                status = cachedDirect.SetTo(GetVolume()->ToBlock(indirect)
                                                + index / runsPerBlock);
                                if (status != B_OK)
                                        return status;

                                block_run* runs = (block_run*)cachedDirect.Block();

                                do {
                                        // are we finished yet?
                                        if (runs[index % runsPerBlock].IsZero())
                                                return B_OK;

                                        status = _CheckAllocated(runs[index % runsPerBlock],
                                                "double indirect->runs->run");
                                        if (status != B_OK)
                                                return status;

                                        Control().stats.double_indirect_block_runs++;
                                        Control().stats.blocks_in_double_indirect
                                                += runs[index % runsPerBlock].Length();
                                } while ((++index % runsPerBlock) != 0);
                        }

                        Control().stats.double_indirect_array_blocks++;
                }
        }

        return B_OK;
}


status_t
CheckVisitor::_CheckAllocated(block_run run, const char* type)
{
        BlockAllocator& allocator = GetVolume()->Allocator();

        // make sure the block run is valid
        if (!allocator.IsValidBlockRun(run, type)) {
                Control().errors |= BFS_INVALID_BLOCK_RUN;
                return B_OK;
        }

        status_t status;

        off_t start = GetVolume()->ToBlock(run);
        off_t end = start + run.Length();

        // check if the run is allocated in the block bitmap on disk
        off_t block = start;

        while (block < end) {
                off_t firstMissing;
                status = allocator.CheckBlocks(block, end - block, true, &firstMissing);
                if (status == B_OK)
                        break;
                else if (status != B_BAD_DATA)
                        return status;

                off_t afterLastMissing;
                status = allocator.CheckBlocks(firstMissing, end - firstMissing, false,
                        &afterLastMissing);
                if (status == B_OK)
                        afterLastMissing = end;
                else if (status != B_BAD_DATA)
                        return status;

                PRINT(("%s: block_run(%" B_PRId32 ", %" B_PRIu16 ", %" B_PRIu16 ")"
                        ": blocks %" B_PRIdOFF " - %" B_PRIdOFF " are not allocated!\n",
                        type, run.AllocationGroup(), run.Start(),
                        run.Length(), firstMissing, afterLastMissing - 1));

                Control().stats.missing += afterLastMissing - firstMissing;

                block = afterLastMissing;
        }

        // set bits in check bitmap, while checking if they're already set
        off_t firstSet = -1;

        for (block = start; block < end; block++) {
                if (_CheckBitmapIsUsedAt(block)) {
                        if (firstSet == -1) {
                                firstSet = block;
                                Control().errors |= BFS_BLOCKS_ALREADY_SET;
                        }
                        Control().stats.already_set++;
                } else {
                        if (firstSet != -1) {
                                FATAL(("%s: block_run(%d, %u, %u): blocks %" B_PRIdOFF
                                        " - %" B_PRIdOFF " are already set!\n", type,
                                        (int)run.AllocationGroup(), run.Start(), run.Length(),
                                        firstSet, block - 1));
                                firstSet = -1;
                        }
                        _SetCheckBitmapAt(block);
                }
        }

        return B_OK;
}


status_t
CheckVisitor::_PrepareIndices()
{
        int32 count = 0;

        for (int32 i = 0; i < Indices().CountItems(); i++) {
                check_index* index = Indices().Array()[i];
                Vnode vnode(GetVolume(), index->run);
                Inode* inode;
                status_t status = vnode.Get(&inode);
                if (status != B_OK) {
                        FATAL(("check: Could not open index at %" B_PRIdOFF "\n",
                                GetVolume()->ToBlock(index->run)));
                        return status;
                }

                BPlusTree* tree = inode->Tree();
                if (tree == NULL) {
                        // TODO: We can't yet repair those
                        continue;
                }

                status = tree->MakeEmpty();
                if (status != B_OK)
                        return status;

                index->inode = inode;
                vnode.Keep();
                count++;
        }

        return count == 0 ? B_ENTRY_NOT_FOUND : B_OK;
}


void
CheckVisitor::_FreeIndices()
{
        for (int32 i = 0; i < Indices().CountItems(); i++) {
                check_index* index = Indices().Array()[i];
                if (index->inode != NULL) {
                        put_vnode(GetVolume()->FSVolume(),
                                GetVolume()->ToVnode(index->inode->BlockRun()));
                }
                delete index;
        }
        Indices().MakeEmpty();
}


status_t
CheckVisitor::_AddInodeToIndex(Inode* inode)
{
        Transaction transaction(GetVolume(), inode->BlockNumber());

        for (int32 i = 0; i < Indices().CountItems(); i++) {
                check_index* index = Indices().Array()[i];
                if (index->inode == NULL)
                        continue;

                index->inode->WriteLockInTransaction(transaction);

                BPlusTree* tree = index->inode->Tree();
                if (tree == NULL)
                        return B_ERROR;

                status_t status = B_OK;

                if (!strcmp(index->name, "name")) {
                        if (inode->InNameIndex()) {
                                char name[B_FILE_NAME_LENGTH];
                                if (inode->GetName(name, B_FILE_NAME_LENGTH) != B_OK)
                                        return B_ERROR;

                                status = tree->Insert(transaction, name, inode->ID());
                        }
                } else if (!strcmp(index->name, "last_modified")) {
                        if (inode->InLastModifiedIndex()) {
                                status = tree->Insert(transaction, inode->OldLastModified(),
                                        inode->ID());
                        }
                } else if (!strcmp(index->name, "size")) {
                        if (inode->InSizeIndex())
                                status = tree->Insert(transaction, inode->Size(), inode->ID());
                } else {
                        uint8 key[MAX_INDEX_KEY_LENGTH];
                        size_t keyLength = sizeof(key);
                        if (inode->ReadAttribute(index->name, B_ANY_TYPE, 0, key,
                                        &keyLength) == B_OK) {
                                status = tree->Insert(transaction, key, keyLength, inode->ID());
                        }
                }

                if (status != B_OK)
                        return status;
        }

        return transaction.Done();
}