nonymous 0.7.0

/// ```sh
/// $ yt-dlp -o '%(id)s.%(ext)s' 'https://www.youtube.com/watch?v=FtutLA63Cp8'
///                                # width:height        fps
/// $ ffmpeg -i FtutLA63Cp8.* -vf scale=48:18,hue=s=0 -r 20 -c:v rawvideo -pix_fmt rgb32 input.rgb
///                                             # listen host/port  path  width height fps
/// $ cargo run --features alloc --example badapple 127.0.0.1 5353 input.rgb 48 18 20
/// ```
/// ```sh
/// $ dig -p 5353 @127.0.0.1 {00000..99999}.test
/// ```
use std::{
    collections::HashMap,
    fs::read,
    net::IpAddr,
    path::PathBuf,
    time::{Duration, Instant},
};

use arrayvec::ArrayVec;
use clap::Parser;
use eyre::bail;
use nonymous::{
    core::{Rcode, Type},
    emit::Buffer,
    fmt::Plain,
    name::NameBuf,
    server::response,
    view::{Message, View},
};
use tokio::net::UdpSocket;

const BYTES_PER_PIXEL: usize = 4;
const VIDEO_DURATION: Duration = Duration::from_secs(60 * 3 + 40);
const MAX_PLAYERS_PER_IP: usize = 3;
const MAX_PLAYERS: usize = 100;

#[derive(Parser)]
struct Args {
    host: String,
    port: u16,
    rgb_path: PathBuf,
    width: usize,
    height: usize,
    fps: usize,
    #[arg(long)]
    debug: bool,
}

#[derive(Clone, Debug, Eq, Hash, PartialEq)]
struct PlayerId(IpAddr, NameBuf);
#[derive(Debug)]
struct PlayerInfo {
    start_time: Instant,
}

#[tokio::main]
async fn main() -> eyre::Result<()> {
    let args = Args::parse();
    let rgb = read(args.rgb_path)?;
    let socket = UdpSocket::bind((args.host, args.port)).await?;
    let mut players: HashMap<PlayerId, PlayerInfo> = HashMap::default();

    loop {
        // remove expired players
        players.retain(|_id, player| player.start_time.elapsed() < VIDEO_DURATION);

        let mut query_buf: ArrayVec<u8, 512> = ArrayVec::new();
        query_buf.resize_zero(query_buf.capacity());
        let (len, remote_addr) = socket.recv_from(&mut query_buf).await?;
        query_buf.resize_zero(len);

        let mut response_buf: ArrayVec<u8, 4096> = ArrayVec::new();
        let Some((query, response)) = response(&query_buf, &mut response_buf, true)? else {
            continue;
        };
        eprintln!(";;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;");
        eprintln!(";; <<< {remote_addr} ({} bytes)", query_buf.len());
        if args.debug {
            eprintln!("{}", Plain(&query));
        }

        let response = response.rcode(Rcode::ServFail);
        let handle_query = || -> eyre::Result<()> {
            let ip = remote_addr.ip();
            let Some(question) = query.qd().next() else {
                bail!("query has no question");
            };
            if question.qtype() != Type::A {
                bail!("expected query for A record");
            }
            let qname = question.qname();
            let Some(base_name) = qname.parent() else {
                bail!("qname has no parent");
            };
            let id = PlayerId(ip, base_name.into());
            if !players.contains_key(&id) {
                if players.len() > MAX_PLAYERS {
                    // reject the new player
                    bail!("too many players! try again later");
                }
                if players.keys().filter(|id| id.0 == ip).count() > MAX_PLAYERS_PER_IP {
                    // remove an arbitrary player for `ip`
                    let remove_id = players
                        .keys()
                        .find(|id| id.0 == ip)
                        .expect("guaranteed by containing branch")
                        .clone();
                    players.remove(&remove_id);
                }
                players.insert(
                    id.clone(),
                    PlayerInfo {
                        start_time: Instant::now(),
                    },
                );
            }
            let player = players.get(&id).expect("guaranteed by branch above");
            let t = player.start_time.elapsed();
            let frame_number = (t.as_secs_f64() * args.fps as f64) as usize;
            eprintln!(";; === players: {players:?}");
            eprintln!(";; === player id {id:?}, t = {t:?}, frame number {frame_number}");

            let mut response = response
                .question_with_name(&question.qname(), question.qtype(), question.qclass())?
                .into_ar();
            let offset = frame_number * BYTES_PER_PIXEL * args.width * args.height;
            let len = BYTES_PER_PIXEL * args.width * args.height;
            if rgb.len() < offset + len {
                bail!("out of bounds");
            }
            let rgb = &rgb[offset..][..len];
            for line in rgb.chunks(BYTES_PER_PIXEL * args.width) {
                let len = args.width.try_into()?;
                let mut record = response
                    .record_with_name(
                        &NameBuf::from_dotted(".")?.view(),
                        Type::TXT,
                        question.qclass(),
                    )?
                    .push_rdata(std::slice::from_ref(&len))?;
                for pixel in line.chunks(BYTES_PER_PIXEL) {
                    let [r, g, b, _a] = pixel else { unreachable!() };
                    let r = *r as f64 / 255.0;
                    let g = *g as f64 / 255.0;
                    let b = *b as f64 / 255.0;
                    // Rec. 709 gamma decoding
                    let linear_r = rec_709_gamma_decode(r);
                    let linear_g = rec_709_gamma_decode(g);
                    let linear_b = rec_709_gamma_decode(b);
                    // linear luminance (grayscale) using Rec. 709 coefficients
                    // (not luma, because the coordinates are linear)
                    let linear_luminance =
                        0.2126 * linear_r + 0.7152 * linear_g + 0.0722 * linear_b;
                    // Rec. 709 gamma encoding
                    if rec_709_gamma_encode(linear_luminance) > 0.5 {
                        record = record.push_rdata(b" ")?;
                    } else {
                        record = record.push_rdata(b"#")?;
                    }
                }
                response = record.finish()?;
            }
            response.rcode(Rcode::NoError).finish()?.finish();
            Ok(())
        };
        if let Err(error) = handle_query() {
            eprintln!(";; error: {error:?}");
        }

        eprintln!(";; >>> {remote_addr} ({} bytes)", response_buf.len());
        if args.debug {
            eprintln!("{}", Plain(&Message::view(&response_buf, ..)?.0));
        }
        socket.send_to(&response_buf, remote_addr).await?;
    }
}

fn rec_709_gamma_decode(v: f64) -> f64 {
    if v < 0.081 {
        v / 4.5
    } else {
        ((v + 0.099) / 1.099).powf(1. / 0.45)
    }
}

fn rec_709_gamma_encode(l: f64) -> f64 {
    if l < 0.018 {
        4.5 * l
    } else {
        1.099 * l.powf(0.45) - 0.099
    }
}