module openssl

import io
import net
import time
import os

// SSLConn is the current connection
pub struct SSLConn {
pub:
    config SSLConnectConfig
pub mut:
    sslctx   &C.SSL_CTX = unsafe { nil }
    ssl      &C.SSL     = unsafe { nil }
    handle   int
    duration time.Duration

    owns_socket bool
    // last_write_sent reports the most recent write_ptr's progress for retry
    // decisions: 0 = provably nothing was sent (safe to replay), or -1 = the
    // count is indeterminate because a failed/retryable write may have already
    // flushed complete records to the peer (TLS cannot prove zero). On full
    // success it equals the bytes written.
    last_write_sent int
}

@[params]
pub struct SSLConnectConfig {
pub:
    verify   string // the path to a rootca.pem file, containing trusted CA certificate(s)
    cert     string // the path to a cert.pem file, containing client certificate(s) for the request
    cert_key string // the path to a key.pem file, containing private keys for the client certificate(s)
    validate bool   // set this to true, if you want to stop requests, when their certificates are found to be invalid

    in_memory_verification bool // if true, verify, cert, and cert_key are read from memory, not from a file

    alpn_protocols []string // the list of ALPN protocols to advertise, e.g. ['h2', 'http/1.1']; empty means no ALPN extension is sent
}

// new_ssl_conn instance an new SSLCon struct
pub fn new_ssl_conn(config SSLConnectConfig) !&SSLConn {
    $if trace_ssl ? {
        eprintln(@METHOD)
    }
    mut conn := &SSLConn{
        config: config
        sslctx: unsafe { nil }
        ssl:    unsafe { nil }
        handle: 0
    }
    conn.init() or { return err }
    return conn
}

// Select operation
enum Select {
    read
    write
    except
}

fn ssl_timeout_deadline(timeout time.Duration) time.Time {
    if timeout <= 0 || timeout == net.infinite_timeout {
        return time.unix(0)
    }
    return time.now().add(timeout)
}

fn ssl_remaining_timeout(deadline time.Time) time.Duration {
    if deadline.unix() == 0 {
        return net.infinite_timeout
    }
    return deadline - time.now()
}

// close closes the ssl connection and does cleanup
pub fn (mut s SSLConn) close() ! {
    s.shutdown()!
}

// negotiated_alpn returns the ALPN protocol selected during the TLS
// handshake (e.g. 'h2' or 'http/1.1'), or an empty string if no protocol
// was negotiated.
pub fn (s &SSLConn) negotiated_alpn() string {
    mut data := &u8(unsafe { nil })
    mut length := u32(0)
    C.v_net_openssl_get0_alpn_selected(voidptr(s.ssl), voidptr(&data), &length)
    if length == 0 || data == unsafe { nil } {
        return ''
    }
    // data points into OpenSSL-owned memory and is not NUL-terminated; copy it.
    return unsafe { data.vbytes(int(length)).bytestr() }
}

// shutdown closes the ssl connection and does cleanup
pub fn (mut s SSLConn) shutdown() ! {
    $if trace_ssl ? {
        eprintln(@METHOD)
    }

    if s.ssl != 0 {
        deadline := ssl_timeout_deadline(s.duration)
        for {
            mut res := C.SSL_shutdown(voidptr(s.ssl))
            if res == 1 {
                break
            }

            err_res := ssl_error(res, s.ssl) or {
                break // We break to free rest of resources
            }
            if err_res == .ssl_error_want_read {
                s.wait_for_read(ssl_remaining_timeout(deadline))!
                continue
            } else if err_res == .ssl_error_want_write {
                s.wait_for_write(ssl_remaining_timeout(deadline))!
                continue
            }
            if s.ssl != 0 {
                unsafe { C.SSL_free(voidptr(s.ssl)) }
            }
            if s.sslctx != 0 {
                C.SSL_CTX_free(s.sslctx)
            }
            return error('net.openssl Could not connect using SSL. (${err_res}),err')
        }
        C.SSL_free(voidptr(s.ssl))
    }
    if s.sslctx != 0 {
        C.SSL_CTX_free(s.sslctx)
    }
    if s.owns_socket {
        net.shutdown(s.handle)
        net.close(s.handle)!
    }
}

