snap_control/crpc_api/
api_service.rs1use std::{
17 sync::Arc,
18 time::{Instant, SystemTime},
19};
20
21use axum::{
22 Extension, Router,
23 extract::{ConnectInfo, State},
24};
25use scion_sdk_axum_connect_rpc::{
26 error::{CrpcError, CrpcErrorCode},
27 extractor::ConnectRpc,
28};
29use scion_sdk_token_validator::validator::Token;
30use snap_tokens::snap_token::SnapTokenClaims;
31use x25519_dalek::PublicKey;
32
33use crate::{
34 crpc_api::api_service::model::{SnapDataPlaneResolver, SnapTunIdentityRegistry},
35 protobuf::anapaya::snap::v1::api_service::{
36 GetSnapDataPlaneRequest, GetSnapDataPlaneResponse, RegisterSnapTunIdentityRequest,
37 RegisterSnapTunIdentityResponse,
38 },
39};
40
41pub mod model {
43 use std::{
44 net::{IpAddr, SocketAddr},
45 time::{Duration, Instant},
46 };
47
48 use axum::http::StatusCode;
49 use x25519_dalek::PublicKey;
50
51 pub trait SnapDataPlaneResolver: Send + Sync {
53 fn get_data_plane_address(
55 &self,
56 endhost_ip: IpAddr,
57 ) -> Result<SnapDataPlane, (StatusCode, anyhow::Error)>;
58 }
59
60 pub struct SnapDataPlane {
62 pub address: SocketAddr,
65 pub snap_tun_control_address: Option<SocketAddr>,
69 pub snap_static_x25519: Option<PublicKey>,
72 }
73
74 pub trait SnapTunIdentityRegistry: Send + Sync {
76 fn register(
86 &self,
87 now: Instant,
88 key: &str,
90 initiator_identity: [u8; 32],
92 psk_share: Option<[u8; 32]>,
94 lifetime: Duration,
97 ) -> bool;
98 }
99}
100
101pub(crate) mod convert {
102 use std::net::AddrParseError;
103
104 use x25519_dalek::PublicKey;
105
106 use crate::{
107 crpc_api::api_service::model::SnapDataPlane,
108 protobuf::anapaya::snap::v1::api_service as rpc,
109 };
110
111 #[derive(thiserror::Error, Debug)]
113 pub enum ConvertError {
114 #[error("failed to parse data plane address: {0}")]
115 ParseAddr(AddrParseError),
116 #[error("failed to parse server control address: {0}")]
117 ParseSnapTunControlAddr(AddrParseError),
118 #[error("server static identity is not 32 bytes")]
119 InvalidServerStaticIdentityLength,
120 }
121
122 impl TryFrom<rpc::GetSnapDataPlaneResponse> for SnapDataPlane {
124 type Error = ConvertError;
125 fn try_from(value: rpc::GetSnapDataPlaneResponse) -> Result<Self, Self::Error> {
126 let snap_tun_control_address = value
127 .snap_tun_control_address
128 .map(|address| {
129 address
130 .parse()
131 .map_err(ConvertError::ParseSnapTunControlAddr)
132 })
133 .transpose()?;
134 let snap_static_x25519 = value
135 .snap_static_x25519
136 .map(|key| {
137 TryInto::<[u8; 32]>::try_into(key.as_slice())
138 .map_err(|_| ConvertError::InvalidServerStaticIdentityLength)
139 .map(PublicKey::from)
140 })
141 .transpose()?;
142 Ok(SnapDataPlane {
143 address: value.address.parse().map_err(ConvertError::ParseAddr)?,
144 snap_tun_control_address,
145 snap_static_x25519,
146 })
147 }
148 }
149}
150
151pub(crate) const SERVICE_PATH: &str = "/anapaya.snap.v1.SnapControl";
152pub(crate) const GET_SNAP_DATA_PLANE_ADDRESS: &str = "/GetSnapDataPlaneAddress";
153pub(crate) const REGISTER_SNAPTUN_IDENTITY: &str = "/RegisterSnapTunIdentity";
154
155pub fn nest_snap_control_api(
157 router: axum::Router,
158 snap_resolver: Arc<dyn SnapDataPlaneResolver>,
159 identity_registrar: Arc<dyn SnapTunIdentityRegistry>,
160) -> axum::Router {
161 router.nest(
162 SERVICE_PATH,
163 Router::new()
164 .route(
165 GET_SNAP_DATA_PLANE_ADDRESS,
166 axum::routing::post(get_snap_data_plane_address_handler),
167 )
168 .with_state(snap_resolver)
169 .route(
170 REGISTER_SNAPTUN_IDENTITY,
171 axum::routing::post(register_snaptun_identity_handler),
172 )
173 .with_state(identity_registrar),
174 )
175}
176
177async fn get_snap_data_plane_address_handler(
178 State(rendezvous_hasher): State<Arc<dyn SnapDataPlaneResolver>>,
179 _snap_token: Extension<SnapTokenClaims>,
180 ConnectInfo(addr): ConnectInfo<std::net::SocketAddr>,
181 ConnectRpc(_request): ConnectRpc<GetSnapDataPlaneRequest>,
182) -> Result<ConnectRpc<GetSnapDataPlaneResponse>, CrpcError> {
183 let addr = rendezvous_hasher.get_data_plane_address(addr.ip())?;
184 Ok(ConnectRpc(GetSnapDataPlaneResponse {
185 address: addr.address.to_string(),
186 snap_tun_control_address: addr
187 .snap_tun_control_address
188 .map(|address| address.to_string()),
189 snap_static_x25519: addr.snap_static_x25519.map(|key| key.to_bytes().to_vec()),
190 }))
191}
192
193async fn register_snaptun_identity_handler(
194 State(identity_registry): State<Arc<dyn SnapTunIdentityRegistry>>,
195 snap_token: Extension<SnapTokenClaims>,
196 ConnectInfo(_): ConnectInfo<std::net::SocketAddr>,
197 ConnectRpc(request): ConnectRpc<RegisterSnapTunIdentityRequest>,
198) -> Result<ConnectRpc<RegisterSnapTunIdentityResponse>, CrpcError> {
199 let now = SystemTime::now();
200 let lifetime = snap_token.0.exp_time().duration_since(now).map_err(|_| {
201 CrpcError::new(
202 CrpcErrorCode::InvalidArgument,
203 "expiration time is in the past".to_string(),
204 )
205 })?;
206
207 let initiator_identity = {
208 let key_bytes: [u8; 32] = request
209 .initiator_static_x25519
210 .as_slice()
211 .try_into()
212 .map_err(|_| {
213 CrpcError::new(
214 CrpcErrorCode::InvalidArgument,
215 "initiator identity is not 32 bytes".to_string(),
216 )
217 })?;
218 PublicKey::from(key_bytes)
219 };
220
221 let psk_share: Option<[u8; 32]> = if request.psk_share.as_slice() == [0u8; 32] {
222 None
223 } else {
224 Some(request.psk_share.as_slice().try_into().map_err(|_| {
225 CrpcError::new(
226 CrpcErrorCode::InvalidArgument,
227 "psk share is not 32 bytes".to_string(),
228 )
229 })?)
230 };
231
232 let pssid = &snap_token.pssid.to_string();
233 if !identity_registry.register(
234 Instant::now(),
235 pssid,
236 *initiator_identity.as_bytes(),
237 psk_share,
238 lifetime,
239 ) {
240 tracing::info!(pssid, "re-registered identity");
241 }
242 Ok(ConnectRpc(RegisterSnapTunIdentityResponse {
243 psk_share: [0u8; 32].to_vec(),
245 }))
246}