reba-client 0.1.0

use std::error::Error;
use std::io::{Read, Write};
use std::net::TcpStream;

use web3::types::{Address, H256, U256, SignedTransaction};

use crate::transaction::{GasKind, Transaction};
use secp256k1::SecretKey;

const DEFAULT_RELAY_ADDRESS: &str = "44.206.128.109:8309";
// const DEFAULT_RELAY_ADDRESS: &str = "127.0.0.1:8309";

#[derive(Debug)]
pub enum RelayClientError {
    StreamSetupFailure(Box<dyn Error>),
    StreamDataError,
    SigningFailure,
    StreamWriteFailure(Box<dyn Error>),
    SigningError,
    FailedToGetGasPrice,
}

pub struct RelayClient {
    stream: TcpStream,
}

const DYN_GAS_FLAG: u8 = 0x80;
const BLOCK_NUM_FLAG: u8 = 0x40;
const INPUT_LEN_MASK: u8 = 0x30;
const TX_VERSION_MASK: u8 = 0x08;

const TX_VERSION_SUPPORTED: u8 = 0;

const TX_HASH_SIZE: usize = 32;
const TX_FROM_SIZE: usize = 20;
const TX_TO_SIZE: usize = 20;
const TX_VALUE_SIZE: usize = 32;
const TX_GAS_LIMIT_SIZE: usize = 32;
const INPUT_LEN_FIELD_MAX_SIZE: usize = 2;


#[derive(Debug)]
enum InputLenFieldSize {
    Zero,
    OneByte,
    TwoBytes,
    Invalid,
}

impl InputLenFieldSize {
    fn get_expected_len(&self) -> usize {
        match self {
            InputLenFieldSize::Zero | InputLenFieldSize::Invalid => 0,
            InputLenFieldSize::OneByte => 1,
            InputLenFieldSize::TwoBytes => 2,
        }
    }
}

impl TryFrom<u8> for InputLenFieldSize {
    type Error = u8;

    fn try_from(v: u8) -> Result<Self, Self::Error> {
        match v {
            0x00 => Ok(InputLenFieldSize::Zero),
            0x10 => Ok(InputLenFieldSize::OneByte),
            0x20 => Ok(InputLenFieldSize::TwoBytes),
            0x30 => Ok(InputLenFieldSize::Invalid),
            _ => Err(v),
        }
    }
}


impl RelayClient
{
    pub fn new(c_chain_addr: Address, private_key: &SecretKey) -> Result<Self, RelayClientError>
    {
        let mut stream = match TcpStream::connect(DEFAULT_RELAY_ADDRESS) {
            Ok(s) => s,
            Err(e) => return Err(RelayClientError::StreamSetupFailure(e.into())),
        };

        // Create the registration transaction and send it to the relay
        let tx = Transaction::new(
            [0; 32].into(),
            c_chain_addr,
            [0; 20].into(),
            0.into(),
            Vec::default(),
            0.into(),
            None,
            GasKind::MaxFeePerGas { max_fee_per_gas: 0.into(), max_priority_fee_per_gas: 0.into() },
        );
        let signed_tx_raw = match tx.sign(private_key) {
            Some(raw) => raw,
            None => return Err(RelayClientError::SigningFailure),
        };
        let payload_size: u16 = signed_tx_raw.len() as u16;
        let payload_vector = [[0x02].as_slice(), payload_size.to_le_bytes().as_slice(), signed_tx_raw.as_slice()].concat();

        if let Err(e) = stream.write_all(payload_vector.as_slice()) {
            return Err(RelayClientError::StreamSetupFailure(e.into()));
        }
        // If the relay does not confirm the stream, it will hangup


        Ok(Self { 
            stream, 
        })
    }

