snap_control/api/
crpc.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 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    api::crpc::model::{SnapDataPlaneResolver, SnapTunIdentityRegistry},
35    proto::anapaya::snap::v1::{
36        GetSnapDataPlaneAddressRequest, GetSnapDataPlaneAddressResponse,
37        RegisterSnapTunIdentityRequest, 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 snap_tokens::AnyClaims;
50    use url::Url;
51    use x25519_dalek::PublicKey;
52
53    /// SNAP data plane discovery trait.
54    pub trait SnapDataPlaneResolver: Send + Sync {
55        /// Get the SNAP data plane address for a given endhost IP address.
56        fn get_data_plane_address(
57            &self,
58            endhost_ip: IpAddr,
59        ) -> Result<SnapDataPlane, (StatusCode, anyhow::Error)>;
60    }
61
62    /// SnapDataPlane resolution response.
63    pub struct SnapDataPlane {
64        /// The SNAP data plane address according to the rendezvous hashing that must be used by
65        /// the client.
66        pub address: SocketAddr,
67        /// XXX(uniquefine): Make this required once all servers have been updated.
68        /// The address (host:port) of the SNAPtun control plane API. This can be the same
69        /// as the data plane address.
70        pub snap_tun_control_address: Option<Url>,
71        /// XXX(uniquefine): Make this required once all servers have been updated.
72        /// The static identity of the snaptun-ng server.
73        pub snap_static_x25519: Option<PublicKey>,
74    }
75
76    /// Trait for registering a static identity for a snaptun connection.
77    pub trait SnapTunIdentityRegistry: Send + Sync {
78        /// Register a static identity for a snaptun connection.
79        ///
80        /// For now, PSK share is ignored.
81        ///
82        /// # Return value
83        ///
84        /// Returns true if the registration is new, otherwise false.
85        ///
86        /// Eventually, might return PSK share of the server.
87        fn register(
88            &self,
89            now: Instant,
90            // The key under which this identity is stored (at most one is allowed per identity)
91            key: &str,
92            // The static identity of the client.
93            initiator_identity: [u8; 32],
94            // The PSK share used to establish a shared secret with the server.
95            psk_share: Option<[u8; 32]>,
96            // The lifetime the registered identity is valid for.
97            // Usually this is determined by the expiration of the SNAP token.
98            lifetime: Duration,
99            // The verified token claims associated with this registration.
100            claims: &AnyClaims,
101        ) -> anyhow::Result<bool>;
102
103        /// Removes registrations whose authorization lifetime has expired.
104        ///
105        /// Implementations must evict expired registrations here. The control
106        /// plane server calls this periodically to enforce authorization
107        /// liveness for active SNAP identities.
108        fn remove_expired(&self, now: Instant);
109    }
110}
111
112pub(crate) mod convert {
113    use std::net::{AddrParseError, SocketAddr};
114
115    use url::Url;
116    use x25519_dalek::PublicKey;
117
118    use crate::{api::crpc::model::SnapDataPlane, proto::anapaya::snap::v1 as rpc};
119
120    /// This error is returned when converting a GetSnapDataPlaneAddressResponse to a SnapDataPlane.
121    #[derive(thiserror::Error, Debug)]
122    pub enum ConvertError {
123        #[error("failed to parse data plane address: {0}")]
124        ParseAddr(AddrParseError),
125        #[error("failed to parse server control address: {0}")]
126        ParseSnapTunControlAddr(AddrParseError),
127        #[error("server static identity is not 32 bytes")]
128        InvalidServerStaticIdentityLength,
129    }
130
131    // Protobuf to Model
132    impl TryFrom<rpc::GetSnapDataPlaneAddressResponse> for SnapDataPlane {
133        type Error = ConvertError;
134        fn try_from(value: rpc::GetSnapDataPlaneAddressResponse) -> Result<Self, Self::Error> {
135            let snap_tun_control_address = value
136                .snap_tun_control_address
137                .map(|address| {
138                    // Try to parse the address as a URL first.
139                    if let Ok(url) = Url::parse(&address) {
140                        return Ok(url);
141                    }
142                    match address.parse::<SocketAddr>() {
143                        Ok(addr) => {
144                            let mut u = Url::parse("http://.").unwrap();
145                            let _ = u.set_ip_host(addr.ip());
146                            let _ = u.set_port(Some(addr.port()));
147                            Ok(u)
148                        }
149                        Err(e) => Err(ConvertError::ParseSnapTunControlAddr(e)),
150                    }
151                })
152                .transpose()?;
153            let snap_static_x25519 = value
154                .snap_static_x25519
155                .map(|key| {
156                    TryInto::<[u8; 32]>::try_into(key.as_slice())
157                        .map_err(|_| ConvertError::InvalidServerStaticIdentityLength)
158                        .map(PublicKey::from)
159                })
160                .transpose()?;
161            Ok(SnapDataPlane {
162                address: value.address.parse().map_err(ConvertError::ParseAddr)?,
163                snap_tun_control_address,
164                snap_static_x25519,
165            })
166        }
167    }
168}
169
170pub(crate) const SERVICE_PATH: &str = "/anapaya.snap.v1.SnapControl";
171pub(crate) const GET_SNAP_DATA_PLANE_ADDRESS: &str = "/GetSnapDataPlaneAddress";
172pub(crate) const REGISTER_SNAPTUN_IDENTITY: &str = "/RegisterSnapTunIdentity";
173
174/// Nests the SNAP control API routes into the provided `base_router`.
175pub fn nest_crpc_api(
176    router: axum::Router,
177    snap_resolver: Arc<dyn SnapDataPlaneResolver>,
178    identity_registrar: Arc<dyn SnapTunIdentityRegistry>,
179) -> axum::Router {
180    router.nest(
181        SERVICE_PATH,
182        Router::new()
183            .route(
184                GET_SNAP_DATA_PLANE_ADDRESS,
185                axum::routing::post(get_snap_data_plane_address_handler),
186            )
187            .with_state(snap_resolver)
188            .route(
189                REGISTER_SNAPTUN_IDENTITY,
190                axum::routing::post(register_snaptun_identity_handler),
191            )
192            .with_state(identity_registrar),
193    )
194}
195
196async fn get_snap_data_plane_address_handler(
197    State(rendezvous_hasher): State<Arc<dyn SnapDataPlaneResolver>>,
198    _snap_token: Extension<AnyClaims>,
199    ConnectInfo(addr): ConnectInfo<std::net::SocketAddr>,
200    ConnectRpc(_request): ConnectRpc<GetSnapDataPlaneAddressRequest>,
201) -> Result<ConnectRpc<GetSnapDataPlaneAddressResponse>, CrpcError> {
202    let addr = rendezvous_hasher.get_data_plane_address(addr.ip())?;
203    Ok(ConnectRpc(GetSnapDataPlaneAddressResponse {
204        address: addr.address.to_string(),
205        snap_tun_control_address: addr
206            .snap_tun_control_address
207            .map(|address| address.to_string()),
208        snap_static_x25519: addr.snap_static_x25519.map(|key| key.to_bytes().to_vec()),
209    }))
210}
211
212async fn register_snaptun_identity_handler(
213    State(identity_registry): State<Arc<dyn SnapTunIdentityRegistry>>,
214    snap_token: Extension<AnyClaims>,
215    ConnectInfo(_): ConnectInfo<std::net::SocketAddr>,
216    ConnectRpc(request): ConnectRpc<RegisterSnapTunIdentityRequest>,
217) -> Result<ConnectRpc<RegisterSnapTunIdentityResponse>, CrpcError> {
218    let now = SystemTime::now();
219    let lifetime = snap_token.0.exp_time().duration_since(now).map_err(|_| {
220        CrpcError::new(
221            CrpcErrorCode::InvalidArgument,
222            "expiration time is in the past".to_string(),
223        )
224    })?;
225
226    let initiator_identity = {
227        let key_bytes: [u8; 32] = request
228            .initiator_static_x25519
229            .as_slice()
230            .try_into()
231            .map_err(|_| {
232                CrpcError::new(
233                    CrpcErrorCode::InvalidArgument,
234                    "initiator identity is not 32 bytes".to_string(),
235                )
236            })?;
237        PublicKey::from(key_bytes)
238    };
239
240    let psk_share: Option<[u8; 32]> = if request.psk_share.as_slice() == [0u8; 32] {
241        None
242    } else {
243        Some(request.psk_share.as_slice().try_into().map_err(|_| {
244            CrpcError::new(
245                CrpcErrorCode::InvalidArgument,
246                "psk share is not 32 bytes".to_string(),
247            )
248        })?)
249    };
250
251    let key = snap_token.jti();
252    if !identity_registry
253        .register(
254            Instant::now(),
255            &key,
256            *initiator_identity.as_bytes(),
257            psk_share,
258            lifetime,
259            &snap_token,
260        )
261        .map_err(|err| CrpcError::new(CrpcErrorCode::InvalidArgument, err.to_string()))?
262    {
263        tracing::info!(key, "re-registered identity");
264    }
265    Ok(ConnectRpc(RegisterSnapTunIdentityResponse {
266        // XXX(uniquefine): PSK is not yet supported.
267        psk_share: [0u8; 32].to_vec(),
268    }))
269}
snap_control/api/crpc.rs

snap_control/api/
crpc.rs