fn (mut s SSLConn) init() ! {
    $if trace_ssl ? {
        eprintln(@METHOD)
    }
    s.sslctx = unsafe { C.SSL_CTX_new(C.SSLv23_client_method()) }
    if s.sslctx == 0 {
        return error('net.openssl Could not get ssl context')
    }

    if s.config.validate {
        C.SSL_CTX_set_verify_depth(s.sslctx, 4)
        C.SSL_CTX_set_options(s.sslctx, C.SSL_OP_NO_COMPRESSION)
    }

    s.ssl = unsafe { &C.SSL(C.SSL_new(s.sslctx)) }
    if s.ssl == 0 {
        return error('net.openssl Could not create OpenSSL instance')
    }

    mut res := 0

    // Advertise ALPN protocols (e.g. ['h2', 'http/1.1']) when requested.
    // OpenSSL expects the length-prefixed wire format: each protocol is a
    // single length byte followed by that many bytes of protocol name.
    if s.config.alpn_protocols.len > 0 {
        mut wire := []u8{cap: 64}
        for proto in s.config.alpn_protocols {
            if proto.len == 0 || proto.len > 255 {
                return error('net.openssl SSLConn.init, invalid ALPN protocol "${proto}"')
            }
            wire << u8(proto.len)
            wire << proto.bytes()
        }
        // Returns 0 on success (opposite of most OpenSSL calls); non-zero also
        // means ALPN is unavailable on OpenSSL versions older than 1.0.2.
        if C.v_net_openssl_set_alpn_protos(voidptr(s.ssl), wire.data, u32(wire.len)) != 0 {
            return error('net.openssl SSLConn.init, failed to set ALPN protocols (requires OpenSSL >= 1.0.2)')
        }
    }

    if s.config.validate {
        mut verify := s.config.verify
        mut cert := s.config.cert
        mut cert_key := s.config.cert_key
        if s.config.in_memory_verification {
            now := time.now().unix().str()
            verify = os.temp_dir() + '/v_verify' + now
            cert = os.temp_dir() + '/v_cert' + now
            cert_key = os.temp_dir() + '/v_cert_key' + now
            if s.config.verify != '' {
                os.write_file(verify, s.config.verify)!
            }
            if s.config.cert != '' {
                os.write_file(cert, s.config.cert)!
            }
            if s.config.cert_key != '' {
                os.write_file(cert_key, s.config.cert_key)!
            }
        }
        if s.config.verify != '' {
            res = C.SSL_CTX_load_verify_locations(voidptr(s.sslctx), &char(verify.str), 0)
            if s.config.validate && res != 1 {
                return error('net.openssl SSLConn.init, SSL_CTX_load_verify_locations failed')
            }
        }
        if s.config.cert != '' {
            res = C.SSL_CTX_use_certificate_file(voidptr(s.sslctx), &char(cert.str),
                C.SSL_FILETYPE_PEM)
            if s.config.validate && res != 1 {
                return error('net.openssl SSLConn.init, SSL_CTX_use_certificate_file failed, res: ${res}')
            }
        }
        if s.config.cert_key != '' {
            res = C.SSL_CTX_use_PrivateKey_file(voidptr(s.sslctx), &char(cert_key.str),
                C.SSL_FILETYPE_PEM)
            if s.config.validate && res != 1 {
                return error('net.openssl SSLConn.init, SSL_CTX_use_PrivateKey_file failed, res: ${res}')
            }
        }

        preferred_ciphers := 'HIGH:!aNULL:!kRSA:!PSK:!SRP:!MD5:!RC4'
        res = C.SSL_set_cipher_list(voidptr(s.ssl), &char(preferred_ciphers.str))
        if s.config.validate && res != 1 {
            println('net.openssl: set cipher failed')
        }
    }
}

