/**
 * Licensed to the University Corporation for Advanced Internet
 * Development, Inc. (UCAID) under one or more contributor license
 * agreements. See the NOTICE file distributed with this work for
 * additional information regarding copyright ownership.
 *
 * UCAID licenses this file to you 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.
 */

/**
 * PKIXTrustEngine.cpp
 * 
 * Shibboleth-specific PKIX-validation TrustEngine.
 */

#include "internal.h"
#include "metadata/MetadataExt.h"
#include "security/PKIXTrustEngine.h"

#include <boost/shared_ptr.hpp>
#include <saml/saml2/metadata/Metadata.h>
#include <saml/saml2/metadata/MetadataCredentialCriteria.h>
#include <saml/saml2/metadata/ObservableMetadataProvider.h>
#include <xmltooling/XMLToolingConfig.h>
#include <xmltooling/security/AbstractPKIXTrustEngine.h>
#include <xmltooling/security/KeyInfoResolver.h>
#include <xmltooling/security/X509Credential.h>
#include <xmltooling/util/Threads.h>

using namespace shibsp;
using namespace opensaml::saml2md;
using namespace xmlsignature;
using namespace xmltooling;
using namespace boost;
using namespace std;

namespace shibsp {

    class SHIBSP_DLLLOCAL PKIXTrustEngine : public AbstractPKIXTrustEngine, public ObservableMetadataProvider::Observer
    {
    public:
        PKIXTrustEngine(const DOMElement* e=nullptr) : AbstractPKIXTrustEngine(e), m_credLock(RWLock::create()) {
        }
        virtual ~PKIXTrustEngine() {
            for (map<const ObservableMetadataProvider*,credmap_t>::iterator i=m_credentialMap.begin(); i!=m_credentialMap.end(); ++i) {
                i->first->removeObserver(this);
            }
        }
        
        AbstractPKIXTrustEngine::PKIXValidationInfoIterator* getPKIXValidationInfoIterator(
            const CredentialResolver& pkixSource, CredentialCriteria* criteria=nullptr
            ) const;

        void onEvent(const ObservableMetadataProvider& metadata) const {
            // Destroy credentials we cached from this provider.
            m_credLock->wrlock();
            SharedLock locker(m_credLock, false);
            m_credentialMap[&metadata].clear();
        }

        const KeyInfoResolver* getKeyInfoResolver() const {
            return m_keyInfoResolver ? m_keyInfoResolver : XMLToolingConfig::getConfig().getKeyInfoResolver();
        }

    private:
        friend class SHIBSP_DLLLOCAL MetadataPKIXIterator;
        scoped_ptr<RWLock> m_credLock;
        typedef map< const KeyAuthority*,vector< boost::shared_ptr<X509Credential> > > credmap_t;
        mutable map<const ObservableMetadataProvider*,credmap_t> m_credentialMap;
    };
    
    SHIBSP_DLLLOCAL PluginManager<TrustEngine,string,const DOMElement*>::Factory PKIXTrustEngineFactory;

    TrustEngine* SHIBSP_DLLLOCAL PKIXTrustEngineFactory(const DOMElement* const & e)
    {
        return new PKIXTrustEngine(e);
    }

    class SHIBSP_DLLLOCAL MetadataPKIXIterator : public AbstractPKIXTrustEngine::PKIXValidationInfoIterator
    {
    public:
        MetadataPKIXIterator(const PKIXTrustEngine& engine, const MetadataProvider& pkixSource, MetadataCredentialCriteria& criteria);

        virtual ~MetadataPKIXIterator() {
            if (m_caching)
                m_engine.m_credLock->unlock();
        }

        bool next();

        int getVerificationDepth() const {
            pair<bool,int> vd = m_current->getVerifyDepth();
            return vd.first ? vd.second : 1;
        }
        
        const vector<XSECCryptoX509*>& getTrustAnchors() const {
            return m_certs;
        }

        const vector<XSECCryptoX509CRL*>& getCRLs() const {
            return m_crls;
        }
    
    private:
        void populate();
        bool m_caching;
        const PKIXTrustEngine& m_engine;
        map<const ObservableMetadataProvider*,PKIXTrustEngine::credmap_t>::iterator m_credCache;
        const XMLObject* m_obj;
        const Extensions* m_extBlock;
        const KeyAuthority* m_current;
        vector<XMLObject*>::const_iterator m_iter;
        vector<XSECCryptoX509*> m_certs;
        vector<XSECCryptoX509CRL*> m_crls;
        vector< boost::shared_ptr<X509Credential> > m_ownedCreds;
    };
};

void shibsp::registerPKIXTrustEngine()
{
    XMLToolingConfig::getConfig().TrustEngineManager.registerFactory(SHIBBOLETH_PKIX_TRUSTENGINE, PKIXTrustEngineFactory);
}

AbstractPKIXTrustEngine::PKIXValidationInfoIterator* PKIXTrustEngine::getPKIXValidationInfoIterator(
    const CredentialResolver& pkixSource, CredentialCriteria* criteria
    ) const
{
    // Make sure these are metadata objects.
    const MetadataProvider& metadata = dynamic_cast<const MetadataProvider&>(pkixSource);
    MetadataCredentialCriteria* metacrit = dynamic_cast<MetadataCredentialCriteria*>(criteria);
    if (!metacrit)
        throw MetadataException("Cannot obtain PKIX information without a MetadataCredentialCriteria object.");

    return new MetadataPKIXIterator(*this, metadata,*metacrit);
}

