1use crate::models::{
7 AnalyzeRequest, AnalyzeResponse, AnonymizeRequest, AnonymizeResponse, EntityResult,
8 ErrorResponse, HealthResponse, TokenInfo,
9};
10use axum::{
11 extract::State,
12 http::StatusCode,
13 response::{IntoResponse, Response},
14 Json,
15};
16use redact_core::{AnalyzerEngine, AnonymizerConfig, EntityType};
17use std::sync::Arc;
18
19#[derive(Clone)]
21pub struct AppState {
22 pub engine: Arc<AnalyzerEngine>,
23}
24
25#[derive(Debug)]
27pub struct ApiError {
28 status: StatusCode,
29 message: String,
30}
31
32impl ApiError {
33 pub fn new(status: StatusCode, message: impl Into<String>) -> Self {
34 Self {
35 status,
36 message: message.into(),
37 }
38 }
39
40 pub fn bad_request(message: impl Into<String>) -> Self {
41 Self::new(StatusCode::BAD_REQUEST, message)
42 }
43
44 pub fn internal_error(message: impl Into<String>) -> Self {
45 Self::new(StatusCode::INTERNAL_SERVER_ERROR, message)
46 }
47}
48
49impl IntoResponse for ApiError {
50 fn into_response(self) -> Response {
51 let body = Json(ErrorResponse::new("error", self.message));
52 (self.status, body).into_response()
53 }
54}
55
56impl From<anyhow::Error> for ApiError {
57 fn from(err: anyhow::Error) -> Self {
58 ApiError::internal_error(err.to_string())
59 }
60}
61
62pub async fn health(State(state): State<AppState>) -> Json<HealthResponse> {
64 Json(HealthResponse {
65 status: "healthy".to_string(),
66 version: redact_core::VERSION.to_string(),
67 recognizers: state.engine.recognizer_registry().recognizers().len(),
68 })
69}
70
71pub async fn analyze(
73 State(state): State<AppState>,
74 Json(request): Json<AnalyzeRequest>,
75) -> Result<Json<AnalyzeResponse>, ApiError> {
76 if request.text.is_empty() {
78 return Err(ApiError::bad_request("Text cannot be empty"));
79 }
80
81 let entity_types: Option<Vec<EntityType>> = request.entities.as_ref().map(|entities| {
83 entities
84 .iter()
85 .map(|e| EntityType::from(e.clone()))
86 .collect()
87 });
88
89 let result = if let Some(entities) = entity_types.as_ref() {
91 state
92 .engine
93 .analyze_with_entities(&request.text, entities, Some(&request.language))
94 .map_err(ApiError::from)?
95 } else {
96 state
97 .engine
98 .analyze(&request.text, Some(&request.language))
99 .map_err(ApiError::from)?
100 };
101
102 let mut results: Vec<EntityResult> = result
104 .detected_entities
105 .into_iter()
106 .filter(|e| {
107 if let Some(min_score) = request.min_score {
108 e.score >= min_score
109 } else {
110 true
111 }
112 })
113 .map(EntityResult::from)
114 .collect();
115
116 results.sort_by_key(|r| r.start);
118
119 Ok(Json(AnalyzeResponse {
120 original_text: None,
121 results,
122 metadata: result.metadata.into(),
123 }))
124}
125
126pub async fn anonymize(
128 State(state): State<AppState>,
129 Json(request): Json<AnonymizeRequest>,
130) -> Result<Json<AnonymizeResponse>, ApiError> {
131 if request.text.is_empty() {
133 return Err(ApiError::bad_request("Text cannot be empty"));
134 }
135
136 if request.config.strategy == redact_core::AnonymizationStrategy::Encrypt
138 && request.config.encryption_key.is_none()
139 {
140 return Err(ApiError::bad_request(
141 "Encryption key required for encrypt strategy",
142 ));
143 }
144
145 let mask_char = request.config.mask_char.chars().next().unwrap_or('*');
147
148 let core_config = AnonymizerConfig {
149 strategy: request.config.strategy,
150 mask_char,
151 mask_start_chars: request.config.mask_start_chars,
152 mask_end_chars: request.config.mask_end_chars,
153 preserve_format: request.config.preserve_format,
154 encryption_key: request.config.encryption_key,
155 hash_salt: request.config.hash_salt,
156 };
157
158 let entity_types: Option<Vec<EntityType>> = request.entities.as_ref().map(|entities| {
160 entities
161 .iter()
162 .map(|e| EntityType::from(e.clone()))
163 .collect()
164 });
165
166 let result = if let Some(entities) = entity_types.as_ref() {
168 let analysis = state
170 .engine
171 .analyze_with_entities(&request.text, entities, Some(&request.language))
172 .map_err(ApiError::from)?;
173
174 let anonymized = state
176 .engine
177 .anonymizer_registry()
178 .anonymize(
179 &request.text,
180 analysis.detected_entities.clone(),
181 &core_config,
182 )
183 .map_err(ApiError::from)?;
184
185 (analysis.detected_entities, anonymized, analysis.metadata)
186 } else {
187 let analysis = state
188 .engine
189 .analyze_and_anonymize(&request.text, Some(&request.language), &core_config)
190 .map_err(ApiError::from)?;
191
192 let anonymized = analysis
193 .anonymized
194 .ok_or_else(|| ApiError::internal_error("Anonymization failed"))?;
195
196 (analysis.detected_entities, anonymized, analysis.metadata)
197 };
198
199 let (detected_entities, anonymized, metadata) = result;
200
201 let results: Vec<EntityResult> = detected_entities
203 .into_iter()
204 .map(EntityResult::from)
205 .collect();
206
207 let tokens: Option<Vec<TokenInfo>> = anonymized
208 .tokens
209 .map(|tokens| tokens.into_iter().map(TokenInfo::from).collect());
210
211 Ok(Json(AnonymizeResponse {
212 text: anonymized.text,
213 results,
214 tokens,
215 metadata: metadata.into(),
216 }))
217}
218
219#[cfg(test)]
220mod tests {
221 use super::*;
222 use crate::models::AnonymizationConfig;
223
224 fn create_test_state() -> AppState {
225 AppState {
226 engine: Arc::new(AnalyzerEngine::new()),
227 }
228 }
229
230 #[tokio::test]
231 async fn test_health() {
232 let state = create_test_state();
233 let response = health(State(state)).await;
234
235 assert_eq!(response.status, "healthy");
236 assert!(!response.version.is_empty());
237 }
238
239 #[tokio::test]
240 async fn test_analyze() {
241 let state = create_test_state();
242 let request = AnalyzeRequest {
243 text: "Email: john@example.com".to_string(),
244 language: "en".to_string(),
245 entities: None,
246 min_score: None,
247 };
248
249 let response = analyze(State(state), Json(request)).await.unwrap();
250
251 assert!(!response.results.is_empty());
252 assert_eq!(response.results[0].entity_type, "EMAIL_ADDRESS");
253 }
254
255 #[tokio::test]
256 async fn test_anonymize() {
257 let state = create_test_state();
258 let request = AnonymizeRequest {
259 text: "Email: john@example.com".to_string(),
260 language: "en".to_string(),
261 config: AnonymizationConfig::default(),
262 entities: None,
263 };
264
265 let response = anonymize(State(state), Json(request)).await.unwrap();
266
267 assert!(response.text.contains("[EMAIL_ADDRESS]"));
268 assert!(!response.results.is_empty());
269 }
270}