From 2a96c5f985f905d0d7d988fbccfeb1dc008e9f07 Mon Sep 17 00:00:00 2001
From: Frank Denis <github@pureftpd.org>
Date: Mon, 20 Apr 2020 15:44:42 +0200
Subject: [PATCH] dafuq

---
 src/dns.rs      | 45 +++++++++++++++++++++++++++++++++++++++++++++
 src/main.rs     |  3 ++-
 src/resolver.rs | 15 ++++++++++++++-
 3 files changed, 61 insertions(+), 2 deletions(-)

diff --git a/src/dns.rs b/src/dns.rs
index 3c41721..e767f9d 100644
--- a/src/dns.rs
+++ b/src/dns.rs
@@ -2,6 +2,7 @@ use crate::dnscrypt_certs::*;
 use crate::errors::*;
 
 use byteorder::{BigEndian, ByteOrder, WriteBytesExt};
+use std::net::{Ipv4Addr, Ipv6Addr};
 use std::sync::Arc;
 
 pub const DNS_MAX_HOSTNAME_SIZE: usize = 256;
@@ -620,3 +621,47 @@ pub fn serve_blocked_response(client_packet: Vec<u8>) -> Result<Vec<u8>, Error>
     packet.extend_from_slice(hinfo_rdata);
     Ok(packet)
 }
+
+pub fn serve_a_response(client_packet: Vec<u8>, ip: Ipv4Addr) -> Result<Vec<u8>, Error> {
+    ensure!(client_packet.len() >= DNS_HEADER_SIZE, "Short packet");
+    ensure!(qdcount(&client_packet) == 1, "No question");
+    ensure!(
+        !is_response(&client_packet),
+        "Question expected, but got a response instead"
+    );
+    let offset = skip_name(&client_packet, DNS_HEADER_SIZE)?;
+    let mut packet = client_packet;
+    ensure!(packet.len() - offset >= 4, "Short packet");
+    packet.truncate(offset + 4);
+    an_ns_ar_count_clear(&mut packet);
+    authoritative_response(&mut packet);
+    ancount_inc(&mut packet)?;
+    packet.write_u16::<BigEndian>(0xc000 + DNS_HEADER_SIZE as u16)?;
+    packet.write_u16::<BigEndian>(DNS_TYPE_A)?;
+    packet.write_u16::<BigEndian>(DNS_CLASS_INET)?;
+    packet.write_u32::<BigEndian>(60)?;
+    packet.extend_from_slice(&ip.octets());
+    Ok(packet)
+}
+
+pub fn serve_aaaa_response(client_packet: Vec<u8>, ip: Ipv6Addr) -> Result<Vec<u8>, Error> {
+    ensure!(client_packet.len() >= DNS_HEADER_SIZE, "Short packet");
+    ensure!(qdcount(&client_packet) == 1, "No question");
+    ensure!(
+        !is_response(&client_packet),
+        "Question expected, but got a response instead"
+    );
+    let offset = skip_name(&client_packet, DNS_HEADER_SIZE)?;
+    let mut packet = client_packet;
+    ensure!(packet.len() - offset >= 4, "Short packet");
+    packet.truncate(offset + 4);
+    an_ns_ar_count_clear(&mut packet);
+    authoritative_response(&mut packet);
+    ancount_inc(&mut packet)?;
+    packet.write_u16::<BigEndian>(0xc000 + DNS_HEADER_SIZE as u16)?;
+    packet.write_u16::<BigEndian>(DNS_TYPE_AAAA)?;
+    packet.write_u16::<BigEndian>(DNS_CLASS_INET)?;
+    packet.write_u32::<BigEndian>(60)?;
+    packet.extend_from_slice(&ip.octets());
+    Ok(packet)
+}
diff --git a/src/main.rs b/src/main.rs
index 6d06248..9050b90 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -224,7 +224,8 @@ async fn handle_client_query(
             Some(token) => ensure!(tokens.contains(&token), "Access token not found"),
         }
     }
-    let response = resolver::get_cached_response_or_resolve(&globals, &mut packet).await?;
+    let response =
+        resolver::get_cached_response_or_resolve(&globals, &client_ctx, &mut packet).await?;
     encrypt_and_respond_to_query(
         globals,
         client_ctx,
diff --git a/src/resolver.rs b/src/resolver.rs
index f0758a9..1f447aa 100644
--- a/src/resolver.rs
+++ b/src/resolver.rs
@@ -2,12 +2,13 @@ use crate::cache::*;
 use crate::dns::{self, *};
 use crate::errors::*;
 use crate::globals::*;
+use crate::ClientCtx;
 
 use byteorder::{BigEndian, ByteOrder};
 use rand::prelude::*;
 use siphasher::sip128::Hasher128;
 use std::hash::Hasher;
-use std::net::{Ipv4Addr, Ipv6Addr, SocketAddr, SocketAddrV4, SocketAddrV6};
+use std::net::{IpAddr, Ipv4Addr, Ipv6Addr, SocketAddr, SocketAddrV4, SocketAddrV6};
 use tokio::net::{TcpStream, UdpSocket};
 use tokio::prelude::*;
 
@@ -164,9 +165,21 @@ pub async fn resolve(
 
 pub async fn get_cached_response_or_resolve(
     globals: &Globals,
+    client_ctx: &ClientCtx,
     mut packet: &mut Vec<u8>,
 ) -> Result<Vec<u8>, Error> {
     let packet_qname = dns::qname(&packet)?;
+    if &packet_qname == b"my.ip" {
+        let client_ip = match client_ctx {
+            ClientCtx::Udp(u) => u.client_addr,
+            ClientCtx::Tcp(t) => t.client_connection.peer_addr()?,
+        }
+        .ip();
+        return match client_ip {
+            IpAddr::V4(ip) => serve_a_response(packet.to_vec(), ip),
+            IpAddr::V6(ip) => serve_aaaa_response(packet.to_vec(), ip),
+        };
+    }
     if let Some(blacklist) = &globals.blacklist {
         if blacklist.find(&packet_qname) {
             #[cfg(feature = "metrics")]