/*
 * securestream.cpp - combines a ByteStream with TLS and SASL
 * Copyright (C) 2004  Justin Karneges
 *
 * This library is free software; you can redistribute it and/or
 * modify it under the terms of the GNU Lesser General Public
 * License as published by the Free Software Foundation; either
 * version 2.1 of the License, or (at your option) any later version.
 *
 * This library 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
 * Lesser General Public License for more details.
 *
 * You should have received a copy of the GNU Lesser General Public
 * License along with this library; if not, write to the Free Software
 * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA
 *
 */

/*
  Note: SecureStream depends on the underlying security layers to signal
    plain-to-encrypted results immediately (as opposed to waiting for the
    event loop) so that the user cannot add/remove security layers during
    this conversion moment.  QCA::TLS and QCA::SASL behave as expected,
    but future layers might not.
*/

#include "securestream.h"

#include <qpointer.h>
#include <QList>
#include <qtimer.h>

#ifdef USE_TLSHANDLER
#include "xmpp.h"
#endif
#include "compressionhandler.h"

//----------------------------------------------------------------------------
// LayerTracker
//----------------------------------------------------------------------------
class LayerTracker
{
public:
    struct Item
    {
        int plain;
        int encoded;
    };

    LayerTracker();

    void reset();
    void addPlain(int plain);
    void specifyEncoded(int encoded, int plain);
    int finished(int encoded);

    int p;
    QList<Item> list;
};

LayerTracker::LayerTracker()
{
    p = 0;
}

void LayerTracker::reset()
{
    p = 0;
    list.clear();
}

void LayerTracker::addPlain(int plain)
{
    p += plain;
}

void LayerTracker::specifyEncoded(int encoded, int plain)
{
    // can't specify more bytes than we have
    if(plain > p)
        plain = p;
    p -= plain;
    Item i;
    i.plain = plain;
    i.encoded = encoded;
    list += i;
}

int LayerTracker::finished(int encoded)
{
    int plain = 0;
    for(QList<Item>::Iterator it = list.begin(); it != list.end();) {
        Item &i = *it;

        // not enough?
        if(encoded < i.encoded) {
            i.encoded -= encoded;
            break;
        }

        encoded -= i.encoded;
        plain += i.plain;
        it = list.erase(it);
    }
    return plain;
}

//----------------------------------------------------------------------------
// SecureStream
//----------------------------------------------------------------------------
class SecureLayer : public QObject
{
    Q_OBJECT
public:
    enum { TLS, SASL, TLSH, Compression };
    int type;
    union {
        QCA::TLS *tls;
        QCA::SASL *sasl;
#ifdef USE_TLSHANDLER
        XMPP::TLSHandler *tlsHandler;
#endif
        CompressionHandler *compressionHandler;
    } p;
    LayerTracker layer;
    bool tls_done;
    int prebytes;

    SecureLayer(QCA::TLS *t)
    {
        type = TLS;
        p.tls = t;
        init();
        connect(p.tls, SIGNAL(handshaken()), SLOT(tls_handshaken()));
        connect(p.tls, SIGNAL(readyRead()), SLOT(tls_readyRead()));
        connect(p.tls, SIGNAL(readyReadOutgoing(int)), SLOT(tls_readyReadOutgoing(int)));
        connect(p.tls, SIGNAL(closed()), SLOT(tls_closed()));
        connect(p.tls, SIGNAL(error(int)), SLOT(tls_error(int)));
    }

    SecureLayer(QCA::SASL *s)
    {
        type = SASL;
        p.sasl = s;
        init();
        connect(p.sasl, SIGNAL(readyRead()), SLOT(sasl_readyRead()));
        connect(p.sasl, SIGNAL(readyReadOutgoing()), SLOT(sasl_readyReadOutgoing()));
        connect(p.sasl, SIGNAL(error()), SLOT(sasl_error()));
    }

