nmWTAI-Platform/Src/nmNum/nmData/nmNumericalResultPersistence.cpp

#include "nmNumericalResultPersistence.h"

#include <QCryptographicHash>
#include <QDataStream>
#include <QDir>
#include <QFile>
#include <QFileInfo>
#include <QtGlobal>

#include <float.h>
#include <new>

#include <vtkNew.h>
#include <vtkUnstructuredGrid.h>
#include <vtkXMLUnstructuredGridReader.h>
#include <vtkXMLUnstructuredGridWriter.h>

namespace {

const char g_aWellHistoryMagic[8] = { 'N', 'M', 'W', 'H', 'I', 'S', '3', 0 };
const quint32 g_nWellHistoryFormatVersion = 1;
const quint32 g_nEndianMarker = 0x01020304u;
const quint32 g_nMaximumCurveSeries = 64u;
const quint64 g_nMaximumCurvePoints = 50000000ull;
const qint64 g_nStreamChunkBytes = 4 * 1024 * 1024;

bool setError(QString* pError, const QString& sError)
{
    if(pError != NULL)
    {
        *pError = sError;
    }
    return false;
}

bool writeAll(QFile& oFile, const char* pData, quint64 nBytes)
{
    quint64 nWritten = 0;
    while(nWritten < nBytes)
    {
        const qint64 nChunk = static_cast<qint64>(qMin(
                static_cast<quint64>(g_nStreamChunkBytes),
                nBytes - nWritten));
        const qint64 nResult = oFile.write(pData + nWritten, nChunk);
        if(nResult != nChunk)
        {
            return false;
        }
        nWritten += static_cast<quint64>(nResult);
    }
    return true;
}

bool readAll(QFile& oFile, char* pData, quint64 nBytes)
{
    quint64 nRead = 0;
    while(nRead < nBytes)
    {
        const qint64 nChunk = static_cast<qint64>(qMin(
                static_cast<quint64>(g_nStreamChunkBytes),
                nBytes - nRead));
        const qint64 nResult = oFile.read(pData + nRead, nChunk);
        if(nResult != nChunk)
        {
            return false;
        }
        nRead += static_cast<quint64>(nResult);
    }
    return true;
}

template<typename T>
bool writePod(QFile& oFile, const T& oValue)
{
    return writeAll(oFile, reinterpret_cast<const char*>(&oValue),
                    static_cast<quint64>(sizeof(T)));
}

template<typename T>
bool readPod(QFile& oFile, T& oValue)
{
    return readAll(oFile, reinterpret_cast<char*>(&oValue),
                   static_cast<quint64>(sizeof(T)));
}

bool safeMultiply(quint64 nLeft, quint64 nRight, quint64& nProduct)
{
    if(nLeft != 0 && nRight > (~static_cast<quint64>(0)) / nLeft)
    {
        return false;
    }
    nProduct = nLeft * nRight;
    return true;
}

bool writeCurve(QFile& oFile, const QVector<QVector<double> >& vecCurve)
{
    const quint32 nSeriesCount = static_cast<quint32>(vecCurve.size());
    if(!writePod(oFile, nSeriesCount))
    {
        return false;
    }
    for(int nSeries = 0; nSeries < vecCurve.size(); ++nSeries)
    {
        const QVector<double>& vecValues = vecCurve[nSeries];
        const quint64 nPointCount = static_cast<quint64>(vecValues.size());
        quint64 nBytes = 0;
        if(!safeMultiply(nPointCount, sizeof(double), nBytes) ||
           !writePod(oFile, nPointCount) ||
           (nBytes > 0 && !writeAll(oFile,
                    reinterpret_cast<const char*>(vecValues.constData()),
                    nBytes)))
        {
            return false;
        }
    }
    return true;
}

bool readCurve(QFile& oFile,
               QVector<QVector<double> >& vecCurve,
               quint64& nTotalPointCount)
{
    quint32 nSeriesCount = 0;
    if(!readPod(oFile, nSeriesCount) || nSeriesCount > g_nMaximumCurveSeries)
    {
        return false;
    }

    QVector<QVector<double> > vecLoaded;
    vecLoaded.resize(static_cast<int>(nSeriesCount));
    for(quint32 nSeries = 0; nSeries < nSeriesCount; ++nSeries)
    {
        quint64 nPointCount = 0;
        quint64 nBytes = 0;
        if(!readPod(oFile, nPointCount) ||
           nPointCount > g_nMaximumCurvePoints ||
           nTotalPointCount > g_nMaximumCurvePoints - nPointCount ||
           !safeMultiply(nPointCount, sizeof(double), nBytes) ||
           nBytes > static_cast<quint64>(oFile.bytesAvailable()) ||
           nPointCount > static_cast<quint64>(INT_MAX))
        {
            return false;
        }
        nTotalPointCount += nPointCount;
        vecLoaded[static_cast<int>(nSeries)].resize(
                static_cast<int>(nPointCount));
        if(nBytes > 0 && !readAll(oFile,
                reinterpret_cast<char*>(
                    vecLoaded[static_cast<int>(nSeries)].data()), nBytes))
        {
            return false;
        }
    }
    vecCurve.swap(vecLoaded);
    return true;
}

bool hasExactMagic(const char* pActual, const char* pExpected, int nSize)
{
    for(int nIndex = 0; nIndex < nSize; ++nIndex)
    {
        if(pActual[nIndex] != pExpected[nIndex])
        {
            return false;
        }
    }
    return true;
}

}

