#![no_std]
#![no_main]

use aya_bpf::{
    bindings::xdp_action,
    macros::{map, xdp},
    maps::Array,
    programs::XdpContext,
};

use responder_ebpf::util::*;
use responder_ebpf::bindings::tcphdr;
use responder_common::*;

use core::mem;

const TCP_HDR_LEN: usize = mem::size_of::<tcphdr>();
const IPPROTO_TCP: u8 = 0x06;


#[inline(always)]
fn parse_tcphdr(ctx: &XdpContext, cursor: &mut usize) -> Option<*mut tcphdr> {
    let tcp = ptr_at_mut::<tcphdr>(&ctx, *cursor);
    if tcp.is_some() {
        *cursor += TCP_HDR_LEN;
    }
    tcp
}

#[xdp(name="responder")]
pub fn responder(ctx: XdpContext) -> u32 {
    match try_responder(ctx) {
        Ok(ret) => ret,
        Err(_) => xdp_action::XDP_ABORTED,
    }
}

#[inline(always)]
unsafe fn bounce_tcp(_ctx: &XdpContext, tcp: *mut tcphdr) {
    mem::swap(&mut (*tcp).source, &mut (*tcp).dest);
    mem::swap(&mut (*tcp).ack_seq, &mut (*tcp).seq); // Swap to keep checksum the same as much as possible
    let (ack_seq, o) = u32::from_be((*tcp).ack_seq).overflowing_add(1); // If overflow: 1's complement sum is unchanged
    (*tcp).ack_seq = u32::to_be(ack_seq);
    (*tcp).set_ack(1);
    (*tcp).check = !(ones_complement_add_u16(!(*tcp).check, (!o as u16) + (1 << 4)));
}

#[map(name = "FILTER_MAP")]
static FILTER_MAP: Array<bloom_filter::ChunkType> =
    Array::<bloom_filter::ChunkType>::with_max_entries(bloom_filter::MAP_SIZE as u32, 0);

#[inline(always)]
unsafe fn matches_filter(_ctx: &XdpContext, daddr: IpAddr) -> bool {
    match daddr {
        IpAddr::V4(daddr) => {
            let mut res = false;
            for hash_offset in 0..bloom_filter::HASH_COUNT {
                res = true;
                let hash = bloom_filter::hash(daddr, hash_offset);
                let map_i = hash >> (bloom_filter::ADDRESS_BITS_CHUNK as u32);
                let chunk_i = hash & (bloom_filter::ADDRESS_MASK_CHUNK as u32);
                // info!(ctx, "{:ipv4} {} {} {}",daddr, hash, map_i, chunk_i);
                let test = if let Some(b) = FILTER_MAP.get(map_i as u32) {
                    let word_i = chunk_i & (bloom_filter::ADDRESS_MASK_WORD as u32);
                    let chunk_i = (chunk_i as usize) >> bloom_filter::ADDRESS_BITS_WORD;
                    // info!(ctx, "{} [{}]", word_i, b[chunk_i]);
                    (b[chunk_i] >> (bloom_filter::ADDRESS_MASK_WORD as u32 - word_i)) & 1 == 1
                } else {
                    false
                };
                if !test {
                    return false
                }
            }
            res
        }
        IpAddr::V6(_daddr) => {
            false // TODO
        }
    }
}

fn try_responder(ctx: XdpContext) -> Result<xdp_action::Type, xdp_action::Type> {
    let mut hdr_cursor = 0usize;

    let (eth, ip) = unsafe {
        parse_routing(&ctx, &mut hdr_cursor)
            .ok_or(xdp_action::XDP_PASS)?
    };

    let protocol = unsafe { l3_get_protocol(&ip) };
    let daddr = unsafe { l3_get_daddr(&ip) };

    if is_local(daddr) {
        // info!(&ctx, "local: pass");
        return Ok(xdp_action::XDP_PASS);
    }

    if unsafe { !matches_filter(&ctx, daddr) } {
        return Ok(xdp_action::XDP_DROP);
    }

    if protocol != IPPROTO_TCP {
        return Ok(xdp_action::XDP_PASS);
    }

    let tcp = parse_tcphdr(&ctx, &mut hdr_cursor).ok_or(xdp_action::XDP_PASS)?;
    let tcp_syn = unsafe { (*tcp).syn() };
    let tcp_ack = unsafe { (*tcp).ack() };

    // match daddr {
    //     IpAddr::V4(ip) => info!(&ctx, "Received packet with matching daddr: {:ipv4}", ip),
    //     IpAddr::V6(ip) => unsafe { info!(&ctx, "Received packet with matching daddr: {:ipv6}", ip.in6_u.u6_addr8) }
    // }
    // info!(&ctx, "and tcp with syn: {}, ack: {}", tcp_syn, tcp_ack);

    if tcp_syn == 0 || tcp_ack != 0 {
        return Ok(xdp_action::XDP_PASS);
    }

    unsafe {
        bounce_routing(&ctx,eth, ip);
        bounce_tcp(&ctx, tcp);
    }

    Ok(xdp_action::XDP_PASS)
}

#[panic_handler]
fn panic(_info: &core::panic::PanicInfo) -> ! {
    unsafe { core::hint::unreachable_unchecked() }
}