snap_control/
client.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 client for the SNAP control plane API.
15
16use std::{net::SocketAddr, ops::Deref, sync::Arc};
17
18use async_trait::async_trait;
19use endhost_api_client::client::CrpcEndhostApiClient;
20use scion_sdk_reqwest_connect_rpc::{client::CrpcClientError, token_source::TokenSource};
21use snap_tun::client::SnapTunControlPlaneClient;
22use url::Url;
23use x25519_dalek::PublicKey;
24
25use crate::{
26    crpc_api::api_service::{GET_SNAP_DATA_PLANE_ADDRESS, REGISTER_SNAPTUN_IDENTITY, SERVICE_PATH},
27    protobuf::anapaya::snap::v1::api_service as proto,
28};
29
30/// Re-export the endhost API client and the reqwest connect RPC cllient.
31pub mod re_export {
32    pub use endhost_api_client::client::{CrpcEndhostApiClient, EndhostApiClient};
33    pub use scion_sdk_reqwest_connect_rpc::{client::CrpcClientError, token_source::*};
34}
35
36/// SNAP data plane address response.
37#[derive(Debug, Clone, PartialEq, Eq)]
38pub struct GetDataPlaneAddressResponse {
39    /// The UDP endpoint (host:port) of the SNAP data plane.
40    pub address: SocketAddr,
41    /// The URL of the SNAPtun control plane API. This can be the same as the data plane address.
42    /// XXX(uniquefine): Make this required once all servers have been updated.
43    pub snap_tun_control_address: Option<Url>,
44    /// The static identity of the snaptun-ng server.
45    /// XXX(uniquefine): Make this required once all servers have been updated.
46    pub snap_static_x25519: Option<PublicKey>,
47}
48
49/// SNAP control plane API trait.
50#[async_trait]
51pub trait ControlPlaneApi: Send + Sync {
52    /// Get the SNAP data plane address.
53    async fn get_data_plane_address(&self) -> Result<GetDataPlaneAddressResponse, CrpcClientError>;
54
55    /// Register a static identity for a snaptun connection.
56    async fn register_snaptun_identity(
57        &self,
58        initiator_identity: PublicKey,
59        psk_share: Option<[u8; 32]>,
60    ) -> Result<Option<[u8; 32]>, CrpcClientError>;
61}
62
63/// Connect RPC client for the SNAP control plane API.
64pub struct CrpcSnapControlClient {
65    client: CrpcEndhostApiClient,
66}
67
68impl Deref for CrpcSnapControlClient {
69    type Target = CrpcEndhostApiClient;
70
71    fn deref(&self) -> &Self::Target {
72        &self.client
73    }
74}
75
76impl CrpcSnapControlClient {
77    /// Creates a new client with default settings
78    pub fn new(base_url: &Url) -> anyhow::Result<Self> {
79        let client = CrpcEndhostApiClient::new(base_url)?;
80        Ok(Self { client })
81    }
82
83    /// Creates a new client with the provided `reqwest::Client`.
84    pub fn new_with_client(base_url: &Url, client: reqwest::Client) -> anyhow::Result<Self> {
85        Ok(Self {
86            client: CrpcEndhostApiClient::new_with_client(base_url, client)?,
87        })
88    }
89
90    /// Uses the provided token source for authentication.
91    pub fn use_token_source(&mut self, token_source: Arc<dyn TokenSource>) -> &mut Self {
92        self.client.use_token_source(token_source);
93        self
94    }
95}
96
97#[async_trait]
98impl ControlPlaneApi for CrpcSnapControlClient {
99    async fn get_data_plane_address(&self) -> Result<GetDataPlaneAddressResponse, CrpcClientError> {
100        let res: proto::GetSnapDataPlaneResponse = self
101            .client
102            .unary_request::<proto::GetSnapDataPlaneRequest, proto::GetSnapDataPlaneResponse>(
103                &format!("{SERVICE_PATH}{GET_SNAP_DATA_PLANE_ADDRESS}"),
104                proto::GetSnapDataPlaneRequest::default(),
105            )
106            .await?;
107        let address = res.address.parse().map_err(|e: std::net::AddrParseError| {
108            CrpcClientError::DecodeError {
109                context: "parsing data plane address".into(),
110                source: Some(e.into()),
111                body: None,
112            }
113        })?;
114
115        let snap_tun_control_address = res
116            .snap_tun_control_address
117            .map(|address| {
118                // Try to parse the address as a URL first.
119                if let Ok(url) = Url::parse(&address) {
120                    return Ok(url);
121                }
122                match address.parse::<SocketAddr>() {
123                    Ok(addr) => {
124                        let mut u = Url::parse("http://.").unwrap();
125                        let _ = u.set_ip_host(addr.ip());
126                        let _ = u.set_port(Some(addr.port()));
127                        Ok(u)
128                    }
129                    Err(e) => {
130                        Err(CrpcClientError::DecodeError {
131                            context: "parsing server control address".into(),
132                            source: Some(e.into()),
133                            body: None,
134                        })
135                    }
136                }
137            })
138            .transpose()?;
139        let snap_static_x25519 = res
140            .snap_static_x25519
141            .map(|key| {
142                let key_bytes: [u8; 32] =
143                    key.as_slice()
144                        .try_into()
145                        .map_err(|e: std::array::TryFromSliceError| {
146                            CrpcClientError::DecodeError {
147                                context: "server static identity is not 32 bytes".into(),
148                                source: Some(e.into()),
149                                body: None,
150                            }
151                        })?;
152                Ok::<_, CrpcClientError>(PublicKey::from(key_bytes))
153            })
154            .transpose()?;
155        Ok(GetDataPlaneAddressResponse {
156            address,
157            snap_tun_control_address,
158            snap_static_x25519,
159        })
160    }
161
162    async fn register_snaptun_identity(
163        &self,
164        initiator_identity: PublicKey,
165        psk_share: Option<[u8; 32]>,
166    ) -> Result<Option<[u8; 32]>, CrpcClientError> {
167        let res = self.client.unary_request::<proto::RegisterSnapTunIdentityRequest, proto::RegisterSnapTunIdentityResponse>(
168            &format!("{SERVICE_PATH}{REGISTER_SNAPTUN_IDENTITY}"),
169            proto::RegisterSnapTunIdentityRequest { initiator_static_x25519: initiator_identity.to_bytes().to_vec(), psk_share: psk_share.unwrap_or([0u8;32]).to_vec() },
170        ).await?;
171        let psk_share = if res.psk_share.as_slice() == [0u8; 32] {
172            None
173        } else {
174            Some(res.psk_share.as_slice().try_into().map_err(
175                |e: std::array::TryFromSliceError| {
176                    CrpcClientError::DecodeError {
177                        context: "psk share is not 32 bytes".into(),
178                        source: Some(e.into()),
179                        body: None,
180                    }
181                },
182            )?)
183        };
184        Ok(psk_share)
185    }
186}
187
188#[async_trait]
189impl SnapTunControlPlaneClient for CrpcSnapControlClient {
190    async fn register_identity(
191        &self,
192        initiator_identity: PublicKey,
193        psk_share: Option<[u8; 32]>,
194    ) -> Result<Option<[u8; 32]>, CrpcClientError> {
195        self.register_snaptun_identity(initiator_identity, psk_share)
196            .await
197    }
198}
snap_control/client.rs

snap_control/
client.rs
