From 87694c6218b5e2187cea645cb144a4ca54b519cc Mon Sep 17 00:00:00 2001
From: ospab <ospab@ospab.host>
Date: Fri, 19 Jun 2026 17:31:43 +0300
Subject: [PATCH] Add update version targeting and fix dns prober

---
 ostp/src/dns_prober.rs | 25 +++++++++++++++++++------
 ostp/src/main.rs       |  4 ++++
 scripts/install.sh     | 27 +++++++++++++++++++++++++--
 3 files changed, 48 insertions(+), 8 deletions(-)

diff --git a/ostp/src/dns_prober.rs b/ostp/src/dns_prober.rs
index d505c6f..c1229af 100644
--- a/ostp/src/dns_prober.rs
+++ b/ostp/src/dns_prober.rs
@@ -99,18 +99,27 @@ const PUBLIC_DNS_SERVERS: &[(&str, &str)] = &[
 ];
 
 pub async fn run_prober() {
-    println!("Starting DNS resolver prober to find the fastest server for DNS Transport...");
+    println!("Enter your OSTP DNS Tunnel domain (e.g., tunnel.example.com):");
+    let mut target_domain = String::new();
+    std::io::stdin().read_line(&mut target_domain).unwrap();
+    let target_domain = target_domain.trim();
+    if target_domain.is_empty() {
+        println!("Domain cannot be empty. Exiting prober.");
+        return;
+    }
+
+    println!("\nStarting DNS resolver prober for domain: {}", target_domain);
     println!("{:<15} | {:<15} | {:<10}", "Name", "IP Address", "Latency");
     println!("{:-<15}-+-{:-<15}-+-{:-<10}", "", "", "");
 
     let mut best_server = "8.8.8.8";
     let mut best_latency = Duration::from_secs(10);
     
-    // We send a random TXT query to test DNS resolution time
+    // Send a real OSTP ping packet encoded as a domain
+    let payload = b"PING";
+    let encoded_domain = ostp_core::dns::encode_payload_to_domain(payload, target_domain);
+
     let mut rng = rand::thread_rng();
-    let id: u16 = rng.gen();
-    let packet = DnsPacket::new_query(id, "example.com", DnsRecordType::TXT);
-    let payload = packet.encode();
 
     for (name, ip) in PUBLIC_DNS_SERVERS {
         let sock = match tokio::net::UdpSocket::bind("0.0.0.0:0").await {
@@ -122,8 +131,12 @@ pub async fn run_prober() {
             continue;
         }
 
+        let id: u16 = rng.gen();
+        let packet = DnsPacket::new_query(id, &encoded_domain, DnsRecordType::TXT);
+        let payload_bytes = packet.encode();
+
         let start = Instant::now();
-        if sock.send(&payload).await.is_ok() {
+        if sock.send(&payload_bytes).await.is_ok() {
             let mut buf = [0u8; 512];
             match tokio::time::timeout(Duration::from_secs(2), sock.recv(&mut buf)).await {
                 Ok(Ok(_)) => {
diff --git a/ostp/src/main.rs b/ostp/src/main.rs
index 97e63a4..5781f92 100644
--- a/ostp/src/main.rs
+++ b/ostp/src/main.rs
@@ -55,6 +55,10 @@ struct Args {
     #[arg(long, help_heading = "Common Commands")]
     update: bool,
 
+    /// Specify a target version for the update command (e.g., -v 0.2.98)
+    #[arg(short = 'v', long = "version", help_heading = "Common Commands")]
+    target_version: Option<String>,
+
     /// Import a share link (ostp://...) into the configuration file and exit
     #[arg(long, help_heading = "Client Commands")]
     import: Option<String>,
diff --git a/scripts/install.sh b/scripts/install.sh
index bd335a8..2834610 100644
--- a/scripts/install.sh
+++ b/scripts/install.sh
@@ -85,10 +85,33 @@ esac
 
 echo "Platform: linux/$ARCH"
 
+# ── Parse arguments ────────────────────────────────────────────────────
+TARGET_VERSION=""
+while [[ $# -gt 0 ]]; do
+  case $1 in
+    -v|--version)
+      TARGET_VERSION="$2"
+      shift 2
+      ;;
+    *)
+      shift
+      ;;
+  esac
+done
+
 # ── Download binary ──────────────────────────────────────────────────
 
-echo "Fetching latest release..."
-LATEST_RELEASE=$(curl -s "https://api.github.com/repos/${GITHUB_REPO}/releases/latest" | grep '"tag_name":' | sed -E 's/.*"([^"]+)".*/\1/')
+if [ -n "$TARGET_VERSION" ]; then
+    LATEST_RELEASE="$TARGET_VERSION"
+    # Ensure it starts with 'v' if it's supposed to
+    if [[ ! "$LATEST_RELEASE" =~ ^v ]]; then
+        LATEST_RELEASE="v$LATEST_RELEASE"
+    fi
+    echo "Fetching requested release $LATEST_RELEASE..."
+else
+    echo "Fetching latest release..."
+    LATEST_RELEASE=$(curl -s "https://api.github.com/repos/${GITHUB_REPO}/releases/latest" | grep '"tag_name":' | sed -E 's/.*"([^"]+)".*/\1/')
+fi
 
 if [ -z "$LATEST_RELEASE" ] || [[ "$LATEST_RELEASE" == *"null"* ]]; then
     echo "[notice] Could not determine latest release automatically."