// connect to server using OpenSSL
pub fn (mut s SSLConn) connect(mut tcp_conn net.TcpConn, hostname string) ! {
    $if trace_ssl ? {
        eprintln('${@METHOD} hostname: ${hostname}')
    }
    s.handle = tcp_conn.sock.handle
    s.duration = tcp_conn.read_timeout()
    mut res := C.SSL_set_tlsext_host_name(voidptr(s.ssl), voidptr(hostname.str))
    if res != 1 {
        return error('net.openssl SSLConn.connect, could not set host name')
    }
    if C.SSL_set_fd(voidptr(s.ssl), tcp_conn.sock.handle) != 1 {
        return error('net.openssl SSLConn.connect, could not assign ssl to socket.')
    }
    s.complete_connect()!
}

// dial opens an ssl connection on hostname:port
pub fn (mut s SSLConn) dial(hostname string, port int) ! {
    $if trace_ssl ? {
        eprintln('${@METHOD} hostname: ${hostname} | port: ${port}')
    }
    mut tcp_conn := net.dial_tcp('${hostname}:${port}') or { return err }
    mut connected := false
    defer {
        if !connected {
            tcp_conn.close() or {}
            if s.ssl != 0 {
                unsafe { C.SSL_free(voidptr(s.ssl)) }
                s.ssl = unsafe { nil }
            }
            if s.sslctx != 0 {
                C.SSL_CTX_free(s.sslctx)
                s.sslctx = unsafe { nil }
            }
            s.handle = 0
            s.owns_socket = false
        }
    }
    s.owns_socket = true
    s.connect(mut tcp_conn, hostname) or { return err }
    connected = true
}

fn (mut s SSLConn) complete_connect() ! {
    $if trace_ssl ? {
        eprintln(@METHOD)
    }

    deadline := ssl_timeout_deadline(s.duration)
    for {
        mut res := C.SSL_connect(voidptr(s.ssl))
        if res == 1 {
            break
        }

        err_res := ssl_error(res, s.ssl)!
        if err_res == .ssl_error_want_read {
            s.wait_for_read(ssl_remaining_timeout(deadline))!
            continue
        }
        if err_res == .ssl_error_want_write {
            s.wait_for_write(ssl_remaining_timeout(deadline))!
            continue
        }
        return error('net.openssl SSLConn.complete_connect, could not connect using SSL. (${err_res}),err')
    }

    if s.config.validate {
        mut pcert := &C.X509(unsafe { nil })
        for {
            mut res := C.SSL_do_handshake(voidptr(s.ssl))
            if res == 1 {
                break
            }

            err_res := ssl_error(res, s.ssl)!
            if err_res == .ssl_error_want_read {
                s.wait_for_read(ssl_remaining_timeout(deadline))!
                continue
            } else if err_res == .ssl_error_want_write {
                s.wait_for_write(ssl_remaining_timeout(deadline))!
                continue
            }
            return error('net.openssl SSLConn.complete_connect, could not validate SSL certificate. (${err_res}),err')
        }
        pcert = C.v_net_openssl_get1_peer_certificate(s.ssl)
        defer {
            if pcert != 0 {
                C.X509_free(pcert)
            }
        }
        res := C.SSL_get_verify_result(voidptr(s.ssl))
        if res != C.X509_V_OK {
            return error('net.openssl SSLConn.complete_connect, failed SSL handshake (OpenSSL SSL_get_verify_result = ${res})')
        }
    }
}

// read_timeout returns the current SSL read timeout.
pub fn (s &SSLConn) read_timeout() time.Duration {
    return s.duration
}

// set_read_timeout sets the read timeout used for subsequent SSL reads.
// A value of 0, or net.infinite_timeout means "wait forever".
pub fn (mut s SSLConn) set_read_timeout(timeout time.Duration) {
    s.duration = timeout
}

