IngestTasksScheduler.java

Go to the documentation of this file.

/*
 * Autopsy Forensic Browser
 *
 * Copyright 2012-2018 Basis Technology Corp.
 * Contact: carrier <at> sleuthkit <dot> org
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package org.sleuthkit.autopsy.ingest;

import java.util.ArrayList;
import java.util.Collection;
import java.util.Collections;
import java.util.Comparator;
import java.util.Deque;
import java.util.Iterator;
import java.util.LinkedList;
import java.util.List;
import java.util.TreeSet;
import java.util.concurrent.BlockingDeque;
import java.util.concurrent.LinkedBlockingDeque;
import java.util.logging.Level;
import java.util.regex.Matcher;
import java.util.regex.Pattern;
import javax.annotation.concurrent.GuardedBy;
import javax.annotation.concurrent.ThreadSafe;
import org.sleuthkit.autopsy.coreutils.Logger;
import org.sleuthkit.datamodel.AbstractFile;
import org.sleuthkit.datamodel.Content;
import org.sleuthkit.datamodel.FileSystem;
import org.sleuthkit.datamodel.TskCoreException;
import org.sleuthkit.datamodel.TskData;

@ThreadSafe
final class IngestTasksScheduler {

    private static final int FAT_NTFS_FLAGS = TskData.TSK_FS_TYPE_ENUM.TSK_FS_TYPE_FAT12.getValue() | TskData.TSK_FS_TYPE_ENUM.TSK_FS_TYPE_FAT16.getValue() | TskData.TSK_FS_TYPE_ENUM.TSK_FS_TYPE_FAT32.getValue() | TskData.TSK_FS_TYPE_ENUM.TSK_FS_TYPE_NTFS.getValue();
    private static final Logger logger = Logger.getLogger(IngestTasksScheduler.class.getName());
    @GuardedBy("IngestTasksScheduler.this")
    private static IngestTasksScheduler instance;
    private final IngestTaskTrackingQueue dataSourceIngestThreadQueue;
    @GuardedBy("this")
    private final TreeSet<FileIngestTask> rootFileTaskQueue;
    @GuardedBy("this")
    private final Deque<FileIngestTask> pendingFileTaskQueue;
    private final IngestTaskTrackingQueue fileIngestThreadsQueue;

    synchronized static IngestTasksScheduler getInstance() {
        if (IngestTasksScheduler.instance == null) {
            IngestTasksScheduler.instance = new IngestTasksScheduler();
        }
        return IngestTasksScheduler.instance;
    }

    private IngestTasksScheduler() {
        this.dataSourceIngestThreadQueue = new IngestTaskTrackingQueue();
        this.rootFileTaskQueue = new TreeSet<>(new RootDirectoryTaskComparator());
        this.pendingFileTaskQueue = new LinkedList<>();
        this.fileIngestThreadsQueue = new IngestTaskTrackingQueue();
    }

    BlockingIngestTaskQueue getDataSourceIngestTaskQueue() {
        return this.dataSourceIngestThreadQueue;
    }

    BlockingIngestTaskQueue getFileIngestTaskQueue() {
        return this.fileIngestThreadsQueue;
    }

    synchronized void scheduleIngestTasks(DataSourceIngestJob job) {
        if (!job.isCancelled()) {
            /*
             * Scheduling of both the data source ingest task and the initial
             * file ingest tasks for a job must be an atomic operation.
             * Otherwise, the data source task might be completed before the
             * file tasks are scheduled, resulting in a potential false positive
             * when another thread checks whether or not all the tasks for the
             * job are completed.
             */
            this.scheduleDataSourceIngestTask(job);
            this.scheduleFileIngestTasks(job, Collections.emptyList());
        }
    }

    synchronized void scheduleDataSourceIngestTask(DataSourceIngestJob job) {
        if (!job.isCancelled()) {
            DataSourceIngestTask task = new DataSourceIngestTask(job);
            try {
                this.dataSourceIngestThreadQueue.putLast(task);
            } catch (InterruptedException ex) {
                IngestTasksScheduler.logger.log(Level.INFO, String.format("Ingest tasks scheduler interrupted while blocked adding a task to the data source level ingest task queue (jobId={%d)", job.getId()), ex);
                Thread.currentThread().interrupt();
            }
        }
    }

    synchronized void scheduleFileIngestTasks(DataSourceIngestJob job, Collection<AbstractFile> files) {
        if (!job.isCancelled()) {
            Collection<AbstractFile> candidateFiles;
            if (files.isEmpty()) {
                candidateFiles = getTopLevelFiles(job.getDataSource());
            } else {
                candidateFiles = files;
            }
            for (AbstractFile file : candidateFiles) {
                FileIngestTask task = new FileIngestTask(job, file);
                if (IngestTasksScheduler.shouldEnqueueFileTask(task)) {
                    this.rootFileTaskQueue.add(task);
                }
            }
            shuffleFileTaskQueues();
        }
    }

    synchronized void fastTrackFileIngestTasks(DataSourceIngestJob job, Collection<AbstractFile> files) {
        if (!job.isCancelled()) {
            /*
             * Put the files directly into the queue for the file ingest
             * threads, if they pass the file filter for the job. The files are
             * added to the queue for the ingest threads BEFORE the other queued
             * tasks because the use case for this method is scheduling new
             * carved or derived files from a higher priority task that is
             * already in progress.
             */
            for (AbstractFile file : files) {
                FileIngestTask fileTask = new FileIngestTask(job, file);
                if (shouldEnqueueFileTask(fileTask)) {
                    try {
                        this.fileIngestThreadsQueue.putFirst(fileTask);
                    } catch (InterruptedException ex) {
                        IngestTasksScheduler.logger.log(Level.INFO, String.format("Ingest tasks scheduler interrupted while scheduling file level ingest tasks (jobId={%d)", job.getId()), ex);
                        Thread.currentThread().interrupt();
                        return;
                    }
                }
            }
        }
    }

    synchronized void notifyTaskCompleted(DataSourceIngestTask task) {
        this.dataSourceIngestThreadQueue.taskCompleted(task);
    }

    synchronized void notifyTaskCompleted(FileIngestTask task) {
        this.fileIngestThreadsQueue.taskCompleted(task);
        shuffleFileTaskQueues();
    }

    synchronized boolean tasksForJobAreCompleted(DataSourceIngestJob job) {
        long jobId = job.getId();
        return !(this.dataSourceIngestThreadQueue.hasTasksForJob(jobId)
                || hasTasksForJob(this.rootFileTaskQueue, jobId)
                || hasTasksForJob(this.pendingFileTaskQueue, jobId)
                || this.fileIngestThreadsQueue.hasTasksForJob(jobId));
    }

    synchronized void cancelPendingTasksForIngestJob(DataSourceIngestJob job) {
        long jobId = job.getId();
        IngestTasksScheduler.removeTasksForJob(this.rootFileTaskQueue, jobId);
        IngestTasksScheduler.removeTasksForJob(this.pendingFileTaskQueue, jobId);
    }

    private static List<AbstractFile> getTopLevelFiles(Content dataSource) {
        List<AbstractFile> topLevelFiles = new ArrayList<>();
        Collection<AbstractFile> rootObjects = dataSource.accept(new GetRootDirectoryVisitor());
        if (rootObjects.isEmpty() && dataSource instanceof AbstractFile) {
            // The data source is itself a file to be processed.
            topLevelFiles.add((AbstractFile) dataSource);
        } else {
            for (AbstractFile root : rootObjects) {
                List<Content> children;
                try {
                    children = root.getChildren();
                    if (children.isEmpty()) {
                        // Add the root object itself, it could be an unallocated
                        // space file, or a child of a volume or an image.
                        topLevelFiles.add(root);
                    } else {
                        // The root object is a file system root directory, get
                        // the files within it.
                        for (Content child : children) {
                            if (child instanceof AbstractFile) {
                                topLevelFiles.add((AbstractFile) child);
                            }
                        }
                    }
                } catch (TskCoreException ex) {
                    logger.log(Level.WARNING, "Could not get children of root to enqueue: " + root.getId() + ": " + root.getName(), ex); //NON-NLS
                }
            }
        }
        return topLevelFiles;
    }

    synchronized private void shuffleFileTaskQueues() {
        while (this.fileIngestThreadsQueue.isEmpty()) {
            /*
             * If the pending file task queue is empty, move the highest
             * priority root file task, if there is one, into it.
             */
            if (this.pendingFileTaskQueue.isEmpty()) {
                final FileIngestTask rootTask = this.rootFileTaskQueue.pollFirst();
                if (rootTask != null) {
                    this.pendingFileTaskQueue.addLast(rootTask);
                }
            }

            /*
             * Try to move the next task from the pending task queue into the
             * queue for the file ingest threads, if it passes the filter for
             * the job.
             */
            final FileIngestTask pendingTask = this.pendingFileTaskQueue.pollFirst();
            if (pendingTask == null) {
                return;
            }
            if (shouldEnqueueFileTask(pendingTask)) {
                try {
                    /*
                     * The task is added to the queue for the ingest threads
                     * AFTER the higher priority tasks that preceded it.
                     */
                    this.fileIngestThreadsQueue.putLast(pendingTask);
                } catch (InterruptedException ex) {
                    IngestTasksScheduler.logger.log(Level.INFO, "Ingest tasks scheduler interrupted while blocked adding a task to the file level ingest task queue", ex);
                    Thread.currentThread().interrupt();
                    return;
                }
            }

            /*
             * If the task that was just queued for the file ingest threads has
             * children, try to queue tasks for the children. Each child task
             * will go into either the directory queue if it has children of its
             * own, or into the queue for the file ingest threads, if it passes
             * the filter for the job.
             */
            final AbstractFile file = pendingTask.getFile();
            try {
                for (Content child : file.getChildren()) {
                    if (child instanceof AbstractFile) {
                        AbstractFile childFile = (AbstractFile) child;
                        FileIngestTask childTask = new FileIngestTask(pendingTask.getIngestJob(), childFile);
                        if (childFile.hasChildren()) {
                            this.pendingFileTaskQueue.add(childTask);
                        } else if (shouldEnqueueFileTask(childTask)) {
                            try {
                                this.fileIngestThreadsQueue.putLast(childTask);
                            } catch (InterruptedException ex) {
                                IngestTasksScheduler.logger.log(Level.INFO, "Ingest tasks scheduler interrupted while blocked adding a task to the file level ingest task queue", ex);
                                Thread.currentThread().interrupt();
                                return;
                            }
                        }
                    }
                }
            } catch (TskCoreException ex) {
                logger.log(Level.SEVERE, String.format("Error getting the children of %s (objId=%d)", file.getName(), file.getId()), ex);  //NON-NLS
            }
        }
    }

    private static boolean shouldEnqueueFileTask(final FileIngestTask task) {
        final AbstractFile file = task.getFile();

        // Skip the task if the file is actually the pseudo-file for the parent
        // or current directory.
        String fileName = file.getName();

        if (fileName.equals(".") || fileName.equals("..")) {
            return false;
        }

        /*
         * Ensures that all directories, files which are members of the ingest
         * file filter, and unallocated blocks (when processUnallocated is
         * enabled) all continue to be processed. AbstractFiles which do not
         * meet one of these criteria will be skipped.
         *
         * An additional check to see if unallocated space should be processed
         * is part of the FilesSet.fileIsMemberOf() method.
         *
         * This code may need to be updated when
         * TSK_DB_FILES_TYPE_ENUM.UNUSED_BLOCKS comes into use by Autopsy.
         */
        if (!file.isDir() && !shouldBeCarved(task) && !fileAcceptedByFilter(task)) {
            return false;
        }

        // Skip the task if the file is one of a select group of special, large
        // NTFS or FAT file system files.
        if (file instanceof org.sleuthkit.datamodel.File) {
            final org.sleuthkit.datamodel.File f = (org.sleuthkit.datamodel.File) file;

            // Get the type of the file system, if any, that owns the file.
            TskData.TSK_FS_TYPE_ENUM fsType = TskData.TSK_FS_TYPE_ENUM.TSK_FS_TYPE_UNSUPP;
            try {
                FileSystem fs = f.getFileSystem();
                if (fs != null) {
                    fsType = fs.getFsType();
                }
            } catch (TskCoreException ex) {
                logger.log(Level.SEVERE, "Error querying file system for " + f, ex); //NON-NLS
            }

            // If the file system is not NTFS or FAT, don't skip the file.
            if ((fsType.getValue() & FAT_NTFS_FLAGS) == 0) {
                return true;
            }

            // Find out whether the file is in a root directory. 
            boolean isInRootDir = false;
            try {
                AbstractFile parent = f.getParentDirectory();
                isInRootDir = parent.isRoot();
            } catch (TskCoreException ex) {
                logger.log(Level.WARNING, "Error querying parent directory for" + f.getName(), ex); //NON-NLS
            }

            // If the file is in the root directory of an NTFS or FAT file 
            // system, check its meta-address and check its name for the '$'
            // character and a ':' character (not a default attribute).
            if (isInRootDir && f.getMetaAddr() < 32) {
                String name = f.getName();
                if (name.length() > 0 && name.charAt(0) == '$' && name.contains(":")) {
                    return false;
                }
            }
        }

        return true;
    }

    private static boolean shouldBeCarved(final FileIngestTask task) {
        return task.getIngestJob().shouldProcessUnallocatedSpace() && task.getFile().getType().equals(TskData.TSK_DB_FILES_TYPE_ENUM.UNALLOC_BLOCKS);
    }

    private static boolean fileAcceptedByFilter(final FileIngestTask task) {
        return !(task.getIngestJob().getFileIngestFilter().fileIsMemberOf(task.getFile()) == null);
    }

    synchronized private static boolean hasTasksForJob(Collection<? extends IngestTask> tasks, long jobId) {
        for (IngestTask task : tasks) {
            if (task.getIngestJob().getId() == jobId) {
                return true;
            }
        }
        return false;
    }

    private static void removeTasksForJob(Collection<? extends IngestTask> tasks, long jobId) {
        Iterator<? extends IngestTask> iterator = tasks.iterator();
        while (iterator.hasNext()) {
            IngestTask task = iterator.next();
            if (task.getIngestJob().getId() == jobId) {
                iterator.remove();
            }
        }
    }

    private static int countTasksForJob(Collection<? extends IngestTask> queue, long jobId) {
        int count = 0;
        for (IngestTask task : queue) {
            if (task.getIngestJob().getId() == jobId) {
                count++;
            }
        }
        return count;
    }

    synchronized IngestJobTasksSnapshot getTasksSnapshotForJob(long jobId) {
        return new IngestJobTasksSnapshot(jobId);
    }

    private static class RootDirectoryTaskComparator implements Comparator<FileIngestTask> {

        @Override
        public int compare(FileIngestTask q1, FileIngestTask q2) {
            AbstractFilePriority.Priority p1 = AbstractFilePriority.getPriority(q1.getFile());
            AbstractFilePriority.Priority p2 = AbstractFilePriority.getPriority(q2.getFile());
            if (p1 == p2) {
                return (int) (q2.getFile().getId() - q1.getFile().getId());
            } else {
                return p2.ordinal() - p1.ordinal();
            }
        }

        private static class AbstractFilePriority {

            private AbstractFilePriority() {
            }

            enum Priority {

                LAST, LOW, MEDIUM, HIGH
            }

            static final List<Pattern> LAST_PRI_PATHS = new ArrayList<>();

            static final List<Pattern> LOW_PRI_PATHS = new ArrayList<>();

            static final List<Pattern> MEDIUM_PRI_PATHS = new ArrayList<>();

            static final List<Pattern> HIGH_PRI_PATHS = new ArrayList<>();

            /*
             * prioritize root directory folders based on the assumption that we
             * are looking for user content. Other types of investigations may
             * want different priorities.
             */
            static /*
             * prioritize root directory folders based on the assumption that we
             * are looking for user content. Other types of investigations may
             * want different priorities.
             */ {
                // these files have no structure, so they go last
                //unalloc files are handled as virtual files in getPriority()
                //LAST_PRI_PATHS.schedule(Pattern.compile("^\\$Unalloc", Pattern.CASE_INSENSITIVE));
                //LAST_PRI_PATHS.schedule(Pattern.compile("^\\Unalloc", Pattern.CASE_INSENSITIVE));
                LAST_PRI_PATHS.add(Pattern.compile("^pagefile", Pattern.CASE_INSENSITIVE));
                LAST_PRI_PATHS.add(Pattern.compile("^hiberfil", Pattern.CASE_INSENSITIVE));
                // orphan files are often corrupt and windows does not typically have
                // user content, so put them towards the bottom
                LOW_PRI_PATHS.add(Pattern.compile("^\\$OrphanFiles", Pattern.CASE_INSENSITIVE));
                LOW_PRI_PATHS.add(Pattern.compile("^Windows", Pattern.CASE_INSENSITIVE));
                // all other files go into the medium category too
                MEDIUM_PRI_PATHS.add(Pattern.compile("^Program Files", Pattern.CASE_INSENSITIVE));
                // user content is top priority
                HIGH_PRI_PATHS.add(Pattern.compile("^Users", Pattern.CASE_INSENSITIVE));
                HIGH_PRI_PATHS.add(Pattern.compile("^Documents and Settings", Pattern.CASE_INSENSITIVE));
                HIGH_PRI_PATHS.add(Pattern.compile("^home", Pattern.CASE_INSENSITIVE));
                HIGH_PRI_PATHS.add(Pattern.compile("^ProgramData", Pattern.CASE_INSENSITIVE));
            }

            static AbstractFilePriority.Priority getPriority(final AbstractFile abstractFile) {
                if (!abstractFile.getType().equals(TskData.TSK_DB_FILES_TYPE_ENUM.FS)) {
                    //quickly filter out unstructured content
                    //non-fs virtual files and dirs, such as representing unalloc space
                    return AbstractFilePriority.Priority.LAST;
                }
                //determine the fs files priority by name
                final String path = abstractFile.getName();
                if (path == null) {
                    return AbstractFilePriority.Priority.MEDIUM;
                }
                for (Pattern p : HIGH_PRI_PATHS) {
                    Matcher m = p.matcher(path);
                    if (m.find()) {
                        return AbstractFilePriority.Priority.HIGH;
                    }
                }
                for (Pattern p : MEDIUM_PRI_PATHS) {
                    Matcher m = p.matcher(path);
                    if (m.find()) {
                        return AbstractFilePriority.Priority.MEDIUM;
                    }
                }
                for (Pattern p : LOW_PRI_PATHS) {
                    Matcher m = p.matcher(path);
                    if (m.find()) {
                        return AbstractFilePriority.Priority.LOW;
                    }
                }
                for (Pattern p : LAST_PRI_PATHS) {
                    Matcher m = p.matcher(path);
                    if (m.find()) {
                        return AbstractFilePriority.Priority.LAST;
                    }
                }
                //default is medium
                return AbstractFilePriority.Priority.MEDIUM;
            }
        }
    }

    @ThreadSafe
    private class IngestTaskTrackingQueue implements BlockingIngestTaskQueue {

        private final BlockingDeque<IngestTask> taskQueue = new LinkedBlockingDeque<>();
        @GuardedBy("this")
        private final List<IngestTask> queuedTasks = new LinkedList<>();
        @GuardedBy("this")
        private final List<IngestTask> tasksInProgress = new LinkedList<>();

        void putFirst(IngestTask task) throws InterruptedException {
            synchronized (this) {
                this.queuedTasks.add(task);
            }
            try {
                this.taskQueue.putFirst(task);
            } catch (InterruptedException ex) {
                synchronized (this) {
                    this.queuedTasks.remove(task);
                }
                throw ex;
            }
        }

        void putLast(IngestTask task) throws InterruptedException {
            synchronized (this) {
                this.queuedTasks.add(task);
            }
            try {
                this.taskQueue.putLast(task);
            } catch (InterruptedException ex) {
                synchronized (this) {
                    this.queuedTasks.remove(task);
                }
                throw ex;
            }
        }

        @Override
        public IngestTask getNextTask() throws InterruptedException {
            IngestTask task = taskQueue.takeFirst();
            synchronized (this) {
                this.queuedTasks.remove(task);
                this.tasksInProgress.add(task);
            }
            return task;
        }

        boolean isEmpty() {
            synchronized (this) {
                return this.queuedTasks.isEmpty();
            }
        }

        void taskCompleted(IngestTask task) {
            synchronized (this) {
                this.tasksInProgress.remove(task);
            }
        }

        boolean hasTasksForJob(long jobId) {
            synchronized (this) {
                return IngestTasksScheduler.hasTasksForJob(this.queuedTasks, jobId) || IngestTasksScheduler.hasTasksForJob(this.tasksInProgress, jobId);
            }
        }

        int countQueuedTasksForJob(long jobId) {
            synchronized (this) {
                return IngestTasksScheduler.countTasksForJob(this.queuedTasks, jobId);
            }
        }

        int countRunningTasksForJob(long jobId) {
            synchronized (this) {
                return IngestTasksScheduler.countTasksForJob(this.tasksInProgress, jobId);
            }
        }

    }

    class IngestJobTasksSnapshot {

        private final long jobId;
        private final long dsQueueSize;
        private final long rootQueueSize;
        private final long dirQueueSize;
        private final long fileQueueSize;
        private final long runningListSize;

        IngestJobTasksSnapshot(long jobId) {
            this.jobId = jobId;
            this.dsQueueSize = IngestTasksScheduler.this.dataSourceIngestThreadQueue.countQueuedTasksForJob(jobId);
            this.rootQueueSize = countTasksForJob(IngestTasksScheduler.this.rootFileTaskQueue, jobId);
            this.dirQueueSize = countTasksForJob(IngestTasksScheduler.this.pendingFileTaskQueue, jobId);
            this.fileQueueSize = IngestTasksScheduler.this.fileIngestThreadsQueue.countQueuedTasksForJob(jobId);;
            this.runningListSize = IngestTasksScheduler.this.dataSourceIngestThreadQueue.countRunningTasksForJob(jobId) + IngestTasksScheduler.this.fileIngestThreadsQueue.countRunningTasksForJob(jobId);
        }

        long getJobId() {
            return jobId;
        }

        long getRootQueueSize() {
            return rootQueueSize;
        }

        long getDirectoryTasksQueueSize() {
            return dirQueueSize;
        }

        long getFileQueueSize() {
            return fileQueueSize;
        }

        long getDsQueueSize() {
            return dsQueueSize;
        }

        long getRunningListSize() {
            return runningListSize;
        }

    }

}