package chat.vojo.proxy.core.upstream

import chat.vojo.proxy.core.ProxyConfig
import chat.vojo.proxy.core.crypto.Aead
import chat.vojo.proxy.core.crypto.ShadowsocksCrypto
import chat.vojo.proxy.core.crypto.SsInputStream
import chat.vojo.proxy.core.crypto.SsOutputStream
import chat.vojo.proxy.core.crypto.ssUdpOpen
import chat.vojo.proxy.core.crypto.ssUdpSeal
import chat.vojo.proxy.core.net.Destination
import chat.vojo.proxy.core.net.parseSocksAddress
import java.net.DatagramPacket
import java.net.DatagramSocket
import java.net.InetAddress
import java.net.InetSocketAddress

/** Upstream that tunnels through a Shadowsocks AEAD server (SIP004). */
class ShadowsocksUpstream(
    private val cfg: ProxyConfig,
    private val protector: SocketProtector,
) : Upstream {

    private val aead: Aead = Aead.of(cfg.cipher)
    private val masterKey: ByteArray = ShadowsocksCrypto.deriveMasterKey(cfg.password, cfg.cipher.keySize)

    override fun connectTcp(dest: Destination): UpstreamConn {
        val socket = protector.connectedSocket(cfg.host, cfg.port)
        try {
            val ssOut = SsOutputStream(socket.getOutputStream(), aead, masterKey)
            val ssIn = SsInputStream(socket.getInputStream(), aead, masterKey)
            // First plaintext bytes are the target address header.
            ssOut.write(dest.raw)
            ssOut.flush()
            return UpstreamConn(socket, ssIn, ssOut)
        } catch (e: Throwable) {
            runCatching { socket.close() }
            throw e
        }
    }

    override fun openUdpAssociation(): UdpAssociation = ShadowsocksUdpAssociation(cfg, protector, aead, masterKey)
}

/**
 * Shadowsocks UDP relay. One socket to the SS server carries every destination;
 * each datagram is independently `[salt][AEAD([ADDR][DATA], nonce=0)]`.
 */
private class ShadowsocksUdpAssociation(
    cfg: ProxyConfig,
    protector: SocketProtector,
    private val aead: Aead,
    private val masterKey: ByteArray,
) : UdpAssociation {

    private val server = InetSocketAddress(InetAddress.getByName(cfg.host), cfg.port)
    private val udp = DatagramSocket().also { protector.protect(it) }

    override fun send(dest: Destination, data: ByteArray, off: Int, len: Int) {
        val plain = ByteArray(dest.raw.size + len)
        System.arraycopy(dest.raw, 0, plain, 0, dest.raw.size)
        System.arraycopy(data, off, plain, dest.raw.size, len)
        val sealed = ssUdpSeal(aead, masterKey, plain, 0, plain.size)
        udp.send(DatagramPacket(sealed, sealed.size, server))
    }

    override fun receive(): UdpReply? {
        val buf = ByteArray(64 * 1024)
        // Skip undecryptable/malformed datagrams (spoofed or corrupt) instead of
        // letting one bad packet kill the reply loop.
        while (true) {
            val dp = DatagramPacket(buf, buf.size)
            udp.receive(dp)
            try {
                val plain = ssUdpOpen(aead, masterKey, buf, 0, dp.length)
                val (dest, headerLen) = parseSocksAddress(plain, 0)
                return UdpReply(dest, plain.copyOfRange(headerLen, plain.size))
            } catch (e: Exception) {
                continue
            }
        }
    }

    override fun close() {
        runCatching { udp.close() }
    }
}