// addr retrieves the local ip address and port number for this connection
pub fn (s &SSLConn) addr() !net.Addr {
    return net.addr_from_socket_handle(s.handle)
}

// peer_addr retrieves the ip address and port number used by the peer
pub fn (s &SSLConn) peer_addr() !net.Addr {
    return net.peer_addr_from_socket_handle(s.handle)
}

pub fn (mut s SSLConn) socket_read_into_ptr(buf_ptr &u8, len int) !int {
    mut res := 0
    $if trace_ssl ? {
        defer(fn) {
            if len > 0 {
                eprintln('${@METHOD} res: ${res}: buf_ptr: ${voidptr(buf_ptr):x}, len: ${len}, hex: ${unsafe { buf_ptr.vbytes(len).hex() }} data: `${unsafe { buf_ptr.vstring_with_len(len) }}`')
            }
        }
    }

    deadline := ssl_timeout_deadline(s.duration)
    // s.wait_for_read(deadline - time.now())!
    for {
        res = C.SSL_read(voidptr(s.ssl), buf_ptr, len)
        if res > 0 {
            return res
        } else if res == 0 {
            $if trace_ssl ? {
                eprintln('${@METHOD} ---> res: io.Eof')
            }
            return io.Eof{}
        } else {
            err_res := ssl_error(res, s.ssl)!
            match err_res {
                .ssl_error_want_read {
                    s.wait_for_read(ssl_remaining_timeout(deadline)) or {
                        $if trace_ssl ? {
                            eprintln('${@METHOD} ---> res: ${err} .ssl_error_want_read')
                        }
                        return err
                    }
                }
                .ssl_error_want_write {
                    s.wait_for_write(ssl_remaining_timeout(deadline)) or {
                        $if trace_ssl ? {
                            eprintln('${@METHOD} ---> res: ${err} .ssl_error_want_write')
                        }
                        return err
                    }
                }
                .ssl_error_zero_return {
                    $if trace_ssl ? {
                        eprintln('${@METHOD} ---> res: 0 .ssl_error_zero_return')
                    }
                    return 0
                }
                else {
                    $if trace_ssl ? {
                        eprintln('${@METHOD} ---> res: could not read, err_res: ${err_res}')
                    }
                    return error('net.openssl SSLConn.socket_read_into_ptr, could not read. (${err_res})')
                }
            }
        }
    }

    // Dead code, to satisfy the compiler
    return error('net.openssl SSLConn.socket_read_into_ptr, unknown error')
}

pub fn (mut s SSLConn) read(mut buffer []u8) !int {
    $if trace_ssl ? {
        eprintln('${@METHOD} buffer.len: ${buffer.len}')
    }
    return s.socket_read_into_ptr(&u8(buffer.data), buffer.len)
}

// write_ptr writes `len` bytes from `bytes` to the ssl connection
pub fn (mut s SSLConn) write_ptr(bytes &u8, len int) !int {
    mut total_sent := 0
    $if trace_ssl ? {
        defer(fn) {
            eprintln('${@METHOD} total_sent: ${total_sent}, bytes: ${voidptr(bytes):x}, len: ${len}, hex: ${unsafe { bytes.vbytes(len).hex() }}, data:-=-=-=-\n${unsafe { bytes.vstring_with_len(len) }}\n-=-=-=-')
        }
    }

    s.last_write_sent = 0
    deadline := ssl_timeout_deadline(s.duration)
    unsafe {
        mut ptr_base := bytes
        for total_sent < len {
            ptr := ptr_base + total_sent
            remaining := len - total_sent
            mut sent := C.SSL_write(voidptr(s.ssl), ptr, remaining)
            if sent <= 0 {
                // SSL_write did not fully complete: OpenSSL may already have
                // flushed one or more complete records (which the peer can
                // decrypt and act on) before returning a retryable error, so the
                // sent count is no longer provable. Mark it indeterminate; a
                // later full success below resets it to the exact length.
                s.last_write_sent = -1
                err_res := ssl_error(sent, s.ssl)!
                if err_res == .ssl_error_want_read {
                    s.wait_for_read(ssl_remaining_timeout(deadline))!
                    continue
                } else if err_res == .ssl_error_want_write {
                    s.wait_for_write(ssl_remaining_timeout(deadline))!
                    continue
                } else if err_res == .ssl_error_zero_return {
                    $if trace_ssl ? {
                        eprintln('${@METHOD} ---> res: ssl write on closed connection .ssl_error_zero_return')
                    }
                    return error('net.openssl SSLConn.write_ptr, ssl write on closed connection') // TODO: error_with_code close
                }
                $if trace_ssl ? {
                    eprintln('${@METHOD} ---> res: could not write SSL, err_res: ${err_res}')
                }
                return error_with_code('net.openssl SSLConn.write_ptr, could not write. (${err_res}),err',
                    int(err_res))
            }
            total_sent += sent
            s.last_write_sent = total_sent
        }
    }
    return total_sent
}

