Add Support for GitHub Personal Access Tokens

src/auth.rs

@@ -0,0 +1,69 @@
+use std::io;
+
+use anyhow::{bail, format_err, Context};
+use serde::{Deserialize, Serialize};
+use toml_edit::Document;
+
+use crate::config::write_if_not_exists;
+use crate::home::Home;
+
+pub static MANIFEST_FILE_NAME: &str = "auth.toml";
+
+static DEFAULT_MANIFEST: &str = r#"
+# This file is for auth tokens managed by Aftman, a cross-platform toolchain manager.
+# For more information, see https://github.com/LPGhatguy/aftman
+
+# github = "token"
+"#;
+
+#[derive(Debug, Serialize, Deserialize)]
+pub struct AuthManifest {
+    pub github: Option<String>,
+}
+
+impl AuthManifest {
+    /// Create an empty global auth manifest if there isn't one already.
+    pub fn init(home: &Home) -> anyhow::Result<()> {
+        let base_dir = home.path();
+        fs_err::create_dir_all(&base_dir)?;
+
+        let manifest_path = base_dir.join(MANIFEST_FILE_NAME);
+        write_if_not_exists(&manifest_path, DEFAULT_MANIFEST.trim())?;
+
+        Ok(())
+    }
+
+    /// Try to load an auth.toml from a directory
+    pub fn load(home: &Home) -> anyhow::Result<Option<AuthManifest>> {
+        let file_path = home.path().join(MANIFEST_FILE_NAME);
+
+        let contents = match fs_err::read(&file_path) {
+            Ok(contents) => contents,
+            Err(err) => {
+                if err.kind() == io::ErrorKind::NotFound {
+                    return Ok(None);
+                }
+
+                bail!(err);
+            }
+        };
+
+        let manifest: AuthManifest = toml::from_slice(&contents)
+            .with_context(|| format_err!("Invalid auth.toml at {}", file_path.display()))?;
+
+        Ok(Some(manifest))
+    }
+
+    fn add_token(home: &Home, token_type: &str, token: &str) -> anyhow::Result<()> {
+        let manifest_path = home.path().join(MANIFEST_FILE_NAME);
+        let content = fs_err::read_to_string(&manifest_path)?;
+        let mut document: Document = content.parse()?;
+        document[token_type.as_ref()] = toml_edit::value(token.to_string());
+
+        fs_err::write(&manifest_path, document.to_string())?;
+
+        log::info!("A {token_type} token has been added globally.");
+
+        Ok(())
+    }
+}

src/main.rs

@@ -1,3 +1,4 @@
+mod auth;
 mod cli;
 mod config;
 mod dirs;
@@ -19,6 +20,7 @@ use std::env::{consts::EXE_SUFFIX, current_dir, current_exe};
 use anyhow::{bail, format_err, Context};
 use structopt::StructOpt;
 
+use crate::auth::AuthManifest;
 use crate::cli::Args;
 use crate::home::Home;
 use crate::manifest::Manifest;
@@ -63,6 +65,7 @@ fn run() -> anyhow::Result<()> {
     }
 
     Manifest::init_global(&home)?;
+    AuthManifest::init(&home)?;
     system_path::init(&home)?;
 
     Args::from_args().run(&home, tool_storage)

src/tool_source/github.rs

@@ -3,11 +3,12 @@ use std::io::{Cursor, Read, Seek};
 use anyhow::Context;
 use reqwest::{
     blocking::Client,
-    header::{ACCEPT, USER_AGENT},
+    header::{ACCEPT, AUTHORIZATION, USER_AGENT},
 };
 use semver::Version;
 use serde::{Deserialize, Serialize};
 
+use crate::auth::AuthManifest;
 use crate::tool_id::ToolId;
 use crate::tool_name::ToolName;
 use crate::tool_source::Asset;
@@ -18,20 +19,24 @@ const APP_NAME: &str = "LPGhatguy/aftman";
 
 pub struct GitHubSource {
     client: Client,
+    token: Option<String>,
 }
 
 impl GitHubSource {
-    pub fn new() -> Self {
+    pub fn new(auth: Option<AuthManifest>) -> Self {
         Self {
             client: Client::new(),
+            token: auth.map(|t| t.github).flatten()
         }
     }
 
     pub fn get_all_releases(&self, name: &ToolName) -> anyhow::Result<Vec<Release>> {
         let url = format!("https://api.github.com/repos/{}/releases", name);
-        let builder = self.client.get(&url).header(USER_AGENT, APP_NAME);
+        let mut builder = self.client.get(&url).header(USER_AGENT, APP_NAME);
 
-        // TODO: Authorization
+        if let Some(token) = &self.token {
+            builder = builder.header(AUTHORIZATION, format!("token {}", token));
+        }
 
         let response_body = builder.send()?.text()?;
 
@@ -78,13 +83,15 @@ impl GitHubSource {
     }
 
     pub fn download_asset(&self, url: &str) -> anyhow::Result<impl Read + Seek> {
-        let builder = self
+        let mut builder = self
             .client
             .get(url)
             .header(USER_AGENT, APP_NAME)
             .header(ACCEPT, "application/octet-stream");
 
-        // TODO: Authorization
+        if let Some(token) = &self.token {
+            builder = builder.header(AUTHORIZATION, format!("token {}", token));
+        }
 
         let response = builder.send()?;
         let body = response.bytes()?.to_vec();

src/tool_storage.rs

@@ -11,6 +11,7 @@ use anyhow::{bail, Context};
 use fs_err::File;
 use once_cell::unsync::OnceCell;
 
+use crate::auth::AuthManifest;
 use crate::home::Home;
 use crate::manifest::Manifest;
 use crate::tool_alias::ToolAlias;
@@ -24,6 +25,7 @@ pub struct ToolStorage {
     pub storage_dir: PathBuf,
     pub bin_dir: PathBuf,
     home: Home,
+    auth: Option<AuthManifest>,
     github: OnceCell<GitHubSource>,
 }
 
@@ -35,10 +37,13 @@ impl ToolStorage {
         let bin_dir = home.path().join("bin");
         fs_err::create_dir_all(&bin_dir)?;
 
+        let auth = AuthManifest::load(home)?;
+
         Ok(Self {
             storage_dir,
             bin_dir,
             home: home.clone(),
+            auth,
             github: OnceCell::new(),
         })
     }
@@ -154,7 +159,7 @@ impl ToolStorage {
         log::info!("Installing tool: {}", spec);
 
         log::debug!("Fetching GitHub releases...");
-        let github = self.github.get_or_init(GitHubSource::new);
+        let github = self.github.get_or_init(|| GitHubSource::new(self.auth));
         let mut releases = github.get_all_releases(spec.name())?;
         releases.sort_by(|a, b| a.version.cmp(&b.version).reverse());
 
@@ -232,7 +237,7 @@ impl ToolStorage {
         log::info!("Installing tool: {id}");
 
         log::debug!("Fetching GitHub release...");
-        let github = self.github.get_or_init(GitHubSource::new);
+        let github = self.github.get_or_init(|| GitHubSource::new(self.auth));
         let release = github.get_release(id)?;
 
         let mut compatible_assets = self.get_compatible_assets(&release);