CWE-ou-minkata/Sources/Plasma/FeatureLib/pfSecurePreloader/pfSecurePreloader.cpp

/*==LICENSE==*

CyanWorlds.com Engine - MMOG client, server and tools
Copyright (C) 2011  Cyan Worlds, Inc.

This program is free software: you can redistribute it and/or modify
it under the terms of the GNU General Public License as published by
the Free Software Foundation, either version 3 of the License, or
(at your option) any later version.

This program is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
GNU General Public License for more details.

You should have received a copy of the GNU General Public License
along with this program.  If not, see <http://www.gnu.org/licenses/>.

You can contact Cyan Worlds, Inc. by email legal@cyan.com
 or by snail mail at:
      Cyan Worlds, Inc.
      14617 N Newport Hwy
      Mead, WA   99021

*==LICENSE==*/
#include "hsSTLStream.h"
#include "hsResMgr.h"
#include "plgDispatch.h"
#include "pnUtils/pnUtils.h"
#include "pnNetBase/pnNetBase.h"
#include "pnAsyncCore/pnAsyncCore.h"
#include "pnNetCli/pnNetCli.h"
#include "plNetGameLib/plNetGameLib.h"
#include "plFile/plFileUtils.h"
#include "plFile/plStreamSource.h"
#include "plNetCommon/plNetCommon.h"
#include "plProgressMgr/plProgressMgr.h"
#include "plMessage/plPreloaderMsg.h"
#include "plMessage/plNetCommMsgs.h"
#include "pfSecurePreloader.h"

#include "plNetClientComm/plNetClientComm.h"

extern  hsBool  gDataServerLocal;


// Max number of concurrent file downloads
static const unsigned kMaxConcurrency   = 1;

pfSecurePreloader * pfSecurePreloader::fInstance;

///////////////////////////////////////////////////////////////////////////////
// Callback routines for the network code

// Called when a file's info is retrieved from the server
static void DefaultFileListRequestCallback(ENetError result, void* param, const NetCliAuthFileInfo infoArr[], unsigned infoCount)
{
    bool success = !IS_NET_ERROR(result);
    
    std::vector<std::wstring> filenames;
    std::vector<UInt32> sizes;
    if (success)
    {
        filenames.reserve(infoCount);
        sizes.reserve(infoCount);
        for (unsigned curFile = 0; curFile < infoCount; curFile++)
        {
            filenames.push_back(infoArr[curFile].filename);
            sizes.push_back(infoArr[curFile].filesize);
        }
    }
    ((pfSecurePreloader*)param)->RequestFinished(filenames, sizes, success);
}

// Called when a file download is either finished, or failed
static void DefaultFileRequestCallback(ENetError result, void* param, const wchar filename[], hsStream* stream)
{
    // Retry download unless shutting down or file not found
    switch (result) {
        case kNetSuccess:
            ((pfSecurePreloader*)param)->FinishedDownload(filename, true);
        break;
        
        case kNetErrFileNotFound:
        case kNetErrRemoteShutdown:
            ((pfSecurePreloader*)param)->FinishedDownload(filename, false);
        break;
        
        default:
            stream->Rewind();
            NetCliAuthFileRequest(
                filename,
                stream, 
                &DefaultFileRequestCallback,
                param
            );
        break;
    }
}


///////////////////////////////////////////////////////////////////////////////
// Our custom stream for writing directly to disk securely, and updating the
//  progress bar. Does NOT support reading (cause it doesn't need to)
class Direct2DiskStream : public hsUNIXStream
{
protected:
    wchar *         fWriteFileName;

    pfSecurePreloader* fPreloader;

public:
    Direct2DiskStream(pfSecurePreloader* preloader);
    ~Direct2DiskStream();

    virtual hsBool Open(const char* name, const char* mode = "wb");
    virtual hsBool Open(const wchar* name, const wchar* mode = L"wb");
    virtual hsBool Close();
    virtual UInt32 Read(UInt32 byteCount, void* buffer);
    virtual UInt32 Write(UInt32 byteCount, const void* buffer);
};


Direct2DiskStream::Direct2DiskStream(pfSecurePreloader* preloader) :
fWriteFileName(nil),
fPreloader(preloader)
{}

