src/docker.rs - master - n0dwis/dns-proxy-server

src/docker.rs
198 строк5 KB
Dmitry Vakulenko
Syntax fix
27 дек 2025, 19:57
27 дек 2025, 19:577e58a58
use std::{collections::HashMap, sync::Arc};

use bollard::{
    Docker,
    query_parameters::{
        EventsOptions, EventsOptionsBuilder, InspectContainerOptions, ListContainersOptions,
    },
    secret::ContainerInspectResponse,
};
use log::{debug, error};
use tokio::sync::Mutex;
use tokio_stream::StreamExt;

pub struct Resolver {
    ips: Mutex<HashMap<String, String>>,
}

impl Resolver {
    fn new() -> Self {
        Resolver {
            ips: Mutex::new(HashMap::new()),
        }
    }

    pub async fn get_ip<'a>(&'a self, name: &str) -> Option<String> {
        debug!("Requested IP for {}", name);
        self.ips
            .lock()
            .await
            .get(name.strip_suffix('.').unwrap_or(name))
            .cloned()
    }

    async fn add_container_ips(&self, container: ContainerInspectResponse) {
        let net_settings = match container.network_settings {
            Some(n) => n,
            None => return,
        };

        let networks = match net_settings.networks {
            Some(n) => n,
            None => return,
        };

        // let name = match container.name {
        //     Some(n) => n,
        //     None => {
        //         warn!("Container has no name");
        //         return;
        //     }
        // };

        // let unslashed_name: &str = name.strip_prefix('/').unwrap_or(&name);

        for (_, n) in networks {
            if let (Some(names), Some(ip)) = (&n.dns_names, &n.ip_address) {
                // debug!("Adding name {}", unslashed_name);
                // self.add_ip(unslashed_name, ip).await;

                for n in names {
                    debug!("Adding name {} with IP {}", n, ip);
                    self.add_ip(&(n.clone()), ip).await;
                }
            };
        }
    }

    async fn remove_container_ips(&self, container: ContainerInspectResponse) {
        // if let Some(name) = container.name {
        //     let unslashed_name: &str = name.strip_prefix('/').unwrap_or(&name);
        //     debug!("Removing name {}", unslashed_name);
        //     self.remove_ip(unslashed_name).await;
        // }

        let net_settings = match container.network_settings {
            Some(n) => n,
            None => return,
        };

        let networks = match net_settings.networks {
            Some(n) => n,
            None => return,
        };

        for (_, n) in networks {
            if let Some(names) = n.dns_names {
                for name in names {
                    debug!("Removing name {}", name);
                    self.remove_ip(name.as_str()).await;
                }
            }
        }
    }

    async fn add_ip(&self, name: &str, ip: &str) {
        self.ips
            .lock()
            .await
            .insert(name.to_string(), ip.to_string());
    }

    async fn remove_ip(&self, name: &str) {
        self.ips.lock().await.remove(name);
    }
}

pub async fn start() -> anyhow::Result<Arc<Resolver>> {
    let docker = Docker::connect_with_socket_defaults()?;

    let mut res = Arc::new(Resolver::new());

    tokio::spawn(events_listener(docker.clone(), res.clone()));

    init_contaiter_ips(&docker, &mut res).await?;

    Ok(res)
}

async fn init_contaiter_ips(docker: &Docker, resolver: &mut Arc<Resolver>) -> anyhow::Result<()> {
    for c in docker
        .list_containers(Option::<ListContainersOptions>::None)
        .await?
    {
        let ext_container = docker
            .inspect_container(&c.id.unwrap(), Option::<InspectContainerOptions>::None)
            .await?;

        resolver.add_container_ips(ext_container).await;
    }

    Ok(())
}

async fn events_listener(docker: Docker, resolver: Arc<Resolver>) -> anyhow::Result<()> {
    let filter = make_events_filter();
    let mut stream = docker.events(Some(filter));

    while let Some(res_evt) = stream.next().await {
        let evt = match res_evt {
            Ok(evt) => evt,
            Err(err) => {
                error!("Error while receiving event: {}", err);
                continue;
            }
        };

        let action = match evt.action {
            Some(action) => action,
            None => continue,
        };

        let container_id = match evt.actor {
            Some(actor) => actor.id.unwrap(),
            _ => continue,
        };

        let containter_res = docker
            .inspect_container(
                container_id.as_str(),
                Option::<InspectContainerOptions>::None,
            )
            .await;

        let container = match containter_res {
            Ok(c) => c,
            Err(err) => {
                error!("Error while inspecting container: {}", err);
                continue;
            }
        };

        match action.as_str() {
            "start" => {
                resolver.add_container_ips(container).await;
            }
            "die" | "destroy" => {
                resolver.remove_container_ips(container).await;
            }
            _ => {}
        }
    }
    Ok(())
}

fn make_events_filter() -> EventsOptions {
    let mut filter = HashMap::new();
    filter.insert("type".to_string(), vec!["container".to_string()]);
    filter.insert(
        "event".to_string(),
        vec![
            "start".to_string(),
            "die".to_string(),
            "destroy".to_string(),
        ],
    );

    EventsOptionsBuilder::new().filters(&filter).build()
}