lib.rs

use std::{
    fmt,
    net::{Ipv4Addr, Ipv6Addr},
    str::FromStr,
};

#[derive(Clone, Debug, PartialEq)]
pub enum Types {
    IPv4,
    IPv6,
    Alias,
    Text,
}

#[derive(Clone, Debug, PartialEq)]
pub enum Content {
    Address(String),
    Data(String),
}

impl FromStr for Types {
    type Err = &'static str;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        match s {
            "ipv4" => Ok(Types::IPv4),
            "ipv6" => Ok(Types::IPv6),
            "alias" => Ok(Types::Alias),
            "txt" => Ok(Types::Text),
            _ => Err("invalid type"),
        }
    }
}

impl fmt::Display for Types {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        let s = match self {
            Types::IPv4 => "ipv4",
            Types::IPv6 => "ipv6",
            Types::Alias => "alias",
            Types::Text => "txt",
        };
        write!(f, "{}", s)
    }
}

#[derive(Clone, Debug, PartialEq)]
pub struct Request {
    pub name: String,
    pub record: Types,
}

#[derive(Clone, Debug, PartialEq)]
pub struct Response {
    pub name: String,
    pub record: Types,
    pub ttl: u64,
    pub len: u32,
    pub content: Content,
}

impl ToString for Request {
    fn to_string(&self) -> String { format!("NAME {}\nTYPE {}", self.name, self.record) }
}

impl ToString for Response {
    fn to_string(&self) -> String {
        match &self.content {
            Content::Address(addr) => format!("NAME {}\nTYPE {}\nTTL {}\nLEN {}\nADDR {}", self.name, self.record, self.ttl, self.len, addr),
            Content::Data(data) => format!("NAME {}\nTYPE {}\nTTL {}\nLEN {}\nDATA {}", self.name, self.record, self.ttl, self.len, data),
        }
    }
}

impl Request {
    pub fn to_bytes(&self) -> Vec<u8> { self.to_string().into_bytes() }
}

impl Response {
    pub fn to_bytes(&self) -> Vec<u8> { self.to_string().into_bytes() }

    // testing response
    pub fn new(r: &Request) -> Self {
        Response {
            ttl: 10000,
            name: r.name.clone(),
            record: r.record.clone(),
            len: r.name.chars().count() as u32,
            content: Content::Address("127.0.0.1:7878".to_string()),
        }
    }
}

impl FromStr for Request {
    type Err = &'static str;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        let mut name = None;
        let mut record = None;

        for line in s.lines() {
            let mut parts = line.splitn(2, " ");
            let key = parts.next().ok_or("missing key")?;
            let value = parts.next().ok_or("missing value")?;

            match key {
                "NAME" => name = Some(value.to_string()),
                "TYPE" => record = Some(value.parse()?),
                _ => return Err("unexpected key"),
            }
        }

        Ok(Request {
            name: name.ok_or("missing name")?,
            record: record.ok_or("missing type")?,
        })
    }
}

impl FromStr for Response {
    type Err = &'static str;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        let mut name = None;
        let mut record = None;
        let mut ttl = None;
        let mut len = None;
        let mut content = (None, None);

        for line in s.lines() {
            let mut parts = line.splitn(2, " ");
            let key = parts.next().ok_or("missing key")?;
            let value = parts.next().ok_or("missing value")?;

            match key {
                "NAME" => name = Some(value.to_string()),
                "TYPE" => record = Some(value.parse()?),
                "TTL" => ttl = Some(value.parse().map_err(|_| "invalid TTL")?),
                "LEN" => len = Some(value.parse().map_err(|_| "invalid LEN")?),
                "ADDR" => content.0 = Some(value.to_string()),
                "DATA" => content.1 = Some(value.to_string()),
                _ => return Err("unexpected key"),
            }
        }

        let kind: Types = record.ok_or("missing TYPE")?;

        match kind {
            Types::IPv4 => {
                let addr = content.0.clone().ok_or("missing ADDR")?;
                if addr.parse::<Ipv4Addr>().is_err() {
                    return Err("invalid IPv4 address");
                }
            }
            Types::IPv6 => {
                let addr = content.0.clone().ok_or("missing ADDR")?;
                if addr.parse::<Ipv6Addr>().is_err() {
                    return Err("invalid IPv6 address");
                }
            }
            _ => {}
        };

        match content.0 {
            Some(addr) => Ok(Response {
                record: kind,
                content: Content::Address(addr),
                name: name.ok_or("missing NAME")?,
                ttl: ttl.ok_or("missing TTL")?,
                len: len.ok_or("missing LEN")?,
            }),
            None => match content.1 {
                Some(data) => Ok(Response {
                    record: kind,
                    content: Content::Data(data),
                    name: name.ok_or("missing NAME")?,
                    ttl: ttl.ok_or("missing TTL")?,
                    len: len.ok_or("missing LEN")?,
                }),
                None => return Err("unexpected key"),
            },
        }
    }
}

impl From<Vec<u8>> for Request {
    fn from(bytes: Vec<u8>) -> Self {
        let s = String::from_utf8(bytes).expect("Invalid UTF-8 data");
        s.parse().unwrap()
    }
}