// write writes data from `bytes` to the ssl connection
pub fn (mut s SSLConn) write(bytes []u8) !int {
    return s.write_ptr(&u8(bytes.data), bytes.len)!
}

// write_string writes a string to the ssl connection
pub fn (mut s SSLConn) write_string(str string) !int {
    $if trace_ssl ? {
        eprintln('${@METHOD} str: ${str}')
    }
    return s.write_ptr(str.str, str.len)
}

// Select waits for an io operation (specified by parameter `test`) to be available
fn select(handle int, test Select, timeout time.Duration) !bool {
    $if trace_ssl ? {
        eprintln('${@METHOD} handle: ${handle}, timeout: ${timeout}')
    }
    set := C.fd_set{}
    C.FD_ZERO(&set)
    C.FD_SET(handle, &set)

    is_infinite := timeout <= 0 || timeout == net.infinite_timeout
    deadline := ssl_timeout_deadline(timeout)
    mut remaining_time := if is_infinite { i64(0) } else { timeout.milliseconds() }
    for is_infinite || remaining_time > 0 {
        seconds := remaining_time / 1000
        microseconds := (remaining_time % 1000) * 1000

        tt := C.timeval{
            tv_sec:  u64(seconds)
            tv_usec: u64(microseconds)
        }
        timeval_timeout := if is_infinite {
            &C.timeval(unsafe { nil })
        } else {
            &tt
        }

        mut res := -1
        match test {
            .read {
                res = net.socket_error(C.select(handle + 1, &set, C.NULL, C.NULL, timeval_timeout))!
            }
            .write {
                res = net.socket_error(C.select(handle + 1, C.NULL, &set, C.NULL, timeval_timeout))!
            }
            .except {
                res = net.socket_error(C.select(handle + 1, C.NULL, C.NULL, &set, timeval_timeout))!
            }
        }

        if res < 0 {
            if C.errno == C.EINTR {
                // errno is 4, Spurious wakeup from signal, keep waiting
                if !is_infinite {
                    remaining_time = ssl_remaining_timeout(deadline).milliseconds()
                }
                continue
            }
            cerr := C.errno
            return error_with_code('net.openssl select, failed: ${res}', cerr)
        } else if res == 0 {
            return net.err_timed_out
        }

        res = C.FD_ISSET(handle, &set)
        $if trace_ssl ? {
            eprintln('${@METHOD} ---> res: ${res}')
        }
        return res != 0
    }

    return net.err_timed_out
}

// wait_for wraps the common wait code
fn wait_for(handle int, what Select, timeout time.Duration) ! {
    ready := select(handle, what, timeout)!
    if ready {
        return
    }

    return net.err_timed_out
}

// wait_for_write waits for a write io operation to be available
fn (mut s SSLConn) wait_for_write(timeout time.Duration) ! {
    return wait_for(s.handle, .write, timeout)
}

// wait_for_read waits for a read io operation to be available
fn (mut s SSLConn) wait_for_read(timeout time.Duration) ! {
    return wait_for(s.handle, .read, timeout)
}