1use async_trait::async_trait;
2use chrono::{DateTime, Utc};
3use sqlx::{Error, PgPool};
4use uuid::Uuid;
5
6use crate::model::{AccessToken, DeviceAuthorization};
7
8#[async_trait]
9pub trait OAuthRepository: Send + Sync + Clone + 'static {
10    async fn create_device_authorization(
11        &self,
12        device_code: &str,
13        user_code: &str,
14        client_id: &str,
15        expires_at: DateTime<Utc>,
16    ) -> Result<DeviceAuthorization, Error>;
17
18    async fn get_device_authorization_by_device_code(
19        &self,
20        device_code: &str,
21    ) -> Result<Option<DeviceAuthorization>, Error>;
22
23    async fn get_device_authorization_by_user_code(
24        &self,
25        user_code: &str,
26    ) -> Result<Option<DeviceAuthorization>, Error>;
27
28    async fn expire_device_authorization(&self, id: Uuid) -> Result<(), Error>;
29
30    async fn authorize_device(
31        &self,
32        user_code: &str,
33        user_id: Uuid,
34    ) -> Result<Option<DeviceAuthorization>, Error>;
35
36    async fn deny_device(
37        &self,
38        user_code: &str,
39        user_id: Uuid,
40    ) -> Result<Option<DeviceAuthorization>, Error>;
41
42    async fn create_access_token(
43        &self,
44        user_id: Uuid,
45        client_id: &str,
46        token_hash: &str,
47    ) -> Result<AccessToken, Error>;
48
49    async fn get_access_token_by_hash(
50        &self,
51        token_hash: &str,
52    ) -> Result<Option<AccessToken>, Error>;
53
54    async fn touch_access_token(&self, id: Uuid) -> Result<(), Error>;
55
56    async fn delete_access_token(&self, id: Uuid) -> Result<(), Error>;
57}
58
59#[derive(Debug, Clone)]
60pub struct OAuthRepositoryImpl {
61    pool: PgPool,
62}
63
64impl OAuthRepositoryImpl {
65    pub fn new(pool: PgPool) -> Self {
66        Self { pool }
67    }
68}
69
70#[async_trait]
71impl OAuthRepository for OAuthRepositoryImpl {
72    async fn create_device_authorization(
73        &self,
74        device_code: &str,
75        user_code: &str,
76        client_id: &str,
77        expires_at: DateTime<Utc>,
78    ) -> Result<DeviceAuthorization, Error> {
79        let device_auth = sqlx::query_as::<_, DeviceAuthorization>(
80            r#"
81            INSERT INTO device_authorizations (device_code, user_code, client_id, expires_at)
82            VALUES ($1, $2, $3, $4)
83            RETURNING id, device_code, user_code, client_id, user_id, status, expires_at, created_at
84            "#,
85        )
86        .bind(device_code)
87        .bind(user_code)
88        .bind(client_id)
89        .bind(expires_at)
90        .fetch_one(&self.pool)
91        .await?;
92
93        Ok(device_auth)
94    }
95
96    async fn get_device_authorization_by_device_code(
97        &self,
98        device_code: &str,
99    ) -> Result<Option<DeviceAuthorization>, Error> {
100        let device_auth = sqlx::query_as::<_, DeviceAuthorization>(
101            r#"
102            SELECT id, device_code, user_code, client_id, user_id, status, expires_at, created_at
103            FROM device_authorizations
104            WHERE device_code = $1
105            "#,
106        )
107        .bind(device_code)
108        .fetch_optional(&self.pool)
109        .await?;
110
111        Ok(device_auth)
112    }
113
114    async fn get_device_authorization_by_user_code(
115        &self,
116        user_code: &str,
117    ) -> Result<Option<DeviceAuthorization>, Error> {
118        let device_auth = sqlx::query_as::<_, DeviceAuthorization>(
119            r#"
120            SELECT id, device_code, user_code, client_id, user_id, status, expires_at, created_at
121            FROM device_authorizations
122            WHERE user_code = $1
123            "#,
124        )
125        .bind(user_code)
126        .fetch_optional(&self.pool)
127        .await?;
128
129        Ok(device_auth)
130    }
131
132    async fn expire_device_authorization(&self, id: Uuid) -> Result<(), Error> {
133        sqlx::query(
134            r#"
135            UPDATE device_authorizations
136            SET status = 'expired'
137            WHERE id = $1
138            "#,
139        )
140        .bind(id)
141        .execute(&self.pool)
142        .await?;
143
144        Ok(())
145    }
146
147    async fn authorize_device(
148        &self,
149        user_code: &str,
150        user_id: Uuid,
151    ) -> Result<Option<DeviceAuthorization>, Error> {
152        let device_auth = sqlx::query_as::<_, DeviceAuthorization>(
153            r#"
154            UPDATE device_authorizations
155            SET status = 'authorized', user_id = $2
156            WHERE user_code = $1 AND status = 'pending' AND expires_at > NOW()
157            RETURNING id, device_code, user_code, client_id, user_id, status, expires_at, created_at
158            "#,
159        )
160        .bind(user_code)
161        .bind(user_id)
162        .fetch_optional(&self.pool)
163        .await?;
164
165        Ok(device_auth)
166    }
167
168    async fn deny_device(
169        &self,
170        user_code: &str,
171        user_id: Uuid,
172    ) -> Result<Option<DeviceAuthorization>, Error> {
173        let device_auth = sqlx::query_as::<_, DeviceAuthorization>(
174            r#"
175            UPDATE device_authorizations
176            SET status = 'denied', user_id = $2
177            WHERE user_code = $1 AND status = 'pending' AND expires_at > NOW()
178            RETURNING id, device_code, user_code, client_id, user_id, status, expires_at, created_at
179            "#,
180        )
181        .bind(user_code)
182        .bind(user_id)
183        .fetch_optional(&self.pool)
184        .await?;
185
186        Ok(device_auth)
187    }
188
189    async fn create_access_token(
190        &self,
191        user_id: Uuid,
192        client_id: &str,
193        token_hash: &str,
194    ) -> Result<AccessToken, Error> {
195        let token = sqlx::query_as::<_, AccessToken>(
196            r#"
197            INSERT INTO access_tokens (user_id, client_id, token_hash)
198            VALUES ($1, $2, $3)
199            RETURNING id, user_id, client_id, token_hash, created_at, last_used_at
200            "#,
201        )
202        .bind(user_id)
203        .bind(client_id)
204        .bind(token_hash)
205        .fetch_one(&self.pool)
206        .await?;
207
208        Ok(token)
209    }
210
211    async fn get_access_token_by_hash(
212        &self,
213        token_hash: &str,
214    ) -> Result<Option<AccessToken>, Error> {
215        let token = sqlx::query_as::<_, AccessToken>(
216            r#"
217            SELECT id, user_id, client_id, token_hash, created_at, last_used_at
218            FROM access_tokens
219            WHERE token_hash = $1
220            "#,
221        )
222        .bind(token_hash)
223        .fetch_optional(&self.pool)
224        .await?;
225
226        Ok(token)
227    }
228
229    async fn touch_access_token(&self, id: Uuid) -> Result<(), Error> {
230        sqlx::query("UPDATE access_tokens SET last_used_at = NOW() WHERE id = $1")
231            .bind(id)
232            .execute(&self.pool)
233            .await?;
234
235        Ok(())
236    }
237
238    async fn delete_access_token(&self, id: Uuid) -> Result<(), Error> {
239        sqlx::query("DELETE FROM access_tokens WHERE id = $1")
240            .bind(id)
241            .execute(&self.pool)
242            .await?;
243
244        Ok(())
245    }
246}
247