mod weighted;
use std::{net::IpAddr, path::PathBuf, time::Duration};
use pharos_corefile::{Cursor, Error};
use pharos_proto::{Message, NameRef, NameWirePart, Rcode, Type};
use pharos_server::{Context, Next, Outcome, Plugin};
#[derive(Debug)]
pub struct LoadBalance {
pub weight_file: Option<PathBuf>,
pub reload: Duration,
weights: Option<weighted::Weights>,
weighted_policy: bool,
prefer: Vec<Subnet>,
}
impl LoadBalance {
pub(crate) fn parse(cursor: &mut Cursor<'_>) -> Result<Self, Error> {
let mut lb = Self {
weight_file: None,
reload: Duration::from_secs(30),
weights: None,
weighted_policy: false,
prefer: Vec::new(),
};
while cursor.next() {
let args = cursor.remaining_args();
match args.first().map(String::as_str) {
None | Some("round_robin") if args.len() <= 1 => lb.weighted_policy = false,
Some("weighted") if args.len() == 2 => {
lb.weighted_policy = true;
lb.prefer.clear();
lb.reload = Duration::from_secs(30);
lb.weight_file = Some(PathBuf::from(&args[1]));
while cursor.next_block() {
if cursor.val() != "reload" {
return Err(cursor.error("unknown weighted property"));
}
let args = cursor.remaining_args();
if args.len() != 1 {
return Err(cursor.arg_err());
}
lb.reload = crate::health::parse_duration(&args[0])
.ok_or_else(|| cursor.error("invalid reload duration"))?;
}
}
_ => return Err(cursor.arg_err()),
}
while cursor.next_block() {
if cursor.val() != "prefer" {
return Err(cursor.error("unknown loadbalance property"));
}
for cidr in cursor.remaining_args() {
lb.prefer
.push(Subnet::parse(&cidr).ok_or_else(|| cursor.error("invalid CIDR"))?);
}
}
}
lb.weights = lb
.weight_file
.as_ref()
.map(|path| weighted::Weights::new(path.clone(), lb.reload));
Ok(lb)
}
pub(crate) fn register_hooks(&self, lifecycle: &mut pharos_server::Lifecycle) {
if let Some(weights) = &self.weights {
weights.register_hooks(lifecycle);
}
}
}
impl Plugin for LoadBalance {
async fn serve<N: Next + Send>(&self, cx: &mut Context<'_>, next: N) -> Outcome {
let outcome = next.run(cx).await;
if outcome == Outcome::Response {
shuffle(
cx.response,
self.weights.as_ref().filter(|_| self.weighted_policy),
&self.prefer,
);
}
outcome
}
}
#[derive(Clone, Copy, Default)]
struct Record {
start: u16,
end: u16,
destination: u16,
sequence: u16,
group: u8,
weight: u8,
priority: usize,
}
const EMPTY_RECORD: Record = Record {
start: 0,
end: 0,
destination: 0,
sequence: 0,
group: 0,
weight: 0,
priority: 0,
};
struct Scratch {
records: [Record; 5957],
order: [Record; 5957],
original: [u8; 65535],
pointers: [Pointer; 32768],
seen: [u8; 8192],
direct: [u8; 8192],
labels: [u8; 8192],
}
#[derive(Clone, Copy, Default)]
struct Pointer {
source: u16,
destination: u16,
target: u16,
original_target: u16,
}
#[derive(Debug)]
struct Subnet {
address: IpAddr,
prefix: u32,
}
impl Subnet {
fn parse(value: &str) -> Option<Self> {
let (address, prefix) = value.split_once('/')?;
let mut address: IpAddr = address.parse().ok()?;
if prefix.is_empty() || !prefix.bytes().all(|byte| byte.is_ascii_digit()) {
return None;
}
let mut prefix: u32 = prefix.parse().ok()?;
if prefix > if address.is_ipv4() { 32 } else { 128 } {
return None;
}
if let IpAddr::V6(ip) = address
&& prefix >= 96
&& let Some(ip) = ip.to_ipv4_mapped()
{
address = IpAddr::V4(ip);
prefix -= 96;
}
Some(Self { address, prefix })
}
fn contains(&self, rr: &pharos_proto::Record<'_>) -> bool {
let address = match (rr.rtype, rr.rdata) {
(Type::A, [a, b, c, d]) => IpAddr::V4(std::net::Ipv4Addr::new(*a, *b, *c, *d)),
(Type::AAAA, bytes) if bytes.len() == 16 => IpAddr::V6(std::net::Ipv6Addr::from(
<[u8; 16]>::try_from(bytes).unwrap(),
))
.to_canonical(),
_ => return false,
};
match (self.address, address) {
(IpAddr::V4(network), IpAddr::V4(ip)) => {
let mask = u32::MAX.checked_shl(32 - self.prefix).unwrap_or(0);
u32::from(network) & mask == u32::from(ip) & mask
}
(IpAddr::V6(network), IpAddr::V6(ip)) => {
let mask = u128::MAX.checked_shl(128 - self.prefix).unwrap_or(0);
u128::from(network) & mask == u128::from(ip) & mask
}
_ => false,
}
}
}
thread_local! {
static SCRATCH: std::cell::RefCell<Scratch> = const { std::cell::RefCell::new(Scratch {
records: [EMPTY_RECORD; 5957],
order: [EMPTY_RECORD; 5957],
original: [0; 65535],
pointers: [Pointer { source: 0, destination: 0, target: 0, original_target: 0 }; 32768],
seen: [0; 8192],
direct: [0; 8192],
labels: [0; 8192],
}) };
static RANDOM: std::cell::Cell<u64> = {
let mut seed = [0; 8];
let _ = getrandom::getrandom(&mut seed);
std::cell::Cell::new(u64::from_ne_bytes(seed).max(1))
};
}
fn random(limit: usize) -> usize {
RANDOM.with(|random| {
let mut value = random.get();
value ^= value << 13;
value ^= value >> 7;
value ^= value << 17;
random.set(value);
(value % limit as u64) as usize
})
}
fn relocated(records: &[Record], offset: usize) -> usize {
let index = records.partition_point(|record| usize::from(record.start) <= offset);
if index == 0 {
return offset;
}
let record = records[index - 1];
if offset < usize::from(record.end) {
usize::from(record.destination) + offset - usize::from(record.start)
} else {
offset
}
}
fn bit(bits: &[u8], offset: usize) -> bool {
bits[offset / 8] & (1 << (offset % 8)) != 0
}
fn mark(bits: &mut [u8], offset: usize) {
bits[offset / 8] |= 1 << (offset % 8);
}
struct PointerCollector<'a> {
records: &'a [Record],
pointers: &'a mut [Pointer],
seen: &'a mut [u8],
direct: &'a mut [u8],
labels: &'a mut [u8],
count: usize,
}
impl PointerCollector<'_> {
fn visit(&mut self, name: NameRef<'_>) -> bool {
let mut previous = None;
let mut pointer_limit = usize::MAX;
for part in name.wire_parts() {
let (offset, len) = match part {
NameWirePart::Label { offset, len } => (offset, len),
NameWirePart::Root { offset } => (offset, 1),
NameWirePart::Pointer { offset, .. } => (offset, 2),
};
let destination = relocated(self.records, offset);
if destination + len > pointer_limit
|| previous
.is_some_and(|(old_end, new_end)| old_end == offset && new_end != destination)
|| (0..len).any(|i| relocated(self.records, offset + i) != destination + i)
{
return false;
}
if let NameWirePart::Pointer { target, .. } = part {
if offset >= name.offset() && offset < name.offset() + name.wire_len() {
mark(self.direct, offset);
}
let target_after = relocated(self.records, target);
if target_after > 0x3fff || target_after >= destination {
return false;
}
if !bit(self.seen, offset) {
if self.count == self.pointers.len() {
return false;
}
mark(self.seen, offset);
self.pointers[self.count] = Pointer {
source: offset as u16,
destination: destination as u16,
target: target_after as u16,
original_target: target as u16,
};
self.count += 1;
}
pointer_limit = destination;
previous = None;
} else {
for byte in offset..offset + len {
mark(self.labels, byte);
}
previous = Some((offset + len, destination + len));
}
}
true
}
fn safe(&mut self) -> bool {
if !self.pointers[..self.count].iter().all(|pointer| {
let source = usize::from(pointer.source);
pointer.target == pointer.original_target
|| (bit(self.direct, source)
&& !bit(self.labels, source)
&& !bit(self.labels, source + 1))
}) {
return false;
}
let pointers = &mut self.pointers[..self.count];
pointers.sort_unstable_by_key(|pointer| pointer.source);
pointers.windows(2).all(|pair| {
usize::from(pair[0].source) + 1 != usize::from(pair[1].source)
|| pair[0].target.to_be_bytes()[1]
== 0xc0 | (pair[1].target.to_be_bytes()[0] & 0x3f)
})
}
}
fn shuffle(wire: &mut [u8], weights: Option<&weighted::Weights>, prefer: &[Subnet]) {
if wire.len() > usize::from(u16::MAX) {
return;
}
let weight_data = weights.map(weighted::Weights::load);
SCRATCH.with(|scratch| {
let Ok(message) = Message::parse(wire) else {
return;
};
if message.rcode() != Rcode::NOERROR
|| message
.questions()
.next()
.is_none_or(|q| matches!(q.qtype, Type::AXFR | Type::IXFR))
|| message.additionals().any(|rr| rr.rtype == Type::TSIG)
{
return;
}
let weighted_query = message
.questions()
.next()
.is_some_and(|q| matches!(q.qtype, Type::A | Type::AAAA | Type::SRV));
if weights.is_some() && !weighted_query && prefer.is_empty() {
return;
}
let mut scratch = scratch.borrow_mut();
let Scratch {
records,
order,
original,
pointers,
seen,
direct,
labels,
} = &mut *scratch;
let mut sections = [(0usize, 0usize); 3];
let mut count = 0;
for (section, iter) in [
message.answers(),
message.authorities(),
message.additionals(),
]
.into_iter()
.enumerate()
{
let first = count;
for rr in iter {
if count == records.len() {
return;
}
records[count] = Record {
start: (rr.ttl_offset - 4 - rr.name.wire_len()) as u16,
end: (rr.rdata_offset + rr.rdata.len()) as u16,
destination: (rr.ttl_offset - 4 - rr.name.wire_len()) as u16,
sequence: 0,
group: match rr.rtype {
Type::CNAME => 0,
Type::A | Type::AAAA => 2,
Type::MX => 3,
_ => 1,
},
weight: weight_data
.as_ref()
.map_or(1, |data| weighted::weight(data, &rr)),
priority: prefer
.iter()
.position(|subnet| subnet.contains(&rr))
.unwrap_or(prefer.len()),
};
count += 1;
}
sections[section] = (first, count);
}
order[..count].copy_from_slice(&records[..count]);
for (section, (first, end)) in sections.into_iter().enumerate() {
if end <= first + 1 || (weights.is_some() && section == 1) {
continue;
}
let section_start = usize::from(records[first].start);
let items = &mut order[first..end];
if weights.is_some() && prefer.is_empty() && !items.iter().any(|item| item.group == 2) {
continue;
}
items.sort_unstable_by_key(|item| (item.group, item.start));
for group in [2, 3] {
let first = items
.iter()
.position(|item| item.group == group)
.unwrap_or(items.len());
let end = items
.iter()
.rposition(|item| item.group == group)
.map_or(first, |i| i + 1);
if weights.is_some() {
if weighted_query && group == 2 && first < end {
let total: usize = items[first..end]
.iter()
.map(|item| usize::from(item.weight))
.sum();
let mut choice = random(total);
for i in first..end {
if choice < usize::from(items[i].weight) {
items.swap(first, i);
break;
}
choice -= usize::from(items[i].weight);
}
}
} else {
for i in first..end.saturating_sub(1) {
items.swap(i, i + random(end - i));
}
}
}
if !prefer.is_empty() && section != 1 {
for (sequence, item) in items.iter_mut().enumerate() {
item.sequence = sequence as u16;
}
let key = |record: &Record| {
let group = match record.group {
0 => 0,
2 => 1,
3 => 2,
_ => 3,
};
(
group,
if group == 1 { record.priority } else { 0 },
record.sequence,
)
};
items.sort_unstable_by_key(key);
}
let mut dest = section_start;
for item in items.iter() {
let len = usize::from(item.end - item.start);
let index = records[..count]
.binary_search_by_key(&item.start, |record| record.start)
.expect("record came from this Message");
records[index].destination = dest as u16;
dest += len;
}
}
seen.fill(0);
direct.fill(0);
labels.fill(0);
let pointer_count = {
let mut collector = PointerCollector {
records: &records[..count],
pointers,
seen,
direct,
labels,
count: 0,
};
for question in message.questions() {
if !collector.visit(question.name) {
return;
}
}
for rr in message
.answers()
.chain(message.authorities())
.chain(message.additionals())
{
if !collector.visit(rr.name) {
return;
}
let Ok(names) = rr.compressible_names() else {
return;
};
for name in names.into_iter().flatten() {
if !collector.visit(name) {
return;
}
}
}
if !collector.safe() {
return;
}
collector.count
};
original[..wire.len()].copy_from_slice(wire);
for item in &order[..count] {
let index = records[..count]
.binary_search_by_key(&item.start, |record| record.start)
.expect("record came from this Message");
let destination = usize::from(records[index].destination);
let len = usize::from(item.end - item.start);
wire[destination..destination + len]
.copy_from_slice(&original[usize::from(item.start)..usize::from(item.end)]);
}
for pointer in &pointers[..pointer_count] {
let destination = usize::from(pointer.destination);
wire[destination..destination + 2]
.copy_from_slice(&(0xc000 | pointer.target).to_be_bytes());
}
});
}