Skip to content

Commit 7822ce9

Browse files
kvinwangkvinwang
committed
teepod: Add script setting the user net firewall
1 parent 047a246 commit 7822ce9

File tree

4 files changed

+160
-5
lines changed

4 files changed

+160
-5
lines changed

‎teepod/src/app/qemu.rs‎

Lines changed: 11 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,9 @@ use crate::{
33
app::Manifest,
44
config::{CvmConfig, GatewayConfig, Networking, ProcessNote},
55
};
6+
use std::os::unix::fs::PermissionsExt;
67
use std::{
8+
fs::Permissions,
79
ops::Deref,
810
path::{Path, PathBuf},
911
process::Command,
@@ -196,6 +198,10 @@ impl VmConfig {
196198
if !hda_path.exists() {
197199
create_hd(&hda_path, self.image.hda.as_ref(), &disk_size)?;
198200
}
201+
if !cfg.user.is_empty() {
202+
fs_err::set_permissions(&hda_path, Permissions::from_mode(0o660))?;
203+
}
204+
199205
if !shared_dir.exists() {
200206
fs::create_dir_all(&shared_dir)?;
201207
}
@@ -340,8 +346,11 @@ impl VmConfig {
340346
cmd_args.splice(0..0, ["taskset", "-c", &cpus].into_iter().map(|s| s.into()));
341347
}
342348

343-
if cfg.sudo {
344-
cmd_args.insert(0, "sudo".into());
349+
if !cfg.user.is_empty() {
350+
cmd_args.splice(
351+
0..0,
352+
["sudo", "-u", &cfg.user].into_iter().map(|s| s.into()),
353+
);
345354
}
346355

347356
let command = cmd_args.remove(0);

‎teepod/src/config.rs‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -100,7 +100,7 @@ pub struct CvmConfig {
100100
/// GPU configuration
101101
pub gpu: GpuConfig,
102102
/// Use sudo to run the VM
103-
pub sudo: bool,
103+
pub user: String,
104104
}
105105

106106
#[derive(Debug, Clone, Deserialize)]

‎teepod/src/setup-user.sh‎

Lines changed: 146 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,146 @@
1+
#!/bin/bash
2+
3+
# Dstack currently runs the CVMs with qemu user networking. To prevent the VM to access 127.0.0.1, we need to
4+
# run the VM as a different user and setup the iptables rules to DROP the traffic to 127.0.0.1.
5+
6+
# This script creates a sandbox user for running VMs with restricted network access.
7+
# It sets up iptables rules to prevent the user from accessing localhost (127.0.0.1),
8+
# with optional exceptions for specific ports.
9+
#
10+
# When deploying a new dstack instance, you can follow the following steps:
11+
# 1. Create a new sandbox user and setup the firewall rules:
12+
# ```
13+
# sudo ./setup-user.sh dstack-prd1 -g $(id -gn)
14+
# ```
15+
#
16+
# If you want allow the VM to access some ports such as local KMS or Tproxy, you can use the following command:
17+
# ```
18+
# sudo ./setup-user.sh dstack-prd1 -g $(id -gn) --allow-tcp 8080 --allow-udp 3000
19+
# ```
20+
# 2. Edit the user in the `[cvm]` section of the teepod.yaml file:
21+
# ```
22+
# [cvm]
23+
# user = "dstack-prd1"
24+
# ```
25+
#
26+
27+
# Default values
28+
USERNAME=""
29+
GROUP_NAME=""
30+
NO_FW=false
31+
ALLOWED_TCP_PORTS=""
32+
ALLOWED_UDP_PORTS=""
33+
34+
# Parse arguments
35+
while [[ $# -gt 0 ]]; do
36+
case "$1" in
37+
--no-fw)
38+
NO_FW=true
39+
shift
40+
;;
41+
--allow-tcp)
42+
ALLOWED_TCP_PORTS="$ALLOWED_TCP_PORTS $2"
43+
shift
44+
shift
45+
;;
46+
--allow-udp)
47+
ALLOWED_UDP_PORTS="$ALLOWED_UDP_PORTS $2"
48+
shift
49+
shift
50+
;;
51+
-g | --group)
52+
GROUP_NAME="$2"
53+
shift
54+
shift
55+
;;
56+
-h | --help)
57+
echo "Usage: $0 <username> [--ufw] [-g|--group] [--no-fw] [--allow-tcp <port> --allow-udp <port>]"
58+
echo "Options:"
59+
echo " --no-fw Do not setup/clear firewall rules"
60+
echo " --allow-tcp Allow the specified TCP port to be accessed"
61+
echo " --allow-udp Allow the specified UDP port to be accessed"
62+
echo " -g, --group Add the user to the specified group"
63+
echo " -h, --help Show this help message"
64+
exit 0
65+
;;
66+
*)
67+
if [[ -z "$USERNAME" ]]; then
68+
USERNAME="$1"
69+
else
70+
echo "Error: Unknown argument '$1'"
71+
echo "Use '$0 --help' for usage information"
72+
exit 1
73+
fi
74+
shift
75+
;;
76+
esac
77+
done
78+
79+
# Check if username is provided
80+
if [[ -z "$USERNAME" ]]; then
81+
echo "Error: Username is required"
82+
echo "Usage: $0 <username> [--ufw] [--no-fw] [--allow <port>]"
83+
exit 1
84+
fi
85+
86+
CHAIN_NAME="DSTACK_SANDBOX_${USERNAME}"
87+
88+
# Create the user if it doesn't exist
89+
if ! id -u $USERNAME >/dev/null 2>&1; then
90+
echo "Creating user $USERNAME"
91+
adduser --disabled-password --gecos '' $USERNAME
92+
fi
93+
94+
usermod -aG kvm $USERNAME
95+
96+
# Add the user to specified group
97+
if [ -n "$GROUP_NAME" ]; then
98+
echo "Adding user $USERNAME to group $GROUP_NAME"
99+
usermod -aG $GROUP_NAME $USERNAME
100+
fi
101+
102+
if iptables -L $CHAIN_NAME >/dev/null 2>&1; then
103+
echo "Removing existing firewall rules"
104+
iptables -D OUTPUT -o lo -m owner --uid-owner $USERNAME -j $CHAIN_NAME 2>/dev/null || true
105+
iptables -F $CHAIN_NAME 2>/dev/null || true
106+
iptables -X $CHAIN_NAME 2>/dev/null || true
107+
echo "Removed iptables chain $CHAIN_NAME"
108+
fi
109+
110+
if [ "$NO_FW" = true ]; then
111+
echo "Skipping firewall rules setup"
112+
exit 0
113+
fi
114+
115+
# Set up firewall rules
116+
# Use iptables with a dedicated chain
117+
echo "Setting up iptables firewall rules with custom chain"
118+
119+
# Create or flush the custom chain
120+
if ! iptables -L $CHAIN_NAME >/dev/null 2>&1; then
121+
iptables -N $CHAIN_NAME
122+
else
123+
iptables -F $CHAIN_NAME
124+
fi
125+
126+
# Add rules to allow specific ports
127+
for port in $ALLOWED_TCP_PORTS; do
128+
echo "Adding exception for TCP port $port"
129+
iptables -A $CHAIN_NAME -o lo -d 127.0.0.1 -p tcp --dport $port -j ACCEPT
130+
iptables -A $CHAIN_NAME -o lo -d 127.0.0.1 -p tcp --sport $port -j ACCEPT
131+
done
132+
for port in $ALLOWED_UDP_PORTS; do
133+
echo "Adding exception for UDP port $port"
134+
iptables -A $CHAIN_NAME -o lo -d 127.0.0.1 -p udp --dport $port -j ACCEPT
135+
iptables -A $CHAIN_NAME -o lo -d 127.0.0.1 -p udp --sport $port -j ACCEPT
136+
done
137+
138+
# Add final DROP rule for all other traffic to localhost
139+
iptables -A $CHAIN_NAME -o lo -d 127.0.0.1 -j DROP
140+
141+
# Ensure our chain is referenced from the OUTPUT chain
142+
if ! iptables -C OUTPUT -o lo -m owner --uid-owner $USERNAME -j $CHAIN_NAME 2>/dev/null; then
143+
iptables -I OUTPUT -o lo -m owner --uid-owner $USERNAME -j $CHAIN_NAME
144+
fi
145+
146+
echo "Setup completed for user $USERNAME"

‎teepod/teepod.toml‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -28,8 +28,8 @@ max_allocable_vcpu = 20
2828
max_allocable_memory_in_mb = 100_000 # MB
2929
# Enable QMP socket
3030
qmp_socket = true
31-
# Use sudo to run the VM
32-
sudo = false
31+
# The user to run the VM as. If empty, the VM will be run as the current user.
32+
user = ""
3333

3434
[cvm.port_mapping]
3535
enabled = false

0 commit comments

Comments
 (0)