// EDNS(0) option codes and helpers (RFC 6891, RFC 7871, RFC 5001, etc.)
// Extended DNS option codes used with OPT pseudo-RRs.

// Standard EDNS option codes

///|
pub let edns_opt_nsid : UInt16 = 3 // Name Server Identifier (RFC 5001)

///|
pub let edns_opt_dau : UInt16 = 5 // DNSSEC Algorithm Understood

///|
pub let edns_opt_dhu : UInt16 = 6 // DS Hash Understood

///|
pub let edns_opt_n3u : UInt16 = 7 // NSEC3 Hash Understood

///|
pub let edns_opt_client_subnet : UInt16 = 8 // Client Subnet (RFC 7871)

///|
pub let edns_opt_expire : UInt16 = 9 // Expire

///|
pub let edns_opt_cookie : UInt16 = 10 // DNS Cookie (RFC 7873)

///|
pub let edns_opt_tcp_keepalive : UInt16 = 11 // TCP Keepalive (RFC 7828)

///|
pub let edns_opt_padding : UInt16 = 12 // Padding (RFC 7830)

///|
pub let edns_opt_chain : UInt16 = 13 // Chain

///|
pub let edns_opt_key_tag : UInt16 = 14 // Key Tag (RFC 8145)

///|
pub let edns_opt_extended_error : UInt16 = 15 // Extended DNS Error (RFC 8914)

///|
pub let edns_opt_client_tag : UInt16 = 16 // Client Tag

///|
pub let edns_opt_server_tag : UInt16 = 17 // Server Tag

// Get the name of an EDNS option code

///|
pub fn edns_option_name(code : UInt16) -> String {
  if code == edns_opt_nsid {
    "NSID"
  } else if code == edns_opt_dau {
    "DAU"
  } else if code == edns_opt_dhu {
    "DHU"
  } else if code == edns_opt_n3u {
    "N3U"
  } else if code == edns_opt_client_subnet {
    "CLIENT-SUBNET"
  } else if code == edns_opt_expire {
    "EXPIRE"
  } else if code == edns_opt_cookie {
    "COOKIE"
  } else if code == edns_opt_tcp_keepalive {
    "TCP-KEEPALIVE"
  } else if code == edns_opt_padding {
    "PADDING"
  } else if code == edns_opt_chain {
    "CHAIN"
  } else if code == edns_opt_key_tag {
    "KEY-TAG"
  } else if code == edns_opt_extended_error {
    "EXTENDED-ERROR"
  } else if code == edns_opt_client_tag {
    "CLIENT-TAG"
  } else if code == edns_opt_server_tag {
    "SERVER-TAG"
  } else {
    "OPT" + code.to_string()
  }
}

// Extended DNS Error codes (RFC 8914)

///|
pub let ede_other : UInt16 = 0

///|
pub let ede_unsupported_dnskey_algorithm : UInt16 = 1

///|
pub let ede_unsupported_ds_digest : UInt16 = 2

///|
pub let ede_stale_answer : UInt16 = 3

///|
pub let ede_forged_answer : UInt16 = 4

///|
pub let ede_dnssec_indeterminate : UInt16 = 5

///|
pub let ede_dnssec_bogus : UInt16 = 6

///|
pub let ede_signature_expired : UInt16 = 7

///|
pub let ede_signature_not_yet_valid : UInt16 = 8

///|
pub let ede_dnskey_missing : UInt16 = 9

///|
pub let ede_rrsigs_missing : UInt16 = 10

///|
pub let ede_no_zone_key : UInt16 = 11

///|
pub let ede_nsec_missing : UInt16 = 12

///|
pub let ede_cached_error : UInt16 = 13

///|
pub let ede_not_ready : UInt16 = 14

///|
pub let ede_blocked : UInt16 = 15

///|
pub let ede_censored : UInt16 = 16

///|
pub let ede_filtered : UInt16 = 17

///|
pub let ede_prohibited : UInt16 = 18

///|
pub let ede_stale_nxdomain : UInt16 = 19

///|
pub let ede_not_authoritative : UInt16 = 20

///|
pub let ede_not_supported : UInt16 = 21

///|
pub let ede_no_reachable_authority : UInt16 = 22

///|
pub let ede_network_error : UInt16 = 23

///|
pub let ede_invalid_data : UInt16 = 24

///|
pub let ede_signature_expired_before : UInt16 = 25

///|
pub let ede_too_early : UInt16 = 26

///|
pub let ede_unsupported_nsec3_iterations : UInt16 = 27

///|
pub let ede_unable_to_conform : UInt16 = 28

// Get the name of an extended DNS error code

