az_gmail_code/
client.rs

1use crate::config::GmailCodeQuery;
2use crate::model::{ExtractedGmailCode, GmailListMessagesResponse, GmailMessage};
3use crate::parser::{collect_message_body_candidates, extract_verification_code};
4use crate::{GmailCodeConfig, GmailCodeError, GmailCodeResult};
5use reqwest::Url;
6use reqwest::blocking::{Client, RequestBuilder, Response};
7use serde::de::DeserializeOwned;
8
9/// Blocking Gmail API client for reading verification codes from an authorized mailbox.
10#[derive(Debug, Clone)]
11pub struct GmailCodeClient {
12    base_url: Url,
13    user_id: String,
14    access_token: String,
15    client: Client,
16}
17
18impl GmailCodeClient {
19    /// Creates a client with default Gmail API settings and the required OAuth access token.
20    pub fn new(access_token: impl Into<String>) -> GmailCodeResult<Self> {
21        Self::with_config(GmailCodeConfig::builder(access_token).build()?)
22    }
23
24    /// Creates a client from explicit configuration.
25    pub fn with_config(config: GmailCodeConfig) -> GmailCodeResult<Self> {
26        config.validate()?;
27        let base_url = Url::parse(&config.base_url)
28            .map_err(|_| GmailCodeError::InvalidBaseUrl(config.base_url.clone()))?;
29
30        let mut builder = Client::builder()
31            .connect_timeout(config.connect_timeout)
32            .timeout(config.request_timeout);
33        if let Some(user_agent) = &config.user_agent {
34            builder = builder.user_agent(user_agent);
35        }
36
37        Ok(Self {
38            base_url,
39            user_id: config.user_id,
40            access_token: config.access_token,
41            client: builder.build()?,
42        })
43    }
44
45    /// Searches Gmail messages using the query object and returns matching message ids.
46    pub fn list_message_ids(&self, query: &GmailCodeQuery) -> GmailCodeResult<Vec<String>> {
47        let response = self.list_messages(query)?;
48        Ok(response
49            .messages
50            .into_iter()
51            .map(|message| message.id)
52            .collect())
53    }
54
55    /// Fetches one Gmail message by id with `format=full`.
56    pub fn get_message(&self, message_id: impl AsRef<str>) -> GmailCodeResult<GmailMessage> {
57        let message_id = message_id.as_ref();
58        let path = format!(
59            "users/{}/messages/{}",
60            urlencoding::encode(&self.user_id),
61            urlencoding::encode(message_id)
62        );
63        let response = self
64            .authenticated_get(&path)?
65            .query(&[("format", "full")])
66            .send()?;
67        Self::read_json(response)
68    }
69
70    /// Finds the latest verification code from messages matching the query.
71    ///
72    /// Gmail returns list results newest-first for typical searches, so this method
73    /// inspects messages in listed order and returns the first code found.
74    pub fn find_latest_code(
75        &self,
76        query: GmailCodeQuery,
77    ) -> GmailCodeResult<Option<ExtractedGmailCode>> {
78        let response = self.list_messages(&query)?;
79        for summary in response.messages {
80            let message = self.get_message(&summary.id)?;
81            if let Some(code) = extract_code_from_message(&message)? {
82                return Ok(Some(code));
83            }
84        }
85        Ok(None)
86    }
87
88    fn list_messages(&self, query: &GmailCodeQuery) -> GmailCodeResult<GmailListMessagesResponse> {
89        let path = format!("users/{}/messages", urlencoding::encode(&self.user_id));
90        let gmail_q = query.gmail_q();
91        let mut request = self
92            .authenticated_get(&path)?
93            .query(&[("maxResults", query.max_results.to_string())]);
94
95        if !gmail_q.is_empty() {
96            request = request.query(&[("q", gmail_q)]);
97        }
98        for label_id in &query.label_ids {
99            if !label_id.trim().is_empty() {
100                request = request.query(&[("labelIds", label_id)]);
101            }
102        }
103
104        Self::read_json(request.send()?)
105    }
106
107    fn authenticated_get(&self, path: &str) -> GmailCodeResult<RequestBuilder> {
108        Ok(self
109            .client
110            .get(self.join_url(path)?)
111            .bearer_auth(&self.access_token))
112    }
113
114    fn read_json<T: DeserializeOwned>(response: Response) -> GmailCodeResult<T> {
115        let response = Self::ensure_success(response)?;
116        let bytes = response.bytes()?;
117        Ok(serde_json::from_slice(bytes.as_ref())?)
118    }
119
120    fn ensure_success(response: Response) -> GmailCodeResult<Response> {
121        let status = response.status();
122        if status.is_success() {
123            return Ok(response);
124        }
125
126        let url = response.url().to_string();
127        let body = match response.bytes() {
128            Ok(bytes) => String::from_utf8_lossy(bytes.as_ref()).into_owned(),
129            Err(error) => return Err(GmailCodeError::Transport(error)),
130        };
131        Err(GmailCodeError::HttpStatus {
132            url,
133            status: status.as_u16(),
134            body,
135        })
136    }
137
138    fn join_url(&self, path: &str) -> GmailCodeResult<Url> {
139        self.base_url
140            .join(path)
141            .map_err(|_| GmailCodeError::InvalidPath(path.to_owned()))
142    }
143}
144
145fn extract_code_from_message(
146    message: &GmailMessage,
147) -> GmailCodeResult<Option<ExtractedGmailCode>> {
148    for candidate in collect_message_body_candidates(message)? {
149        if let Some(code) = extract_verification_code(&candidate.text) {
150            return Ok(Some(ExtractedGmailCode {
151                code,
152                message_id: message.id.clone(),
153                thread_id: message.thread_id.clone(),
154                from: message.header("From").map(ToOwned::to_owned),
155                subject: message.header("Subject").map(ToOwned::to_owned),
156                source_mime_type: candidate.mime_type,
157            }));
158        }
159    }
160    Ok(None)
161}
162
163#[cfg(test)]
164mod tests {
165    use super::GmailCodeClient;
166    use crate::{GmailCodeConfig, GmailCodeQuery};
167    use base64::Engine;
168    use base64::engine::general_purpose::URL_SAFE_NO_PAD;
169    use std::io::{Read, Write};
170    use std::net::TcpListener;
171    use std::sync::{Arc, Mutex};
172    use std::thread;
173    use std::time::Duration;
174
175    #[test]
176    fn find_latest_code_searches_and_fetches_message() {
177        let server = MockGmailServer::spawn(vec![
178            MockResponse::json(
179                200,
180                r#"{"messages":[{"id":"msg-1","threadId":"th-1"}],"resultSizeEstimate":1}"#,
181            ),
182            MockResponse::json(
183                200,
184                &format!(
185                    r#"{{
186                        "id":"msg-1",
187                        "threadId":"th-1",
188                        "payload":{{
189                            "mimeType":"multipart/alternative",
190                            "headers":[
191                                {{"name":"From","value":"security@example.com"}},
192                                {{"name":"Subject","value":"Login verification"}}
193                            ],
194                            "parts":[
195                                {{
196                                    "partId":"1",
197                                    "mimeType":"text/plain",
198                                    "body":{{"data":"{}"}}
199                                }}
200                            ]
201                        }}
202                    }}"#,
203                    encode("Your verification code: 246810")
204                ),
205            ),
206        ]);
207        let client = test_client(server.base_url());
208
209        let code = client
210            .find_latest_code(
211                GmailCodeQuery::new()
212                    .from("security@example.com")
213                    .subject("Login verification")
214                    .newer_than("10m")
215                    .unread(true),
216            )
217            .expect("code lookup")
218            .expect("code should exist");
219
220        assert_eq!(code.code, "246810");
221        assert_eq!(code.message_id, "msg-1");
222        assert_eq!(code.from.as_deref(), Some("security@example.com"));
223        assert_eq!(code.subject.as_deref(), Some("Login verification"));
224
225        let requests = server.requests();
226        assert_eq!(requests.len(), 2);
227        assert!(requests[0].starts_with("GET /gmail/v1/users/me/messages?"));
228        assert!(requests[0].contains("maxResults=10"));
229        assert!(requests[0].contains("labelIds=INBOX"));
230        assert!(
231            requests[0]
232                .to_ascii_lowercase()
233                .contains("authorization: bearer test-token")
234        );
235        assert!(requests[1].starts_with("GET /gmail/v1/users/me/messages/msg-1?format=full"));
236    }
237
238    #[test]
239    fn http_status_keeps_response_body() {
240        let server = MockGmailServer::spawn(vec![MockResponse::text(401, "bad token")]);
241        let client = test_client(server.base_url());
242
243        let error = client
244            .list_message_ids(&GmailCodeQuery::new())
245            .expect_err("401 should fail");
246
247        assert!(error.to_string().contains("HTTP 401"));
248        assert!(error.to_string().contains("bad token"));
249    }
250
251    fn test_client(base_url: String) -> GmailCodeClient {
252        let config = GmailCodeConfig::builder("test-token")
253            .base_url(base_url)
254            .request_timeout(Duration::from_secs(3))
255            .connect_timeout(Duration::from_secs(3))
256            .build()
257            .expect("config should build");
258        GmailCodeClient::with_config(config).expect("client should build")
259    }
260
261    fn encode(value: &str) -> String {
262        URL_SAFE_NO_PAD.encode(value.as_bytes())
263    }
264
265    #[derive(Debug, Clone)]
266    struct MockResponse {
267        status: u16,
268        content_type: &'static str,
269        body: String,
270    }
271
272    impl MockResponse {
273        fn json(status: u16, body: &str) -> Self {
274            Self {
275                status,
276                content_type: "application/json",
277                body: body.to_owned(),
278            }
279        }
280
281        fn text(status: u16, body: &str) -> Self {
282            Self {
283                status,
284                content_type: "text/plain",
285                body: body.to_owned(),
286            }
287        }
288    }
289
290    struct MockGmailServer {
291        base_url: String,
292        requests: Arc<Mutex<Vec<String>>>,
293        handle: thread::JoinHandle<()>,
294    }
295
296    impl MockGmailServer {
297        fn spawn(responses: Vec<MockResponse>) -> Self {
298            let listener = TcpListener::bind("127.0.0.1:0").expect("listener should bind");
299            let address = listener.local_addr().expect("address should exist");
300            let requests = Arc::new(Mutex::new(Vec::new()));
301            let thread_requests = Arc::clone(&requests);
302            let handle = thread::spawn(move || {
303                for response in responses {
304                    let (mut stream, _) = listener.accept().expect("connection should arrive");
305                    stream
306                        .set_read_timeout(Some(Duration::from_secs(2)))
307                        .expect("timeout should set");
308
309                    let request = read_http_request(&mut stream);
310                    thread_requests.lock().expect("requests lock").push(request);
311
312                    let payload = format!(
313                        "HTTP/1.1 {} OK\r\nContent-Length: {}\r\nContent-Type: {}\r\nConnection: close\r\n\r\n{}",
314                        response.status,
315                        response.body.len(),
316                        response.content_type,
317                        response.body
318                    );
319                    stream
320                        .write_all(payload.as_bytes())
321                        .expect("response should write");
322                }
323            });
324
325            Self {
326                base_url: format!("http://{address}/gmail/v1/"),
327                requests,
328                handle,
329            }
330        }
331
332        fn base_url(&self) -> String {
333            self.base_url.clone()
334        }
335
336        fn requests(self) -> Vec<String> {
337            self.handle.join().expect("server should finish");
338            Arc::try_unwrap(self.requests)
339                .expect("server requests should be owned")
340                .into_inner()
341                .expect("requests lock")
342        }
343    }
344
345    fn read_http_request(stream: &mut std::net::TcpStream) -> String {
346        let mut buffer = Vec::new();
347        let mut chunk = [0u8; 1024];
348        loop {
349            let read = stream.read(&mut chunk).expect("request should read");
350            if read == 0 {
351                break;
352            }
353            buffer.extend_from_slice(&chunk[..read]);
354            if buffer.windows(4).any(|window| window == b"\r\n\r\n") {
355                break;
356            }
357        }
358        String::from_utf8_lossy(&buffer).into_owned()
359    }
360}
az_gmail_code/client.rs

az_gmail_code/
client.rs