    SecureLayer(CompressionHandler *t)
    {
        t->setParent(this); // automatically clean up CompressionHandler when SecureLayer is destroyed
        type = Compression;
        p.compressionHandler = t;
        init();
        connect(p.compressionHandler, SIGNAL(readyRead()), SLOT(compressionHandler_readyRead()));
        connect(p.compressionHandler, SIGNAL(readyReadOutgoing()), SLOT(compressionHandler_readyReadOutgoing()));
        connect(p.compressionHandler, SIGNAL(error()), SLOT(compressionHandler_error()));
    }

#ifdef USE_TLSHANDLER
    SecureLayer(XMPP::TLSHandler *t)
    {
        type = TLSH;
        p.tlsHandler = t;
        init();
        connect(p.tlsHandler, SIGNAL(success()), SLOT(tlsHandler_success()));
        connect(p.tlsHandler, SIGNAL(fail()), SLOT(tlsHandler_fail()));
        connect(p.tlsHandler, SIGNAL(closed()), SLOT(tlsHandler_closed()));
        connect(p.tlsHandler, SIGNAL(readyRead(QByteArray)), SLOT(tlsHandler_readyRead(QByteArray)));
        connect(p.tlsHandler, SIGNAL(readyReadOutgoing(QByteArray,int)), SLOT(tlsHandler_readyReadOutgoing(QByteArray,int)));
    }
#endif

    void init()
    {
        tls_done = false;
        prebytes = 0;
    }

    void write(const QByteArray &a)
    {
        layer.addPlain(a.size());
        switch(type) {
            case TLS:  { p.tls->write(a); break; }
            case SASL: { p.sasl->write(a); break; }
#ifdef USE_TLSHANDLER
            case TLSH: { p.tlsHandler->write(a); break; }
#endif
            case Compression: { p.compressionHandler->write(a); break; }
        }
    }

    void writeIncoming(const QByteArray &a)
    {
        switch(type) {
            case TLS:  { p.tls->writeIncoming(a); break; }
            case SASL: { p.sasl->writeIncoming(a); break; }
#ifdef USE_TLSHANDLER
            case TLSH: { p.tlsHandler->writeIncoming(a); break; }
#endif
            case Compression: { p.compressionHandler->writeIncoming(a); break; }
        }
    }

    int finished(int plain)
    {
        int written = 0;

        // deal with prebytes (bytes sent prior to this security layer)
        if(prebytes > 0) {
            if(prebytes >= plain) {
                written += plain;
                prebytes -= plain;
                plain = 0;
            }
            else {
                written += prebytes;
                plain -= prebytes;
                prebytes = 0;
            }
        }

        // put remainder into the layer tracker
        if(type == SASL || tls_done)
            written += layer.finished(plain);

        return written;
    }

signals:
    void tlsHandshaken();
    void tlsClosed(const QByteArray &);
    void readyRead(const QByteArray &);
    void needWrite(const QByteArray &);
    void error(int);

private slots:
    void tls_handshaken()
    {
        tls_done = true;
        tlsHandshaken();
    }

    void tls_readyRead()
    {
        QByteArray a = p.tls->read();
        readyRead(a);
    }

    void tls_readyReadOutgoing(int plainBytes)
    {
        QByteArray a = p.tls->readOutgoing();
        if(tls_done)
            layer.specifyEncoded(a.size(), plainBytes);
        needWrite(a);
    }

    void tls_closed()
    {
        QByteArray a = p.tls->readUnprocessed();
        tlsClosed(a);
    }

    void tls_error(int x)
    {
        error(x);
    }

    void sasl_readyRead()
    {
        QByteArray a = p.sasl->read();
        readyRead(a);
    }

    void sasl_readyReadOutgoing()
    {
        int plainBytes;
        QByteArray a = p.sasl->readOutgoing(&plainBytes);
        layer.specifyEncoded(a.size(), plainBytes);
        needWrite(a);
    }

    void sasl_error()
    {
        error(p.sasl->errorCode());
    }

    void compressionHandler_readyRead()
    {
        QByteArray a = p.compressionHandler->read();
        readyRead(a);
    }

    void compressionHandler_readyReadOutgoing()
    {
        int plainBytes;
        QByteArray a = p.compressionHandler->readOutgoing(&plainBytes);
        layer.specifyEncoded(a.size(), plainBytes);
        needWrite(a);
    }

    void compressionHandler_error()
    {
        error(p.compressionHandler->errorCode());
    }

#ifdef USE_TLSHANDLER
    void tlsHandler_success()
    {
        tls_done = true;
        tlsHandshaken();
    }

    void tlsHandler_fail()
    {
        error(0);
    }

