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::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
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 x25519_dalek::PublicKey;
50
51    /// SNAP data plane discovery trait.
52    pub trait SnapDataPlaneResolver: Send + Sync {
53        /// Get the SNAP data plane address for a given endhost IP address.
54        fn get_data_plane_address(
55            &self,
56            endhost_ip: IpAddr,
57        ) -> Result<SnapDataPlane, (StatusCode, anyhow::Error)>;
58    }
59
60    /// SnapDataPlane resolution response.
61    pub struct SnapDataPlane {
62        /// The SNAP data plane address according to the rendezvous hashing that must be used by
63        /// the client.
64        pub address: SocketAddr,
65        /// XXX(uniquefine): Make this required once all servers have been updated.
66        /// The address (host:port) of the SNAPtun control plane API. This can be the same
67        /// as the data plane address.
68        pub snap_tun_control_address: Option<SocketAddr>,
69        /// XXX(uniquefine): Make this required once all servers have been updated.
70        /// The static identity of the snaptun-ng server.
71        pub snap_static_x25519: Option<PublicKey>,
72    }
73
74    /// Trait for registering a static identity for a snaptun connection.
75    pub trait SnapTunIdentityRegistry: Send + Sync {
76        /// Register a static identity for a snaptun connection.
77        ///
78        /// For now, PSK share is ignored.
79        ///
80        /// # Return value
81        ///
82        /// Returns true if the registration is new, otherwise false.
83        ///
84        /// Eventually, might return PSK share of the server.
85        fn register(
86            &self,
87            now: Instant,
88            // The key under which this identity is stored (at most one is allowed per identity)
89            key: &str,
90            // The static identity of the client.
91            initiator_identity: [u8; 32],
92            // The PSK share used to establish a shared secret with the server.
93            psk_share: Option<[u8; 32]>,
94            // The lifetime the registered identity is valid for.
95            // Usually this is determined by the expiration of the SNAP token.
96            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    /// This error is returned when converting a GetSnapDataPlaneResponse to a SnapDataPlane.
112    #[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    // Protobuf to Model
123    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
155/// Nests the SNAP control API routes into the provided `base_router`.
156pub 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        // XXX(uniquefine): PSK is not yet supported.
244        psk_share: [0u8; 32].to_vec(),
245    }))
246}