Skip to main content

snap_control/crpc_api/
api_service.rs

1// Copyright 2025 Anapaya Systems
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//   http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14//! Connect RPC API endpoint definitions and endpoint handlers.
15
16use 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::AnyClaims;
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
41/// SNAP control plane API models.
42pub mod model {
43    use std::{
44        net::{IpAddr, SocketAddr},
45        time::{Duration, Instant},
46    };
47
48    use axum::http::StatusCode;
49    use url::Url;
50    use x25519_dalek::PublicKey;
51
52    /// SNAP data plane discovery trait.
53    pub trait SnapDataPlaneResolver: Send + Sync {
54        /// Get the SNAP data plane address for a given endhost IP address.
55        fn get_data_plane_address(
56            &self,
57            endhost_ip: IpAddr,
58        ) -> Result<SnapDataPlane, (StatusCode, anyhow::Error)>;
59    }
60
61    /// SnapDataPlane resolution response.
62    pub struct SnapDataPlane {
63        /// The SNAP data plane address according to the rendezvous hashing that must be used by
64        /// the client.
65        pub address: SocketAddr,
66        /// XXX(uniquefine): Make this required once all servers have been updated.
67        /// The address (host:port) of the SNAPtun control plane API. This can be the same
68        /// as the data plane address.
69        pub snap_tun_control_address: Option<Url>,
70        /// XXX(uniquefine): Make this required once all servers have been updated.
71        /// The static identity of the snaptun-ng server.
72        pub snap_static_x25519: Option<PublicKey>,
73    }
74
75    /// Trait for registering a static identity for a snaptun connection.
76    pub trait SnapTunIdentityRegistry: Send + Sync {
77        /// Register a static identity for a snaptun connection.
78        ///
79        /// For now, PSK share is ignored.
80        ///
81        /// # Return value
82        ///
83        /// Returns true if the registration is new, otherwise false.
84        ///
85        /// Eventually, might return PSK share of the server.
86        fn register(
87            &self,
88            now: Instant,
89            // The key under which this identity is stored (at most one is allowed per identity)
90            key: &str,
91            // The static identity of the client.
92            initiator_identity: [u8; 32],
93            // The PSK share used to establish a shared secret with the server.
94            psk_share: Option<[u8; 32]>,
95            // The lifetime the registered identity is valid for.
96            // Usually this is determined by the expiration of the SNAP token.
97            lifetime: Duration,
98        ) -> bool;
99    }
100}
101
102pub(crate) mod convert {
103    use std::net::{AddrParseError, SocketAddr};
104
105    use url::Url;
106    use x25519_dalek::PublicKey;
107
108    use crate::{
109        crpc_api::api_service::model::SnapDataPlane,
110        protobuf::anapaya::snap::v1::api_service as rpc,
111    };
112
113    /// This error is returned when converting a GetSnapDataPlaneResponse to a SnapDataPlane.
114    #[derive(thiserror::Error, Debug)]
115    pub enum ConvertError {
116        #[error("failed to parse data plane address: {0}")]
117        ParseAddr(AddrParseError),
118        #[error("failed to parse server control address: {0}")]
119        ParseSnapTunControlAddr(AddrParseError),
120        #[error("server static identity is not 32 bytes")]
121        InvalidServerStaticIdentityLength,
122    }
123
124    // Protobuf to Model
125    impl TryFrom<rpc::GetSnapDataPlaneResponse> for SnapDataPlane {
126        type Error = ConvertError;
127        fn try_from(value: rpc::GetSnapDataPlaneResponse) -> Result<Self, Self::Error> {
128            let snap_tun_control_address = value
129                .snap_tun_control_address
130                .map(|address| {
131                    // Try to parse the address as a URL first.
132                    if let Ok(url) = Url::parse(&address) {
133                        return Ok(url);
134                    }
135                    match address.parse::<SocketAddr>() {
136                        Ok(addr) => {
137                            let mut u = Url::parse("http://.").unwrap();
138                            let _ = u.set_ip_host(addr.ip());
139                            let _ = u.set_port(Some(addr.port()));
140                            Ok(u)
141                        }
142                        Err(e) => Err(ConvertError::ParseSnapTunControlAddr(e)),
143                    }
144                })
145                .transpose()?;
146            let snap_static_x25519 = value
147                .snap_static_x25519
148                .map(|key| {
149                    TryInto::<[u8; 32]>::try_into(key.as_slice())
150                        .map_err(|_| ConvertError::InvalidServerStaticIdentityLength)
151                        .map(PublicKey::from)
152                })
153                .transpose()?;
154            Ok(SnapDataPlane {
155                address: value.address.parse().map_err(ConvertError::ParseAddr)?,
156                snap_tun_control_address,
157                snap_static_x25519,
158            })
159        }
160    }
161}
162
163pub(crate) const SERVICE_PATH: &str = "/anapaya.snap.v1.SnapControl";
164pub(crate) const GET_SNAP_DATA_PLANE_ADDRESS: &str = "/GetSnapDataPlaneAddress";
165pub(crate) const REGISTER_SNAPTUN_IDENTITY: &str = "/RegisterSnapTunIdentity";
166
167/// Nests the SNAP control API routes into the provided `base_router`.
168pub fn nest_snap_control_api(
169    router: axum::Router,
170    snap_resolver: Arc<dyn SnapDataPlaneResolver>,
171    identity_registrar: Arc<dyn SnapTunIdentityRegistry>,
172) -> axum::Router {
173    router.nest(
174        SERVICE_PATH,
175        Router::new()
176            .route(
177                GET_SNAP_DATA_PLANE_ADDRESS,
178                axum::routing::post(get_snap_data_plane_address_handler),
179            )
180            .with_state(snap_resolver)
181            .route(
182                REGISTER_SNAPTUN_IDENTITY,
183                axum::routing::post(register_snaptun_identity_handler),
184            )
185            .with_state(identity_registrar),
186    )
187}
188
189async fn get_snap_data_plane_address_handler(
190    State(rendezvous_hasher): State<Arc<dyn SnapDataPlaneResolver>>,
191    _snap_token: Extension<AnyClaims>,
192    ConnectInfo(addr): ConnectInfo<std::net::SocketAddr>,
193    ConnectRpc(_request): ConnectRpc<GetSnapDataPlaneRequest>,
194) -> Result<ConnectRpc<GetSnapDataPlaneResponse>, CrpcError> {
195    let addr = rendezvous_hasher.get_data_plane_address(addr.ip())?;
196    Ok(ConnectRpc(GetSnapDataPlaneResponse {
197        address: addr.address.to_string(),
198        snap_tun_control_address: addr
199            .snap_tun_control_address
200            .map(|address| address.to_string()),
201        snap_static_x25519: addr.snap_static_x25519.map(|key| key.to_bytes().to_vec()),
202    }))
203}
204
205async fn register_snaptun_identity_handler(
206    State(identity_registry): State<Arc<dyn SnapTunIdentityRegistry>>,
207    snap_token: Extension<AnyClaims>,
208    ConnectInfo(_): ConnectInfo<std::net::SocketAddr>,
209    ConnectRpc(request): ConnectRpc<RegisterSnapTunIdentityRequest>,
210) -> Result<ConnectRpc<RegisterSnapTunIdentityResponse>, CrpcError> {
211    let now = SystemTime::now();
212    let lifetime = snap_token.0.exp_time().duration_since(now).map_err(|_| {
213        CrpcError::new(
214            CrpcErrorCode::InvalidArgument,
215            "expiration time is in the past".to_string(),
216        )
217    })?;
218
219    let initiator_identity = {
220        let key_bytes: [u8; 32] = request
221            .initiator_static_x25519
222            .as_slice()
223            .try_into()
224            .map_err(|_| {
225                CrpcError::new(
226                    CrpcErrorCode::InvalidArgument,
227                    "initiator identity is not 32 bytes".to_string(),
228                )
229            })?;
230        PublicKey::from(key_bytes)
231    };
232
233    let psk_share: Option<[u8; 32]> = if request.psk_share.as_slice() == [0u8; 32] {
234        None
235    } else {
236        Some(request.psk_share.as_slice().try_into().map_err(|_| {
237            CrpcError::new(
238                CrpcErrorCode::InvalidArgument,
239                "psk share is not 32 bytes".to_string(),
240            )
241        })?)
242    };
243
244    let key = &snap_token.jti();
245    if !identity_registry.register(
246        Instant::now(),
247        key,
248        *initiator_identity.as_bytes(),
249        psk_share,
250        lifetime,
251    ) {
252        tracing::info!(key, "re-registered identity");
253    }
254    Ok(ConnectRpc(RegisterSnapTunIdentityResponse {
255        // XXX(uniquefine): PSK is not yet supported.
256        psk_share: [0u8; 32].to_vec(),
257    }))
258}