tcpdns.rs

use std::io::{Read, Write, Result as IoResult};
use std::net::{UdpSocket, TcpStream};

const HOST: &str = "127.0.0.1";
const PORT: u16 = 5354;
const UPSTREAM: &str = "114.114.114.114";
const UPSTREAM_PORT: u16 = 53;

fn be_bytes_to_u16(b: &[u8;2]) -> u16 {
    return ((b[0] as u16) << 8) | (b[1] as u16);
}

fn handle_dns(query: &[u8]) -> IoResult<Vec<u8>>{
    match TcpStream::connect(format!("{}:{}", UPSTREAM, UPSTREAM_PORT)) {
        Ok(mut stream) => {
            stream.write(&[0, query.len() as u8])?;
            stream.write(query)?;
            let mut len = [0; 2];
            stream.read(&mut len)?;
            let len = be_bytes_to_u16(&len);
            let mut result = vec![0u8;len as usize];
            stream.read(&mut result)?;
            Ok(result)
        }
        Err(err) => Err(err)
    }
}

fn main() {
    let socket = UdpSocket::bind(format!("{}:{}", HOST, PORT)).expect("failed to bind address");
    loop {
        let mut buf = [0; 512];
        let (n, src) = socket.recv_from(&mut buf).expect("failed to recv");
        let buf = &mut buf[..n];
        // parse dns query id
        println!("{:}", be_bytes_to_u16(&[buf[0], buf[1]]));

        let result = handle_dns(&buf);
        if result.is_ok() {
            socket.send_to(&result.unwrap(), &src).expect("failed to send");
        }
    };
}