shared/models/server/
firewall.rs

1use compact_str::ToCompactString;
2use garde::Validate;
3use serde::{Deserialize, Serialize};
4use sqlx::Row;
5use std::collections::HashSet;
6use utoipa::ToSchema;
7
8#[derive(ToSchema, Serialize, Deserialize, Clone, Copy, PartialEq, Eq)]
9#[serde(rename_all = "snake_case")]
10pub enum FirewallRuleAction {
11    Allow,
12    Deny,
13}
14
15#[derive(ToSchema, Serialize, Deserialize, Clone, Copy, PartialEq, Eq, Hash)]
16#[serde(rename_all = "snake_case")]
17pub enum FirewallRuleProtocol {
18    Tcp,
19    Udp,
20}
21
22#[derive(ToSchema, Validate, Serialize, Deserialize, Clone)]
23pub struct FirewallRule {
24    #[garde(skip)]
25    pub action: FirewallRuleAction,
26    #[garde(skip)]
27    #[serde(default)]
28    pub protocols: HashSet<FirewallRuleProtocol>,
29    #[garde(skip)]
30    #[schema(value_type = Vec<String>)]
31    #[serde(default)]
32    pub sources: Vec<cidr::IpCidr>,
33    #[garde(inner(length(min = 1, max = 1024), inner(range(min = 1))))]
34    #[schema(min_items = 1, max_items = 1024)]
35    #[serde(default)]
36    pub ports: Option<Vec<u16>>,
37    #[garde(inner(length(min = 1, max = 512)))]
38    #[schema(min_length = 1, max_length = 512)]
39    #[serde(default)]
40    pub source_file: Option<String>,
41}
42
43impl From<FirewallRule> for wings_api::FirewallRule {
44    fn from(value: FirewallRule) -> Self {
45        Self {
46            action: match value.action {
47                FirewallRuleAction::Allow => wings_api::FirewallRuleAction::Allow,
48                FirewallRuleAction::Deny => wings_api::FirewallRuleAction::Deny,
49            },
50            protocols: value
51                .protocols
52                .into_iter()
53                .map(|protocol| match protocol {
54                    FirewallRuleProtocol::Tcp => wings_api::FirewallRuleProtocol::Tcp,
55                    FirewallRuleProtocol::Udp => wings_api::FirewallRuleProtocol::Udp,
56                })
57                .collect(),
58            sources: value
59                .sources
60                .iter()
61                .map(|source| source.to_compact_string())
62                .collect(),
63            ports: value
64                .ports
65                .map(|ports| ports.into_iter().map(|port| port as u32).collect()),
66            source_file: value.source_file.map(Into::into),
67        }
68    }
69}
70
71pub(super) async fn fetch_raw_rules(
72    database: &crate::database::Database,
73    server_uuid: uuid::Uuid,
74) -> Result<Option<serde_json::Value>, sqlx::Error> {
75    sqlx::query_scalar(
76        r#"
77        SELECT server_firewalls.rules
78        FROM server_firewalls
79        WHERE server_firewalls.server_uuid = $1
80        "#,
81    )
82    .bind(server_uuid)
83    .fetch_optional(database.read())
84    .await
85}
86
87pub(super) fn decode_rules(
88    rules: Option<serde_json::Value>,
89) -> Result<Vec<FirewallRule>, serde_json::Error> {
90    match rules {
91        Some(rules) => serde_json::from_value(rules),
92        None => Ok(Vec::new()),
93    }
94}
95
96impl super::Server {
97    pub async fn firewall_rules(
98        &self,
99        database: &crate::database::Database,
100    ) -> Result<Vec<FirewallRule>, crate::database::DatabaseError> {
101        Ok(decode_rules(fetch_raw_rules(database, self.uuid).await?)?)
102    }
103
104    #[inline]
105    pub async fn allocation_ports(
106        &self,
107        database: &crate::database::Database,
108    ) -> Result<Vec<u16>, crate::database::DatabaseError> {
109        Self::allocation_ports_by_uuid(database, self.uuid).await
110    }
111
112    pub async fn allocation_ports_by_uuid(
113        database: &crate::database::Database,
114        server_uuid: uuid::Uuid,
115    ) -> Result<Vec<u16>, crate::database::DatabaseError> {
116        let ports: Vec<i32> = sqlx::query_scalar(
117            r#"
118            SELECT node_allocations.port
119            FROM server_allocations
120            JOIN node_allocations ON node_allocations.uuid = server_allocations.allocation_uuid
121            WHERE server_allocations.server_uuid = $1
122            ORDER BY node_allocations.port
123            "#,
124        )
125        .bind(server_uuid)
126        .fetch_all(database.read())
127        .await?;
128
129        Ok(ports.into_iter().map(|port| port as u16).collect())
130    }
131
132    pub async fn allocation_ports_by_uuids(
133        database: &crate::database::Database,
134        server_uuids: &[uuid::Uuid],
135    ) -> Result<std::collections::HashMap<uuid::Uuid, Vec<u16>>, crate::database::DatabaseError>
136    {
137        let mut ports: std::collections::HashMap<uuid::Uuid, Vec<u16>> =
138            std::collections::HashMap::new();
139
140        for row in sqlx::query(
141            r#"
142            SELECT server_allocations.server_uuid, node_allocations.port
143            FROM server_allocations
144            JOIN node_allocations ON node_allocations.uuid = server_allocations.allocation_uuid
145            WHERE server_allocations.server_uuid = ANY($1)
146            ORDER BY node_allocations.port
147            "#,
148        )
149        .bind(server_uuids)
150        .fetch_all(database.read())
151        .await?
152        {
153            ports
154                .entry(row.try_get("server_uuid")?)
155                .or_default()
156                .push(row.try_get::<i32, _>("port")? as u16);
157        }
158
159        Ok(ports)
160    }
161
162    pub async fn set_firewall_rules(
163        &self,
164        database: &crate::database::Database,
165        rules: &[FirewallRule],
166    ) -> Result<(), crate::database::DatabaseError> {
167        sqlx::query(
168            r#"
169            INSERT INTO server_firewalls (server_uuid, rules)
170            VALUES ($1, $2)
171            ON CONFLICT (server_uuid) DO UPDATE SET rules = EXCLUDED.rules
172            "#,
173        )
174        .bind(self.uuid)
175        .bind(serde_json::to_value(rules)?)
176        .execute(database.write())
177        .await?;
178
179        Ok(())
180    }
181}
shared/models/server/firewall.rs

shared/models/server/
firewall.rs