    void tlsHandler_closed()
    {
        tlsClosed(QByteArray());
    }

    void tlsHandler_readyRead(const QByteArray &a)
    {
        readyRead(a);
    }

    void tlsHandler_readyReadOutgoing(const QByteArray &a, int plainBytes)
    {
        if(tls_done)
            layer.specifyEncoded(a.size(), plainBytes);
        needWrite(a);
    }
#endif
};

#include "securestream.moc"

class SecureStream::Private
{
public:
    ByteStream *bs;
    QList<SecureLayer*> layers;
    int pending;
    int errorCode;
    bool active;
    bool topInProgress;

    bool haveTLS() const
    {
        foreach(SecureLayer *s, layers) {
            if(s->type == SecureLayer::TLS
#ifdef USE_TLSHANDLER
            || s->type == SecureLayer::TLSH
#endif
            ) {
                return true;
            }
        }
        return false;
    }

    bool haveSASL() const
    {
        foreach(SecureLayer *s, layers) {
            if(s->type == SecureLayer::SASL)
                return true;
        }
        return false;
    }

    bool haveCompress() const
    {
        foreach(SecureLayer *s, layers) {
            if(s->type == SecureLayer::Compression)
                return true;
        }
        return false;
    }

    void deleteLayers()
    {
        qDeleteAll(layers);
        layers.clear();
    }
};

SecureStream::SecureStream(ByteStream *s)
:ByteStream(0)
{
    d = new Private;

    d->bs = s;
    connect(d->bs, SIGNAL(readyRead()), SLOT(bs_readyRead()));
    connect(d->bs, SIGNAL(bytesWritten(qint64)), SLOT(bs_bytesWritten(qint64)));

    d->pending = 0;
    d->active = true;
    d->topInProgress = false;
    setOpenMode(QIODevice::ReadWrite);
}

SecureStream::~SecureStream()
{
    d->deleteLayers();
    delete d;
}

void SecureStream::linkLayer(QObject *s)
{
    connect(s, SIGNAL(tlsHandshaken()), SLOT(layer_tlsHandshaken()));
    connect(s, SIGNAL(tlsClosed(QByteArray)), SLOT(layer_tlsClosed(QByteArray)));
    connect(s, SIGNAL(readyRead(QByteArray)), SLOT(layer_readyRead(QByteArray)));
    connect(s, SIGNAL(needWrite(QByteArray)), SLOT(layer_needWrite(QByteArray)));
    connect(s, SIGNAL(error(int)), SLOT(layer_error(int)));
}

int SecureStream::calcPrebytes() const
{
    int x = 0;
    foreach(SecureLayer *s, d->layers) {
        x += s->prebytes;
    }
    return (d->pending - x);
}

void SecureStream::startTLSClient(QCA::TLS *t, const QByteArray &spare)
{
    if(!d->active || d->topInProgress || d->haveTLS())
        return;

    SecureLayer *s = new SecureLayer(t);
    s->prebytes = calcPrebytes();
    linkLayer(s);
    d->layers.append(s);
    d->topInProgress = true;

    insertData(spare);
}

void SecureStream::startTLSServer(QCA::TLS *t, const QByteArray &spare)
{
    if(!d->active || d->topInProgress || d->haveTLS())
        return;

    SecureLayer *s = new SecureLayer(t);
    s->prebytes = calcPrebytes();
    linkLayer(s);
    d->layers.append(s);
    d->topInProgress = true;

    insertData(spare);
}

void SecureStream::setLayerCompress(const QByteArray& spare)
{
    if(!d->active || d->topInProgress || d->haveCompress())
        return;

    SecureLayer *s = new SecureLayer(new CompressionHandler());
    s->prebytes = calcPrebytes();
    linkLayer(s);
    d->layers.append(s);

    insertData(spare);
}

void SecureStream::setLayerSASL(QCA::SASL *sasl, const QByteArray &spare)
{
    if(!d->active || d->topInProgress || d->haveSASL())
        return;

    SecureLayer *s = new SecureLayer(sasl);
    s->prebytes = calcPrebytes();
    linkLayer(s);
    d->layers.append(s);

    insertData(spare);
}

