/*******************************************************************************
* Copyright (c) 2014, Jan Koester jan.koester@gmx.net
* All rights reserved.
*
* Redistribution and use in source and binary forms, with or without
* modification, are permitted provided that the following conditions are met:
* Redistributions of source code must retain the above copyright
*      notice, this list of conditions and the following disclaimer.
* Redistributions in binary form must reproduce the above copyright
*      notice, this list of conditions and the following disclaimer in the
*      documentation and/or other materials provided with the distribution.
* Neither the name of the <organization> nor the
*      names of its contributors may be used to endorse or promote products
*      derived from this software without specific prior written permission.
*
* THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
* ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
* WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
* DISCLAIMED. IN NO EVENT SHALL <COPYRIGHT HOLDER> BE LIABLE FOR ANY
* DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
* (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
* LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
* ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
* (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
* SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
*******************************************************************************/

#include <cstring>
#include <vector>

#include "exception.h"
#include "socket.h"

namespace netplus {

// ============================================================================
// SSL Constructors - Initialize TLS state and set socket reference
// ============================================================================

ssl::ssl(const std::map<std::string, CertificateBundle>& certs) : tcp() {
    _tls.cert_map = std::make_shared<const std::map<std::string, CertificateBundle>>(certs);
    if (!certs.empty()) {
        _tls.selected_cert_bundle = const_cast<CertificateBundle*>(&_tls.cert_map->begin()->second);
    }
    _tls.setSocket(this);
    _Type = sockettype::SSL;
}

ssl::ssl(const std::map<std::string, CertificateBundle>& certs, int sock) : tcp(sock) {
    _tls.cert_map = std::make_shared<const std::map<std::string, CertificateBundle>>(certs);
    if (!certs.empty()) {
        _tls.selected_cert_bundle = const_cast<CertificateBundle*>(&_tls.cert_map->begin()->second);
    }
    _tls.setSocket(this);
    _Type = sockettype::SSL;
}

ssl::ssl(const std::map<std::string, CertificateBundle>& certs, 
         const std::string& addr, int port, int maxconnections, int sockopts)
    : tcp(addr, port, maxconnections, sockopts) {
    _tls.cert_map = std::make_shared<const std::map<std::string, CertificateBundle>>(certs);
    if (!certs.empty()) {
        _tls.selected_cert_bundle = const_cast<CertificateBundle*>(&_tls.cert_map->begin()->second);
    }
    _tls.setSocket(this);
    _Type = sockettype::SSL;
}

ssl::ssl(std::shared_ptr<const std::map<std::string, CertificateBundle>> shared_certs) : tcp() {
    _tls.cert_map = std::move(shared_certs);
    if (_tls.cert_map && !_tls.cert_map->empty()) {
        _tls.selected_cert_bundle = const_cast<CertificateBundle*>(&_tls.cert_map->begin()->second);
    }
    _tls.setSocket(this);
    _Type = sockettype::SSL;
}

// ============================================================================
// Accept - Create child SSL socket for accepted connection
// ============================================================================

void ssl::accept(std::unique_ptr<socket>& csock, bool nonblock) {
    // Share cert_map with child — no deep copy
    std::unique_ptr<ssl> cssock(new ssl(_tls.cert_map));
    
    tcp::accept(reinterpret_cast<std::unique_ptr<socket>&>(cssock), nonblock);
    
    // Propagate ALPN select callback and framing callback to child socket
    cssock->_tls.alpn_select_cb = _tls.alpn_select_cb;
    cssock->_framing_cb = _framing_cb;
    // Propagate TLS session cache for abbreviated handshake support
    cssock->_tls.session_cache = _tls.session_cache;
    cssock->_tls.setSocket(cssock.get());
    cssock->_Type = sockettype::SSL;
    
    csock = std::move(cssock);
}

// ============================================================================
// Connect - Initiate TLS client handshake
// ============================================================================

void ssl::connect(const std::string& addr, int port, bool nonblock) {
    // First establish TCP connection
    tcp::connect(addr, port, nonblock);

    if (nonblock)
        setNonBlock();

    // Set up for client mode
    _tls.is_client = true;
    _tls.hostname = addr;
    _tls.setSocket(this);

    // Build and queue ClientHello, set state machine to flush it
    _tls.queueClientHello();
    _tls.hs_state = tls::HsState::CLI_FLUSH_CH;

    // For blocking connections, drive the handshake to completion.
    // For non-blocking, the caller (e.g. epoll event loop) will drive it.
    if (!nonblock) {
        while (!_tls.getHandshakeDone()) {
            // Flush any pending writes first
            while (_tls.hasPendingWrite()) {
                _tls.flush_out();
            }
            // Drive the handshake state machine
            _tls.handshake_after_connect();
        }
    }
}

// ============================================================================
// Close - Reset TLS state and close TCP connection
// ============================================================================

void ssl::close() {
    _tls.resetTLS();
    tcp::close();
}

// ============================================================================
// sendData - Encrypt and send application data via TLS
// ============================================================================

size_t ssl::sendData(buffer& data, int flags) {
    return _tls.sendData(data, flags);
}

// ============================================================================
// recvData - Receive and decrypt TLS application data
// ============================================================================

size_t ssl::recvData(buffer& data, int flags) {
    return _tls.recvData(data, flags);
}

#ifdef Windows
// ============================================================================
// Windows-specific accept using AcceptEx
// ============================================================================

void ssl::accept(LPFN_ACCEPTEX lpfnAcceptEx, std::unique_ptr<socket>& csock) {
    std::unique_ptr<ssl> cssock(new ssl(_tls.cert_map));
    
    tcp::accept(lpfnAcceptEx, reinterpret_cast<std::unique_ptr<socket>&>(cssock));
    
    cssock->_tls.setSocket(cssock.get());
    cssock->_Type = sockettype::SSL;
    
    csock = std::move(cssock);
}

void ssl::prime_read() {
    // SSL uses the underlying TCP transport for raw I/O.
    // Post a WSARecv on the TCP socket so IOCP delivers the ClientHello
    // (and all subsequent TLS records) to the worker thread.
    tcp::prime_read();
}
#endif

} // namespace netplus
