From 4141a62e472a4dc5dfed9022045cb2f5ebd258eb Mon Sep 17 00:00:00 2001
From: Michael Krasnitski <michael.krasnitski@gmail.com>
Date: Thu, 15 Aug 2024 14:57:55 -0400
Subject: [PATCH] Add `DeviceCodes::poll_until_available` method

---
 Cargo.toml              |  2 +-
 examples/device_flow.rs | 25 +------------------------
 src/auth.rs             | 35 +++++++++++++++++++++++++++++++++--
 3 files changed, 35 insertions(+), 27 deletions(-)

diff --git a/Cargo.toml b/Cargo.toml
index 6cd5c728..661ec55c 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -51,7 +51,7 @@ serde_json = "1.0.64"
 serde_path_to_error = "0.1.4"
 serde_urlencoded = "0.7.1"
 snafu = "0.8"
-tokio = { version = "1.17.0", default-features = false, optional = true }
+tokio = { version = "1.17.0", default-features = false, features = ["time"], optional = true }
 tower = { version = "0.4.13", default-features = false, features = ["util", "buffer"] }
 tower-http = { version = "0.5.1", features = ["map-response-body", "trace"] }
 tracing = { version = "0.1.37", features = ["log"], optional = true }
diff --git a/examples/device_flow.rs b/examples/device_flow.rs
index b57111cf..1903fe38 100644
--- a/examples/device_flow.rs
+++ b/examples/device_flow.rs
@@ -1,6 +1,4 @@
-use either::Either;
 use http::header::ACCEPT;
-use std::time::Duration;
 
 #[tokio::main]
 async fn main() -> octocrab::Result<()> {
@@ -17,28 +15,7 @@ async fn main() -> octocrab::Result<()> {
         "Go to {} and enter code {}",
         codes.verification_uri, codes.user_code
     );
-    let mut interval = Duration::from_secs(codes.interval);
-    let mut clock = tokio::time::interval(interval);
-    let auth = loop {
-        clock.tick().await;
-        match codes.poll_once(&crab, &client_id).await? {
-            Either::Left(auth) => break auth,
-            Either::Right(cont) => match cont {
-                octocrab::auth::Continue::SlowDown => {
-                    // We were request to slow down. We add five seconds to the polling
-                    // duration.
-                    interval += Duration::from_secs(5);
-                    clock = tokio::time::interval(interval);
-                    // The first tick happens instantly, so we tick that off immediately.
-                    clock.tick().await;
-                }
-                octocrab::auth::Continue::AuthorizationPending => {
-                    // The user has not clicked authorize yet, but nothing has gone wrong.
-                    // We keep polling.
-                }
-            },
-        }
-    };
+    let auth = codes.poll_until_available(&crab, &client_id).await?;
 
     println!("Authorization succeeded with access to {:?}", auth.scope);
     Ok(())
diff --git a/src/auth.rs b/src/auth.rs
index 494941b7..89521eb7 100644
--- a/src/auth.rs
+++ b/src/auth.rs
@@ -7,7 +7,7 @@ use jsonwebtoken::{Algorithm, EncodingKey, Header};
 use secrecy::{ExposeSecret, SecretString};
 use serde::{Deserialize, Serialize};
 use std::fmt;
-use std::time::SystemTime;
+use std::time::{Duration, SystemTime};
 
 use snafu::*;
 
@@ -204,7 +204,7 @@ pub struct DeviceCodes {
 impl DeviceCodes {
     /// Poll Github to see if authentication codes are available.
     ///
-    /// See `https://docs.github.com/en/developers/apps/building-oauth-apps/authorizing-oauth-apps#response-parameters` for details.
+    /// See `https://docs.github.com/en/developers/apps/building-oauth-apps/authorizing-oauth-apps` for details.
     pub async fn poll_once(
         &self,
         crab: &crate::Octocrab,
@@ -225,6 +225,37 @@ impl DeviceCodes {
             TokenResponse::Continue { error } => Either::Right(error),
         })
     }
+
+    /// Poll Github in a loop until authentication codes become available.
+    #[cfg(feature = "tokio")]
+    pub async fn poll_until_available(
+        &self,
+        crab: &crate::Octocrab,
+        client_id: &SecretString,
+    ) -> Result<OAuth> {
+        let mut interval = Duration::from_secs(self.interval);
+        let mut clock = tokio::time::interval(interval);
+
+        loop {
+            clock.tick().await;
+            match self.poll_once(crab, client_id).await? {
+                Either::Left(auth) => return Ok(auth),
+                Either::Right(cont) => match cont {
+                    Continue::SlowDown => {
+                        // We were requested to slow down, so add five seconds to the polling
+                        // duration.
+                        interval += Duration::from_secs(5);
+                        clock = tokio::time::interval(interval);
+                        // The first tick happens instantly, so we tick that off immediately.
+                        clock.tick().await;
+                    }
+                    Continue::AuthorizationPending => {
+                        // The user has not clicked authorize yet, so we keep polling as normal.
+                    }
+                },
+            }
+        }
+    }
 }
 
 /// See https://docs.github.com/en/developers/apps/building-oauth-apps/authorizing-oauth-apps#input-parameters