MetadataPKIXIterator::MetadataPKIXIterator(
    const PKIXTrustEngine& engine, const MetadataProvider& pkixSource, MetadataCredentialCriteria& criteria
    ) : m_caching(false), m_engine(engine), m_obj(criteria.getRole().getParent()), m_extBlock(nullptr), m_current(nullptr)
{
    // If we can't hook the metadata for changes, then we can't do any caching and the rest of this is academic.
    const ObservableMetadataProvider* observable = dynamic_cast<const ObservableMetadataProvider*>(&pkixSource);
    if (!observable)
        return;

    // While holding read lock, see if this metadata plugin has been seen before.
    m_engine.m_credLock->rdlock();
    m_credCache = m_engine.m_credentialMap.find(observable);
    if (m_credCache == m_engine.m_credentialMap.end()) {

        // We need to elevate the lock and retry.
        m_engine.m_credLock->unlock();
        m_engine.m_credLock->wrlock();
        m_credCache = m_engine.m_credentialMap.find(observable);
        if (m_credCache == m_engine.m_credentialMap.end()) {
            try {
                // It's still brand new, so hook it for cache activation.
                observable->addObserver(&m_engine);

                // Prime the map reference with an empty credential map.
                m_credCache = m_engine.m_credentialMap.insert(make_pair(observable,PKIXTrustEngine::credmap_t())).first;
            }
            catch (std::exception&) {
                // The destructor won't run if we throw here, so we need to unlock.
                m_engine.m_credLock->unlock();
                throw;
            }
            
            // Downgrade the lock.
            // We don't have to recheck because we never erase the master map entry entirely, even on changes.
            m_engine.m_credLock->unlock();
            m_engine.m_credLock->rdlock();
        }
    }
    
    // We've hooked the metadata for changes, and we know we can cache against it.
    m_caching = true;
}


bool MetadataPKIXIterator::next()
{
    // If we had an active block, look for another in the same block.
    if (m_extBlock) {
        // Keep going until we hit the end of the block.
        vector<XMLObject*>::const_iterator end = m_extBlock->getUnknownXMLObjects().end();
        while (m_iter != end) {
            // If we hit a KeyAuthority, remember it and signal.
            if (m_current=dynamic_cast<KeyAuthority*>(*m_iter++)) {
                populate();
                return true;
            }
        }
        
        // If we get here, we hit the end of this Extensions block.
        // Climb a level, if possible.
        m_obj = m_obj->getParent();
        m_current = nullptr;
        m_extBlock = nullptr;
    }

    // If we get here, we try and find an Extensions block.
    while (m_obj) {
        const EntityDescriptor* entity = dynamic_cast<const EntityDescriptor*>(m_obj);
        if (entity) {
            m_extBlock = entity->getExtensions();
        }
        else {
            const EntitiesDescriptor* entities = dynamic_cast<const EntitiesDescriptor*>(m_obj);
            if (entities) {
                m_extBlock = entities->getExtensions();
            }
        }
        
        if (m_extBlock) {
            m_iter = m_extBlock->getUnknownXMLObjects().begin();
            return next();
        }
        
        // Jump a level and try again.
        m_obj = m_obj->getParent();
    }

    return false;
}

void MetadataPKIXIterator::populate()
{
    // Dump anything old.
    m_certs.clear();
    m_crls.clear();
    m_ownedCreds.clear();

    if (m_caching) {
        // We're holding a read lock. Search for "resolved" creds.
        PKIXTrustEngine::credmap_t::iterator cached = m_credCache->second.find(m_current);
        if (cached != m_credCache->second.end()) {
            // Copy over the information.
            for (vector< boost::shared_ptr<X509Credential> >::const_iterator c = cached->second.begin(); c != cached->second.end(); ++c) {
                m_certs.insert(m_certs.end(), (*c)->getEntityCertificateChain().begin(), (*c)->getEntityCertificateChain().end());
                if ((*c)->getCRL())
                    m_crls.push_back((*c)->getCRL());
            }
            return;
        }
    }

    // We're either not caching or didn't find the results we need, so we have to resolve them.
    const vector<KeyInfo*>& keyInfos = m_current->getKeyInfos();
    for (vector<KeyInfo*>::const_iterator k = keyInfos.begin(); k != keyInfos.end(); ++k) {
        boost::shared_ptr<Credential> cred(m_engine.getKeyInfoResolver()->resolve(*k, X509Credential::RESOLVE_CERTS | X509Credential::RESOLVE_CRLS));
        boost::shared_ptr<X509Credential> xcred = boost::dynamic_pointer_cast<X509Credential>(cred);
        if (xcred)
            m_ownedCreds.push_back(xcred);
    }

    // Copy over the new information.
    for (vector< boost::shared_ptr<X509Credential> >::const_iterator c = m_ownedCreds.begin(); c != m_ownedCreds.end(); ++c) {
        m_certs.insert(m_certs.end(), (*c)->getEntityCertificateChain().begin(), (*c)->getEntityCertificateChain().end());
        if ((*c)->getCRL())
            m_crls.push_back((*c)->getCRL());
    }

    // As a last step, if we're caching, try and elevate to a write lock for cache insertion.
    if (m_caching) {
        m_engine.m_credLock->unlock();
        m_engine.m_credLock->wrlock();
        if (m_credCache->second.count(m_current) == 0) {
            // Transfer objects into cache.
            m_credCache->second[m_current] = m_ownedCreds;
            m_ownedCreds.clear();
        }
        m_engine.m_credLock->unlock();
        m_engine.m_credLock->rdlock();

        // In theory we could have lost the objects but that shouldn't be possible
        // since the metadata itself is locked and shouldn't change behind us.
    }
}
