From 661dfc55fd804a669419b95233ed242c0be8d0e5 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Kaan=20Barmore-Gen=C3=A7?= <kaan@bgenc.net>
Date: Thu, 9 Jun 2022 18:29:02 -0700
Subject: [PATCH] Use a single threaded setup & allow configurable TTL (#21)

* use a single threaded executor

* add configurable TTL support
---
 example.toml  |  8 ++++++++
 src/config.rs | 15 ++++++++++++---
 src/main.rs   | 21 +++++++++++++++------
 src/opts.rs   |  2 +-
 4 files changed, 36 insertions(+), 10 deletions(-)

diff --git a/example.toml b/example.toml
index 63f5a98..7f1d9be 100644
--- a/example.toml
+++ b/example.toml
@@ -9,6 +9,12 @@ fqdn = "example.com"
 # otherwise do things that will cause you to be charged money.
 api_key = "xxxxxxxxxxxxxxxxxxxxxxxx"
 
+# The Time To Live value to be used by entries. This can be an integer between
+# 300 and 2592000. It is 300 by default. This is roughly how quickly DNS changes
+# will propagate when updated, you should keep this the minimum so changes to
+# your IP address propagate quickly.
+ttl = 300
+
 # For every domain or subdomain you want to update, create an entry below.
 
 [[entry]]
@@ -24,3 +30,5 @@ types = ["A", "AAAA"]
 # Updates A for some.example.net
 name = "some"
 fqdn = "example.net" # Overrides top level setting
+# Individual entries can override the global TTL
+ttl = 600
diff --git a/src/config.rs b/src/config.rs
index b2b4cf8..66d2e3b 100644
--- a/src/config.rs
+++ b/src/config.rs
@@ -10,28 +10,37 @@ pub struct Entry {
     pub name: String,
     types: Option<Vec<String>>,
     fqdn: Option<String>,
+    ttl: Option<u32>,
 }
 
+fn default_ttl() -> u32 { return 300; }
+
 #[derive(Deserialize, Debug)]
 pub struct Config {
     fqdn: String,
     pub api_key: String,
     pub entry: Vec<Entry>,
+    #[serde(default = "default_ttl")]
+    pub ttl: u32,
 }
 
 const DEFAULT_TYPES: &'static [&'static str] = &["A"];
 
 impl Config {
     pub fn fqdn<'c>(entry: &'c Entry, config: &'c Config) -> &'c str {
-        return entry.fqdn.as_ref().unwrap_or(&config.fqdn).as_str();
+        entry.fqdn.as_ref().unwrap_or(&config.fqdn).as_str()
+    }
+
+    pub fn ttl(entry: &Entry, config: &Config) -> u32 {
+        entry.ttl.unwrap_or(config.ttl)
     }
 
     pub fn types<'e>(entry: &'e Entry) -> Vec<&'e str> {
-        return entry
+        entry
             .types
             .as_ref()
             .and_then(|ts| Some(ts.iter().map(|t| t.as_str()).collect()))
-            .unwrap_or_else(|| DEFAULT_TYPES.to_vec());
+            .unwrap_or_else(|| DEFAULT_TYPES.to_vec())
     }
 }
 
diff --git a/src/main.rs b/src/main.rs
index d4da066..4442608 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -3,7 +3,8 @@ use anyhow;
 use clap::Parser;
 use futures;
 use reqwest::{header, Client, ClientBuilder, StatusCode};
-use std::{collections::HashMap, num::NonZeroU32, sync::Arc, time::Duration};
+use serde::Serialize;
+use std::{num::NonZeroU32, sync::Arc, time::Duration};
 use tokio::{self, task::JoinHandle};
 mod config;
 mod opts;
@@ -43,7 +44,13 @@ async fn get_ip(api_url: &str) -> anyhow::Result<String> {
     Ok(text)
 }
 
-#[tokio::main]
+#[derive(Serialize)]
+pub struct APIPayload {
+    pub rrset_values: Vec<String>,
+    pub rrset_ttl: u32,
+}
+
+#[tokio::main(flavor = "current_thread")]
 async fn main() -> anyhow::Result<()> {
     let opts = opts::Opts::parse();
     let conf = config::load_config(&opts)
@@ -83,11 +90,13 @@ async fn main() -> anyhow::Result<()> {
                 "AAAA" => ipv6.die_with(|error| format!("Needed IPv6 for {}: {}", fqdn, error)),
                 bad_entry_type => die!("Unexpected type in config: {}", bad_entry_type),
             };
-            let mut map = HashMap::new();
-            map.insert("rrset_values", vec![ip]);
-            let req = client.put(url).json(&map);
+            let payload = APIPayload {
+                rrset_values: vec![ip.to_string()],
+                rrset_ttl: Config::ttl(&entry, &conf),
+            };
+            let req = client.put(url).json(&payload);
             let task_governor = governor.clone();
-            let task = tokio::task::spawn(async move {
+            let task = tokio::task::spawn_local(async move {
                 task_governor.until_ready_with_jitter(retry_jitter).await;
                 println!("Updating {}", &fqdn);
                 match req.send().await {
diff --git a/src/opts.rs b/src/opts.rs
index 6b0bb22..85e4e87 100644
--- a/src/opts.rs
+++ b/src/opts.rs
@@ -1,4 +1,4 @@
-use clap::{Parser, ArgEnum};
+use clap::Parser;
 
 /// A tool to automatically update DNS entries on Gandi, using it as a dynamic DNS system.
 #[derive(Parser, Debug)]