Direct2DiskStream::~Direct2DiskStream()
{
    Close();
}

hsBool Direct2DiskStream::Open(const char* name, const char* mode)
{
    wchar* wName = hsStringToWString(name);
    wchar* wMode = hsStringToWString(mode);
    hsBool ret = Open(wName, wMode);
    delete [] wName;
    delete [] wMode;
    return ret;
}

hsBool Direct2DiskStream::Open(const wchar* name, const wchar* mode)
{
    if (0 != wcscmp(mode, L"wb")) {
        hsAssert(0, "Unsupported open mode");
        return false;
    }
    
    fWriteFileName = TRACKED_NEW(wchar[wcslen(name) + 1]);
    wcscpy(fWriteFileName, name);
    
//  LogMsg(kLogPerf, L"Opening disk file %S", fWriteFileName);
    return hsUNIXStream::Open(name, mode);
}

hsBool Direct2DiskStream::Close()
{
    delete [] fWriteFileName;
    fWriteFileName = nil;
    return hsUNIXStream::Close();
}

UInt32 Direct2DiskStream::Read(UInt32 bytes, void* buffer)
{
    hsAssert(0, "not implemented");
    return 0; // we don't read
}

UInt32 Direct2DiskStream::Write(UInt32 bytes, const void* buffer)
{
//  LogMsg(kLogPerf, L"Writing %u bytes to disk file %S", bytes, fWriteFileName);
    fPreloader->UpdateProgressBar(bytes);
    return hsUNIXStream::Write(bytes, buffer);
}


///////////////////////////////////////////////////////////////////////////////
// secure preloader class implementation

// closes and deletes all streams
void pfSecurePreloader::ICleanupStreams()
{
    if (fD2DStreams.size() > 0)
    {
        std::map<std::wstring, hsStream*>::iterator curStream;
        for (curStream = fD2DStreams.begin(); curStream != fD2DStreams.end(); curStream++)
        {
            curStream->second->Close();
            delete curStream->second;
            curStream->second = nil;
        }
        fD2DStreams.clear();
    }
}

// queues a single file to be preloaded (does nothing if already preloaded)
void pfSecurePreloader::RequestSingleFile(std::wstring filename)
{
    fileRequest request;
    ZERO(request);
    request.fType = fileRequest::kSingleFile;
    request.fPath = filename;
    request.fExt = L"";

    fRequests.push_back(request);
}

// queues a group of files to be preloaded (does nothing if already preloaded)
void pfSecurePreloader::RequestFileGroup(std::wstring dir, std::wstring ext)
{
    fileRequest request;
    ZERO(request);
    request.fType = fileRequest::kFileList;
    request.fPath = dir;
    request.fExt = ext;

    fRequests.push_back(request);
}

