如何在Rust中正确格式化JWT以连接到snowflake

jhkqcmku  于 2023-05-29  发布在  其他
关注(0)|答案(1)|浏览(220)

我最近开始了学习Rust的漫长旅程。我以前从来没有使用过低级语言,我经常发现自己处于这样的情况,我是大麦把它放在一起,因为以前很多繁重的工作都是从我身上抽象出来的。
我正在尝试使用JWT连接到我们的Snowflake示例。相关文档可以在this page上找到。
我似乎无法获得有效的JWT进行身份验证。我已经通过snowSql在终端中连接确认了我的公钥和私钥对工作。所以问题几乎可以肯定是我创建JWT的方式。
我有几个问题可能指向解决方案。我创建了一个结构体来打包“sub”字段,我把所有东西都放在那里,因为文档列出了有效载荷头的方式,我不确定这是否应该用.with_issuer()这样的东西来完成。
我模糊地理解了指纹的概念,这似乎是需要的,但在文档中也不像,它看起来像公钥,没有公钥开始。
我也完全不知道encoding_key和encode函数...
没有进一步的到期日,这是我到目前为止拼凑的。

use std::time::Duration;
use reqwest;
use jwt_simple::prelude::*;
use jwt_simple::reexports::serde_json;
use chrono::{DateTime, Utc, Duration as Duration2};
use thiserror::Error;
use jsonwebtoken::{encode, Algorithm, EncodingKey, Header};

#[tokio::main]
async fn main() {

    let client = reqwest::Client::new();

    let account_identifier = String::from("*<Account Identifier>*");
    let user = String::from("RUST");

    //create_jwt_token(&account_identifier, &user);
    let token = create_jwt_token(&account_identifier, &user);
    //dbg!(&token);
    let response = client.post( "https://*<Account Identifier>*.us-east-1.snowflakecomputing.com/api/v2/statements")
        .header("Authorization", "Bearer ".to_owned() + &token.unwrap().to_string())
        .header("X-Snowflake-Authorization-Token-Type", "KEYPAIR_JWT")
        .header("CONTENT-TYPE", "application/json")
        .header("ACCEPT", "application/json")
        .header("User-Agent", "snowflake-rust/0.1")
        .send()
        .await
        .unwrap()
        .text()
        .await;

    println!("{:?}", response);

}

fn create_jwt_token(account_identifier: &str, user: &str)
                    -> Result<String, KeyPairError>
{
    let private_key = get_private_key()?;
    let public_key_fingerprint = RS256PublicKey::from_pem(&get_public_key()?)
        .map_err(KeyPairError::FingerprintGeneration)?;

    let qualified_username = format!("{}.{}", account_identifier, user);
    let issuer = format!("{}:{:?}", qualified_username, public_key_fingerprint);

    let current_time = Utc::now();

    #[derive(Serialize, Deserialize, Debug)]
    struct MyAdditionalData {
        iss: String,
        sub: String,
        iat: i64,
        exp: i64,
    }
    let additional_data = MyAdditionalData {
        iss: issuer,
        sub: qualified_username,
        iat: current_time.timestamp(),
        exp: (current_time + Duration2::seconds(3600)).timestamp(),
    };

    let payload = Claims::with_custom_claims(additional_data,Duration::from_secs(3600).into());

    let encoding_key = EncodingKey::from_rsa_pem(private_key.as_bytes()).expect("something went wrong");
    let token = encode(&Header::new(Algorithm::RS256), &payload, &encoding_key).expect("failed to create token");

    Ok(token)
}

fn get_private_key() -> Result<String, KeyPairError> {
    let path = "./environment_variables/snowflake_private_key_path.txt";
    let private_key = std::fs::read_to_string(path)
        .map_err(|e| KeyPairError::PrivateKeyRead(e, path.into()))?;
    std::fs::read_to_string(&private_key)
        .map_err(|e| KeyPairError::PrivateKeyRead(e, private_key))
}

fn get_public_key() -> Result<String, KeyPairError> {
    let path = "./environment_variables/snowflake_public_key_path.txt";
    let public_key = std::fs::read_to_string(path)
        .map_err(|e| KeyPairError::PublicKeyRead(e, path.into()))?;
    std::fs::read_to_string(&public_key)
        .map_err(|e| KeyPairError::PublicKeyRead(e, public_key))
}

#[derive(thiserror::Error, Debug)]
pub enum KeyPairError {
    #[error("failed to read public key, path: {1}—{0}")]
    PublicKeyRead(std::io::Error, String),
    #[error("failed to read private key, path: {1}—{0}")]
    PrivateKeyRead(std::io::Error, String),
    #[error("failed to generate fingerprint from public key—{0}")]
    FingerprintGeneration(anyhow::Error),
    #[error("failed to generate key pair from private key—{0}")]
    KayPairGeneration(anyhow::Error),
}

我从中得到的错误是:

Ok("{\n  \"code\" : \"390144\",\n  \"message\" : \"JWT token is invalid.\"\n}")

编辑:现在我正在努力获得适当的JWT指纹。这是我在尝试复制python实现后目前正在使用的。它“工作”,但输出错误的指纹。

let sha256_digest = digest(&SHA256, get_public_key().unwrap().as_bytes());
let public_key_fingerprint = format!("SHA256:{}", base64::encode(sha256_digest.as_ref()));
hjqgdpho

hjqgdpho1#

因此,问题原来是我从snowflake-connector crate中提取的用于创建令牌的代码返回了带有下划线的指纹,而不是指纹中的斜线。
我用.replace('_',"/”)修改了指纹;并且能够发送有效的JWT。
最后的结果看起来像这样->

pub(crate) fn create_token(account_identifier: &str, user: &str) -> Result<String, KeyPairError> {
let private_key = get_private_key()?;
let public_key = get_public_key()?;

let mut public_key_fingerprint = RS256PublicKey::from_pem(&public_key)
    .map_err(KeyPairError::FingerprintGeneration)?
    .sha256_thumbprint();
let padding = public_key_fingerprint.len() % 3;
for _ in 0..padding {
    public_key_fingerprint.push('=');
}
let public_key_fingerprint = public_key_fingerprint.replace('_',"/");
let qualified_username  = format!("{account_identifier}.{user}");
let issuer = format!("{qualified_username}.SHA256:{public_key_fingerprint}");
let claims = Claims::create(jwt_simple::prelude::Duration::from_secs(3_600)) //3_600 is 1 hour
    .with_issuer(issuer)
    .with_subject(qualified_username);
let key_pair = RS256KeyPair::from_pem(&private_key)
    .map_err(KeyPairError::KayPairGeneration)?;
key_pair.sign(claims)
    .map_err(KeyPairError::KayPairGeneration)

}

相关问题