nmNumericalFileReference::nmNumericalFileReference()
    : m_nLength(0)
{
}

bool nmNumericalFileReference::isValid() const
{
    QString sNormalizedPath;
    if(!nmNumericalResultPersistence::normalizeRelativePath(
                m_sRelativePath, sNormalizedPath) ||
       sNormalizedPath != m_sRelativePath || m_sSha1.size() != 40)
    {
        return false;
    }
    for(int nIndex = 0; nIndex < m_sSha1.size(); ++nIndex)
    {
        const QChar oChar = m_sSha1[nIndex];
        if(!((oChar >= '0' && oChar <= '9') ||
             (oChar >= 'a' && oChar <= 'f')))
        {
            return false;
        }
    }
    return true;
}

nmNumericalWindowPayloadReferences::nmNumericalWindowPayloadReferences()
    : m_bHasCurrentGrid(false),
      m_bHasSnapshot(false)
{
}

int nmNumericalResultPersistence::projectVersion()
{
    return 3;
}

quint64 nmNumericalResultPersistence::maximumBinaryPayloadBytes()
{
    return 512ull * 1024ull * 1024ull;
}

bool nmNumericalResultPersistence::normalizeRelativePath(
        const QString& sPath,
        QString& sNormalizedPath)
{
    sNormalizedPath.clear();
    if(sPath.isEmpty() || QDir::isAbsolutePath(sPath) ||
       sPath.contains(':'))
    {
        return false;
    }

    QString sForwardPath = sPath;
    sForwardPath.replace('\\', '/');
    const QStringList listParts = sForwardPath.split(
            '/', QString::KeepEmptyParts);
    if(listParts.isEmpty())
    {
        return false;
    }
    for(int nIndex = 0; nIndex < listParts.size(); ++nIndex)
    {
        if(listParts[nIndex].isEmpty() || listParts[nIndex] == "." ||
           listParts[nIndex] == "..")
        {
            return false;
        }
    }

    sNormalizedPath = listParts.join("/");
    return !sNormalizedPath.isEmpty();
}

bool nmNumericalResultPersistence::resolveReferencedPath(
        const QString& sRootDirectory,
        const QString& sRelativePath,
        QString& sAbsolutePath)
{
    QString sBaseCanonical;
    QString sFileCanonical;
    sAbsolutePath.clear();
    QString sNormalizedPath;
    if(!normalizeRelativePath(sRelativePath, sNormalizedPath))
    {
        return false;
    }

    QString sRoot = QDir::cleanPath(QFileInfo(sRootDirectory).absoluteFilePath());
    QString sTarget = QDir::cleanPath(
            QDir(sRoot).absoluteFilePath(sNormalizedPath));
    sRoot.replace('\\', '/');
    sTarget.replace('\\', '/');
    const QString sRootPrefix = sRoot.endsWith('/') ? sRoot : sRoot + "/";
    if(!sTarget.startsWith(sRootPrefix, Qt::CaseInsensitive))
    {
        return false;
    }

    // 已存在文件再核对规范路径，防止目录联接或符号链接逃逸成果根目录。
    sBaseCanonical = QFileInfo(sRoot).canonicalFilePath();
    sFileCanonical = QFileInfo(sTarget).canonicalFilePath();
    if(!sBaseCanonical.isEmpty() && !sFileCanonical.isEmpty())
    {
        QString sCanonicalPrefix = QDir::fromNativeSeparators(sBaseCanonical);
        if(!sCanonicalPrefix.endsWith('/'))
        {
            sCanonicalPrefix += "/";
        }
        if(!QDir::fromNativeSeparators(sFileCanonical).startsWith(
                    sCanonicalPrefix, Qt::CaseInsensitive))
        {
            return false;
        }
    }

    sAbsolutePath = QDir::toNativeSeparators(sTarget);
    return true;
}