// preloads all requested files from the server (does nothing if already preloaded)
void pfSecurePreloader::Start()
{
    if (gDataServerLocal) {
        // using local data, don't do anything
        plPreloaderMsg * msg = TRACKED_NEW plPreloaderMsg();
        msg->fSuccess = true;
        msg->Send();
        return;
    }

    NetCliAuthGetEncryptionKey(fEncryptionKey, 4); // grab the encryption key from the server

    fNetError = false;

    // make sure we are all cleaned up
    ICleanupStreams();
    fTotalDataReceived = 0;

    // update the progress bar for downloading
    if (!fProgressBar)
        fProgressBar = plProgressMgr::GetInstance()->RegisterOperation((hsScalar)(fRequests.size()), "Getting file info...", plProgressMgr::kUpdateText, false, true);
    
    for (unsigned curRequest = 0; curRequest < fRequests.size(); curRequest++)
    {
        fNumInfoRequestsRemaining++; // increment the counter
        if (fRequests[curRequest].fType == fileRequest::kSingleFile)
        {
#ifndef PLASMA_EXTERNAL_RELEASE
            // in internal releases, we can use on-disk files if they exist
            if (plFileUtils::FileExists(fRequests[curRequest].fPath.c_str()))
            {
                fileInfo info;
                info.fOriginalNameAndPath = fRequests[curRequest].fPath;
                info.fSizeInBytes = plFileUtils::GetFileSize(info.fOriginalNameAndPath.c_str());
                info.fDownloading = false;
                info.fDownloaded = false;
                info.fLocal = true;

                // generate garbled name
                wchar_t pathBuffer[MAX_PATH + 1];
                wchar_t filename[arrsize(pathBuffer)];
                GetTempPathW(arrsize(pathBuffer), pathBuffer);
                GetTempFileNameW(pathBuffer, L"CYN", 0, filename);
                info.fGarbledNameAndPath = filename;

                fTotalDataDownload += info.fSizeInBytes;

                fFileInfoMap[info.fOriginalNameAndPath] = info;
            }
            // internal client will still request it, even if it exists locally,
            // so that things get updated properly
#endif // PLASMA_EXTERNAL_RELEASE
            NetCliAuthFileListRequest(
                fRequests[curRequest].fPath.c_str(),
                nil,
                &DefaultFileListRequestCallback,
                (void*)this
            );
        }
        else
        {
#ifndef PLASMA_EXTERNAL_RELEASE
            // in internal releases, we can use on-disk files if they exist
            // Build the search string as "dir\\*.ext"
            wchar searchStr[MAX_PATH];

            PathAddFilename(searchStr, fRequests[curRequest].fPath.c_str(), L"*", arrsize(searchStr));
            PathSetExtension(searchStr, searchStr, fRequests[curRequest].fExt.c_str(), arrsize(searchStr));

            ARRAY(PathFind) paths;
            PathFindFiles(&paths, searchStr, kPathFlagFile); // find all files that match

            // convert it to our little file info array
            PathFind* curFile = paths.Ptr();
            PathFind* lastFile = paths.Term();
            while (curFile != lastFile) {
                fileInfo info;
                info.fOriginalNameAndPath = curFile->name;
                info.fSizeInBytes = (UInt32)curFile->fileLength;
                info.fDownloading = false;
                info.fDownloaded = false;
                info.fLocal = true;

                // generate garbled name
                wchar_t pathBuffer[MAX_PATH + 1];
                wchar_t filename[arrsize(pathBuffer)];
                GetTempPathW(arrsize(pathBuffer), pathBuffer);
                GetTempFileNameW(pathBuffer, L"CYN", 0, filename);
                info.fGarbledNameAndPath = filename;

                fTotalDataDownload += info.fSizeInBytes;

                fFileInfoMap[info.fOriginalNameAndPath] = info;
                curFile++;
            }
#endif // PLASMA_EXTERNAL_RELEASE

            NetCliAuthFileListRequest(
                fRequests[curRequest].fPath.c_str(),
                fRequests[curRequest].fExt.c_str(),
                &DefaultFileListRequestCallback,
                (void*)this
            );
        }
    }
}

// closes all file pointers and cleans up after itself
void pfSecurePreloader::Cleanup()
{
    ICleanupStreams();

    fRequests.clear();
    fFileInfoMap.clear();

    fNumInfoRequestsRemaining = 0;
    fTotalDataDownload = 0;
    fTotalDataReceived = 0;

    DEL(fProgressBar);
    fProgressBar = nil;
}

//============================================================================
void pfSecurePreloader::RequestFinished(const std::vector<std::wstring> & filenames, const std::vector<UInt32> & sizes, bool succeeded)
{
    fNetError |= !succeeded;
    
    if (succeeded)
    {
        unsigned count = 0;
        for (int curFile = 0; curFile < filenames.size(); curFile++)
        {
            if (fFileInfoMap.find(filenames[curFile]) != fFileInfoMap.end())
                continue; // if it is a duplicate, ignore it (the duplicate is probably one we found locally)

            fileInfo info;
            info.fOriginalNameAndPath = filenames[curFile];
            info.fSizeInBytes = sizes[curFile];
            info.fDownloading = false;
            info.fDownloaded = false;
            info.fLocal = false; // if we get here, it was retrieved remotely

            // generate garbled name
            wchar_t pathBuffer[MAX_PATH + 1];
            wchar_t filename[arrsize(pathBuffer)];
            GetTempPathW(arrsize(pathBuffer), pathBuffer);
            GetTempFileNameW(pathBuffer, L"CYN", 0, filename);
            info.fGarbledNameAndPath = filename;

            fTotalDataDownload += info.fSizeInBytes;

            fFileInfoMap[info.fOriginalNameAndPath] = info;
            ++count;
        }
        LogMsg(kLogPerf, "Added %u files to secure download queue", count);
    }
    if (fProgressBar)
        fProgressBar->Increment(1.f);
        
    --fNumInfoRequestsRemaining;    // even if we fail, decrement the counter

    if (succeeded) {
        DEL(fProgressBar);
        fProgressBar = plProgressMgr::GetInstance()->RegisterOperation((hsScalar)(fTotalDataDownload), "Downloading...", plProgressMgr::kUpdateText, false, true);

        // Issue some file download requests (up to kMaxConcurrency)
        IIssueDownloadRequests();
    }
    else {
        IPreloadComplete();
    }   
}

