GitHub backend, refactor a bit

Author: jackTabsCode

Cargo.lock

@@ -3,14 +3,14 @@ name = "drillbit"
 version = "0.1.4"
 edition = "2024"
 description = " A plugin installation tool for Roblox"
-license = "MIT"
-authors = ["Jack Taylor <j@jackt.io>"]
-repository = "https://github.com/jacktabscode/drillbit"
-homepage = "https://github.com/jacktabscode/drillbit"
 readme = "README.md"
+homepage = "https://github.com/jacktabscode/drillbit"
+repository = "https://github.com/jacktabscode/drillbit"
+license = "MIT"
 
 [dependencies]
 anyhow = "1.0.99"
+async-trait = "0.1.86"
 blake3 = "1.8.2"
 env_logger = "0.11.8"
 fs-err = { version = "3.1.1", features = ["tokio"] }

Cargo.toml

@@ -30,6 +30,10 @@ local = "plugins/editor_position.luau" # note that Roblox doesn't load `.luau` f
 # You can also add plugins on the Creator Store
 [plugins.hoarcekat]
 cloud = 4621580428
+
+# You can also use GitHub release artifacts
+[plugins.jest_companion]
+github = "https://github.com/jackTabsCode/jest-companion/releases/download/v0.1.1/plugin.rbxm"
 ```
 
 ## Usage

README.md

@@ -0,0 +1,184 @@
+use anyhow::{Context, Result, bail};
+use async_trait::async_trait;
+use fs_err::tokio as fs;
+use reqwest::Client;
+use serde::Deserialize;
+
+use crate::config::Plugin;
+
+#[async_trait]
+pub trait Backend {
+    fn new() -> Result<Self>
+    where
+        Self: Sized;
+
+    async fn download(&mut self, plugin: &Plugin) -> Result<(Vec<u8>, Option<String>)>;
+
+    fn plugin_id(&self, plugin: &Plugin, key: &str, cwd: &str) -> String;
+}
+
+pub struct LocalBackend;
+
+#[async_trait]
+impl Backend for LocalBackend {
+    fn new() -> Result<Self> {
+        Ok(LocalBackend)
+    }
+
+    async fn download(&mut self, plugin: &Plugin) -> Result<(Vec<u8>, Option<String>)> {
+        let Plugin::Local(path) = plugin else {
+            bail!("LocalBackend can only handle Local plugins")
+        };
+
+        let mut new_ext = None;
+
+        if let Some(ext) = path.extension()
+            && ext == "luau"
+        {
+            new_ext = Some("lua".to_string());
+        }
+
+        let path = path.to_path(".");
+        let data = fs::read(path).await.context("Failed to read plugin")?;
+
+        Ok((data, new_ext))
+    }
+
+    fn plugin_id(&self, plugin: &Plugin, _key: &str, cwd: &str) -> String {
+        let Plugin::Local(path) = plugin else {
+            unreachable!()
+        };
+
+        let filename = path
+            .to_path(".")
+            .file_name()
+            .unwrap()
+            .to_string_lossy()
+            .to_string();
+
+        format!("{}_{}", cwd, filename)
+    }
+}
+
+pub struct CloudBackend {
+    cookie: Option<String>,
+    client: Client,
+}
+
+impl CloudBackend {
+    fn get_cookie(&mut self) -> Result<String> {
+        if self.cookie.is_none() {
+            let cookie = rbx_cookie::get().context("Couldn't get Roblox cookie")?;
+            self.cookie = Some(cookie);
+        }
+        Ok(self.cookie.as_ref().unwrap().clone())
+    }
+}
+
+#[async_trait]
+impl Backend for CloudBackend {
+    fn new() -> Result<Self> {
+        Ok(Self {
+            cookie: None,
+            client: Client::new(),
+        })
+    }
+
+    async fn download(&mut self, plugin: &Plugin) -> Result<(Vec<u8>, Option<String>)> {
+        let Plugin::Cloud(id) = plugin else {
+            unreachable!()
+        };
+
+        let cookie = self.get_cookie()?;
+        let url = format!("https://assetdelivery.roblox.com/v2/asset?id={id}");
+        let res = self
+            .client
+            .get(&url)
+            .header("Cookie", &cookie)
+            .send()
+            .await?;
+
+        if !res.status().is_success() {
+            bail!("Request failed with status: {}", res.status());
+        }
+
+        let asset: AssetResponse = res.json().await?;
+        let download_url = asset
+            .locations
+            .first()
+            .ok_or_else(|| anyhow::anyhow!("No download locations found"))?
+            .location
+            .clone();
+
+        let file_res = self
+            .client
+            .get(download_url)
+            .header("Cookie", &cookie)
+            .send()
+            .await?;
+
+        if !file_res.status().is_success() {
+            bail!("Download failed with status: {}", file_res.status());
+        }
+
+        Ok((file_res.bytes().await?.to_vec(), Some("rbxm".to_string())))
+    }
+
+    fn plugin_id(&self, plugin: &Plugin, key: &str, cwd: &str) -> String {
+        let Plugin::Cloud(id) = plugin else {
+            unreachable!()
+        };
+
+        format!("{}_{}_{}", cwd, key, id)
+    }
+}
+
+pub struct GitHubBackend {
+    client: Client,
+}
+
+#[async_trait]
+impl Backend for GitHubBackend {
+    fn new() -> Result<Self> {
+        Ok(Self {
+            client: Client::new(),
+        })
+    }
+
+    async fn download(&mut self, plugin: &Plugin) -> Result<(Vec<u8>, Option<String>)> {
+        let Plugin::GitHub(url) = plugin else {
+            unimplemented!()
+        };
+
+        let res = self.client.get(url).send().await?;
+
+        if !res.status().is_success() {
+            bail!(
+                "GitHub release download failed with status: {}",
+                res.status()
+            );
+        }
+
+        let ext = url.split('.').next_back().map(|s| s.to_string());
+        Ok((res.bytes().await?.to_vec(), ext))
+    }
+
+    fn plugin_id(&self, plugin: &Plugin, _key: &str, cwd: &str) -> String {
+        let Plugin::GitHub(url) = plugin else {
+            unimplemented!()
+        };
+
+        let filename = url.split('/').next_back().unwrap_or("unknown");
+        format!("{}_{}", cwd, filename)
+    }
+}
+
+#[derive(Deserialize)]
+struct Location {
+    location: String,
+}
+
+#[derive(Deserialize)]
+struct AssetResponse {
+    locations: Vec<Location>,
+}

src/backends.rs

@@ -48,8 +48,9 @@ impl Config {
 }
 
 #[derive(Debug, Deserialize)]
-#[serde(rename_all = "snake_case")]
+#[serde(rename_all = "lowercase")]
 pub enum Plugin {
     Local(RelativePathBuf),
     Cloud(u64),
+    GitHub(String),
 }

src/config.rs

@@ -1,6 +1,6 @@
 use crate::{
+    backends::{Backend, CloudBackend, GitHubBackend, LocalBackend},
     config::{Config, Plugin},
-    web::WebClient,
 };
 use anyhow::{Context, bail};
 use fs_err::tokio as fs;
@@ -12,8 +12,8 @@ use std::{
     path::{Path, PathBuf},
 };
 
+mod backends;
 mod config;
-mod web;
 
 #[tokio::main]
 async fn main() -> anyhow::Result<()> {
@@ -24,29 +24,36 @@ async fn main() -> anyhow::Result<()> {
 
     let config = Config::read().await.context("Failed to read config")?;
 
-    let has_cloud_plugins = config
-        .plugins
-        .values()
-        .any(|plugin| matches!(plugin, Plugin::Cloud(_)));
-    let client = if has_cloud_plugins {
-        Some(WebClient::new()?)
-    } else {
-        None
-    };
-
     let cwd_path = env::current_dir().unwrap();
     let cwd = cwd_path.file_name().unwrap().to_str().unwrap();
 
     let plugins_path = get_plugins_path().context("Failed to get plugins path")?;
 
     let mut existing_plugins = get_existing_hashes(&plugins_path).await?;
 
+    fn get_or_create_backend<T: Backend>(backend: &mut Option<T>) -> anyhow::Result<&mut T> {
+        if backend.is_none() {
+            *backend = Some(T::new()?);
+        }
+        Ok(backend.as_mut().unwrap())
+    }
+
+    let mut local_backend: Option<LocalBackend> = None;
+    let mut github_backend: Option<GitHubBackend> = None;
+    let mut cloud_backend: Option<CloudBackend> = None;
+
     for (key, plugin) in config.plugins {
-        let id = plugin_id(&plugin, &key, cwd);
-        let mut path = plugins_path.join(&id);
+        let backend: &mut dyn Backend = match &plugin {
+            Plugin::Local(_) => get_or_create_backend(&mut local_backend)?,
+            Plugin::GitHub(_) => get_or_create_backend(&mut github_backend)?,
+            Plugin::Cloud(_) => get_or_create_backend(&mut cloud_backend)?,
+        };
 
+        let id = backend.plugin_id(&plugin, &key, cwd);
         info!("Reading \"{key}\"...");
-        let (data, ext) = read_plugin(client.as_ref(), &plugin).await?;
+        let (data, ext) = backend.download(&plugin).await?;
+
+        let mut path = plugins_path.join(&id);
 
         if let Some(ext) = ext {
             path.set_extension(ext);
@@ -110,50 +117,3 @@ async fn get_existing_hashes(
 
     Ok(existing_plugins)
 }
-
-fn plugin_id(plugin: &Plugin, key: &str, cwd: &str) -> String {
-    match plugin {
-        Plugin::Local(path) => {
-            let filename = path
-                .to_path(".")
-                .file_name()
-                .unwrap()
-                .to_string_lossy()
-                .to_string();
-            format!("{}_{}", cwd, filename)
-        }
-        Plugin::Cloud(id) => {
-            format!("{}_{}_{}", cwd, key, id)
-        }
-    }
-}
-
-async fn read_plugin(
-    client: Option<&WebClient>,
-    plugin: &Plugin,
-) -> anyhow::Result<(Vec<u8>, Option<String>)> {
-    match plugin {
-        Plugin::Cloud(id) => {
-            let client = client.context("WebClient is required for cloud plugins")?;
-            let data = client
-                .download_plugin(*id)
-                .await
-                .context("Failed to download plugin")?;
-
-            Ok((data, Some("rbxm".to_string())))
-        }
-        Plugin::Local(path) => {
-            let mut new_ext = None;
-
-            if let Some(ext) = path.extension()
-                && ext == "luau"
-            {
-                new_ext = Some("lua".to_string());
-            }
-            let path = path.to_path(".");
-            let data = fs::read(path).await.context("Failed to read plugin")?;
-
-            Ok((data, new_ext))
-        }
-    }
-}