#ifdef USE_TLSHANDLER
void SecureStream::startTLSClient(XMPP::TLSHandler *t, const QString &server, const QByteArray &spare)
{
    if(!d->active || d->topInProgress || d->haveTLS())
        return;

    SecureLayer *s = new SecureLayer(t);
    s->prebytes = calcPrebytes();
    linkLayer(s);
    d->layers.append(s);
    d->topInProgress = true;

    // unlike QCA::TLS, XMPP::TLSHandler has no return value
    s->p.tlsHandler->startClient(server);

    insertData(spare);
}
#endif

void SecureStream::closeTLS()
{
    if (!d->layers.isEmpty()) {
        SecureLayer *s = d->layers.last();
        if(s->type == SecureLayer::TLS) {
            s->p.tls->close();
        }
    }
}

int SecureStream::errorCode() const
{
    return d->errorCode;
}

bool SecureStream::isOpen() const
{
    return d->active;
}

void SecureStream::write(const QByteArray &a)
{
    if(!isOpen()) {
        qDebug("Writing to closed stream!");
        return;
    }

    d->pending += a.size();

    // send to the last layer
    if (!d->layers.isEmpty()) {
        SecureLayer *s = d->layers.last();
        s->write(a);
    }
    else {
        writeRawData(a);
    }
}

qint64 SecureStream::bytesToWrite() const
{
    return d->pending;
}

void SecureStream::bs_readyRead()
{
    QByteArray a = d->bs->readAll();

    // send to the first layer
    if (!d->layers.isEmpty()) {
        SecureLayer *s = d->layers.first();
        s->writeIncoming(a);
    }
    else {
        incomingData(a);
    }
}

void SecureStream::bs_bytesWritten(qint64 bytes)
{
    foreach(SecureLayer *s, d->layers) {
        bytes = s->finished(bytes);
    }

    if(bytes > 0) {
        d->pending -= bytes;
        bytesWritten(bytes);
    }
}

void SecureStream::layer_tlsHandshaken()
{
    d->topInProgress = false;
    tlsHandshaken();
}

void SecureStream::layer_tlsClosed(const QByteArray &)
{
    setOpenMode(QIODevice::NotOpen);
    d->active = false;
    d->deleteLayers();
    tlsClosed();
}

void SecureStream::layer_readyRead(const QByteArray &a)
{
    SecureLayer *s = static_cast<SecureLayer *>(sender());
    QList<SecureLayer*>::Iterator it(d->layers.begin());
    while((*it) != s) {
        Q_ASSERT(it != d->layers.end());
        ++it;
    }
    Q_ASSERT(it != d->layers.end());

    // pass upwards
    ++it;
    if (it != d->layers.end()) {
        s = (*it);
        s->writeIncoming(a);
    }
    else {
        incomingData(a);
    }
}

void SecureStream::layer_needWrite(const QByteArray &a)
{
    SecureLayer *s = static_cast<SecureLayer *>(sender());
    QList<SecureLayer*>::Iterator it(d->layers.begin());
    while((*it) != s) {
        Q_ASSERT(it != d->layers.end());
        ++it;
    }
    Q_ASSERT(it != d->layers.end());

    // pass downwards
    if (it != d->layers.begin()) {
        --it;
        s = (*it);
        s->write(a);
    }
    else {
        writeRawData(a);
    }
}

void SecureStream::layer_error(int x)
{
    SecureLayer *s = static_cast<SecureLayer *>(sender());
    int type = s->type;
    d->errorCode = x;
    setOpenMode(QIODevice::NotOpen);
    d->active = false;
    d->deleteLayers();
    if(type == SecureLayer::TLS)
        setError(ErrTLS);
    else if(type == SecureLayer::SASL)
        setError(ErrSASL);
#ifdef USE_TLSHANDLER
    else if(type == SecureLayer::TLSH)
        setError(ErrTLS);
#endif
}

void SecureStream::insertData(const QByteArray &a)
{
    if(!a.isEmpty()) {
        if (!d->layers.isEmpty()) {
            SecureLayer *s = d->layers.last();
            s->writeIncoming(a);
        }
        else {
            incomingData(a);
        }
    }
}

void SecureStream::writeRawData(const QByteArray &a)
{
    d->bs->write(a);
}

void SecureStream::incomingData(const QByteArray &a)
{
    appendRead(a);
    if(bytesAvailable())
        emit readyRead();
}
