1use sova_core::extend::{named, with_leaked};
4use sova_core::{App, Plugin};
5
6#[derive(Clone)]
8pub struct Shield {
9 frame: Option<&'static str>,
10 content_type: Option<&'static str>,
11 referrer: Option<&'static str>,
12 coop: Option<&'static str>,
13 corp: Option<&'static str>,
14 dns_prefetch: Option<&'static str>,
15 download_options: Option<&'static str>,
16 permitted_cross_domain: Option<&'static str>,
17 xss_protection: Option<&'static str>,
18 csp: Option<&'static str>,
19 csp_explicit: bool,
20}
21
22impl Shield {
23 pub fn new() -> Self {
24 Self {
25 frame: Some("SAMEORIGIN"),
26 content_type: Some("nosniff"),
27 referrer: Some("no-referrer"),
28 coop: Some("same-origin"),
29 corp: Some("same-origin"),
30 dns_prefetch: Some("off"),
31 download_options: Some("noopen"),
32 permitted_cross_domain: Some("none"),
33 xss_protection: Some("0"),
34 csp: None,
35 csp_explicit: false,
36 }
37 }
38
39 pub fn frame(mut self, value: &'static str) -> Self {
40 self.frame = Some(value);
41 self
42 }
43
44 pub fn frame_off(mut self) -> Self {
45 self.frame = None;
46 self
47 }
48
49 pub fn content_type(mut self, value: &'static str) -> Self {
50 self.content_type = Some(value);
51 self
52 }
53
54 pub fn content_type_off(mut self) -> Self {
55 self.content_type = None;
56 self
57 }
58
59 pub fn referrer(mut self, value: &'static str) -> Self {
60 self.referrer = Some(value);
61 self
62 }
63
64 pub fn referrer_off(mut self) -> Self {
65 self.referrer = None;
66 self
67 }
68
69 pub fn cross_origin_opener(mut self, value: &'static str) -> Self {
70 self.coop = Some(value);
71 self
72 }
73
74 pub fn cross_origin_opener_off(mut self) -> Self {
75 self.coop = None;
76 self
77 }
78
79 pub fn cross_origin_resource(mut self, value: &'static str) -> Self {
80 self.corp = Some(value);
81 self
82 }
83
84 pub fn cross_origin_resource_off(mut self) -> Self {
85 self.corp = None;
86 self
87 }
88
89 pub fn dns_prefetch(mut self, value: &'static str) -> Self {
90 self.dns_prefetch = Some(value);
91 self
92 }
93
94 pub fn dns_prefetch_off(mut self) -> Self {
95 self.dns_prefetch = None;
96 self
97 }
98
99 pub fn download_options(mut self, value: &'static str) -> Self {
100 self.download_options = Some(value);
101 self
102 }
103
104 pub fn download_options_off(mut self) -> Self {
105 self.download_options = None;
106 self
107 }
108
109 pub fn permitted_cross_domain(mut self, value: &'static str) -> Self {
110 self.permitted_cross_domain = Some(value);
111 self
112 }
113
114 pub fn permitted_cross_domain_off(mut self) -> Self {
115 self.permitted_cross_domain = None;
116 self
117 }
118
119 pub fn xss_protection(mut self, value: &'static str) -> Self {
120 self.xss_protection = Some(value);
121 self
122 }
123
124 pub fn xss_protection_off(mut self) -> Self {
125 self.xss_protection = None;
126 self
127 }
128
129 pub fn csp(mut self, policy: &'static str) -> Self {
131 self.csp = Some(policy);
132 self.csp_explicit = true;
133 self
134 }
135
136 pub fn csp_off(mut self) -> Self {
137 self.csp = None;
138 self.csp_explicit = true;
139 self
140 }
141}
142
143impl Default for Shield {
144 fn default() -> Self {
145 Self::new()
146 }
147}
148
149impl Plugin for Shield {
150 fn id(&self) -> &'static str {
151 "shield"
152 }
153
154 fn meta(&self) -> sova_core::PluginMeta {
155 sova_core::PluginMeta::new("Shield")
156 .description("Baseline security response headers (helmet-style)")
157 .version(env!("CARGO_PKG_VERSION"))
158 }
159
160 fn install(mut self, app: &mut App) {
161 if let Some(doc) = app.config_doc() {
162 if let Some(section) = doc.section("shield") {
163 if !self.csp_explicit {
164 if let Some(p) = section.get("csp").and_then(|v| v.as_str()) {
165 self.csp = Some(Box::leak(p.to_string().into_boxed_str()));
166 }
167 }
168 if let Some(v) = section.get("frame").and_then(|v| v.as_str()) {
169 self.frame = Some(Box::leak(v.to_string().into_boxed_str()));
170 }
171 }
172 }
173 app.use_middleware(named(
174 "shield",
175 with_leaked(self, |shield, req, next| async move {
176 let mut res = next(req).await;
177 if let Some(v) = shield.frame {
178 res = res.header("x-frame-options", v);
179 }
180 if let Some(v) = shield.content_type {
181 res = res.header("x-content-type-options", v);
182 }
183 if let Some(v) = shield.referrer {
184 res = res.header("referrer-policy", v);
185 }
186 if let Some(v) = shield.coop {
187 res = res.header("cross-origin-opener-policy", v);
188 }
189 if let Some(v) = shield.corp {
190 res = res.header("cross-origin-resource-policy", v);
191 }
192 if let Some(v) = shield.dns_prefetch {
193 res = res.header("x-dns-prefetch-control", v);
194 }
195 if let Some(v) = shield.download_options {
196 res = res.header("x-download-options", v);
197 }
198 if let Some(v) = shield.permitted_cross_domain {
199 res = res.header("x-permitted-cross-domain-policies", v);
200 }
201 if let Some(v) = shield.xss_protection {
202 res = res.header("x-xss-protection", v);
203 }
204 if let Some(v) = shield.csp {
205 res = res.header("content-security-policy", v);
206 }
207 res
208 }),
209 ));
210 }
211}
212
213#[cfg(test)]
214mod tests {
215 use super::*;
216 use http::Method;
217 use sova_core::{Request, Response};
218
219 #[tokio::test]
220 async fn default_headers() {
221 let mut app = App::new();
222 app.install(Shield::default());
223 app.get("/", |_r: Request| async { Response::text("ok") });
224 let res = app.handle_request(Method::GET, "/", "").await;
225 assert_eq!(
226 res.headers()
227 .get("x-frame-options")
228 .and_then(|v| v.to_str().ok()),
229 Some("SAMEORIGIN")
230 );
231 assert_eq!(
232 res.headers()
233 .get("x-content-type-options")
234 .and_then(|v| v.to_str().ok()),
235 Some("nosniff")
236 );
237 assert_eq!(
238 res.headers()
239 .get("referrer-policy")
240 .and_then(|v| v.to_str().ok()),
241 Some("no-referrer")
242 );
243 assert_eq!(
244 res.headers()
245 .get("cross-origin-opener-policy")
246 .and_then(|v| v.to_str().ok()),
247 Some("same-origin")
248 );
249 assert_eq!(
250 res.headers()
251 .get("x-xss-protection")
252 .and_then(|v| v.to_str().ok()),
253 Some("0")
254 );
255 assert!(res.headers().get("content-security-policy").is_none());
256 }
257}