rusty_feeder/exchange/upbit/
auth.rs

1/*
2 * Upbit JWT authentication implementation
3 * Provides JWT token generation for private WebSocket and REST API endpoints
4 */
5
6use anyhow::{Result, anyhow};
7use hmac::Hmac;
8use jsonwebtoken::{Algorithm, EncodingKey, Header, encode};
9use serde::{Deserialize, Serialize};
10use sha2::{Digest, Sha256, Sha512};
11use smartstring::{SmartString, alias::String};
12use uuid::Uuid;
13
14/// Type alias for HMAC-SHA256
15type HmacSha256 = Hmac<Sha256>;
16
17/// JWT claims for Upbit authentication
18#[derive(Debug, Serialize, Deserialize)]
19pub struct UpbitJwtClaims {
20    /// Access key (API key)
21    pub access_key: String,
22
23    /// Nonce - unique identifier for each request
24    pub nonce: String,
25
26    /// Query hash for GET/DELETE requests with parameters (optional)
27    #[serde(skip_serializing_if = "Option::is_none")]
28    pub query_hash: Option<String>,
29
30    /// Body hash for POST/PUT requests with body (optional)
31    #[serde(skip_serializing_if = "Option::is_none")]
32    pub body_hash: Option<String>,
33}
34
35/// Upbit authentication configuration
36#[derive(Debug, Clone)]
37pub struct UpbitAuthConfig {
38    /// API access key
39    pub access_key: String,
40
41    /// API secret key
42    pub secret_key: String,
43}
44
45impl UpbitAuthConfig {
46    /// Create a new authentication configuration
47    #[must_use]
48    pub const fn new(access_key: String, secret_key: String) -> Self {
49        Self {
50            access_key,
51            secret_key,
52        }
53    }
54}
55
56/// Upbit JWT authentication handler
57#[derive(Debug)]
58pub struct UpbitAuth {
59    /// Authentication configuration
60    config: UpbitAuthConfig,
61}
62
63impl UpbitAuth {
64    /// Create a new authentication handler
65    #[must_use]
66    pub const fn new(config: UpbitAuthConfig) -> Self {
67        Self { config }
68    }
69
70    /// Generate JWT token for WebSocket authentication
71    ///
72    /// # Returns
73    /// JWT token String for use in Authorization header
74    pub fn generate_websocket_jwt(&self) -> Result<String> {
75        // Create claims with nonce
76        let claims = UpbitJwtClaims {
77            access_key: self.config.access_key.clone(),
78            nonce: String::from(Uuid::new_v4().to_string()),
79            query_hash: None,
80            body_hash: None,
81        };
82
83        self.encode_jwt(&claims)
84    }
85
86    /// Generate JWT token for REST API GET/DELETE requests
87    ///
88    /// # Parameters
89    /// - `query_params`: Query parameters as key-value pairs (optional)
90    ///
91    /// # Returns
92    /// JWT token String for use in Authorization header
93    pub fn generate_rest_jwt_get(&self, query_params: Option<&[(&str, &str)]>) -> Result<String> {
94        let mut claims = UpbitJwtClaims {
95            access_key: self.config.access_key.clone(),
96            nonce: String::from(Uuid::new_v4().to_string()),
97            query_hash: None,
98            body_hash: None,
99        };
100
101        // If query parameters exist, calculate query hash
102        if let Some(params) = query_params
103            && !params.is_empty()
104        {
105            let query_string = self.build_sorted_query_string(params);
106            let query_hash = self.calculate_sha512_hash(&query_string);
107            claims.query_hash = Some(query_hash);
108        }
109
110        self.encode_jwt(&claims)
111    }
112
113    /// Generate JWT token for REST API POST/PUT requests
114    ///
115    /// # Parameters
116    /// - `body`: JSON request body as String
117    ///
118    /// # Returns
119    /// JWT token String for use in Authorization header
120    pub fn generate_rest_jwt_post(&self, body: &str) -> Result<String> {
121        let claims = UpbitJwtClaims {
122            access_key: self.config.access_key.clone(),
123            nonce: String::from(Uuid::new_v4().to_string()),
124            query_hash: None,
125            body_hash: Some(self.calculate_sha512_hash(body)),
126        };
127
128        self.encode_jwt(&claims)
129    }
130
131    /// Build sorted query String from parameters
132    ///
133    /// # Parameters
134    /// - `params`: Query parameters as key-value pairs
135    ///
136    /// # Returns
137    /// Sorted query String (e.g., "market=KRW-BTC&state=wait")
138    fn build_sorted_query_string(&self, params: &[(&str, &str)]) -> String {
139        let mut params_vec: Vec<(&str, &str)> = params.to_vec();
140        params_vec.sort_by(|a, b| a.0.cmp(b.0)); // Sort by key alphabetically
141
142        params_vec
143            .iter()
144            .map(|(k, v)| format!("{k}={v}").into())
145            .collect::<Vec<String>>()
146            .join("&")
147            .into()
148    }
149
150    /// Calculate SHA512 hash of a String
151    ///
152    /// # Parameters
153    /// - `data`: String to hash
154    ///
155    /// # Returns
156    /// Hexadecimal representation of the hash
157    fn calculate_sha512_hash(&self, data: &str) -> String {
158        let mut hasher = Sha512::new();
159        hasher.update(data.as_bytes());
160        hex::encode(hasher.finalize()).into()
161    }
162
163    /// Encode JWT with the given claims
164    ///
165    /// # Parameters
166    /// - `claims`: JWT claims to encode
167    ///
168    /// # Returns
169    /// Encoded JWT token String
170    fn encode_jwt(&self, claims: &UpbitJwtClaims) -> Result<String> {
171        // Create header with HS256 algorithm
172        let header = Header::new(Algorithm::HS256);
173
174        // Create encoding key from secret
175        let encoding_key = EncodingKey::from_secret(self.config.secret_key.as_bytes());
176
177        // Encode the JWT
178        encode(&header, claims, &encoding_key)
179            .map(SmartString::from)
180            .map_err(|e| anyhow!("Failed to encode JWT: {}", e))
181    }
182}
183
184// Migrated to rusty-common WebSocket auth infrastructure
185/// Helper function to add JWT authentication to WebSocket connection request
186/// Uses rusty-common WebSocket auth system
187pub fn add_jwt_to_websocket_request(
188    headers: &mut rusty_common::collections::FxHashMap<String, String>,
189    jwt_token: &str,
190) -> Result<()> {
191    // Add Authorization header with Bearer token using rusty-common WebSocket auth
192    headers.insert("Authorization".into(), format!("Bearer {jwt_token}").into());
193
194    Ok(())
195}
196
197#[cfg(test)]
198mod tests {
199    use super::*;
200
201    #[test]
202    fn test_jwt_generation() {
203        // Test configuration
204        let config = UpbitAuthConfig::new("test_access_key".into(), "test_secret_key".into());
205
206        let auth = UpbitAuth::new(config);
207
208        // Test WebSocket JWT generation
209        let ws_jwt = auth.generate_websocket_jwt().unwrap();
210        assert!(!ws_jwt.is_empty());
211
212        // Test REST GET JWT generation without params
213        let get_jwt = auth.generate_rest_jwt_get(None).unwrap();
214        assert!(!get_jwt.is_empty());
215
216        // Test REST GET JWT generation with params
217        let params = vec![("market", "KRW-BTC"), ("state", "wait")];
218        let get_jwt_with_params = auth.generate_rest_jwt_get(Some(&params)).unwrap();
219        assert!(!get_jwt_with_params.is_empty());
220
221        // Test REST POST JWT generation
222        let body = r#"{"market":"KRW-BTC","side":"bid","volume":"0.01","price":"20000000","ord_type":"limit"}"#;
223        let post_jwt = auth.generate_rest_jwt_post(body).unwrap();
224        assert!(!post_jwt.is_empty());
225    }
226
227    #[test]
228    fn test_query_string_sorting() {
229        let config = UpbitAuthConfig::new("test_access_key".into(), "test_secret_key".into());
230
231        let auth = UpbitAuth::new(config);
232
233        // Test query String sorting
234        let params = vec![("state", "wait"), ("market", "KRW-BTC"), ("uuids", "uuid1")];
235        let sorted = auth.build_sorted_query_string(&params);
236        assert_eq!(sorted, "market=KRW-BTC&state=wait&uuids=uuid1");
237    }
238
239    #[test]
240    fn test_sha512_hash() {
241        let config = UpbitAuthConfig::new("test_access_key".into(), "test_secret_key".into());
242
243        let auth = UpbitAuth::new(config);
244
245        // Test SHA512 hash calculation
246        let data = "test_data";
247        let hash = auth.calculate_sha512_hash(data);
248
249        // SHA512 hash should be 128 characters (64 bytes in hex)
250        assert_eq!(hash.len(), 128);
251
252        // Known hash for "test_data"
253        let expected_hash = "66f993ac0f65e1a301f4924cc38fc4da1d6ac7c53c3e75e535735788263d97f50a323cec0b033e28a95a4549a5d39a45f85d6f6fb7b02b75e3d5de5e9727a3ea";
254        assert_eq!(hash, expected_hash);
255    }
256}
rusty_feeder/exchange/upbit/auth.rs

rusty_feeder/exchange/upbit/
auth.rs