bool nmNumericalResultPersistence::buildFileReference(
        const QString& sAbsolutePath,
        const QString& sRelativePath,
        nmNumericalFileReference& oReference)
{
    QString sNormalizedPath;
    if(!normalizeRelativePath(sRelativePath, sNormalizedPath))
    {
        return false;
    }

    QFile oFile(sAbsolutePath);
    if(!oFile.open(QIODevice::ReadOnly))
    {
        return false;
    }
    QCryptographicHash oHash(QCryptographicHash::Sha1);
    while(!oFile.atEnd())
    {
        const QByteArray baChunk = oFile.read(g_nStreamChunkBytes);
        if(baChunk.isEmpty() && oFile.error() != QFile::NoError)
        {
            oFile.close();
            return false;
        }
        oHash.addData(baChunk);
    }
    const qint64 nSize = oFile.size();
    oFile.close();
    if(nSize < 0)
    {
        return false;
    }

    oReference.m_sRelativePath = sNormalizedPath;
    oReference.m_nLength = static_cast<quint64>(nSize);
    oReference.m_sSha1 = QString::fromLatin1(oHash.result().toHex());
    return oReference.isValid();
}

bool nmNumericalResultPersistence::validateFileReference(
        const QString& sRootDirectory,
        const nmNumericalFileReference& oReference,
        QString* pAbsolutePath,
        QString* pError)
{
    if(pError != NULL)
    {
        pError->clear();
    }
    if(!oReference.isValid())
    {
        return setError(pError, "Numerical file reference is invalid.");
    }

    QString sFilePath;
    if(!resolveReferencedPath(sRootDirectory,
                              oReference.m_sRelativePath, sFilePath))
    {
        return setError(pError, "Numerical file reference leaves its root directory.");
    }
    const QFileInfo oInfo(sFilePath);
    if(!oInfo.exists() || !oInfo.isFile() || oInfo.size() < 0 ||
       static_cast<quint64>(oInfo.size()) != oReference.m_nLength)
    {
        return setError(pError, "Numerical file length does not match its manifest.");
    }

    nmNumericalFileReference oActual;
    if(!buildFileReference(sFilePath, oReference.m_sRelativePath, oActual) ||
       oActual.m_nLength != oReference.m_nLength ||
       oActual.m_sSha1 != oReference.m_sSha1)
    {
        return setError(pError, "Numerical file SHA-1 does not match its manifest.");
    }
    if(pAbsolutePath != NULL)
    {
        *pAbsolutePath = sFilePath;
    }
    return true;
}

bool nmNumericalResultPersistence::writeWellHistory(
        const QString& sFilePath,
        const nmNumericalWellHistoryData& oHistory,
        QString* pError)
{
    if(pError != NULL)
    {
        pError->clear();
    }
    QFile oFile(sFilePath);
    if(!oFile.open(QIODevice::WriteOnly | QIODevice::Truncate))
    {
        return setError(pError, "Cannot create numerical well history file.");
    }

    const bool bWritten = writeAll(oFile, g_aWellHistoryMagic, 8) &&
            writePod(oFile, g_nWellHistoryFormatVersion) &&
            writePod(oFile, g_nEndianMarker) &&
            writeCurve(oFile, oHistory.m_vecPressure) &&
            writeCurve(oFile, oHistory.m_vecLogLog) &&
            writeCurve(oFile, oHistory.m_vecSemiLog) &&
            oFile.flush();
    oFile.close();
    if(!bWritten)
    {
        QFile::remove(sFilePath);
        return setError(pError, "Cannot write complete numerical well history file.");
    }
    return true;
}