//============================================================================
void pfSecurePreloader::IIssueDownloadRequests () {

    std::map<std::wstring, fileInfo>::iterator curFile;
    for (curFile = fFileInfoMap.begin(); curFile != fFileInfoMap.end(); curFile++)
    {
        // Skip files already downloaded or currently downloading
        if (curFile->second.fDownloaded || curFile->second.fDownloading)
            continue;
            
        std::wstring filename = curFile->second.fOriginalNameAndPath;
#ifndef PLASMA_EXTERNAL_RELEASE
        // in internal releases, we can use on-disk files if they exist
        if (plFileUtils::FileExists(filename.c_str()))
        {
            // don't bother streaming, just make the secure stream using the local file

            // a local key overrides the server-downloaded key
            UInt32 localKey[4];
            bool hasLocalKey = plFileUtils::GetSecureEncryptionKey(filename.c_str(), localKey, arrsize(localKey));
            hsStream* stream = nil;
            if (hasLocalKey)
                stream = plSecureStream::OpenSecureFile(filename.c_str(), 0, localKey);
            else
                stream = plSecureStream::OpenSecureFile(filename.c_str(), 0, fEncryptionKey);

            // add it to the stream source
            bool added = plStreamSource::GetInstance()->InsertFile(filename.c_str(), stream);
            if (!added)
                DEL(stream); // wasn't added, so nuke our local copy

            // and make sure the vars are set up right
            curFile->second.fDownloaded = true;
            curFile->second.fLocal = true;
        }
        else
#endif
        {
            // Enforce concurrency limit
            if (fNumDownloadRequestsRemaining >= kMaxConcurrency)
                break;

            curFile->second.fDownloading = true;
            curFile->second.fDownloaded = false;
            curFile->second.fLocal = false;

            // create and setup the stream
            Direct2DiskStream* fileStream = TRACKED_NEW Direct2DiskStream(this);
            fileStream->Open(curFile->second.fGarbledNameAndPath.c_str(), L"wb");
            fD2DStreams[filename] = (hsStream*)fileStream;

            // request the file from the server
            LogMsg(kLogPerf, L"Requesting secure file:%s", filename.c_str());
            ++fNumDownloadRequestsRemaining;
            NetCliAuthFileRequest(
                filename.c_str(),
                (hsStream*)fileStream, 
                &DefaultFileRequestCallback,
                this
            );
        }
    }
    
    if (!fNumDownloadRequestsRemaining)
        IPreloadComplete();
}

void pfSecurePreloader::UpdateProgressBar(UInt32 bytesReceived)
{
    fTotalDataReceived += bytesReceived;
    if (fTotalDataReceived > fTotalDataDownload)
        fTotalDataReceived = fTotalDataDownload; // shouldn't happen... but just in case

    if (fProgressBar)
        fProgressBar->Increment((hsScalar)bytesReceived);
}

