Implement s3:// protocol

For those that want to pull from s3

Signed-off-by: Eric Curtin <[email protected]>

Author: ericcurtin

examples/run/run.cpp

@@ -65,6 +65,23 @@ static int printe(const char * fmt, ...) {
     return ret;
 }
 
+static std::string strftime_fmt(const char * fmt, const std::tm & tm) {
+    // Estimate the size of the output buffer
+    std::string buffer;
+    buffer.resize(128);
+
+    // Try to format the string
+    size_t len;
+    while ((len = std::strftime(buffer.data(), buffer.size(), fmt, &tm)) == 0) {
+        // If the buffer was too small, double its size and try again
+        buffer.resize(buffer.size() * 2);
+    }
+
+    buffer.resize(len);
+
+    return buffer;
+}
+
 class Opt {
   public:
     int init(int argc, const char ** argv) {
@@ -698,6 +715,49 @@ class LlamaData {
         return download(url, bn, true);
     }
 
+    int s3_dl(const std::string & model, const std::string & bn) {
+        // Extract bucket and key from S3 URI
+        const std::string prefix = "s3://";
+        const size_t      pos    = model.find(prefix);
+        if (pos != 0) {
+            return 1;
+        }
+
+        const std::string path      = model.substr(prefix.length());
+        const size_t      slash_pos = path.find('/');
+        if (slash_pos == std::string::npos) {
+            return 1;
+        }
+
+        const std::string bucket = path.substr(0, slash_pos);
+        const std::string key    = path.substr(slash_pos + 1);
+
+        // Get AWS credentials from environment, used env vars, open to change
+        const char * access_key = std::getenv("AWS_ACCESS_KEY_ID");
+        const char * secret_key = std::getenv("AWS_SECRET_ACCESS_KEY");
+        if (!access_key || !secret_key) {
+            printe("AWS credentials not found in environment\n");
+            return 1;
+        }
+
+        // Generate AWS Signature Version 4 headers
+        // (Implementation requires HMAC-SHA256 and date handling)
+        // Get current timestamp
+        const time_t                   now     = time(nullptr);
+        const tm                       tm      = *gmtime(&now);
+        std::string                    date    = strftime_fmt("%Y%m%d", tm);
+        const std::vector<std::string> headers = {
+            "Authorization: AWS4-HMAC-SHA256 Credential=" + std::string(access_key) + "/" + date +
+                "/us-east-1/s3/aws4_request",
+            "x-amz-content-sha256: UNSIGNED-PAYLOAD", "x-amz-date: 20240130T000000Z"
+        };
+
+        // Construct S3 endpoint URL
+        const std::string url = "https://" + bucket + ".s3.amazonaws.com/" + key;
+
+        return download(url, bn, true, headers);
+    }
+
     std::string basename(const std::string & path) {
         const size_t pos = path.find_last_of("/\\");
         if (pos == std::string::npos) {
@@ -738,6 +798,9 @@ class LlamaData {
             rm_until_substring(model_, "github:");
             rm_until_substring(model_, "://");
             ret = github_dl(model_, bn);
+        } else if (string_starts_with(model_, "s3://")) {
+            rm_until_substring(model_, "://");
+            ret = s3_dl(model_, bn);
         } else {  // ollama:// or nothing
             rm_until_substring(model_, "ollama.com/library/");
             rm_until_substring(model_, "://");