bool nmNumericalResultPersistence::readWellHistory(
        const QString& sFilePath,
        nmNumericalWellHistoryData& oHistory,
        QString* pError)
{
    if(pError != NULL)
    {
        pError->clear();
    }
    const QFileInfo oInfo(sFilePath);
    if(!oInfo.exists() || oInfo.size() < 0 ||
       static_cast<quint64>(oInfo.size()) > maximumBinaryPayloadBytes())
    {
        return setError(pError, "Numerical well history payload is missing or too large.");
    }

    try
    {
        QFile oFile(sFilePath);
        if(!oFile.open(QIODevice::ReadOnly))
        {
            return setError(pError, "Cannot open numerical well history file.");
        }
        char aMagic[8] = { 0 };
        quint32 nFormatVersion = 0;
        quint32 nEndianMarker = 0;
        quint64 nTotalPointCount = 0;
        nmNumericalWellHistoryData oLoaded;
        const bool bRead = readAll(oFile, aMagic, 8) &&
                hasExactMagic(aMagic, g_aWellHistoryMagic, 8) &&
                readPod(oFile, nFormatVersion) &&
                nFormatVersion == g_nWellHistoryFormatVersion &&
                readPod(oFile, nEndianMarker) &&
                nEndianMarker == g_nEndianMarker &&
                readCurve(oFile, oLoaded.m_vecPressure, nTotalPointCount) &&
                readCurve(oFile, oLoaded.m_vecLogLog, nTotalPointCount) &&
                readCurve(oFile, oLoaded.m_vecSemiLog, nTotalPointCount) &&
                oFile.atEnd();
        oFile.close();
        if(!bRead)
        {
            return setError(pError, "Numerical well history file is damaged.");
        }
        oHistory = oLoaded;
        return true;
    }
    catch(const std::bad_alloc&)
    {
        return setError(pError, "Not enough memory to load numerical well history.");
    }
}

bool nmNumericalResultPersistence::writeGrid(
        vtkUnstructuredGrid* pGrid,
        const QString& sFilePath,
        QString* pError)
{
    if(pError != NULL)
    {
        pError->clear();
    }
    if(pGrid == NULL || pGrid->GetNumberOfPoints() <= 0 ||
       pGrid->GetNumberOfCells() <= 0)
    {
        return setError(pError, "Numerical grid is empty.");
    }

    vtkNew<vtkXMLUnstructuredGridWriter> pWriter;
    pWriter->SetInputData(pGrid);
    pWriter->SetFileName(sFilePath.toLocal8Bit().constData());
    pWriter->SetDataModeToBinary();
    if(pWriter->Write() == 0 || !QFileInfo(sFilePath).isFile() ||
       QFileInfo(sFilePath).size() <= 0)
    {
        QFile::remove(sFilePath);
        return setError(pError, "Cannot write numerical VTU grid.");
    }
    return true;
}

bool nmNumericalResultPersistence::readGrid(
        const QString& sFilePath,
        vtkSmartPointer<vtkUnstructuredGrid>& pGrid,
        QString* pError)
{
    if(pError != NULL)
    {
        pError->clear();
    }
    pGrid = NULL;
    if(!QFileInfo(sFilePath).isFile())
    {
        return setError(pError, "Numerical VTU grid is missing.");
    }

    try
    {
        vtkNew<vtkXMLUnstructuredGridReader> pReader;
        pReader->SetFileName(sFilePath.toLocal8Bit().constData());
        pReader->Update();
        vtkUnstructuredGrid* pOutput = pReader->GetOutput();
        if(pOutput == NULL || pOutput->GetNumberOfPoints() <= 0 ||
           pOutput->GetNumberOfCells() <= 0)
        {
            return setError(pError, "Numerical VTU grid is invalid.");
        }
        pGrid = pOutput;
        return pGrid != NULL;
    }
    catch(const std::bad_alloc&)
    {
        pGrid = NULL;
        return setError(pError, "Not enough memory to load numerical VTU grid.");
    }
}