arrow_integration_testing/flight_server_scenarios/
auth_basic_proto.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: See the NOTICE file distributed with this work for additional information regarding copyright ownership
3// SPDX-FileContributor: Licensed to the Apache Software Foundation (ASF) under one or more contributor license agreements
4
5//! Basic auth test for the Flight server.
6
7use std::pin::Pin;
8use std::sync::Arc;
9
10use arrow_flight::{
11    Action, ActionType, BasicAuth, Criteria, Empty, FlightData, FlightDescriptor, FlightInfo,
12    HandshakeRequest, HandshakeResponse, PollInfo, PutResult, SchemaResult, Ticket,
13    flight_service_server::FlightService, flight_service_server::FlightServiceServer,
14};
15use futures::{Stream, StreamExt, channel::mpsc, sink::SinkExt};
16use tokio::sync::Mutex;
17use tonic::{Request, Response, Status, Streaming, metadata::MetadataMap, transport::Server};
18type TonicStream<T> = Pin<Box<dyn Stream<Item = T> + Send + Sync + 'static>>;
19
20type Error = Box<dyn std::error::Error + Send + Sync + 'static>;
21type Result<T = (), E = Error> = std::result::Result<T, E>;
22
23use prost::Message;
24
25use crate::{AUTH_PASSWORD, AUTH_USERNAME};
26
27/// Run a scenario that tests basic auth.
28pub async fn scenario_setup(port: u16) -> Result {
29    let service = AuthBasicProtoScenarioImpl {
30        username: AUTH_USERNAME.into(),
31        password: AUTH_PASSWORD.into(),
32        peer_identity: Arc::new(Mutex::new(None)),
33    };
34    let addr = super::listen_on(port).await?;
35    let svc = FlightServiceServer::new(service);
36
37    let server = Server::builder().add_service(svc).serve(addr);
38
39    // NOTE: Log output used in tests to signal server is ready
40    println!("Server listening on localhost:{}", addr.port());
41    server.await?;
42    Ok(())
43}
44
45/// Scenario for testing basic auth.
46#[derive(Clone)]
47pub struct AuthBasicProtoScenarioImpl {
48    username: Arc<str>,
49    password: Arc<str>,
50    #[expect(dead_code)]
51    peer_identity: Arc<Mutex<Option<String>>>,
52}
53
54impl AuthBasicProtoScenarioImpl {
55    fn check_auth(&self, metadata: &MetadataMap) -> Result<GrpcServerCallContext, Status> {
56        let token = metadata
57            .get_bin("auth-token-bin")
58            .and_then(|v| v.to_bytes().ok())
59            .and_then(|b| String::from_utf8(b.to_vec()).ok());
60        self.is_valid(token)
61    }
62
63    fn is_valid(&self, token: Option<String>) -> Result<GrpcServerCallContext, Status> {
64        match token {
65            Some(t) if t == *self.username => Ok(GrpcServerCallContext {
66                peer_identity: self.username.to_string(),
67            }),
68            _ => Err(Status::unauthenticated("Invalid token")),
69        }
70    }
71}
72
73struct GrpcServerCallContext {
74    peer_identity: String,
75}
76
77impl GrpcServerCallContext {
78    pub fn peer_identity(&self) -> &str {
79        &self.peer_identity
80    }
81}
82
83#[tonic::async_trait]
84impl FlightService for AuthBasicProtoScenarioImpl {
85    type HandshakeStream = TonicStream<Result<HandshakeResponse, Status>>;
86    type ListFlightsStream = TonicStream<Result<FlightInfo, Status>>;
87    type DoGetStream = TonicStream<Result<FlightData, Status>>;
88    type DoPutStream = TonicStream<Result<PutResult, Status>>;
89    type DoActionStream = TonicStream<Result<arrow_flight::Result, Status>>;
90    type ListActionsStream = TonicStream<Result<ActionType, Status>>;
91    type DoExchangeStream = TonicStream<Result<FlightData, Status>>;
92
93    async fn get_schema(
94        &self,
95        request: Request<FlightDescriptor>,
96    ) -> Result<Response<SchemaResult>, Status> {
97        self.check_auth(request.metadata())?;
98        Err(Status::unimplemented("Not yet implemented"))
99    }
100
101    async fn do_get(
102        &self,
103        request: Request<Ticket>,
104    ) -> Result<Response<Self::DoGetStream>, Status> {
105        self.check_auth(request.metadata())?;
106        Err(Status::unimplemented("Not yet implemented"))
107    }
108
109    async fn handshake(
110        &self,
111        request: Request<Streaming<HandshakeRequest>>,
112    ) -> Result<Response<Self::HandshakeStream>, Status> {
113        let (tx, rx) = mpsc::channel(10);
114
115        tokio::spawn({
116            let username = self.username.clone();
117            let password = self.password.clone();
118
119            async move {
120                let requests = request.into_inner();
121
122                requests
123                    .for_each(move |req| {
124                        let mut tx = tx.clone();
125                        let req = req.expect("Error reading handshake request");
126                        let HandshakeRequest { payload, .. } = req;
127
128                        let auth =
129                            BasicAuth::decode(&*payload).expect("Error parsing handshake request");
130
131                        let resp = if *auth.username == *username && *auth.password == *password {
132                            Ok(HandshakeResponse {
133                                payload: username.as_bytes().to_vec().into(),
134                                ..HandshakeResponse::default()
135                            })
136                        } else {
137                            Err(Status::unauthenticated(format!(
138                                "Don't know user {}",
139                                auth.username
140                            )))
141                        };
142
143                        async move {
144                            tx.send(resp)
145                                .await
146                                .expect("Error sending handshake response");
147                        }
148                    })
149                    .await;
150            }
151        });
152
153        Ok(Response::new(Box::pin(rx)))
154    }
155
156    async fn list_flights(
157        &self,
158        request: Request<Criteria>,
159    ) -> Result<Response<Self::ListFlightsStream>, Status> {
160        self.check_auth(request.metadata())?;
161        Err(Status::unimplemented("Not yet implemented"))
162    }
163
164    async fn get_flight_info(
165        &self,
166        request: Request<FlightDescriptor>,
167    ) -> Result<Response<FlightInfo>, Status> {
168        self.check_auth(request.metadata())?;
169        Err(Status::unimplemented("Not yet implemented"))
170    }
171
172    async fn poll_flight_info(
173        &self,
174        request: Request<FlightDescriptor>,
175    ) -> Result<Response<PollInfo>, Status> {
176        self.check_auth(request.metadata())?;
177        Err(Status::unimplemented("Not yet implemented"))
178    }
179
180    async fn do_put(
181        &self,
182        request: Request<Streaming<FlightData>>,
183    ) -> Result<Response<Self::DoPutStream>, Status> {
184        let metadata = request.metadata();
185        self.check_auth(metadata)?;
186        Err(Status::unimplemented("Not yet implemented"))
187    }
188
189    async fn do_action(
190        &self,
191        request: Request<Action>,
192    ) -> Result<Response<Self::DoActionStream>, Status> {
193        let flight_context = self.check_auth(request.metadata())?;
194        // Respond with the authenticated username.
195        let buf = flight_context.peer_identity().as_bytes().to_vec().into();
196        let result = arrow_flight::Result { body: buf };
197        let output = futures::stream::once(async { Ok(result) });
198        Ok(Response::new(Box::pin(output) as Self::DoActionStream))
199    }
200
201    async fn list_actions(
202        &self,
203        request: Request<Empty>,
204    ) -> Result<Response<Self::ListActionsStream>, Status> {
205        self.check_auth(request.metadata())?;
206        Err(Status::unimplemented("Not yet implemented"))
207    }
208
209    async fn do_exchange(
210        &self,
211        request: Request<Streaming<FlightData>>,
212    ) -> Result<Response<Self::DoExchangeStream>, Status> {
213        let metadata = request.metadata();
214        self.check_auth(metadata)?;
215        Err(Status::unimplemented("Not yet implemented"))
216    }
217}
arrow_integration_testing/flight_server_scenarios/auth_basic_proto.rs

arrow_integration_testing/flight_server_scenarios/
auth_basic_proto.rs
