zeek/src/ScriptCoverageManager.cc

#include "zeek/ScriptCoverageManager.h"

#include <sys/stat.h>
#include <algorithm>
#include <cstdio>
#include <cstring>
#include <fstream>
#include <sstream>
#include <utility>

#include "zeek/Desc.h"
#include "zeek/Reporter.h"
#include "zeek/Type.h"
#include "zeek/script_opt/ScriptOpt.h"

using namespace std;

namespace zeek::detail {

ScriptCoverageManager::ScriptCoverageManager() : ignoring(0), delim('\t') {}

void ScriptCoverageManager::AddStmt(Stmt* s) {
    if ( ignoring != 0 || analysis_options.gen_ZAM )
        return;

    stmts.emplace_back(NewRef{}, s);
}

void ScriptCoverageManager::AddFunction(IDPtr func_id, StmtPtr body) {
    if ( analysis_options.gen_ZAM )
        return;

    func_instances.emplace_back(func_id, body);
}

bool ScriptCoverageManager::ReadStats() {
    char* bf = getenv("ZEEK_PROFILER_FILE");

    if ( ! bf )
        return false;

    std::ifstream ifs;
    ifs.open(bf, std::ifstream::in);

    if ( ! ifs )
        return false;

    std::stringstream ss;
    ss << ifs.rdbuf();
    std::string file_contents = ss.str();
    ss.clear();

    std::vector<std::string> lines;
    util::tokenize_string(file_contents, "\n", &lines);
    string delimiter;
    delimiter = delim;

    for ( const auto& line : lines ) {
        if ( line.empty() )
            continue;

        std::vector<std::string> line_components;
        util::tokenize_string(line, delimiter, &line_components);

        if ( line_components.size() != 3 ) {
            fprintf(stderr, "invalid ZEEK_PROFILER_FILE line: %s\n", line.data());
            continue;
        }

        std::string& cnt = line_components[0];
        std::string& location = line_components[1];
        std::string& desc = line_components[2];

        pair<string, string> location_desc(std::move(location), std::move(desc));
        uint64_t count;
        util::atoi_n(cnt.size(), cnt.c_str(), nullptr, 10, count);
        usage_map.emplace(std::move(location_desc), count);
    }

    return true;
}

bool ScriptCoverageManager::WriteStats() {
    char* bf = getenv("ZEEK_PROFILER_FILE");

    if ( ! bf )
        return false;

    util::SafeDirname dirname{bf};

    if ( ! util::detail::ensure_intermediate_dirs(dirname.result.data()) ) {
        reporter->Error("Failed to open ZEEK_PROFILER_FILE destination '%s' for writing", bf);
        return false;
    }

    FILE* f;
    const char* p = strstr(bf, "XXXXXX");

    if ( p && ! p[6] ) {
        mode_t old_umask = umask(S_IXUSR | S_IRWXO | S_IRWXG);
        int fd = mkstemp(bf);
        umask(old_umask);

        if ( fd == -1 ) {
            reporter->Error("Failed to generate unique file name from ZEEK_PROFILER_FILE: %s", bf);
            return false;
        }
        f = fdopen(fd, "w");
    }
    else {
        f = fopen(bf, "w");
    }

    if ( ! f ) {
        reporter->Error("Failed to open ZEEK_PROFILER_FILE destination '%s' for writing", bf);
        return false;
    }

    for ( auto& s : stmts ) {
        ODesc desc_info;
        s->Describe(&desc_info);
        TrackUsage(s, desc_info.Description(), s->GetAccessCount());
    }

    for ( auto& [func, body] : func_instances ) {
        auto ft = func->GetType<FuncType>();
        auto desc = ft->FlavorString() + " " + func->Name() + " BODY";

        TrackUsage(body, desc, body->GetAccessCount());
    }

    for ( auto& [location_info, cnt] : usage_map )
        Report(f, cnt, location_info.first, location_info.second);

    fclose(f);
    return true;
}

void ScriptCoverageManager::TrackUsage(const ObjPtr& obj, std::string desc, uint64_t cnt) {
    ODesc location_info;
    obj->GetLocationInfo()->Describe(&location_info);

    static canonicalize_desc cd{delim};
    for_each(desc.begin(), desc.end(), cd);

    pair<string, string> location_desc(location_info.Description(), desc);

    if ( usage_map.find(location_desc) != usage_map.end() )
        usage_map[location_desc] += cnt;
    else
        usage_map[location_desc] = cnt;
}

void ScriptCoverageManager::Report(FILE* f, uint64_t cnt, std::string loc, std::string desc) {
    fprintf(f, "%" PRIu64 "%c%s%c%s\n", cnt, delim, loc.c_str(), delim, desc.c_str());
}

} // namespace zeek::detail