Skip to main content

pingap_discovery/
lib.rs

1// Copyright 2024-2025 Tree xie.
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7// http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use hickory_resolver::ResolveError;
16use pingap_core::NotificationSender;
17use snafu::Snafu;
18use std::sync::Arc;
19
20pub static LOG_TARGET: &str = "pingap::discovery";
21
22#[derive(Debug, Snafu)]
23pub enum Error {
24    #[snafu(display("Io error {source}, {content}"))]
25    Io {
26        source: std::io::Error,
27        content: String,
28    },
29    #[snafu(display("Resolve error {source}"))]
30    Resolve { source: ResolveError },
31    #[snafu(display("{message}"))]
32    Invalid { message: String },
33    #[snafu(display("Docker error {source}"))]
34    Docker { source: bollard::errors::Error },
35}
36impl From<Error> for pingora::BError {
37    fn from(value: Error) -> Self {
38        pingora::Error::because(
39            pingora::ErrorType::HTTPStatus(500),
40            value.to_string(),
41            pingora::Error::new(pingora::ErrorType::InternalError),
42        )
43    }
44}
45
46pub type Result<T, E = Error> = std::result::Result<T, E>;
47
48pub(crate) type Addr = (String, String, usize);
49
50/// Formats a list of address strings into a vector of structured address tuples.
51///
52/// # Arguments
53///
54/// * `addrs` - A slice of strings containing addresses in the format "host:port weight" or "host:port" or "host"
55/// * `tls` - A boolean indicating whether to use TLS default port (443) or HTTP default port (80)
56///
57/// # Returns
58///
59/// Returns a vector of tuples containing (host, port, weight), where:
60/// * host is the hostname or IP address
61/// * port is either specified in the address or defaults to 443/80 based on TLS setting
62/// * weight is either specified after the address or defaults to 1
63pub(crate) fn format_addrs(addrs: &[String], tls: bool) -> Vec<Addr> {
64    let mut new_addrs = vec![];
65    for addr in addrs.iter() {
66        // get the weight of address
67        let arr: Vec<_> = addr.split(' ').collect();
68        let weight = if arr.len() == 2 {
69            arr[1].parse::<usize>().unwrap_or(1)
70        } else {
71            1
72        };
73        // split ip and port
74        // the port will use default value if none
75        if let Some((host, port)) = arr[0].split_once(':') {
76            new_addrs.push((host.to_string(), port.to_string(), weight));
77        } else {
78            let port = if tls {
79                "443".to_string()
80            } else {
81                "80".to_string()
82            };
83            new_addrs.push((arr[0].to_string(), port, weight));
84        }
85    }
86    new_addrs
87}
88
89pub const DNS_DISCOVERY: &str = "dns";
90pub const DOCKER_DISCOVERY: &str = "docker";
91pub const STATIC_DISCOVERY: &str = "static";
92pub const TRANSPARENT_DISCOVERY: &str = "transparent";
93
94#[derive(Default)]
95pub struct Discovery {
96    addr: Vec<String>,
97    tls: bool,
98    ipv4_only: bool,
99    dns_server: Option<String>,
100    dns_domain: Option<String>,
101    dns_search: Option<String>,
102    sender: Option<Arc<NotificationSender>>,
103}
104
105impl Discovery {
106    pub fn new(addr: Vec<String>) -> Self {
107        Self {
108            addr,
109            tls: false,
110            ipv4_only: false,
111            dns_server: None,
112            dns_domain: None,
113            dns_search: None,
114            sender: None,
115        }
116    }
117    pub fn with_sender(
118        mut self,
119        sender: Option<Arc<NotificationSender>>,
120    ) -> Self {
121        self.sender = sender;
122        self
123    }
124    pub fn with_tls(mut self, tls: bool) -> Self {
125        self.tls = tls;
126        self
127    }
128    pub fn with_ipv4_only(mut self, ipv4_only: bool) -> Self {
129        self.ipv4_only = ipv4_only;
130        self
131    }
132    pub fn with_dns_server(mut self, dns_server: String) -> Self {
133        if dns_server.is_empty() {
134            self.dns_server = None;
135        } else {
136            self.dns_server = Some(dns_server);
137        }
138        self
139    }
140    pub fn with_domain(mut self, domain: String) -> Self {
141        self.dns_domain = Some(domain);
142        self
143    }
144    pub fn with_search(mut self, search: String) -> Self {
145        self.dns_search = Some(search);
146        self
147    }
148}
149
150mod common;
151mod dns;
152mod docker;
153pub use common::{is_static_discovery, new_static_discovery};
154pub use dns::{is_dns_discovery, new_dns_discover_backends};
155pub use docker::{is_docker_discovery, new_docker_discover_backends};
156
157#[cfg(test)]
158mod tests {
159    use super::format_addrs;
160    use pretty_assertions::assert_eq;
161
162    #[test]
163    fn test_format_addrs() {
164        let addrs = format_addrs(&["127.0.0.1:8080".to_string()], false);
165        assert_eq!(format!("{:?}", addrs), r#"[("127.0.0.1", "8080", 1)]"#);
166
167        let addrs = format_addrs(&["127.0.0.1".to_string()], false);
168        assert_eq!(format!("{:?}", addrs), r#"[("127.0.0.1", "80", 1)]"#);
169
170        let addrs = format_addrs(&["127.0.0.1".to_string()], true);
171        assert_eq!(format!("{:?}", addrs), r#"[("127.0.0.1", "443", 1)]"#);
172
173        let addrs = format_addrs(&["127.0.0.1 10".to_string()], false);
174        assert_eq!(format!("{:?}", addrs), r#"[("127.0.0.1", "80", 10)]"#);
175    }
176}