/*
 * goostring-format-checker.cc
 *
 * This file is licensed under the GPLv2 or later
 *
 * Clang++ compiler plugin that checks usage of GooString::format-like functions
 *
 * Copyright (C) 2014 Fabio D'Urso <fabiodurso@hotmail.it>
 */

#include <cctype>

#include <clang/Frontend/FrontendPluginRegistry.h>
#include <clang/AST/AST.h>
#include <clang/AST/ASTConsumer.h>
#include <clang/AST/Attr.h>
#include <clang/AST/RecursiveASTVisitor.h>
#include <clang/Frontend/CompilerInstance.h>

using namespace clang;

namespace {

class GooStringFormatCheckerVisitor : public RecursiveASTVisitor<GooStringFormatCheckerVisitor>
{
public:
    explicit GooStringFormatCheckerVisitor(CompilerInstance *compInst);

    bool VisitFunctionDecl(FunctionDecl *funcDecl);
    bool VisitCallExpr(CallExpr *callExpr);

private:
    /* Returns the index of the format argument, or -1 if the function must
     * not be checked */
    int findFormatArgumentIndex(const FunctionDecl *funcDecl) const;

    /* Returns the SourceLocation of the n-th character */
    SourceLocation getLocationOfCharacter(const StringLiteral *strLiteral, unsigned n);

    /* Validates usage of a placeholder and returns the corresponding
     * argument index, or -1 in case of errors */
    int verifyPlaceholder(const CallExpr *callExpr, const SourceLocation &placeholderLocation, std::string &placeholderText, int baseArgIdx) const;