impl From<Vec<u8>> for Response {
    fn from(bytes: Vec<u8>) -> Self {
        let s = String::from_utf8(bytes).expect("Invalid UTF-8 data");
        s.parse().unwrap()
    }
}

main.rs

use async_std::{net::UdpSocket, task};
use dns::{Request, Response};
use std::error::Error;

async fn serve(address: &str) -> Result<(), Box<dyn Error + Send + Sync>> {
    let socket = UdpSocket::bind(address).await?;
    println!("Server listening on {}", socket.local_addr()?);

    let mut buf = [0u8; 1024];
    loop {
        let (n, src) = socket.recv_from(&mut buf).await?;
        println!("Received {} bytes from {}", n, src);

        let response = buf[..n].to_vec();

        task::spawn(async move {
            let socket = UdpSocket::bind("0.0.0.0:0").await.expect("Failed to bind socket");
            let request: Request = response.try_into().expect("Failed to parse request");
            let data = Response::new(&request);

            if let Err(err) = socket.send_to(&data.to_bytes(), &src).await {
                eprintln!("Failed to send response: {}", err);
            }
        });
    }
}

#[async_std::main]
async fn main() -> Result<(), Box<dyn Error + Send + Sync>> {
    serve("127.0.0.1:5533").await?;
    Ok(())
}

#[cfg(test)]
mod tests {
    use std::net::UdpSocket;

    #[test]
    fn mock_request() {
        let socket = UdpSocket::bind("0.0.0.0:0").expect("Failed to bind socket");
        let server_address = "127.0.0.1:5533";

        let request = dns::Request {
            name: "example.com".to_string(),
            record: dns::Types::IPv4,
        };

        socket.send_to(&request.to_bytes(), server_address).expect("Failed to send data");

        println!("Data sent to server");

        let mut buf = [0; 1024];
        let (amt, _) = socket.recv_from(&mut buf).expect("Failed to receive data");

        println!("Received {} bytes from server: \n{}", amt, String::from_utf8_lossy(&buf[..amt]));
    }
}

protocol.ts

import { protocol, net, webContents } from 'electron';
import { join } from 'path';
import { parse } from 'url';
import { createSocket } from 'dgram';

import { ERROR_PROTOCOL, WEBUI_PROTOCOL } from '~/constants/files';

protocol.registerSchemesAsPrivileged([
	{
		scheme: 'browser',
		privileges: {
			bypassCSP: true,
			secure: true,
			standard: true,
			supportFetchAPI: true,
			allowServiceWorkers: true,
			corsEnabled: false
		}
	}
]);

protocol.registerSchemesAsPrivileged([
	{
		scheme: 'meow',
		privileges: {
			secure: true,
			standard: true,
			supportFetchAPI: true,
			allowServiceWorkers: true,
			corsEnabled: true
		}
	}
]);

const parseInput = (s) => {
	const result = {
		name: null,
		record: null,
		ttl: null,
		len: null,
		content: { addr: null, data: null }
	};

	s.split('\n').forEach((line) => {
		let [key, value] = line.split(' ', 2);
		if (!key || !value) {
			throw new Error('missing key or value');
		}

		switch (key) {
			case 'NAME':
				result.name = value;
				break;
			case 'TYPE':
				result.record = parseInt(value);
				break;
			case 'TTL':
				result.ttl = parseInt(value);
				break;
			case 'LEN':
				result.len = parseInt(value);
				break;
			case 'ADDR':
				result.content.addr = value;
				break;
			case 'DATA':
				result.content.data = value;
				break;
			default:
				throw new Error('unexpected key');
		}
	});

	return result;
};

export const registerProtocol = (session: Electron.Session) => {
	session.protocol.handle('meow', async (request) => {
		let url = '';

		const parsed = parse(request.url);
		const client = createSocket('udp4');
		const data = Buffer.from(`NAME ${parsed.hostname}\nTYPE ipv4`);

		const dnsResolutionPromise = new Promise((resolve, reject) => {
			client.on('message', (msg, info) => {
				const dns = parseInput(msg.toString());
				url = `http://${dns.content.addr}${parsed.path || ''}`;
				resolve(url);
				client.close();
			});

			client.send(data, 5533, 'localhost', (error) => {
				if (error) {
					client.close();
					reject(error);
				}
			});
		});

		const resolvedUrl = await dnsResolutionPromise;
		return session.fetch(resolvedUrl);
	});

	session.protocol.registerFileProtocol(ERROR_PROTOCOL, (request, callback: any) => {
		const parsed = parse(request.url);

		if (parsed.hostname === 'network-error') {
			return callback({
				path: join(__dirname, '../static/pages/', `network-error.html`)
			});
		}
	});

	if (process.env.NODE_ENV !== 'development') {
		session.protocol.registerFileProtocol(WEBUI_PROTOCOL, (request, callback: any) => {
			const parsed = parse(request.url);

			if (parsed.path === '/') {
				return callback({
					path: join(__dirname, `${parsed.hostname}.html`)
				});
			}

			callback({ path: join(__dirname, parsed.path) });
		});
	}
};