///|
pub fn extended_error_name(code : UInt16) -> String {
  if code == ede_other {
    "Other"
  } else if code == ede_unsupported_dnskey_algorithm {
    "Unsupported DNSKEY Algorithm"
  } else if code == ede_unsupported_ds_digest {
    "Unsupported DS Digest Type"
  } else if code == ede_stale_answer {
    "Stale Answer"
  } else if code == ede_forged_answer {
    "Forged Answer"
  } else if code == ede_dnssec_indeterminate {
    "DNSSEC Indeterminate"
  } else if code == ede_dnssec_bogus {
    "DNSSEC Bogus"
  } else if code == ede_signature_expired {
    "Signature Expired"
  } else if code == ede_signature_not_yet_valid {
    "Signature Not Yet Valid"
  } else if code == ede_dnskey_missing {
    "DNSKEY Missing"
  } else if code == ede_rrsigs_missing {
    "RRSIGs Missing"
  } else if code == ede_no_zone_key {
    "No Zone Key"
  } else if code == ede_nsec_missing {
    "NSEC Missing"
  } else if code == ede_cached_error {
    "Cached Error"
  } else if code == ede_not_ready {
    "Not Ready"
  } else if code == ede_blocked {
    "Blocked"
  } else if code == ede_censored {
    "Censored"
  } else if code == ede_filtered {
    "Filtered"
  } else if code == ede_prohibited {
    "Prohibited"
  } else if code == ede_stale_nxdomain {
    "Stale NXDomain Answer"
  } else if code == ede_not_authoritative {
    "Not Authoritative"
  } else if code == ede_not_supported {
    "Not Supported"
  } else if code == ede_no_reachable_authority {
    "No Reachable Authority"
  } else if code == ede_network_error {
    "Network Error"
  } else if code == ede_invalid_data {
    "Invalid Data"
  } else {
    "EDE" + code.to_string()
  }
}

///|
fn client_subnet_address_length(source_prefix : Int) -> Int {
  (source_prefix + 7) / 8
}

///|
fn validate_client_subnet_data(data : Array[Byte]) -> Result[Unit, String] {
  if data.length() < 4 {
    return Err("EDNS Client Subnet option is shorter than its 4-octet header")
  }
  let family = (data[0].to_int() << 8) | data[1].to_int()
  let address_bits = if family == 1 {
    32
  } else if family == 2 {
    128
  } else {
    return Err("EDNS Client Subnet family must be 1 (IPv4) or 2 (IPv6)")
  }
  let source_prefix = data[2].to_int()
  let scope_prefix = data[3].to_int()
  if source_prefix > address_bits {
    return Err("EDNS Client Subnet source prefix exceeds address family width")
  }
  if scope_prefix > address_bits {
    return Err("EDNS Client Subnet scope prefix exceeds address family width")
  }
  let address_length = client_subnet_address_length(source_prefix)
  if data.length() != 4 + address_length {
    return Err(
      "EDNS Client Subnet address length does not match the source prefix",
    )
  }
  let used_bits = source_prefix % 8
  if used_bits != 0 && address_length > 0 {
    let host_mask = (1 << (8 - used_bits)) - 1
    if (data[data.length() - 1].to_int() & host_mask) != 0 {
      return Err("EDNS Client Subnet address has non-zero padding bits")
    }
  }
  Ok(())
}

///|
fn validate_cookie_data(data : Array[Byte]) -> Result[Unit, String] {
  let length = data.length()
  // RFC 7873: CLIENT COOKIE is exactly 8 octets. SERVER COOKIE is absent in
  // an initial request or between 8 and 32 octets, for a total of 8 or 16..40.
  if length == 8 || (length >= 16 && length <= 40) {
    Ok(())
  } else {
    Err("EDNS Cookie length must be 8 or between 16 and 40 octets")
  }
}

///|
fn validate_keepalive_data(data : Array[Byte]) -> Result[Unit, String] {
  // RFC 7828 requests carry an empty option; responses may carry one UInt16
  // idle timeout in units of 100 milliseconds.
  if data.length() == 0 || data.length() == 2 {
    Ok(())
  } else {
    Err("EDNS TCP Keepalive option must contain zero or two octets")
  }
}

///|
fn validate_edns_option(option : OptOption) -> Result[Unit, String] {
  if option.opt_code == edns_opt_client_subnet {
    validate_client_subnet_data(option.opt_data)
  } else if option.opt_code == edns_opt_cookie {
    validate_cookie_data(option.opt_data)
  } else if option.opt_code == edns_opt_tcp_keepalive {
    validate_keepalive_data(option.opt_data)
  } else {
    Ok(())
  }
}

// Build a padding option for EDNS (RFC 7830)
// Pads DNS queries to obscure the actual query length

///|
pub fn build_padding_option(target_size : UInt16) -> OptOption {
  // A target smaller than an option header cannot carry padding.  Clamp instead
  // of passing a negative length to Array::make.
  let pad_size = if target_size.to_int() > 4 {
    target_size.to_int() - 4
  } else {
    0
  }
  let data = Array::make(pad_size, (0).to_byte())
  { opt_code: edns_opt_padding, opt_data: data }
}

// Build a validated Client Subnet option (RFC 7871).