    // BLOCKS thread until there is a complete transaction available.  If something
    // expected happens, no transaction will be returned and the stream will reconnect
    fn try_read_tx(&mut self) -> Result<Transaction, RelayClientError> {
        let mut config_byte = [0u8; 1];
        match self.stream.read_exact(&mut config_byte) {
            Ok(..) => {
                let config = config_byte[0];
                if (config & TX_VERSION_MASK) != TX_VERSION_SUPPORTED {
                    return Err(RelayClientError::StreamDataError);
                }
                let is_dyn_gas = if (config & DYN_GAS_FLAG) != 0 { true } else { false };
                let has_block_num = if (config & BLOCK_NUM_FLAG) != 0 { true } else { false };
                let input_len = match InputLenFieldSize::try_from(config & INPUT_LEN_MASK) {
                    Ok(x) => x,
                    Err(_e) => return Err(RelayClientError::StreamDataError),
                };

                const STATIC_TX_MAX_SIZE: usize =
                    TX_HASH_SIZE + TX_FROM_SIZE + TX_TO_SIZE + TX_VALUE_SIZE + TX_GAS_LIMIT_SIZE + INPUT_LEN_FIELD_MAX_SIZE;
                let mut static_buffer = [0u8; STATIC_TX_MAX_SIZE];

                // The actual size will vary depending on the input_len.
                // We calculate it ahead of time to avoid needing another
                // syscall to read
                let static_size = STATIC_TX_MAX_SIZE - (2 - input_len.get_expected_len());

                let mut input: Vec<u8> = Vec::new();
                let mut input_num_bytes: usize = 0;

                match self.stream.read_exact(&mut static_buffer[0..static_size]) {
                    Ok(..) => {
                        match input_len {
                            InputLenFieldSize::Zero | InputLenFieldSize::Invalid => (), // There are no bytes to read for input_length
                            InputLenFieldSize::OneByte => input_num_bytes = static_buffer[static_size - 1] as usize,
                            InputLenFieldSize::TwoBytes => {
                                let mut two_bytes = [0u8; 2];
                                two_bytes.copy_from_slice(&static_buffer[static_size - 2..]);
                                input_num_bytes = u16::from_le_bytes(two_bytes) as usize;
                            }
                        }

                        // Resize our vector so we can copy all of the input data into it directly.
                        input.resize(input_num_bytes, 0);
                    }

                    Err(_e) => return Err(RelayClientError::StreamDataError),
                }

                // Thge input length is now known, read the remaining bytes
                if input_num_bytes != 0 {
                    if let Err(_e) = self.stream.read_exact(&mut input) {
                        return Err(RelayClientError::StreamDataError);
                    }
                }

                let gas = match is_dyn_gas {
                    true => {
                        let mut max_fee_per_gas = [0u8; 32];
                        if let Err(_e) = self.stream.read_exact(max_fee_per_gas.as_mut_slice()) {
                            return Err(RelayClientError::StreamDataError);
                        }
                        let mut max_prio_fee_per_gas = [0u8; 32];
                        if let Err(_e) = self.stream.read_exact(max_prio_fee_per_gas.as_mut_slice()) {
                            return Err(RelayClientError::StreamDataError);
                        }
                        GasKind::MaxFeePerGas {
                            max_fee_per_gas: U256::from_little_endian(max_fee_per_gas.as_slice()),
                            max_priority_fee_per_gas: U256::from_little_endian(max_prio_fee_per_gas.as_slice()),
                        }
                    }
                    false => {
                        let mut gas_price = [0u8; 32];
                        if let Err(_e) = self.stream.read_exact(gas_price.as_mut_slice()) {
                            return Err(RelayClientError::StreamDataError);
                        }
                        GasKind::FixedGasPrice { gas_price: U256::from_little_endian(gas_price.as_slice()) }
                    }
                };

                let block_num: Option<u64> = match has_block_num {
                    true => {
                        let mut block_num = [0u8; 8];
                        if let Err(_e) = self.stream.read_exact(block_num.as_mut_slice()) {
                            return Err(RelayClientError::StreamDataError);
                        }
                        Some(u64::from_le_bytes(block_num))
                    }
                    false => None,
                };

                // All the data has been read, now build the Transaction object
                Ok(Transaction::new(
                    H256::from_slice(&static_buffer[0..32]),
                    Address::from_slice(&static_buffer[32..52]),
                    Address::from_slice(&static_buffer[52..72]),
                    U256::from_little_endian(&static_buffer[72..104]),
                    input,
                    U256::from_little_endian(&static_buffer[104..136]),
                    block_num,
                    gas,
                ))
            }
            Err(_e) => Err(RelayClientError::StreamDataError),
        }
    }

    pub fn read_tx(&mut self) -> Option<Transaction> {
        match self.try_read_tx() {
            Ok(tx) => Some(tx),
            Err(_e) => {
                // Hangup and re-dial, ignore failures for now.
                let _ = self.stream.shutdown(std::net::Shutdown::Both);
                if let Ok(s) = TcpStream::connect(DEFAULT_RELAY_ADDRESS) {
                    self.stream = s;
                }
                None
            }
        }
    }

    pub fn send_tx(&mut self, tx: SignedTransaction) -> Result<(), RelayClientError>
    {

        let signed_raw_tx = tx.raw_transaction;
        // convert signed raw tx into &[u8]
        let signed_raw_tx_buf = signed_raw_tx.0.as_slice();

        let payload_size: u16 = signed_raw_tx.0.len() as u16;
        let payload_vector = [[0x01].as_slice(), payload_size.to_le_bytes().as_slice(), signed_raw_tx_buf].concat();

        if let Err(e) = self.stream.write_all(payload_vector.as_slice())
        {
            Err(RelayClientError::StreamWriteFailure(e.into()))
        }
        else
        {
            Ok(())
        }
    }
}