void pfSecurePreloader::FinishedDownload(std::wstring filename, bool succeeded)
{
    for (;;)
    {
        if (fFileInfoMap.find(filename) == fFileInfoMap.end())
        {
            // file doesn't exist... abort
            succeeded = false;
            break;
        }

        fFileInfoMap[filename].fDownloading = false;

        // close and delete the writer stream (even if we failed)
        fD2DStreams[filename]->Close();
        delete fD2DStreams[filename];
        fD2DStreams.erase(fD2DStreams.find(filename));

        if (succeeded)
        {
            // open a secure stream to that file
            hsStream* stream = plSecureStream::OpenSecureFile(
                fFileInfoMap[filename].fGarbledNameAndPath.c_str(),
                plSecureStream::kRequireEncryption | plSecureStream::kDeleteOnExit, // force delete and encryption
                fEncryptionKey
            );

            bool addedToSource = plStreamSource::GetInstance()->InsertFile(filename.c_str(), stream);
            if (!addedToSource)
                DEL(stream); // cleanup if it wasn't added

            fFileInfoMap[filename].fDownloaded = true;
            break;
        }
        
        // file download failed, clean up after it

        // delete the temporary file
        if (plFileUtils::FileExists(fFileInfoMap[filename].fGarbledNameAndPath.c_str()))
            plFileUtils::RemoveFile(fFileInfoMap[filename].fGarbledNameAndPath.c_str(), true);

        // and remove it from the info map
        fFileInfoMap.erase(fFileInfoMap.find(filename));
        break;
    }
        
    fNetError |= !succeeded;
    --fNumDownloadRequestsRemaining;
    LogMsg(kLogPerf, L"Received secure file:%s, success:%s", filename.c_str(), succeeded ? L"Yep" : L"Nope");

    if (!succeeded)
        IPreloadComplete();
    else
        // Issue some file download requests (up to kMaxConcurrency)
        IIssueDownloadRequests();
}

//============================================================================
void pfSecurePreloader::INotifyAuthReconnected () {

    // The secure file download network protocol will now just pick up downloading
    // where it left off before the reconnect, so no need to reset in-progress files.
    
    /*
    std::map<std::wstring, fileInfo>::iterator curFile;
    for (curFile = fFileInfoMap.begin(); curFile != fFileInfoMap.end(); curFile++) {

        // Reset files that were currently downloading
        if (curFile->second.fDownloading)
            curFile->second.fDownloading = false;
    }

    if (fNumDownloadRequestsRemaining > 0) {

        LogMsg(kLogPerf, L"pfSecurePreloader: Auth reconnected, resetting in-progress file downloads");

        // Issue some file download requests (up to kMaxConcurrency)
        IIssueDownloadRequests();
    }
    */
}

//============================================================================
void pfSecurePreloader::IPreloadComplete () {
    DEL(fProgressBar);
    fProgressBar = nil;
    
    plPreloaderMsg * msg = TRACKED_NEW plPreloaderMsg();
    msg->fSuccess = !fNetError;
    msg->Send();
}

//============================================================================
hsBool pfSecurePreloader::MsgReceive (plMessage * msg) {

    if (plNetCommAuthConnectedMsg * authMsg = plNetCommAuthConnectedMsg::ConvertNoRef(msg)) {
    
        INotifyAuthReconnected();
        return true;
    }
    
    return hsKeyedObject::MsgReceive(msg);
}

//============================================================================
pfSecurePreloader * pfSecurePreloader::GetInstance () {

    if (!fInstance) {
    
        fInstance = NEWZERO(pfSecurePreloader);
        fInstance->RegisterAs(kSecurePreloader_KEY);
    }

    return fInstance;
}

//============================================================================
bool pfSecurePreloader::IsInstanced () {

    return fInstance != nil;
}

//============================================================================
void pfSecurePreloader::Init () {

    if (!fInitialized) {
        
        fInitialized = true;
        plgDispatch::Dispatch()->RegisterForExactType(plNetCommAuthConnectedMsg::Index(), GetKey());
    }
}

//============================================================================
void pfSecurePreloader::Shutdown () {

    if (fInitialized) {
        
        fInitialized = false;
        plgDispatch::Dispatch()->UnRegisterForExactType(plNetCommAuthConnectedMsg::Index(), GetKey());
    }

    if (fInstance) {
    
        fInstance->UnRegister();
        fInstance = nil;
    }
}

//============================================================================
pfSecurePreloader::pfSecurePreloader () {
}

//============================================================================
pfSecurePreloader::~pfSecurePreloader () {

    Cleanup();
}