///|
pub fn build_client_subnet_option_checked(
  family : UInt16,
  source_prefix : UInt16,
  scope_prefix : UInt16,
  address : Array[Byte],
) -> Result[OptOption, String] {
  let address_bits = if family == 1 {
    32
  } else if family == 2 {
    128
  } else {
    return Err("EDNS Client Subnet family must be 1 (IPv4) or 2 (IPv6)")
  }
  if source_prefix.to_int() > address_bits {
    return Err("EDNS Client Subnet source prefix exceeds address family width")
  }
  if scope_prefix.to_int() > address_bits {
    return Err("EDNS Client Subnet scope prefix exceeds address family width")
  }
  let total = 4 + address.length()
  let data = Array::make(total, (0).to_byte())
  data[0] = ((family >> 8) & 0xFF).to_byte()
  data[1] = (family & 0xFF).to_byte()
  data[2] = (source_prefix & 0xFF).to_byte()
  data[3] = (scope_prefix & 0xFF).to_byte()
  for i = 0; i < address.length(); i = i + 1 {
    data[4 + i] = address[i]
  }
  let option = { opt_code: edns_opt_client_subnet, opt_data: data }
  match validate_edns_option(option) {
    Ok(_) => Ok(option)
    Err(error) => Err(error)
  }
}

///|
/// Compatibility constructor for the original API. Invalid input aborts
/// immediately instead of returning an option that could become illegal wire.
/// New code should use `build_client_subnet_option_checked`.
pub fn build_client_subnet_option(
  family : UInt16,
  source_prefix : UInt16,
  scope_prefix : UInt16,
  address : Array[Byte],
) -> OptOption {
  match
    build_client_subnet_option_checked(
      family, source_prefix, scope_prefix, address,
    ) {
    Ok(option) => option
    Err(error) => abort(error)
  }
}

// Build a validated DNS Cookie option (RFC 7873).

///|
pub fn build_cookie_option_checked(
  client_cookie : Array[Byte],
  server_cookie : Array[Byte],
) -> Result[OptOption, String] {
  if client_cookie.length() != 8 {
    return Err("DNS Client Cookie must contain exactly 8 octets")
  }
  if server_cookie.length() != 0 &&
    (server_cookie.length() < 8 || server_cookie.length() > 32) {
    return Err("DNS Server Cookie must be absent or contain 8 to 32 octets")
  }
  let total = client_cookie.length() + server_cookie.length()
  let data = Array::make(total, (0).to_byte())
  for i = 0; i < client_cookie.length(); i = i + 1 {
    data[i] = client_cookie[i]
  }
  for i = 0; i < server_cookie.length(); i = i + 1 {
    data[client_cookie.length() + i] = server_cookie[i]
  }
  Ok({ opt_code: edns_opt_cookie, opt_data: data })
}

///|
/// Compatibility constructor for the original API. Invalid cookie lengths
/// abort immediately; new code should use `build_cookie_option_checked`.
pub fn build_cookie_option(
  client_cookie : Array[Byte],
  server_cookie : Array[Byte],
) -> OptOption {
  match build_cookie_option_checked(client_cookie, server_cookie) {
    Ok(option) => option
    Err(error) => abort(error)
  }
}

// Build a TCP Keepalive response option (RFC 7828).

///|
/// Build the empty TCP Keepalive option sent in a DNS request.
pub fn build_keepalive_request_option() -> OptOption {
  { opt_code: edns_opt_tcp_keepalive, opt_data: [] }
}

///|
pub fn build_keepalive_option_checked(
  timeout_ms : Int,
) -> Result[OptOption, String] {
  if timeout_ms < 0 {
    return Err("EDNS TCP Keepalive timeout cannot be negative")
  }
  if timeout_ms > 6_553_500 {
    return Err("EDNS TCP Keepalive timeout exceeds the UInt16 100ms range")
  }
  // The wire value is an idle timeout in 100ms units. Round a positive value
  // upward so encoding never advertises a timeout shorter than requested.
  let units = if timeout_ms == 0 { 0 } else { (timeout_ms + 99) / 100 }
  let data = Array::make(2, (0).to_byte())
  data[0] = ((units >> 8) & 0xFF).to_byte()
  data[1] = (units & 0xFF).to_byte()
  Ok({ opt_code: edns_opt_tcp_keepalive, opt_data: data })
}

///|
/// Compatibility constructor accepting the original UInt16 millisecond
/// argument. The encoded field is converted to RFC 7828's 100ms units.
pub fn build_keepalive_option(timeout_ms : UInt16) -> OptOption {
  match build_keepalive_option_checked(timeout_ms.to_int()) {
    Ok(option) => option
    Err(error) => abort(error)
  }
}

// Check if an EDNS option is known/standard

///|
pub fn is_known_edns_option(code : UInt16) -> Bool {
  code == edns_opt_nsid ||
  code == edns_opt_dau ||
  code == edns_opt_dhu ||
  code == edns_opt_n3u ||
  (code >= edns_opt_client_subnet && code <= edns_opt_server_tag)
}