    CompilerInstance *compInst;
    DiagnosticsEngine *diag;
    unsigned diag_badFuncZeroArgs;
    unsigned diag_badFuncNonVariadic;
    unsigned diag_badFuncLastArgInvalidType;
    unsigned diag_notStringLiteral;
    unsigned diag_notPlainASCII;
    unsigned diag_wrongOrder;
    unsigned diag_unescapedBracket;
    unsigned diag_unterminatedPlaceholder;
    unsigned diag_unconsumedArgs;
    unsigned diag_missingColon;
    unsigned diag_missingArgNumber;
    unsigned diag_badArgNumber;
    unsigned diag_argumentNotPresent;
    unsigned diag_badPrecision;
    unsigned diag_badType;
    unsigned diag_wrongArgExprType;
};

GooStringFormatCheckerVisitor::GooStringFormatCheckerVisitor(CompilerInstance *compInst) : compInst(compInst)
{
    diag = &compInst->getDiagnostics();

    diag_badFuncZeroArgs = diag->getCustomDiagID(DiagnosticsEngine::Error, "Cannot enforce format string checks on a function that takes no arguments");
    diag_badFuncNonVariadic = diag->getCustomDiagID(DiagnosticsEngine::Error, "Cannot enforce format string checks on a non-variadic function");
    diag_badFuncLastArgInvalidType = diag->getCustomDiagID(DiagnosticsEngine::Error, "Cannot enforce format string checks if the last non-variadic argument is not const char *");
    diag_notStringLiteral = diag->getCustomDiagID(DiagnosticsEngine::Warning, "Format string is not a string literal. Skipping format checks");
    diag_notPlainASCII = diag->getCustomDiagID(DiagnosticsEngine::Warning, "Format string contains non-ASCII or NUL characters. Skipping format checks");
    diag_wrongOrder = diag->getCustomDiagID(DiagnosticsEngine::Error, "Argument %0 must be consumed before argument %1");
    diag_unescapedBracket = diag->getCustomDiagID(DiagnosticsEngine::Error, "Unescaped '}' character");
    diag_unterminatedPlaceholder = diag->getCustomDiagID(DiagnosticsEngine::Error, "Unterminated placeholder");
    diag_unconsumedArgs = diag->getCustomDiagID(DiagnosticsEngine::Warning, "Unconsumed argument(s)");
    diag_missingColon = diag->getCustomDiagID(DiagnosticsEngine::Error, "Invalid placeholder '{%0}': missing colon character");
    diag_missingArgNumber = diag->getCustomDiagID(DiagnosticsEngine::Error, "Invalid placeholder '{%0}': missing <arg> number");
    diag_badArgNumber = diag->getCustomDiagID(DiagnosticsEngine::Error, "Invalid placeholder '{%0}': bad <arg> number");
    diag_argumentNotPresent = diag->getCustomDiagID(DiagnosticsEngine::Error, "Argument for placeholder '{%0}' is not present");
    diag_badPrecision = diag->getCustomDiagID(DiagnosticsEngine::Error, "Invalid placeholder '{%0}': bad <precision> value");
    diag_badType = diag->getCustomDiagID(DiagnosticsEngine::Error, "Invalid placeholder '{%0}': bad <type> specifier");
    diag_wrongArgExprType = diag->getCustomDiagID(DiagnosticsEngine::Error, "Expected %0 for placeholder '{%1}', found %2");
}

bool GooStringFormatCheckerVisitor::VisitFunctionDecl(FunctionDecl *funcDecl)
{
    findFormatArgumentIndex(funcDecl); // Spot misuse of the "gooformat" annotation
    return true;
}

bool GooStringFormatCheckerVisitor::VisitCallExpr(CallExpr *callExpr)
{
    /*** Locate format argument or skip calls that needn't be checked ***/

    const int formatArgIdx = findFormatArgumentIndex(callExpr->getDirectCallee());
    if (formatArgIdx == -1)
        return true;

    /*** Obtain format string value ***/

    const Expr *formatArgExpr = callExpr->getArg(formatArgIdx);
    while (formatArgExpr->getStmtClass() == Stmt::ImplicitCastExprClass) {
        formatArgExpr = static_cast<const ImplicitCastExpr *>(formatArgExpr)->getSubExpr();
    }
    if (formatArgExpr->getStmtClass() != Stmt::StringLiteralClass) {
        diag->Report(formatArgExpr->getExprLoc(), diag_notStringLiteral);
        return true;
    }
    const StringLiteral *formatArgStrLiteral = static_cast<const StringLiteral *>(formatArgExpr);
    if (formatArgStrLiteral->containsNonAsciiOrNull()) {
        diag->Report(formatArgExpr->getExprLoc(), diag_notPlainASCII);
        return true;
    }

    /*** Parse format string and verify arguments ***/

    const std::string format = formatArgStrLiteral->getString().str();

    /* Keeps track of whether we are currently parsing a character contained
     * within '{' ... '}'. If set, current_placeholder contains the contents
     * parsed so far (without brackets) */
    bool in_placeholder = false;
    std::string current_placeholder;

    // Source location of the current placeholder's opening bracket
    SourceLocation placeholderLoc;

    /* Keeps track of the next expected argument number, to check that
     * arguments are first consumed in order (eg {0:d}{2:d}{1:d} is wrong).
     * Note that it's possible to "look back" at already consumed
     * arguments (eg {0:d}{1:d}{0:d} is OK) */
    int nextExpectedArgNum = 0;

    for (unsigned i = 0; i < format.length(); i++) {
        if (in_placeholder) {
            // Have we reached the end of the placeholder?
            if (format[i] == '}') {
                in_placeholder = false;

                // Verifies the placeholder and returns the argument number
                const int foundArgNum = verifyPlaceholder(callExpr, placeholderLoc, current_placeholder, formatArgIdx + 1);

                // If the placeholder wasn't valid, disable argument order checks
                if (foundArgNum == -1) {
                    nextExpectedArgNum = -1;
                }

                // If argument order checks are enabled, let's check!
                if (nextExpectedArgNum != -1) {
                    if (foundArgNum == nextExpectedArgNum) {
                        nextExpectedArgNum++;
                    } else if (foundArgNum > nextExpectedArgNum) {
                        diag->Report(placeholderLoc, diag_wrongOrder) << nextExpectedArgNum << foundArgNum;
                        nextExpectedArgNum = -1; // disable further checks
                    }
                }
            } else {
                current_placeholder += format[i];
            }
        } else if (format[i] == '{') {
            // If we find a '{' then a placeholder is starting...
            in_placeholder = true;
            current_placeholder = "";
            placeholderLoc = getLocationOfCharacter(formatArgStrLiteral, i);

            // ...unless it's followed by another '{' (escape sequence)
            if (i + 1 < format.length() && format[i + 1] == '{') {
                i++; // skip next '{' character
                in_placeholder = false;
            }
        } else if (format[i] == '}') {
            /* If we have found a '}' and we're not in a placeholder,
             * then it *MUST* be followed by another '}' (escape sequence) */
            if (i + 1 >= format.length() || format[i + 1] != '}') {
                diag->Report(getLocationOfCharacter(formatArgStrLiteral, i), diag_unescapedBracket);
            } else {
                i++; // skip next '}' character
            }
        }
    }

    /* If we've reached the end of the format string and in_placeholder is
     * still set, then the last placeholder wasn't terminated properly */
    if (in_placeholder)
        diag->Report(placeholderLoc, diag_unterminatedPlaceholder);

    int unconsumedArgs = callExpr->getNumArgs() - (formatArgIdx + 1 + nextExpectedArgNum);
    if (unconsumedArgs > 0)
        diag->Report(callExpr->getArg(callExpr->getNumArgs() - unconsumedArgs)->getExprLoc(), diag_unconsumedArgs);

    return true;
}

int GooStringFormatCheckerVisitor::findFormatArgumentIndex(const FunctionDecl *funcDecl) const
{
    if (!funcDecl)
        return -1;

    AnnotateAttr *annotation = NULL;
    for (specific_attr_iterator<AnnotateAttr> it = funcDecl->specific_attr_begin<AnnotateAttr>(); it != funcDecl->specific_attr_end<AnnotateAttr>() && !annotation; ++it) {
        if (it->getAnnotation() == "gooformat")
            annotation = *it;
    }

    // If this function hasn't got the "gooformat" annotation on it
    if (!annotation)
        return -1;

    if (funcDecl->getNumParams() == 0) {
        diag->Report(annotation->getLocation(), diag_badFuncZeroArgs);
        return -1;
    }

    if (!funcDecl->isVariadic()) {
        diag->Report(annotation->getLocation(), diag_badFuncNonVariadic);
        return -1;
    }

    // Assume the last non-variadic argument is the format specifier
    const int formatArgIdx = funcDecl->getNumParams() - 1;
    const QualType formatArgType = funcDecl->getParamDecl(formatArgIdx)->getType();
    if (formatArgType.getAsString() != "const char *") {
        diag->Report(annotation->getLocation(), diag_badFuncLastArgInvalidType);
        return -1;
    }

    return formatArgIdx;
}

SourceLocation GooStringFormatCheckerVisitor::getLocationOfCharacter(const StringLiteral *strLiteral, unsigned n)
{
    return strLiteral->getLocationOfByte(n, compInst->getSourceManager(), compInst->getLangOpts(), compInst->getTarget());
}

int GooStringFormatCheckerVisitor::verifyPlaceholder(const CallExpr *callExpr, const SourceLocation &placeholderLocation, std::string &placeholderText, int baseArgIdx) const
{
    // Find the colon that separates the argument number and the format specifier
    const size_t delim = placeholderText.find(':');
    if (delim == std::string::npos) {
        diag->Report(placeholderLocation, diag_missingColon) << placeholderText;
        return -1;
    }
    if (delim == 0) {
        diag->Report(placeholderLocation, diag_missingArgNumber) << placeholderText;
        return -1;
    }
    for (unsigned int i = 0; i < delim; i++) {
        if (!isdigit(placeholderText[i])) {
            diag->Report(placeholderLocation, diag_badArgNumber) << placeholderText;
            return -1;
        }
    }

    // Extract argument number and its actual position in the call's argument list
    const int argNum = atoi(placeholderText.substr(0, delim).c_str());
    const int argIdx = baseArgIdx + argNum;
    if (argIdx >= callExpr->getNumArgs()) {
        diag->Report(placeholderLocation, diag_argumentNotPresent) << placeholderText;
        return argNum;
    }

    // Check and strip width/precision specifiers
    std::string format = placeholderText.substr(delim + 1);
    bool dot_found = false;
    while (isdigit(format[0]) || format[0] == '.') {
        if (format[0] == '.') {
            if (dot_found) {
                diag->Report(placeholderLocation, diag_badPrecision) << placeholderText;
                return argNum;
            }
            dot_found = true;
        }
        format = format.substr(1);
    }

    const Expr *argExpr = callExpr->getArg(argIdx);
    const QualType qualType = argExpr->getType();
    const Type *valueType = qualType->getUnqualifiedDesugaredType();

    if (format == "d" || format == "x" || format == "X" || format == "o" || format == "b" || format == "w") {
        if (!valueType->isSpecificBuiltinType(BuiltinType::Int)) {
            diag->Report(argExpr->getExprLoc(), diag_wrongArgExprType) << "int" << placeholderText << qualType.getAsString();
        }
    } else if (format == "ud" || format == "ux" || format == "uX" || format == "uo" || format == "ub") {
        if (!valueType->isSpecificBuiltinType(BuiltinType::UInt)) {
            diag->Report(argExpr->getExprLoc(), diag_wrongArgExprType) << "unsigned int" << placeholderText << qualType.getAsString();
        }
    } else if (format == "ld" || format == "lx" || format == "lX" || format == "lo" || format == "lb") {
        if (!valueType->isSpecificBuiltinType(BuiltinType::Long)) {
            diag->Report(argExpr->getExprLoc(), diag_wrongArgExprType) << "long" << placeholderText << qualType.getAsString();
        }
    } else if (format == "uld" || format == "ulx" || format == "ulX" || format == "ulo" || format == "ulb") {
        if (!valueType->isSpecificBuiltinType(BuiltinType::ULong)) {
            diag->Report(argExpr->getExprLoc(), diag_wrongArgExprType) << "unsigned long" << placeholderText << qualType.getAsString();
        }
    } else if (format == "lld" || format == "llx" || format == "llX" || format == "llo" || format == "llb") {
        if (!valueType->isSpecificBuiltinType(BuiltinType::LongLong)) {
            diag->Report(argExpr->getExprLoc(), diag_wrongArgExprType) << "long long" << placeholderText << qualType.getAsString();
        }
    } else if (format == "ulld" || format == "ullx" || format == "ullX" || format == "ullo" || format == "ullb") {
        if (!valueType->isSpecificBuiltinType(BuiltinType::ULongLong)) {
            diag->Report(argExpr->getExprLoc(), diag_wrongArgExprType) << "unsigned long long" << placeholderText << qualType.getAsString();
        }
    } else if (format == "f" || format == "g" || format == "gs") {
        if (!valueType->isSpecificBuiltinType(BuiltinType::Double)) {
            diag->Report(argExpr->getExprLoc(), diag_wrongArgExprType) << "float or double" << placeholderText << qualType.getAsString();
        }
    } else if (format == "c") {
        if (!valueType->isSpecificBuiltinType(BuiltinType::UInt) && !valueType->isSpecificBuiltinType(BuiltinType::Int)) {
            diag->Report(argExpr->getExprLoc(), diag_wrongArgExprType) << "char, short or int" << placeholderText << qualType.getAsString();
        }
    } else if (format == "s") {
        if (!valueType->isPointerType() || !valueType->getPointeeType()->getUnqualifiedDesugaredType()->isCharType()) {
            diag->Report(argExpr->getExprLoc(), diag_wrongArgExprType) << "char *" << placeholderText << qualType.getAsString();
        }
    } else if (format == "t") {
        const CXXRecordDecl *pointeeType = valueType->isPointerType() ? valueType->getPointeeType()->getAsCXXRecordDecl() : 0;
        if (pointeeType == 0 || pointeeType->getQualifiedNameAsString() != "GooString") {
            diag->Report(argExpr->getExprLoc(), diag_wrongArgExprType) << "GooString *" << placeholderText << qualType.getAsString();
        }
    } else {
        diag->Report(placeholderLocation, diag_badType) << placeholderText;
        return argNum;
    }

    return argNum;
}

class GooStringFormatCheckerConsumer : public clang::ASTConsumer
{
public:
    GooStringFormatCheckerConsumer(CompilerInstance *compInst) : visitor(compInst) { }

    virtual void HandleTranslationUnit(clang::ASTContext &ctx) { visitor.TraverseDecl(ctx.getTranslationUnitDecl()); }

private:
    GooStringFormatCheckerVisitor visitor;
};

class GooStringFormatCheckerAction : public PluginASTAction
{
protected:
    ASTConsumer *CreateASTConsumer(CompilerInstance &compInst, llvm::StringRef inFile) { return new GooStringFormatCheckerConsumer(&compInst); }

    bool ParseArgs(const CompilerInstance &compInst, const std::vector<std::string> &args)
    {
        if (args.size() != 0) {
            DiagnosticsEngine &D = compInst.getDiagnostics();
            D.Report(D.getCustomDiagID(DiagnosticsEngine::Error, "goostring-format-checker takes no arguments"));
            return false;
        } else {
            return true;
        }
    }
};

}

static FrontendPluginRegistry::Add<GooStringFormatCheckerAction> X("goostring-format-checker", "Checks usage of GooString::format-like functions");