mirror of
https://github.com/open-webui/open-webui
synced 2024-11-07 00:59:52 +00:00
120 lines
3.8 KiB
Python
120 lines
3.8 KiB
Python
import re
|
|
import os
|
|
|
|
from fastapi import Request
|
|
from starlette.middleware.base import BaseHTTPMiddleware
|
|
from typing import Dict
|
|
|
|
class SecurityHeadersMiddleware(BaseHTTPMiddleware):
|
|
async def dispatch(self, request: Request, call_next):
|
|
response = await call_next(request)
|
|
response.headers.update(set_security_headers())
|
|
return response
|
|
|
|
def set_security_headers() -> Dict[str, str]:
|
|
"""
|
|
Sets security headers based on environment variables.
|
|
|
|
This function reads specific environment variables and uses their values
|
|
to set corresponding security headers. The headers that can be set are:
|
|
- cache-control
|
|
- strict-transport-security
|
|
- referrer-policy
|
|
- x-content-type-options
|
|
- x-download-options
|
|
- x-frame-options
|
|
- x-permitted-cross-domain-policies
|
|
|
|
Each environment variable is associated with a specific setter function
|
|
that constructs the header. If the environment variable is set, the
|
|
corresponding header is added to the options dictionary.
|
|
|
|
Returns:
|
|
dict: A dictionary containing the security headers and their values.
|
|
"""
|
|
options = {}
|
|
header_setters = {
|
|
'CACHE_CONTROL': set_cache_control,
|
|
'HSTS': set_hsts,
|
|
'REFERRER_POLICY': set_referrer,
|
|
'XCONTENT_TYPE': set_xcontent_type,
|
|
'XDOWNLOAD_OPTIONS': set_xdownload_options,
|
|
'XFRAME_OPTIONS': set_xframe,
|
|
'XPERMITTED_CROSS_DOMAIN_POLICIES': set_xpermitted_cross_domain_policies,
|
|
}
|
|
|
|
for env_var, setter in header_setters.items():
|
|
value = os.environ.get(env_var, None)
|
|
if value:
|
|
header = setter(value)
|
|
if header:
|
|
options.update(header)
|
|
|
|
return options
|
|
|
|
# Set HTTP Strict Transport Security(HSTS) response header
|
|
def set_hsts(value: str):
|
|
pattern = r'^max-age=(\d+)(;includeSubDomains)?(;preload)?$'
|
|
match = re.match(pattern, value, re.IGNORECASE)
|
|
if not match:
|
|
return 'max-age=31536000;includeSubDomains'
|
|
return {
|
|
'Strict-Transport-Security': value
|
|
}
|
|
|
|
# Set X-Frame-Options response header
|
|
def set_xframe(value: str):
|
|
pattern = r'^(DENY|SAMEORIGIN)$'
|
|
match = re.match(pattern, value, re.IGNORECASE)
|
|
if not match:
|
|
value = 'DENY'
|
|
return {
|
|
"X-Frame-Options": value
|
|
}
|
|
|
|
# Set Referrer-Policy response header
|
|
def set_referrer(value: str):
|
|
pattern = r'^(no-referrer|no-referrer-when-downgrade|origin|origin-when-cross-origin|same-origin|strict-origin|strict-origin-when-cross-origin|unsafe-url)$'
|
|
match = re.match(pattern, value, re.IGNORECASE)
|
|
if not match:
|
|
value = 'no-referrer'
|
|
return {
|
|
'Referrer-Policy': value
|
|
}
|
|
|
|
# Set Cache-Control response header
|
|
def set_cache_control(value: str):
|
|
pattern = r'^(public|private|no-cache|no-store|must-revalidate|proxy-revalidate|max-age=\d+|s-maxage=\d+|no-transform|immutable)(,\s*(public|private|no-cache|no-store|must-revalidate|proxy-revalidate|max-age=\d+|s-maxage=\d+|no-transform|immutable))*$'
|
|
match = re.match(pattern, value, re.IGNORECASE)
|
|
if not match:
|
|
value = 'no-store, max-age=0'
|
|
|
|
return {
|
|
'Cache-Control': value
|
|
}
|
|
|
|
# Set X-Download-Options response header
|
|
def set_xdownload_options(value: str):
|
|
if value != 'noopen':
|
|
value = 'noopen'
|
|
return {
|
|
'X-Download-Options': value
|
|
}
|
|
|
|
# Set X-Content-Type-Options response header
|
|
def set_xcontent_type(value: str):
|
|
if value != 'nosniff':
|
|
value = 'nosniff'
|
|
return {
|
|
'X-Content-Type-Options': value
|
|
}
|
|
|
|
# Set X-Permitted-Cross-Domain-Policies response header
|
|
def set_xpermitted_cross_domain_policies(value: str):
|
|
pattern = r'^(none|master-only|by-content-type|by-ftp-filename)$'
|
|
match = re.match(pattern, value, re.IGNORECASE)
|
|
if not match:
|
|
value = 'none'
|
|
return {
|
|
'X-Permitted-Cross-Domain-Policies': value
|
|
} |