diff --git a/.gitignore b/.gitignore index 1dda66de..805deafc 100644 --- a/.gitignore +++ b/.gitignore @@ -1,4 +1,6 @@ *agent +.beast-local/ +.codex/ beast.log venv/ site/ @@ -7,4 +9,4 @@ site/ ubuntu-bionic-18.04-cloudimg-console.log target/ .vscode/ -vendor/ \ No newline at end of file +vendor/ diff --git a/_examples/bare-docker/beast.toml b/_examples/bare-docker/beast.toml index cfbba35e..10f4f65d 100644 --- a/_examples/bare-docker/beast.toml +++ b/_examples/bare-docker/beast.toml @@ -20,3 +20,10 @@ points = 20 [challenge.env] docker_context = "docker-file" ports = [10005] +default_port = 10005 + +[resource] +cpu_shares = 1024 +memory_limit = 536870912 +pids_limit = 100 +cpuslimit = 0.25 diff --git a/_examples/compose-type/beast.toml b/_examples/compose-type/beast.toml index 3adf7f08..31698fa3 100644 --- a/_examples/compose-type/beast.toml +++ b/_examples/compose-type/beast.toml @@ -11,10 +11,6 @@ type = "web" points = 200 [challenge.env] -# Beast still requires ports/default_port for validation and metadata. -ports = [10020] -default_port = 10020 - - docker_compose = "docker-compose.yml" +default_port_var = "APP_PORT" web_root = "challenge" diff --git a/_examples/compose-type/docker-compose.yml b/_examples/compose-type/docker-compose.yml index 2d3d94e0..c945a976 100644 --- a/_examples/compose-type/docker-compose.yml +++ b/_examples/compose-type/docker-compose.yml @@ -2,7 +2,7 @@ services: app: build: . ports: - - "10020:80" + - "${APP_PORT}:80" environment: MYSQL_HOST: mysql MYSQL_DATABASE: my_db diff --git a/_examples/example.config.toml b/_examples/example.config.toml index f54a1876..8c916360 100644 --- a/_examples/example.config.toml +++ b/_examples/example.config.toml @@ -44,6 +44,9 @@ username = "user1" # Path to private SSH key for interacting with the server. ssh_key_path = "/path/to/your/private/key1" +# Port range for this server (format: START:END) +port_range = "30000:40000" + # Status of remote server to be used # If it is set to false then that remote server will not be used active = false @@ -53,10 +56,13 @@ active = false host = "localhost" # Username to be used for ssh connection (Leave empty for localhost) -username = "user1" +username = "" # Path to private SSH key for interacting with the server. (Leave empty for localhost) -ssh_key_path = "/path/to/your/private/key1" +ssh_key_path = "" + +# Port range for this server (format: START:END) - uses local_host_port_range if empty +port_range = "" # Status of remote server to be used active = true @@ -113,6 +119,17 @@ host = "localhost" port = "5432" sslmode = "prefer" +[redis_config] +host = "localhost" +port = "6379" +password = "" +user = "" + +[instance_config] +default_expiration = 300 +max_extension = 600 +max_instances_per_user = 3 + # The following fields are required only while hosting a competition on beast # This section contains information about the competition to be hosted # Structure of the sections with the acceptable fields are: diff --git a/_examples/instanced-compose/Dockerfile b/_examples/instanced-compose/Dockerfile new file mode 100644 index 00000000..360a032e --- /dev/null +++ b/_examples/instanced-compose/Dockerfile @@ -0,0 +1,12 @@ +FROM php:7.4-apache + +# Install MySQL extension +RUN docker-php-ext-install mysqli pdo pdo_mysql + +# Copy challenge files +COPY challenge/ /var/www/html/ + +# Set permissions +RUN chown -R www-data:www-data /var/www/html + +EXPOSE 80 diff --git a/_examples/instanced-compose/README.md b/_examples/instanced-compose/README.md new file mode 100644 index 00000000..d6240e8d --- /dev/null +++ b/_examples/instanced-compose/README.md @@ -0,0 +1,89 @@ +# Instanced Docker Compose Challenge Example + +This is an example of an **instanced challenge using Docker Compose** - a multi-container challenge where each user gets their own isolated environment with a web server and database. + +## Architecture + +``` +┌──────────────────────────────────────────┐ +│ User's Instanced Environment │ +│ ┌─────────────┐ ┌─────────────┐ │ +│ │ PHP/Apache │ ───▶ │ MySQL │ │ +│ │ (web) │ │ (db) │ │ +│ └─────────────┘ └─────────────┘ │ +│ │ │ +│ ▼ │ +│ Port: 31234 (dynamically assigned) │ +└──────────────────────────────────────────┘ +``` + +## Key Configuration + +In `beast.toml`: + +```toml +[challenge.metadata] +instanced = true +instance_expiration = 600 # 10 minutes + +[challenge.env] +docker_compose = "docker-compose.yml" +default_port = 8080 +``` + +In `docker-compose.yml`, use the `INSTANCE_PORT` environment variable: + +```yaml +services: + web: + ports: + - "${INSTANCE_PORT:-8080}:80" +``` + +## Challenge Details + +This is a SQL injection challenge: + +1. The login form is vulnerable to SQL injection +2. Bypass authentication to login as admin +3. The flag is stored in the `secrets` table + +### Solution + +``` +Username: admin' OR '1'='1' -- +Password: anything +``` + +Or use UNION-based injection to extract data directly. + +## Testing Locally + +```bash +# Build and run locally (for testing) +cd _examples/instanced-compose +docker-compose up -d + +# Access at http://localhost:8080 +``` + +## Usage via Beast API + +```bash +# Spawn your instance +curl -X POST -H "Authorization: Bearer $TOKEN" \ + http://localhost:8080/api/instances/instanced-compose/spawn + +# Response: +# { +# "instance_id": "abc123def456", +# "challenge_name": "instanced-compose", +# "hosted_address": "localhost", +# "port": 31234, +# "expires_at": "2024-01-15T10:40:00Z", +# "ttl_seconds": 600 +# } + +# Access your instance +open http://localhost:31234 +``` diff --git a/_examples/instanced-compose/beast.toml b/_examples/instanced-compose/beast.toml new file mode 100644 index 00000000..b07b2df7 --- /dev/null +++ b/_examples/instanced-compose/beast.toml @@ -0,0 +1,34 @@ +[author] +name = "beast-admin" +email = "admin@beast.local" +ssh_key = "ssh-rsa AAAAB3NzaC1yc2EAAAADAQABAAABAQ" + +[challenge.metadata] +name = "instanced-compose" +flag = "FLAG{c0mp0s3_1nst4nc3s_r0ck!}" +type = "web" +description = "A web challenge with database backend. Each user gets their own isolated environment!" +points = 200 +difficulty = "medium" +tags = ["web", "sql", "instanced"] +maxAttemptLimit = 100 +instanced = true +instance_expiration = 12 + +[[challenge.metadata.hints]] +text = "Check for SQL injection vulnerabilities" +points = 30 + +[[challenge.metadata.hints]] +text = "The admin password might be in the database..." +points = 50 + +[challenge.env] +docker_compose = "docker-compose.yml" +default_port_var = "INSTANCE_PORT" + +[resource] +cpu_shares = 1024 +memory_limit = 536870912 +pids_limit = 100 +cpuslimit = 0.25 diff --git a/_examples/instanced-compose/challenge/index.php b/_examples/instanced-compose/challenge/index.php new file mode 100644 index 00000000..42a0e5a5 --- /dev/null +++ b/_examples/instanced-compose/challenge/index.php @@ -0,0 +1,158 @@ + + + + Secret Vault - Login + + + +
+

Secret Vault

+ + connect_error) { + $error = "Connection failed. Please try again."; + } else { + $username = $_POST['username']; + $password = $_POST['password']; + + // VULNERABLE: SQL Injection! + $query = "SELECT * FROM users WHERE username='$username' AND password='$password'"; + $result = $conn->query($query); + + if ($result && $result->num_rows > 0) { + $row = $result->fetch_assoc(); + + if ($row['role'] === 'admin') { + // Admin login - show secrets + $secrets_query = "SELECT * FROM secrets"; + $secrets_result = $conn->query($secrets_query); + + $success = "Welcome Admin! Here are your secrets:

"; + while ($secret = $secrets_result->fetch_assoc()) { + $success .= "" . htmlspecialchars($secret['secret_name']) . ": " . + htmlspecialchars($secret['secret_value']) . "
"; + } + } else { + $success = "Welcome, " . htmlspecialchars($row['username']) . "! You're logged in as a regular user."; + } + } else { + $error = "Invalid username or password!"; + } + + $conn->close(); + } + } + ?> + + +
+ + + +
+ +
+
+ + +
+
+ + +
+ +
+ + +

Hint: Try logging in as admin to see the secrets!

+
+ + diff --git a/_examples/instanced-compose/docker-compose.yml b/_examples/instanced-compose/docker-compose.yml new file mode 100644 index 00000000..2f51b88c --- /dev/null +++ b/_examples/instanced-compose/docker-compose.yml @@ -0,0 +1,28 @@ +version: '3.8' + +services: + web: + build: + context: . + dockerfile: Dockerfile + ports: + - "${INSTANCE_PORT}:80" + environment: + - DB_HOST=db + - DB_USER=challenge + - DB_PASS=challengepass + - DB_NAME=ctf + depends_on: + - db + restart: unless-stopped + + db: + image: mysql:5.7 + environment: + - MYSQL_ROOT_PASSWORD=rootpass + - MYSQL_DATABASE=ctf + - MYSQL_USER=challenge + - MYSQL_PASSWORD=challengepass + volumes: + - ./init.sql:/docker-entrypoint-initdb.d/init.sql:ro + restart: unless-stopped diff --git a/_examples/instanced-compose/init.sql b/_examples/instanced-compose/init.sql new file mode 100644 index 00000000..f7acd93a --- /dev/null +++ b/_examples/instanced-compose/init.sql @@ -0,0 +1,26 @@ +-- Initialize the CTF database + +USE ctf; + +CREATE TABLE users ( + id INT AUTO_INCREMENT PRIMARY KEY, + username VARCHAR(50) NOT NULL, + password VARCHAR(255) NOT NULL, + role VARCHAR(20) DEFAULT 'user' +); + +CREATE TABLE secrets ( + id INT AUTO_INCREMENT PRIMARY KEY, + secret_name VARCHAR(100) NOT NULL, + secret_value TEXT NOT NULL +); + +-- Insert some users +INSERT INTO users (username, password, role) VALUES + ('guest', 'guest123', 'user'), + ('admin', 'sup3rs3cr3t_4dm1n_p4ss!', 'admin'); + +-- Insert the flag as a secret +INSERT INTO secrets (secret_name, secret_value) VALUES + ('flag', 'FLAG{c0mp0s3_1nst4nc3s_r0ck!}'), + ('admin_note', 'Remember to change the admin password!'); diff --git a/_examples/instanced-service/README.md b/_examples/instanced-service/README.md new file mode 100644 index 00000000..7710cbcf --- /dev/null +++ b/_examples/instanced-service/README.md @@ -0,0 +1,119 @@ +# Instanced Service Challenge Example + +This is an example of an **instanced challenge** - a challenge where each user gets their own dedicated container instance. + +## Key Features + +- **Per-user isolation**: Each user spawns their own container +- **Automatic expiration**: Instances expire after a configurable time (default: 5 minutes) +- **Dynamic port allocation**: Ports are assigned from a configured range (not from the challenge config) + +## Configuration + +In `beast.toml`, the key settings for instanced challenges are: + +```toml +[challenge.metadata] +instanced = true # Enable instancing +instance_expiration = 300 # Optional: override default expiration (in seconds) + +[challenge.env] +# DO NOT specify ports for instanced challenges! +# Instead, use default_port to indicate which container port to expose +default_port = 9999 +``` + +## Global Configuration + +In your Beast `config.toml`, configure the instance settings: + +```toml +[instance_config] +local_host_port_range = "10000-11000" # Host port range for instances +default_expiration = 300 # Default TTL in seconds (5 minutes) +max_extension = 600 # Maximum extension time (10 minutes) +max_instances_per_user = 3 # Max concurrent instances per user +``` + +## API Usage + +### User Endpoints + +1. **Spawn an instance**: + ```bash + curl -X POST -H "Authorization: Bearer $TOKEN" \ + http://localhost:8080/api/instances/instanced-service/spawn + ``` + +2. **Get your instance**: + ```bash + curl -H "Authorization: Bearer $TOKEN" \ + http://localhost:8080/api/instances/instanced-service + ``` + +3. **Get all your instances**: + ```bash + curl -H "Authorization: Bearer $TOKEN" \ + http://localhost:8080/api/instances + ``` + +4. **Extend instance lifetime**: + ```bash + curl -X POST -H "Authorization: Bearer $TOKEN" \ + -d "seconds=300" \ + http://localhost:8080/api/instances/instanced-service/extend + ``` + +5. **Kill your instance**: + ```bash + curl -X DELETE -H "Authorization: Bearer $TOKEN" \ + http://localhost:8080/api/instances/instanced-service + ``` + +### Admin Endpoints + +1. **List all instances**: + ```bash + curl -H "Authorization: Bearer $ADMIN_TOKEN" \ + http://localhost:8080/api/admin/instances + ``` + +2. **Kill any instance**: + ```bash + curl -X DELETE -H "Authorization: Bearer $ADMIN_TOKEN" \ + http://localhost:8080/api/admin/instances/{instance_id} + ``` + +3. **Kill all instances for a challenge**: + ```bash + curl -X DELETE -H "Authorization: Bearer $ADMIN_TOKEN" \ + http://localhost:8080/api/admin/instances/challenge/instanced-service + ``` + +## Response Example + +When spawning an instance, you'll receive: + +```json +{ + "instance_id": "a1b2c3d4e5f6", + "challenge_name": "instanced-service", + "hosted_address": "localhost", + "port": 31234, + "created_at": "2024-01-15T10:30:00Z", + "expires_at": "2024-01-15T10:35:00Z", + "ttl_seconds": 300 +} +``` + +Connect to your instance: +```bash +nc localhost 31234 +``` + +## Challenge Details + +This example is a simple buffer overflow challenge: +- The `vulnerable()` function uses `gets()` which doesn't check bounds +- Overflow the 64-byte buffer to overwrite the return address +- Redirect execution to the `win()` function to get the flag diff --git a/_examples/instanced-service/beast.toml b/_examples/instanced-service/beast.toml new file mode 100644 index 00000000..695a8870 --- /dev/null +++ b/_examples/instanced-service/beast.toml @@ -0,0 +1,38 @@ +[author] +name = "beast-admin" +email = "admin@beast.local" +ssh_key = "ssh-rsa AAAAB3NzaC1yc2EAAAADAQABAAABAQ" + +[challenge.metadata] +name = "instanced-service" +flag = "FLAG{1nst4nc3d_ch4ll3ng3_w0rks!}" +type = "service" +description = "A simple buffer overflow challenge. Each user gets their own instance!" +points = 150 +difficulty = "easy" +tags = ["pwn", "beginner", "instanced"] +maxAttemptLimit = 50 +instanced = true +instance_expiration = 10 + +[[challenge.metadata.hints]] +text = "Have you tried overflowing the buffer?" +points = 25 + +[[challenge.metadata.hints]] +text = "The sample() function looks interesting..." +points = 50 + +[challenge.env] +ports = [9999] +default_port = 9999 +apt_deps = ["gcc", "xinetd"] +setup_scripts = ["setup.sh"] +service_path = "pwn" +base_image = "ubuntu:18.04" + +[resource] +cpu_shares = 1024 +memory_limit = 536870912 +pids_limit = 100 +cpuslimit = 0.25 diff --git a/_examples/instanced-service/pwn_me.c b/_examples/instanced-service/pwn_me.c new file mode 100644 index 00000000..2ac0ae7e --- /dev/null +++ b/_examples/instanced-service/pwn_me.c @@ -0,0 +1,50 @@ +#include +#include +#include +#include + +// Compile with: gcc -o pwn pwn_me.c -fno-stack-protector -no-pie + +void win() { + FILE *fp; + char flag[100]; + + fp = fopen("/challenge/flag.txt", "r"); + if (fp == NULL) { + printf("Error: Could not open flag file!\n"); + return; + } + + if (fgets(flag, sizeof(flag), fp) != NULL) { + printf("Congratulations! Here's your flag: %s\n", flag); + } + + fclose(fp); +} + +void vulnerable() { + char buffer[64]; + + printf("Welcome to the Instanced PWN Challenge!\n"); + printf("Each user gets their own container instance.\n"); + printf("Can you overflow the buffer and call win()?\n\n"); + printf("Enter your payload: "); + fflush(stdout); + + // Vulnerable: no bounds checking! + gets(buffer); + + printf("You entered: %s\n", buffer); + printf("Better luck next time!\n"); +} + +int main() { + // Disable buffering for proper network I/O + setvbuf(stdin, NULL, _IONBF, 0); + setvbuf(stdout, NULL, _IONBF, 0); + setvbuf(stderr, NULL, _IONBF, 0); + + vulnerable(); + + return 0; +} diff --git a/_examples/instanced-service/setup.sh b/_examples/instanced-service/setup.sh new file mode 100644 index 00000000..c8a6f9dc --- /dev/null +++ b/_examples/instanced-service/setup.sh @@ -0,0 +1,16 @@ +#!/bin/bash +set -e + +echo "[*] Setting up instanced-service challenge..." + +# Compile the vulnerable binary +gcc -o pwn pwn_me.c -fno-stack-protector -no-pie -z execstack + +# Make it executable +chmod +x pwn + +# Create flag file +echo "FLAG{1nst4nc3d_ch4ll3ng3_w0rks!}" > flag.txt +chmod 444 flag.txt + +echo "[*] Setup complete!" diff --git a/_examples/service/beast.toml b/_examples/service/beast.toml index e95a4a2d..f0a3e554 100644 --- a/_examples/service/beast.toml +++ b/_examples/service/beast.toml @@ -24,3 +24,10 @@ apt_deps = ["gcc", "socat"] setup_scripts = ["setup.sh"] service_path = "pwn" ports = [10004] +default_port = 10004 + +[resource] +cpu_shares = 1024 +memory_limit = 536870912 +pids_limit = 100 +cpuslimit = 0.25 diff --git a/_examples/simple/beast.toml b/_examples/simple/beast.toml index 661245f3..1a408966 100644 --- a/_examples/simple/beast.toml +++ b/_examples/simple/beast.toml @@ -22,3 +22,10 @@ apt_deps = ["gcc", "socat"] setup_scripts = ["setup.sh"] run_cmd = "socat tcp-l:10005,fork,reuseaddr exec:./pwn" ports = [10005] +default_port = 10005 + +[resource] +cpu_shares = 1024 +memory_limit = 536870912 +pids_limit = 100 +cpuslimit = 0.25 diff --git a/_examples/static-chall/beast.toml b/_examples/static-chall/beast.toml index 6a74ddea..7d9c5560 100644 --- a/_examples/static-chall/beast.toml +++ b/_examples/static-chall/beast.toml @@ -22,4 +22,10 @@ minPoints = 50 tags = ["easy", "web"] [challenge.env] -static_dir = "static" \ No newline at end of file +static_dir = "static" + +[resource] +cpu_shares = 1024 +memory_limit = 536870912 +pids_limit = 100 +cpuslimit = 0.25 diff --git a/_examples/web-php-mysql/beast.toml b/_examples/web-php-mysql/beast.toml index 73ecf752..5e45eb8c 100644 --- a/_examples/web-php-mysql/beast.toml +++ b/_examples/web-php-mysql/beast.toml @@ -25,3 +25,9 @@ setup_scripts = ["setup.sh"] ports = [10004] web_root = "challenge" default_port = 10004 + +[resource] +cpu_shares = 1024 +memory_limit = 536870912 +pids_limit = 100 +cpuslimit = 0.25 diff --git a/_examples/web-php/beast.toml b/_examples/web-php/beast.toml index 018376f8..eb68611e 100644 --- a/_examples/web-php/beast.toml +++ b/_examples/web-php/beast.toml @@ -23,3 +23,9 @@ points = 20 ports = [10002] web_root = "challenge" default_port = 10002 + +[resource] +cpu_shares = 1024 +memory_limit = 536870912 +pids_limit = 100 +cpuslimit = 0.25 diff --git a/_examples/xinetd-service/beast.toml b/_examples/xinetd-service/beast.toml index f4c7bbf7..3ebbbcfc 100644 --- a/_examples/xinetd-service/beast.toml +++ b/_examples/xinetd-service/beast.toml @@ -27,3 +27,10 @@ setup_scripts = ["setup.sh"] xinetd_config = "ctf.xinetd" service_path = "pwn" ports = [10003] +default_port = 10003 + +[resource] +cpu_shares = 1024 +memory_limit = 536870912 +pids_limit = 100 +cpuslimit = 0.25 diff --git a/api/info.go b/api/info.go index c604c7a6..175a52f9 100644 --- a/api/info.go +++ b/api/info.go @@ -29,24 +29,6 @@ var ( graphCacheStale = true ) -// Returns port in use by beast. -// @Summary Returns ports in use by beast by looking in the hack git repository, also returns min and max value of port allowed while specifying in beast challenge config. -// @Description Returns the ports in use by beast, which cannot be used in creating a new challenge.. -// @Tags info -// @Accept json -// @Produce json -// @Param Authorization header string true "Bearer" -// @Success 200 {object} api.PortsInUseResp -// @Router /api/info/ports/used [get] - -func usedPortsInfoHandler(c *gin.Context) { - c.JSON(http.StatusOK, PortsInUseResp{ - MinPortValue: core.ALLOWED_MIN_PORT_VALUE, - MaxPortValue: core.ALLOWED_MAX_PORT_VALUE, - PortsInUse: cfg.USED_PORTS_LIST, - }) -} - func hintHandler(c *gin.Context) { hintIDStr := c.Param("hintID") @@ -157,7 +139,7 @@ func hintHandler(c *gin.Context) { }) return } - + oldScore := user.Score newScore := oldScore - hint.Points if newScore < 0 { @@ -394,7 +376,7 @@ func challengesMetadataHandler(c *gin.Context) { } } - availableChallenges := make([]ChallengeMetadata, len(challenges)) + availableChallenges := make([]ChallengeMetadata, 0, len(challenges)) authHeader := c.GetHeader("Authorization") username, err := coreUtils.GetUser(authHeader) @@ -413,7 +395,7 @@ func challengesMetadataHandler(c *gin.Context) { return } - for index, challenge := range challenges { + for _, challenge := range challenges { if challenge.Status == "Undeployed" && user.Role == core.USER_ROLES["contestant"] { continue } @@ -427,22 +409,24 @@ func challengesMetadataHandler(c *gin.Context) { } challengeTags := make([]string, len(challenge.Tags)) - for index, tags := range challenge.Tags { - challengeTags[index] = tags.TagName + for i, tags := range challenge.Tags { + challengeTags[i] = tags.TagName } - availableChallenges[index] = ChallengeMetadata{ - Name: challenge.Name, - ChallId: challenge.ID, - Tags: challengeTags, - CreatedAt: challenge.CreatedAt, - Points: challenge.Points, - SolvesNumber: totalSolves, - SolveStatus: solveStatus, - Difficulty: challenge.Difficulty, - PreRequisite: strings.Split(challenge.PreReqs, core.DELIMITER), - DeployedStatus: challenge.Status, - } + availableChallenges = append(availableChallenges, ChallengeMetadata{ + Name: challenge.Name, + ChallId: challenge.ID, + Tags: challengeTags, + CreatedAt: challenge.CreatedAt, + Points: challenge.Points, + SolvesNumber: totalSolves, + SolveStatus: solveStatus, + Difficulty: challenge.Difficulty, + PreRequisite: strings.Split(challenge.PreReqs, core.DELIMITER), + DeployedStatus: challenge.Status, + Instanced: challenge.Instanced, + InstanceExpiration: challenge.InstanceExpiration, + }) } c.JSON(http.StatusOK, availableChallenges) @@ -524,7 +508,6 @@ func userInfoHandler(c *gin.Context) { } var user database.User var err error - var parsedUserId uint if userId != "" { id, err := strconv.ParseUint(userId, 10, 64) if err != nil { @@ -533,9 +516,8 @@ func userInfoHandler(c *gin.Context) { }) return } - parsedUserId = uint(id) - user, err = database.QueryUserById(parsedUserId) + user, err = database.QueryUserById(uint(id)) if err != nil { c.JSON(http.StatusInternalServerError, HTTPErrorResp{ Error: "DATABASE ERROR while processing the request.", @@ -552,7 +534,7 @@ func userInfoHandler(c *gin.Context) { } } - challenges, err := database.GetRelatedChallenges(&user) + solvedChallenges, err := database.GetUserSolvedChallenges(user.ID) if err != nil { log.Error(err) c.JSON(http.StatusInternalServerError, HTTPErrorResp{ @@ -560,36 +542,29 @@ func userInfoHandler(c *gin.Context) { }) return } - var resp UserResp - - var challNameString []string - for _, challenge := range challenges { - challNameString = append(challNameString, challenge.Name) - } - userChallenges := make([]ChallengeSolveResp, len(challenges)) - for index, challenge := range challenges { + userChallenges := make([]ChallengeSolveResp, len(solvedChallenges)) + for index, sc := range solvedChallenges { - challengeTags := make([]string, len(challenge.Tags)) - - for index, tags := range challenge.Tags { - challengeTags[index] = tags.TagName + challengeTags := make([]string, len(sc.Tags)) + for i, tag := range sc.Tags { + challengeTags[i] = tag.TagName } challResp := ChallengeSolveResp{ - Id: challenge.ID, - Name: challenge.Name, + Id: sc.ChallengeID, + Name: sc.Name, Tags: challengeTags, - Category: challenge.Type, - SolvedAt: challenge.CreatedAt, - Points: challenge.Points, + Category: sc.Type, + SolvedAt: sc.SolvedAt, + Points: sc.Points, } userChallenges[index] = challResp } var rank int64 if user.Status == 0 { - rank, err = database.GetUserRank(parsedUserId, user.Score, user.UpdatedAt) + rank, err = database.GetUserRank(user.ID, user.Score, user.UpdatedAt) } else { rank = 1e9 } @@ -602,7 +577,7 @@ func userInfoHandler(c *gin.Context) { return } - resp = UserResp{ + resp := UserResp{ Username: user.Username, Id: user.ID, Role: user.Role, diff --git a/api/instance.go b/api/instance.go new file mode 100644 index 00000000..e2adc9ab --- /dev/null +++ b/api/instance.go @@ -0,0 +1,413 @@ +package api + +import ( + "fmt" + "github.com/sdslabs/beastv4/core" + "net/http" + "time" + + "github.com/gin-gonic/gin" + "github.com/sdslabs/beastv4/core/cache" + "github.com/sdslabs/beastv4/core/config" + "github.com/sdslabs/beastv4/core/database" + "github.com/sdslabs/beastv4/core/manager" + coreUtils "github.com/sdslabs/beastv4/core/utils" +) + +type InstanceResponse struct { + InstanceID string `json:"instance_id"` + ChallengeName string `json:"challenge_name"` + HostedAddress string `json:"hosted_address"` + Port uint32 `json:"port"` + CreatedAt time.Time `json:"created_at"` + ExpiresAt time.Time `json:"expires_at"` + TTLSeconds int64 `json:"ttl_seconds"` +} + +type AdminInstanceResponse struct { + InstanceResponse + UserID string `json:"user_id"` + Username string `json:"username"` + ContainerID string `json:"container_id"` + DeploymentType string `json:"deployment_type"` +} + +func instanceToResponse(instance *cache.Instance) InstanceResponse { + ttl := time.Until(instance.ExpiresAt).Seconds() + if ttl < 0 { + ttl = 0 + } + + return InstanceResponse{ + InstanceID: instance.InstanceID, + ChallengeName: instance.ChallengeName, + HostedAddress: instance.ServerDeployed, + Port: instance.Port, + CreatedAt: instance.CreatedAt, + ExpiresAt: instance.ExpiresAt, + TTLSeconds: int64(ttl), + } +} + +func instanceToAdminResponse(instance *cache.Instance) AdminInstanceResponse { + return AdminInstanceResponse{ + InstanceResponse: instanceToResponse(instance), + UserID: instance.UserID, + Username: instance.Username, + ContainerID: instance.ContainerID, + DeploymentType: instance.DeploymentType, + } +} + +func spawnInstanceHandler(ctx *gin.Context) { + challengeName := ctx.Param("challenge_name") + if challengeName == "" { + ctx.JSON(http.StatusBadRequest, HTTPErrorResp{ + Error: "challenge_name is required", + }) + return + } + + username, err := coreUtils.GetUser(ctx.GetHeader("Authorization")) + if err != nil { + ctx.JSON(http.StatusUnauthorized, HTTPPlainResp{ + Message: "Unauthorized", + }) + return + } + + user, err := database.QueryFirstUserEntry("username", username) + if err != nil { + ctx.JSON(http.StatusUnauthorized, HTTPPlainResp{ + Message: "User not found", + }) + return + } + + userID := fmt.Sprintf("%d", user.ID) + + instance, err := manager.SpawnInstance(challengeName, userID, username) + if err != nil { + if instance != nil { + ctx.JSON(http.StatusConflict, HTTPErrorResp{ + Error: err.Error(), + }) + return + } + ctx.JSON(http.StatusBadRequest, HTTPErrorResp{ + Error: err.Error(), + }) + return + } + + ctx.JSON(http.StatusOK, instanceToResponse(instance)) +} + +func getUserInstanceHandler(ctx *gin.Context) { + challengeName := ctx.Param("challenge_name") + if challengeName == "" { + ctx.JSON(http.StatusBadRequest, HTTPErrorResp{ + Error: "challenge_name is required", + }) + return + } + + username, err := coreUtils.GetUser(ctx.GetHeader("Authorization")) + if err != nil { + ctx.JSON(http.StatusUnauthorized, HTTPPlainResp{ + Message: "Unauthorized", + }) + return + } + + user, err := database.QueryFirstUserEntry("username", username) + if err != nil { + ctx.JSON(http.StatusUnauthorized, HTTPPlainResp{ + Message: "User not found", + }) + return + } + + userID := fmt.Sprintf("%d", user.ID) + + instance, err := manager.GetUserInstance(userID, challengeName) + if err != nil { + ctx.JSON(http.StatusNotFound, HTTPErrorResp{ + Error: "No active instance found for this challenge", + }) + return + } + + ctx.JSON(http.StatusOK, instanceToResponse(instance)) +} + +func getUserInstancesHandler(ctx *gin.Context) { + username, err := coreUtils.GetUser(ctx.GetHeader("Authorization")) + if err != nil { + ctx.JSON(http.StatusUnauthorized, HTTPPlainResp{ + Message: "Unauthorized", + }) + return + } + + user, err := database.QueryFirstUserEntry("username", username) + if err != nil { + ctx.JSON(http.StatusUnauthorized, HTTPPlainResp{ + Message: "User not found", + }) + return + } + + userID := fmt.Sprintf("%d", user.ID) + + instances, err := manager.GetUserInstances(userID) + if err != nil { + ctx.JSON(http.StatusInternalServerError, HTTPErrorResp{ + Error: err.Error(), + }) + return + } + + var response []InstanceResponse + for _, instance := range instances { + response = append(response, instanceToResponse(instance)) + } + + if response == nil { + response = []InstanceResponse{} + } + + ctx.JSON(http.StatusOK, response) +} + +func extendInstanceHandler(ctx *gin.Context) { + challengeName := ctx.Param("challenge_name") + if challengeName == "" { + ctx.JSON(http.StatusBadRequest, HTTPErrorResp{ + Error: "challenge_name is required", + }) + return + } + + username, err := coreUtils.GetUser(ctx.GetHeader("Authorization")) + if err != nil { + ctx.JSON(http.StatusUnauthorized, HTTPPlainResp{ + Message: "Unauthorized", + }) + return + } + + user, err := database.QueryFirstUserEntry("username", username) + if err != nil { + ctx.JSON(http.StatusUnauthorized, HTTPPlainResp{ + Message: "User not found", + }) + return + } + + userID := fmt.Sprintf("%d", user.ID) + + instance, err := manager.GetUserInstance(userID, challengeName) + if err != nil { + ctx.JSON(http.StatusNotFound, HTTPErrorResp{ + Error: "No active instance found for this challenge", + }) + return + } + + additionalSeconds := core.DEFAULT_MINIMUM_EXTEND_TIME + if seconds := ctx.PostForm("seconds"); seconds != "" { + var parsedSeconds int64 + _, err := fmt.Sscanf(seconds, "%d", &parsedSeconds) + if err == nil && parsedSeconds > 0 { + additionalSeconds = parsedSeconds + } + } + + maxExtension := config.Cfg.InstanceConfig.MaxExtension + if additionalSeconds > maxExtension { + additionalSeconds = maxExtension + } + + err = manager.ExtendInstance(instance.InstanceID, additionalSeconds) + if err != nil { + ctx.JSON(http.StatusBadRequest, HTTPErrorResp{ + Error: err.Error(), + }) + return + } + + instance, err = manager.GetUserInstance(userID, challengeName) + if err != nil { + ctx.JSON(http.StatusInternalServerError, HTTPErrorResp{ + Error: "Failed to get updated instance", + }) + return + } + + ctx.JSON(http.StatusOK, instanceToResponse(instance)) +} + +func killUserInstanceHandler(ctx *gin.Context) { + challengeName := ctx.Param("challenge_name") + if challengeName == "" { + ctx.JSON(http.StatusBadRequest, HTTPErrorResp{ + Error: "challenge_name is required", + }) + return + } + + username, err := coreUtils.GetUser(ctx.GetHeader("Authorization")) + if err != nil { + ctx.JSON(http.StatusUnauthorized, HTTPPlainResp{ + Message: "Unauthorized", + }) + return + } + + user, err := database.QueryFirstUserEntry("username", username) + if err != nil { + ctx.JSON(http.StatusUnauthorized, HTTPPlainResp{ + Message: "User not found", + }) + return + } + + userID := fmt.Sprintf("%d", user.ID) + + err = manager.KillUserInstance(userID, challengeName) + if err != nil { + ctx.JSON(http.StatusNotFound, HTTPErrorResp{ + Error: err.Error(), + }) + return + } + + ctx.JSON(http.StatusOK, HTTPPlainResp{ + Message: "Instance killed successfully", + }) +} + +func adminGetAllInstancesHandler(ctx *gin.Context) { + instances, err := manager.GetAllInstances() + if err != nil { + ctx.JSON(http.StatusInternalServerError, HTTPErrorResp{ + Error: err.Error(), + }) + return + } + + var response []AdminInstanceResponse + for _, instance := range instances { + response = append(response, instanceToAdminResponse(instance)) + } + + if response == nil { + response = []AdminInstanceResponse{} + } + + ctx.JSON(http.StatusOK, response) +} + +func adminGetInstanceHandler(ctx *gin.Context) { + instanceID := ctx.Param("instance_id") + if instanceID == "" { + ctx.JSON(http.StatusBadRequest, HTTPErrorResp{ + Error: "instance_id is required", + }) + return + } + + instance, err := manager.GetInstance(instanceID) + if err != nil { + ctx.JSON(http.StatusNotFound, HTTPErrorResp{ + Error: "Instance not found", + }) + return + } + + ctx.JSON(http.StatusOK, instanceToAdminResponse(instance)) +} + +func adminKillInstanceHandler(ctx *gin.Context) { + instanceID := ctx.Param("instance_id") + if instanceID == "" { + ctx.JSON(http.StatusBadRequest, HTTPErrorResp{ + Error: "instance_id is required", + }) + return + } + + err := manager.KillInstance(instanceID) + if err != nil { + ctx.JSON(http.StatusNotFound, HTTPErrorResp{ + Error: err.Error(), + }) + return + } + + ctx.JSON(http.StatusOK, HTTPPlainResp{ + Message: "Instance killed successfully", + }) +} + +func adminKillUserInstancesHandler(ctx *gin.Context) { + userID := ctx.Param("user_id") + if userID == "" { + ctx.JSON(http.StatusBadRequest, HTTPErrorResp{ + Error: "user_id is required", + }) + return + } + + instances, err := manager.GetUserInstances(userID) + if err != nil { + ctx.JSON(http.StatusInternalServerError, HTTPErrorResp{ + Error: err.Error(), + }) + return + } + + killedCount := 0 + for _, instance := range instances { + err := manager.KillInstance(instance.InstanceID) + if err == nil { + killedCount++ + } + } + + ctx.JSON(http.StatusOK, HTTPPlainResp{ + Message: fmt.Sprintf("%d instances killed", killedCount), + }) +} + +func adminKillChallengeInstancesHandler(ctx *gin.Context) { + challengeName := ctx.Param("challenge_name") + if challengeName == "" { + ctx.JSON(http.StatusBadRequest, HTTPErrorResp{ + Error: "challenge_name is required", + }) + return + } + + instances, err := manager.GetChallengeInstances(challengeName) + if err != nil { + ctx.JSON(http.StatusInternalServerError, HTTPErrorResp{ + Error: err.Error(), + }) + return + } + + // can be delegated to a coroutine if bottlenecks performance + killedCount := 0 + for _, instance := range instances { + err = manager.KillInstance(instance.InstanceID) + if err == nil { + killedCount++ + } + } + + ctx.JSON(http.StatusOK, HTTPPlainResp{ + Message: fmt.Sprintf("%d instances killed", killedCount), + }) +} diff --git a/api/main.go b/api/main.go index 8cfd5e55..37c932bb 100644 --- a/api/main.go +++ b/api/main.go @@ -4,6 +4,7 @@ import ( "net/http" "github.com/gin-gonic/gin" + "github.com/sdslabs/beastv4/core/cache" log "github.com/sirupsen/logrus" ginSwagger "github.com/swaggo/gin-swagger" swaggerFiles "github.com/swaggo/gin-swagger/swaggerFiles" @@ -67,6 +68,9 @@ func RunBeastApiServer(port, defaultauthorpassword string, autoDeploy, healthPro auth.Init(core.ITERATIONS, core.HASH_LENGTH, core.TIMEPERIOD, core.ISSUER, config.Cfg.JWTSecret, []string{core.USER_ROLES["author"]}, []string{core.USER_ROLES["admin"]}, []string{core.USER_ROLES["contestant"]}) remoteManager.Init() database.Init() + cache.Init() + startDynamicScoreWorker() + go manager.InstanceCleanupProber() // Initialise and start the Hub // Must be started before the Notification Router, since SSE handler has access to SSE Hub diff --git a/api/response.go b/api/response.go index 4038b207..04d0dce8 100644 --- a/api/response.go +++ b/api/response.go @@ -106,16 +106,18 @@ type HintResponse struct { } type ChallengeMetadata struct { - ChallId uint `json:"id" example:"0"` - Name string `json:"name" example:"Web Challenge"` - Tags []string `json:"tags" example:"['pwn','misc']"` - Points uint `json:"points" example:"50"` - Difficulty string `json:"difficulty" example:"easy"` // e.g., "easy", "medium", "hard" - SolvesNumber uint16 `json:"solvesNumber" example:"100"` - SolveStatus bool `json:"solveStatus" example:"True"` // e.g., True: "solved", False: "unsolved" - CreatedAt time.Time `json:"createdAt"` - DeployedStatus string `json:"deployedStatus" example:"deployed"` - PreRequisite []string `json:"preRequisite" example:"['chall1', chall2]"` + ChallId uint `json:"id" example:"0"` + Name string `json:"name" example:"Web Challenge"` + Tags []string `json:"tags" example:"['pwn','misc']"` + Points uint `json:"points" example:"50"` + Difficulty string `json:"difficulty" example:"easy"` + SolvesNumber uint16 `json:"solvesNumber" example:"100"` + SolveStatus bool `json:"solveStatus" example:"True"` + CreatedAt time.Time `json:"createdAt"` + DeployedStatus string `json:"deployedStatus" example:"deployed"` + PreRequisite []string `json:"preRequisite" example:"['chall1', chall2]"` + Instanced bool `json:"instanced" example:"false"` + InstanceExpiration int64 `json:"instanceExpiration" example:"300"` } type Challenge struct { diff --git a/api/router.go b/api/router.go index c395b38c..ef667fab 100644 --- a/api/router.go +++ b/api/router.go @@ -132,6 +132,21 @@ func initGinRouter() *gin.Engine { adminPanelGroup.POST("/freezeLeaderboard", freezeLeaderboardHandler) adminPanelGroup.POST("/unfreezeLeaderboard", unfreezeLeaderboardHandler) adminPanelGroup.GET("/submissions", submissionsHandler) + + adminPanelGroup.GET("/instances", adminGetAllInstancesHandler) + adminPanelGroup.GET("/instances/:instance_id", adminGetInstanceHandler) + adminPanelGroup.DELETE("/instances/:instance_id", adminKillInstanceHandler) + adminPanelGroup.DELETE("/instances/user/:user_id", adminKillUserInstancesHandler) + adminPanelGroup.DELETE("/instances/challenge/:challenge_name", adminKillChallengeInstancesHandler) + } + + instanceGroup := apiGroup.Group("/instances") + { + instanceGroup.GET("", getUserInstancesHandler) + instanceGroup.GET("/:challenge_name", getUserInstanceHandler) + instanceGroup.POST("/:challenge_name/spawn", spawnInstanceHandler) + instanceGroup.POST("/:challenge_name/extend", extendInstanceHandler) + instanceGroup.DELETE("/:challenge_name", killUserInstanceHandler) } } diff --git a/api/submit.go b/api/submit.go index f2d6ee80..c08b4373 100644 --- a/api/submit.go +++ b/api/submit.go @@ -4,6 +4,7 @@ import ( "math" "net/http" "strconv" + "sync" "time" "github.com/gin-gonic/gin" @@ -15,6 +16,11 @@ import ( log "github.com/sirupsen/logrus" ) +var ( + dynamicScoreWorkerOnce sync.Once + dynamicScoreNotify = make(chan struct{}, 1) +) + // Verifies and creates an entry in the database for successful submission of flag for a challenge. // @Summary Verifies and creates an entry in the database for successful submission of flag for a challenge. // @Description Returns success or error response based on the flag submitted. Also, the flag will not be submitted if it was previously submitted @@ -31,6 +37,7 @@ import ( func submitFlagHandler(c *gin.Context) { challId := c.PostForm("chall_id") flag := c.PostForm("flag") + now := time.Now() err, state := coreUtils.CheckTime() if err != nil { @@ -51,53 +58,80 @@ func submitFlagHandler(c *gin.Context) { }) return } - if state == 1 { - username, err := coreUtils.GetUser(c.GetHeader("Authorization")) - if err != nil { - c.JSON(http.StatusUnauthorized, HTTPErrorResp{ - Error: "Unauthorized user", - }) - return - } + if state != 1 { + return + } - if challId == "" { - c.JSON(http.StatusBadRequest, HTTPErrorResp{ - Error: "Id of the challenge is a required parameter to process request.", - }) - return - } + username, err := coreUtils.GetUser(c.GetHeader("Authorization")) + if err != nil { + c.JSON(http.StatusUnauthorized, HTTPErrorResp{ + Error: "Unauthorized user", + }) + return + } - if flag == "" { - c.JSON(http.StatusBadRequest, HTTPErrorResp{ - Error: "Flag for the challenge is a required parameter to process request.", - }) - return - } + if challId == "" { + c.JSON(http.StatusBadRequest, HTTPErrorResp{ + Error: "Id of the challenge is a required parameter to process request.", + }) + return + } - user, err := database.QueryFirstUserEntry("username", username) - if err != nil { - c.JSON(http.StatusUnauthorized, HTTPErrorResp{ - Error: "Unauthorized user", - }) - return - } + if flag == "" { + c.JSON(http.StatusBadRequest, HTTPErrorResp{ + Error: "Flag for the challenge is a required parameter to process request.", + }) + return + } - if user.Status == 1 { - c.JSON(http.StatusUnauthorized, HTTPErrorResp{ - Error: "Banned user", - }) - return - } + user, err := database.QueryFirstUserEntry("username", username) + if err != nil || user.ID == 0 { + c.JSON(http.StatusUnauthorized, HTTPErrorResp{ + Error: "Unauthorized user", + }) + return + } - parsedChallId, err := strconv.Atoi(challId) - if err != nil { - c.JSON(http.StatusInternalServerError, HTTPErrorResp{ - Error: "DATABASE ERROR while processing the request.", - }) - return - } + if user.Status == 1 { + c.JSON(http.StatusUnauthorized, HTTPErrorResp{ + Error: "Banned user", + }) + return + } - chall, err := database.QueryChallengeEntries("id", strconv.Itoa(int(parsedChallId))) + parsedChallId, err := strconv.Atoi(challId) + if err != nil { + c.JSON(http.StatusBadRequest, HTTPErrorResp{ + Error: "Invalid challenge id.", + }) + return + } + + chall, err := database.QueryChallengeEntries("id", strconv.Itoa(parsedChallId)) + if err != nil { + c.JSON(http.StatusInternalServerError, HTTPErrorResp{ + Error: "DATABASE ERROR while processing the request.", + }) + return + } + if len(chall) == 0 { + c.JSON(http.StatusBadRequest, HTTPErrorResp{ + Error: "Challenge not found.", + }) + return + } + + challenge := chall[0] + if challenge.Status != core.DEPLOY_STATUS["deployed"] { + c.JSON(http.StatusOK, FlagSubmitResp{ + Message: "Challenge is unavailable", + Success: false, + }) + return + } + + if challenge.PreReqs != "" { + preReqsStatus, err := database.CheckPreReqsStatus(challenge, user.ID) if err != nil { c.JSON(http.StatusInternalServerError, HTTPErrorResp{ Error: "DATABASE ERROR while processing the request.", @@ -105,240 +139,125 @@ func submitFlagHandler(c *gin.Context) { return } - challenge := chall[0] - if challenge.Status != core.DEPLOY_STATUS["deployed"] { + if !preReqsStatus { c.JSON(http.StatusOK, FlagSubmitResp{ - Message: "Challenge is unavailable", + Message: "You have not solved the prerequisites of this challenge.", Success: false, }) return } + } - if challenge.PreReqs != "" { - preReqsStatus, err := database.CheckPreReqsStatus(challenge, user.ID) - - if err != nil { - c.JSON(http.StatusInternalServerError, HTTPErrorResp{ - Error: "DATABASE ERROR while processing the request.", - }) - return - } - - if !preReqsStatus { - c.JSON(http.StatusOK, FlagSubmitResp{ - Message: "You have not solved the prerequisites of this challenge.", - Success: false, - }) - return - } - } - - if challenge.MaxAttemptLimit > 0 { - previousTries, err := database.GetUserPreviousTries(user.ID, challenge.ID) - if err != nil { - c.JSON(http.StatusInternalServerError, HTTPErrorResp{ - Error: "DATABASE ERROR while processing the request."}) - return - } - - if previousTries >= challenge.MaxAttemptLimit { - c.JSON(http.StatusOK, FlagSubmitResp{ - Message: "You have reached the maximum number of tries for this challenge.", - Success: false, - }) - return - } - } - - // Increase user tries by 1 - err = database.UpdateUserChallengeTries(user.ID, challenge.ID) + attempt, err := database.ReserveSubmissionAttempt(user.ID, challenge.ID, challenge.MaxAttemptLimit, flag, now) + if err != nil { + c.JSON(http.StatusInternalServerError, HTTPErrorResp{ + Error: "DATABASE ERROR while processing the request.", + }) + return + } + switch attempt.Status { + case database.SubmissionAttemptAlreadySolved: + c.JSON(http.StatusOK, FlagSubmitResp{ + Message: "Challenge has already been solved.", + Success: false, + }) + return + case database.SubmissionAttemptMaxAttempts: + c.JSON(http.StatusOK, FlagSubmitResp{ + Message: "You have reached the maximum number of tries for this challenge.", + Success: false, + }) + return + } + isCheating := false + if challenge.DynamicFlag { + validFlags, err := database.QueryDynamicFlagEntries(map[string]interface{}{ + "Name": challenge.Name, + "Flag": flag, + }) if err != nil { c.JSON(http.StatusInternalServerError, HTTPErrorResp{ Error: "DATABASE ERROR while processing the request.", }) return } - solved, err := database.CheckPreviousSubmissions(user.ID, challenge.ID) - if err != nil { - c.JSON(http.StatusInternalServerError, HTTPErrorResp{ - Error: "DATABASE ERROR while processing the request.", - }) - return - } - - if solved { + if len(validFlags) == 0 { c.JSON(http.StatusOK, FlagSubmitResp{ - Message: "Challenge has already been solved.", + Message: "Your flag is incorrect", Success: false, }) return } - // If the challenge is dynamic, then the flag is not stored in the database - var isCheating bool - if challenge.DynamicFlag { - whereMap := map[string]interface{}{ - "Name": challenge.Name, - "Flag": flag, - } - validFlags, err := database.QueryDynamicFlagEntries(whereMap) - if err != nil { - c.JSON(http.StatusInternalServerError, HTTPErrorResp{ - Error: "DATABASE ERROR while processing the request.", - }) - return - } - - wheremap := map[string]interface{}{ - "challenge_id": challenge.ID, - "flag": flag, - } - submissions, err := database.QuerySubmissions(wheremap) - if err != nil { - c.JSON(http.StatusInternalServerError, HTTPErrorResp{ - Error: "DATABASE ERROR while processing the request.", - }) - return - } - - flagInValidFlags := len(validFlags) > 0 - flagInSubmissions := len(submissions) > 0 - - // Case 1: Flag not in validFlags (incorrect flag) - no cheating detection for wrong flags - if !flagInValidFlags { - UserChallengesEntry := database.UserChallenges{ - CreatedAt: time.Now(), - UserID: user.ID, - ChallengeID: challenge.ID, - Solved: false, - Flag: flag, - Cheating: false, - } - err = database.SaveFlagSubmission(&UserChallengesEntry) - if err != nil { - c.JSON(http.StatusInternalServerError, HTTPErrorResp{ - Error: "DATABASE ERROR while processing the request.", - }) - return - } - c.JSON(http.StatusOK, FlagSubmitResp{ - Message: "Your flag is incorrect", - Success: false, - }) - return - } - - // Case 2: Flag in validFlags and in submissions, cheating with valid flag - if flagInValidFlags && flagInSubmissions { - if user.ID != submissions[0].UserID { - subuser, _ := database.QueryUserById(submissions[0].UserID) - msg := "User " + user.Username + " has submitted the flag " + flag + " for challenge " + challenge.Name + " which has already been solved by user " + subuser.Username - go notify.SendNotification(notify.Warning, msg) - isCheating = true - // Continue to end of function with Solved: true, Cheating: true - } else { - c.JSON(http.StatusOK, FlagSubmitResp{ - Message: "You have already solved this challenge", - Success: false, - }) - return - } - } - - // Case 3: Flag in validFlags but not in submissions (solved without cheating) and saved at the end. - - } else { - if challenge.Flag != flag { - UserChallengesEntry := database.UserChallenges{ - CreatedAt: time.Now(), - UserID: user.ID, - ChallengeID: challenge.ID, - Solved: false, - Flag: flag, - } - err = database.SaveFlagSubmission(&UserChallengesEntry) - if err != nil { - c.JSON(http.StatusInternalServerError, HTTPErrorResp{ - Error: "DATABASE ERROR while processing the request.", - }) - return - } - c.JSON(http.StatusOK, FlagSubmitResp{ - Message: "Your flag is incorrect", - Success: false, - }) - return - } - } - challengePoints := challenge.Points - log.Debugf("Dynamic scoring is set to %t", config.Cfg.CompetitionInfo.DynamicScore) - if config.Cfg.CompetitionInfo.DynamicScore { - submissions, err := database.QuerySubmissions(map[string]interface{}{ - "challenge_id": parsedChallId, - }) - if err != nil { - log.Error(err) - } - solvers := len(submissions) - newPoints := dynamicScore(challenge.MaxPoints, challenge.MinPoints, uint(solvers)) - if newPoints != challengePoints { - database.UpdateChallenge(&challenge, map[string]interface{}{ - "Points": newPoints, - }) - log.Debugf("By dynamic scoring the points of challenge %s are changed to %d from %d", challenge.Name, newPoints, challengePoints) - err = updatePointsOfSolvers(submissions, newPoints, challengePoints) - if err != nil { - log.Error(err) - } - challengePoints = newPoints - } - } - oldScore := user.Score - newScore := user.Score + challengePoints - if newScore <= 0 { - newScore = 0 - } - err = database.UpdateUser(&user, map[string]interface{}{"Score": newScore}) + claim, err := database.ClaimDynamicFlag(challenge.ID, user.ID, flag, now) if err != nil { c.JSON(http.StatusInternalServerError, HTTPErrorResp{ Error: "DATABASE ERROR while processing the request.", }) return } - - if len(adminLeaderboardCache) < core.LEADERBOARD_SIZE || - (len(adminLeaderboardCache) > 0 && (newScore >= adminLeaderboardCache[len(adminLeaderboardCache)-1].Score || - oldScore >= adminLeaderboardCache[len(adminLeaderboardCache)-1].Score)) { - leaderboardStale = true - graphCacheStale = true - adminLeaderboardStale = true - } - - UserChallengesEntry := database.UserChallenges{ - CreatedAt: time.Now(), - UserID: user.ID, - ChallengeID: challenge.ID, - Solved: true, - Flag: flag, - Cheating: isCheating, - } - - err = database.SaveFlagSubmission(&UserChallengesEntry) - if err != nil { - c.JSON(http.StatusInternalServerError, HTTPErrorResp{ - Error: "DATABASE ERROR while processing the request.", + if claim.Status == database.DynamicFlagClaimedByOtherUser { + subuser, _ := database.QueryUserById(claim.ClaimedByID) + msg := "User " + user.Username + " has submitted the flag " + flag + " for challenge " + challenge.Name + " which has already been claimed by user " + subuser.Username + go notify.SendNotification(notify.Warning, msg) + if err := database.MarkSubmissionCheating(user.ID, challenge.ID, flag); err != nil { + log.Warnf("failed to mark duplicate dynamic flag submission as cheating: %v", err) + } + c.JSON(http.StatusOK, FlagSubmitResp{ + Message: "This dynamic flag has already been claimed.", + Success: false, }) return } + } else if challenge.Flag != flag { + c.JSON(http.StatusOK, FlagSubmitResp{ + Message: "Your flag is incorrect", + Success: false, + }) + return + } + wonSolveRace, err := database.MarkSubmissionSolved(user.ID, challenge.ID, flag, isCheating, now) + if err != nil { + c.JSON(http.StatusInternalServerError, HTTPErrorResp{ + Error: "DATABASE ERROR while processing the request.", + }) + return + } + if !wonSolveRace { c.JSON(http.StatusOK, FlagSubmitResp{ - Message: "Your flag is correct", - Success: true, + Message: "Challenge has already been solved.", + Success: false, }) + return + } + challengePoints := challenge.Points + if err := database.AwardUserScore(user.ID, int64(challengePoints)); err != nil { + c.JSON(http.StatusInternalServerError, HTTPErrorResp{ + Error: "DATABASE ERROR while processing the request.", + }) return } + + log.Debugf("Dynamic scoring is set to %t", config.Cfg.CompetitionInfo.DynamicScore) + if config.Cfg.CompetitionInfo.DynamicScore { + if err := database.MarkDynamicScoreDirty(challenge.ID, user.ID, now); err != nil { + log.Errorf("failed to mark dynamic score dirty for challenge %s: %v", challenge.Name, err) + } else { + notifyDynamicScoreWorker() + } + } + + leaderboardStale = true + graphCacheStale = true + adminLeaderboardStale = true + + c.JSON(http.StatusOK, FlagSubmitResp{ + Message: "Your flag is correct", + Success: true, + }) } // dynamicScore returns dynamic score of the challenge based on number of solves @@ -350,6 +269,84 @@ func dynamicScore(maxPoints, minPoints, solvers uint) uint { return uint(math.Round(float64(minPoints) + (float64(maxPoints)-float64(minPoints))/divisor)) } +func startDynamicScoreWorker() { + dynamicScoreWorkerOnce.Do(func() { + go func() { + ticker := time.NewTicker(30 * time.Second) + defer ticker.Stop() + + for { + select { + case <-dynamicScoreNotify: + processDirtyDynamicScores() + case <-ticker.C: + processDirtyDynamicScores() + } + } + }() + }) +} + +func notifyDynamicScoreWorker() { + select { + case dynamicScoreNotify <- struct{}{}: + default: + } +} + +func processDirtyDynamicScores() { + dirtyScores, err := database.QueryDirtyDynamicScores(100) + if err != nil { + log.Errorf("failed to query dirty dynamic scores: %v", err) + return + } + + for _, dirty := range dirtyScores { + if err := recomputeDynamicScore(dirty); err != nil { + log.Errorf("failed to recompute dynamic score for challenge %d: %v", dirty.ChallengeID, err) + continue + } + if err := database.ClearDynamicScoreDirty(dirty.ChallengeID, dirty.UpdatedAt); err != nil { + log.Errorf("failed to clear dynamic score dirty marker for challenge %d: %v", dirty.ChallengeID, err) + } + } +} + +func recomputeDynamicScore(dirty database.DynamicScoreDirty) error { + challs, err := database.QueryChallengeEntries("id", strconv.Itoa(int(dirty.ChallengeID))) + if err != nil { + return err + } + if len(challs) == 0 { + return nil + } + + challenge := challs[0] + if !challenge.DynamicFlag { + return nil + } + + solvers, err := database.CountSolvedSubmissionsForChallenge(challenge.ID) + if err != nil { + return err + } + + newPoints := dynamicScore(challenge.MaxPoints, challenge.MinPoints, solvers) + delta := int64(newPoints) - int64(challenge.Points) + if err := database.ApplyDynamicScoreDelta(challenge.ID, newPoints, delta); err != nil { + return err + } + + if delta != 0 { + log.Debugf("By dynamic scoring the points of challenge %s are changed to %d from %d", challenge.Name, newPoints, challenge.Points) + leaderboardStale = true + graphCacheStale = true + adminLeaderboardStale = true + } + + return nil +} + // updatePointsOfSolvers updates the points of solvers, whenever points of challenge changes func updatePointsOfSolvers(submissions []database.UserChallenges, newChallengePointsAfterSolve, oldChallengePointsBeforeSolve uint) error { scoreChanged := false diff --git a/cmd/beast/backup.go b/cmd/beast/backup.go index 14e088af..59b1e2f7 100644 --- a/cmd/beast/backup.go +++ b/cmd/beast/backup.go @@ -1,6 +1,7 @@ package main import ( + "github.com/sdslabs/beastv4/core/cache" "github.com/sdslabs/beastv4/core/database" "github.com/spf13/cobra" ) @@ -12,3 +13,11 @@ var backupDatabase = &cobra.Command{ database.BackupDatabase() }, } + +var backupCache = &cobra.Command{ + Use: "backup-cache", + Short: "Backups the existing cache and remote/staging directories", + Run: func(cmd *cobra.Command, args []string) { + cache.BackupCache() + }, +} diff --git a/cmd/beast/cache.go b/cmd/beast/cache.go new file mode 100644 index 00000000..c3b9086c --- /dev/null +++ b/cmd/beast/cache.go @@ -0,0 +1,30 @@ +package main + +import ( + "github.com/sdslabs/beastv4/core/cache" + log "github.com/sirupsen/logrus" + "github.com/spf13/cobra" +) + +var resetCacheCmd = &cobra.Command{ + Use: "reset-cache", + Short: "Backups the existing cache and cleans up old cache and remote/staging directories", + Run: func(cmd *cobra.Command, args []string) { + cache.BackupAndReset() + }, +} + +var restoreCacheCmd = &cobra.Command{ + Use: "restore-cache", + Short: "Restores the cache, with the backed-up file", + Run: func(cmd *cobra.Command, args []string) { + if RestoreFile != "" { + err := cache.RestoreCache(RestoreFile) + if err != nil { + log.Errorf("Error restoring cache from file %s: %v\n", RestoreFile, err) + } + } else { + log.Fatalf("Restore file not specified.") + } + }, +} diff --git a/cmd/beast/commands.go b/cmd/beast/commands.go index 05f1ecda..6ca5467d 100644 --- a/cmd/beast/commands.go +++ b/cmd/beast/commands.go @@ -113,6 +113,9 @@ func init() { restoreDatabaseCmd.PersistentFlags().StringVarP(&RestoreFile, "restore-file", "r", "", "Backup file to be used for restoration.") + + restoreCacheCmd.PersistentFlags().StringVarP(&RestoreFile, "restore-file", "r", "", "Restore file to be used for restoration.") + rootCmd.AddCommand(versionCmd) rootCmd.AddCommand(initCmd) rootCmd.AddCommand(configCmd) @@ -131,4 +134,7 @@ func init() { rootCmd.AddCommand(resetDatabaseCmd) rootCmd.AddCommand(restoreDatabaseCmd) rootCmd.AddCommand(backupDatabase) + rootCmd.AddCommand(resetCacheCmd) + rootCmd.AddCommand(restoreCacheCmd) + rootCmd.AddCommand(backupCache) } diff --git a/cmd/beast/config.go b/cmd/beast/config.go index 7868640d..2b82f59b 100644 --- a/cmd/beast/config.go +++ b/cmd/beast/config.go @@ -115,16 +115,20 @@ func promptServerDetails(configuration *config.BeastConfig) { var server config.AvailableServer server.Host = utils.PromptString("Enter Host Name, leave empty for localhost") + if server.Host == "" { + server.Host = core.LOCALHOST + } server.Username = utils.PromptString("Enter Username") server.SSHKeyPath = utils.PromptString("Enter SSH Key Path") server.Active = utils.PromptBinary("Enable this server?") - configuration.AvailableServers[server.Username] = server + configuration.AvailableServers[server.Host] = server } } func promptResourceLimits(configuration *config.BeastConfig) { configuration.CPUShares = utils.PromptInt64("Default CPU Share (must be over 6MB):", core.DEFAULT_CPU_SHARE) + configuration.CPUsLimit = utils.PromptFloat32("Default CPU Limit", core.DEFAULT_CPU_LIMIT) configuration.PidsLimit = utils.PromptInt64("Default PIDs Limit:", core.DEFAULT_PIDS_LIMIT) configuration.Memory = utils.PromptInt64("Default Memory Limit:", core.DEFAULT_MEMORY_LIMIT) @@ -192,6 +196,28 @@ func promptNotificationWebhooks(configuration *config.BeastConfig) { } } +func promptCacheConnectionDetails(configuration *config.BeastConfig) { + configuration.RedisConf.User = utils.PromptString("Enter Redis User Name (this user will be created if does not exist)... leaving it empty will default it to beast") + if configuration.RedisConf.User == "" { + configuration.RedisConf.User = "beast" + } + + configuration.RedisConf.Password = utils.PromptSecret(fmt.Sprintf("Enter Redis User %s Password... leaving it empty will default it to beast", configuration.RedisConf.User)) + if configuration.RedisConf.Password == "" { + configuration.RedisConf.Password = "beast" + } + + configuration.RedisConf.Host = utils.PromptString("Enter Redis Host Name, leave empty for localhost") + if configuration.RedisConf.Host == "" { + configuration.RedisConf.Host = core.LOCALHOST + } + + configuration.RedisConf.Port = strconv.FormatInt(utils.PromptInt64("Enter Redis Port", 6379), 10) + + log.Infoln("Setting Redis DB to 0...") + configuration.RedisConf.Db = 0 +} + func promptDatabaseConnectionDetails(configuration *config.BeastConfig) { configuration.PsqlConf.User = utils.PromptString("Enter Postgres User Name (this user will be created if does not exist)... leaving it empty will default it to beast") if configuration.PsqlConf.User == "" { @@ -210,7 +236,7 @@ func promptDatabaseConnectionDetails(configuration *config.BeastConfig) { configuration.PsqlConf.Host = utils.PromptString("Enter Postgres Host Name, leave empty for localhost") if configuration.PsqlConf.Host == "" { - configuration.PsqlConf.Host = "localhost" + configuration.PsqlConf.Host = core.LOCALHOST } configuration.PsqlConf.Port = strconv.FormatInt(utils.PromptInt64("Enter Postgres Port", 5432), 10) configuration.PsqlConf.SslMode = utils.PromptSelection("Enter Postgres SSL Mode", []string{ @@ -227,6 +253,7 @@ func promptBeastConfiguration(configuration *config.BeastConfig) { promptRemoteRepository(configuration) promptCompetitionDetails(configuration) promptNotificationWebhooks(configuration) + promptCacheConnectionDetails(configuration) promptDatabaseConnectionDetails(configuration) } diff --git a/cmd/beast/init.go b/cmd/beast/init.go index 90a0dd6f..9a62b0a3 100644 --- a/cmd/beast/init.go +++ b/cmd/beast/init.go @@ -1,12 +1,21 @@ package main import ( + "context" "database/sql" "errors" "fmt" - "github.com/BurntSushi/toml" + "io" + "net/http" + "os" + "os/exec" + "os/user" + "path/filepath" + "strings" + _ "github.com/jackc/pgx/v5/stdlib" "github.com/lib/pq" + "github.com/redis/go-redis/v9" "github.com/sdslabs/beastv4/core" "github.com/sdslabs/beastv4/core/config" "github.com/sdslabs/beastv4/core/database" @@ -14,13 +23,6 @@ import ( "github.com/sdslabs/beastv4/utils" log "github.com/sirupsen/logrus" "github.com/spf13/cobra" - "io" - "net/http" - "os" - "os/exec" - "os/user" - "path/filepath" - "strings" ) const ( @@ -95,6 +97,63 @@ func installAir() error { return cmd.Run() } +func createBeastRedisUser(cache *redis.Client, configuration *config.RedisConfig) error { + ctx := context.Background() + log.Warnln("Beast expects Redis ACLs to be enabled. If ACLs are not configured, some features may not function correctly.") + + result, err := cache.ACLUsers(ctx).Result() + if err != nil { + return err + } + + for _, user := range result { + if user == configuration.User { + log.Infoln(fmt.Sprintf("Redis user %s already exists", configuration.User)) + break + } + } + + _, err = cache.ACLSetUser(ctx, configuration.User, "on", ">"+configuration.Password, "~beast:*", "+@all").Result() + if err != nil { + return err + } + log.Infoln(fmt.Sprintf("Initialised redis user %s", configuration.User)) + + err = cache.Do(ctx, "acl", "save").Err() + if err != nil { + return fmt.Errorf("error while trying to save the acl file: %s", err.Error()) + } + + return nil +} + +func initCache() error { + log.Infoln("Initializing cache...") + + redisConfig := config.Cfg.RedisConf + var cache *redis.Client + if utils.PromptBinary("Do you use password authentication for the redis default user?") { + cache = redis.NewClient(&redis.Options{ + Addr: fmt.Sprintf("%s:%s", redisConfig.Host, redisConfig.Port), + Username: core.REDIS_DEFAULT_USER, + Password: utils.PromptSecret("Enter default redis user password"), + }) + } else { + cache = redis.NewClient(&redis.Options{ + Addr: fmt.Sprintf("%s:%s", redisConfig.Host, redisConfig.Port), + Username: core.REDIS_DEFAULT_USER, + }) + } + + _, err := cache.Ping(context.Background()).Result() + if err != nil { + return fmt.Errorf("failed to connected to redis: %s", err.Error()) + } + + defer cache.Close() + return createBeastRedisUser(cache, &config.Cfg.RedisConf) +} + func createBeastDbUser(db *sql.DB, configuration *config.PsqlConfig) error { if result := utils.PromptBinary("Create default beast postgres user?"); !result { return errors.New("failed to create database") @@ -125,12 +184,6 @@ func dbUserCheck() (bool, error) { func initDb() error { log.Infoln("Initializing database...") - var configuration config.BeastConfig - _, err := toml.DecodeFile(BEAST_GLOBAL_CONFIG, &configuration) - if err != nil { - return err - } - isPostgres, err := dbUserCheck() if err != nil { return err @@ -166,42 +219,44 @@ func initDb() error { defer db.Close() + configuration := config.Cfg.PsqlConf + var exists int - err = db.QueryRow("SELECT 1 FROM pg_roles WHERE rolname = $1", configuration.PsqlConf.User).Scan(&exists) + err = db.QueryRow("SELECT 1 FROM pg_roles WHERE rolname = $1", configuration.User).Scan(&exists) if errors.Is(err, sql.ErrNoRows) { - if err = createBeastDbUser(db, &configuration.PsqlConf); err != nil { + if err = createBeastDbUser(db, &configuration); err != nil { return err } } else if err != nil { return err } else { - log.Infoln(fmt.Sprintf("User %s already exists", configuration.PsqlConf.User)) + log.Infoln(fmt.Sprintf("User %s already exists", configuration.User)) } - log.Infoln(fmt.Sprintf("Changing password for user %s", configuration.PsqlConf.User)) - query := fmt.Sprintf("ALTER USER %s WITH PASSWORD %s", pq.QuoteIdentifier(configuration.PsqlConf.User), utils.QuoteLiteral(configuration.PsqlConf.Password)) + log.Infoln(fmt.Sprintf("Changing password for user %s", configuration.User)) + query := fmt.Sprintf("ALTER USER %s WITH PASSWORD %s", pq.QuoteIdentifier(configuration.User), utils.QuoteLiteral(configuration.Password)) _, err = db.Exec(query) if err != nil { return err } - err = db.QueryRow("SELECT 1 FROM pg_database WHERE datname = $1", configuration.PsqlConf.Dbname).Scan(&exists) + err = db.QueryRow("SELECT 1 FROM pg_database WHERE datname = $1", configuration.Dbname).Scan(&exists) if errors.Is(err, sql.ErrNoRows) { - if err = createBeastDatabase(db, &configuration.PsqlConf); err != nil { + if err = createBeastDatabase(db, &configuration); err != nil { return err } } else if err != nil { return err } else { - log.Infoln(fmt.Sprintf("Database %s already exists", configuration.PsqlConf.Dbname)) + log.Infoln(fmt.Sprintf("Database %s already exists", configuration.Dbname)) } - _, err = db.Exec(fmt.Sprintf("ALTER DATABASE %s OWNER TO %s", pq.QuoteIdentifier(configuration.PsqlConf.Dbname), pq.QuoteIdentifier(configuration.PsqlConf.User))) + _, err = db.Exec(fmt.Sprintf("ALTER DATABASE %s OWNER TO %s", pq.QuoteIdentifier(configuration.Dbname), pq.QuoteIdentifier(configuration.User))) if err != nil { return err } - log.Infoln(fmt.Sprintf("%s set as owner of database %s", configuration.PsqlConf.User, configuration.PsqlConf.Dbname)) + log.Infoln(fmt.Sprintf("%s set as owner of database %s", configuration.User, configuration.Dbname)) return nil } @@ -273,6 +328,14 @@ func runBeastBootsteps() error { log.Infoln("Successfully installed air for live reloading...") + config.InitConfig() + + if err := initCache(); err != nil { + return err + } + + log.Infoln("Verified redis setup for beast") + if err := initDb(); err != nil { return err } diff --git a/cmd/beast/run.go b/cmd/beast/run.go index d389ebc1..a582c0d5 100644 --- a/cmd/beast/run.go +++ b/cmd/beast/run.go @@ -3,6 +3,7 @@ package main import ( "encoding/json" "fmt" + "github.com/sdslabs/beastv4/core/cache" "math" "os" "os/signal" @@ -55,6 +56,10 @@ func cleanupRunningContainers() { for _, challenge := range challenges { if challenge.Status == core.DEPLOY_STATUS["deployed"] { + if challenge.Instanced { + _ = manager.KillChallengeInstances(challenge.Name) + } + err = manager.UndeployChallenge(challenge.Name) if err != nil { log.Errorln(fmt.Sprintf("Failed to undeploy challenge [Id: %v] %s", challenge.ID, challenge.Name)) @@ -66,6 +71,26 @@ func cleanupRunningContainers() { } } +func cleanupCacheConnections() { + log.Infoln("Cleaning up cache connections...") + + err := cache.BackupCache() + if err != nil { + log.Errorln("Error while backing up cache:", err) + } else { + log.Infoln("Cache backup completed successfully") + } + + log.Infoln("Terminating cache connection...") + + err = cache.Close() + if err != nil { + log.Errorln("Unable to terminate cache connections:", err) + } else { + log.Infoln("Cache connections terminated successfully") + } +} + func cleanupDatabaseConnections() { log.Infoln("Backing up database...") @@ -131,12 +156,14 @@ func cleanup() { stopSseNotificationHub() stopApiScheduler() + cleanupRunningContainers() + stopWorkerQueue() stopRemoteManagers() saveLeaderboardCache() - cleanupRunningContainers() + cleanupCacheConnections() cleanupDatabaseConnections() // - Clean up temporary files: found no files to be cleared as of now diff --git a/core/cache/cache.go b/core/cache/cache.go new file mode 100644 index 00000000..68bfe7ba --- /dev/null +++ b/core/cache/cache.go @@ -0,0 +1,317 @@ +package cache + +import ( + "context" + "fmt" + "os" + "os/exec" + "path/filepath" + "strconv" + "strings" + "sync" + "time" + + "github.com/BurntSushi/toml" + "github.com/redis/go-redis/v9" + "github.com/sdslabs/beastv4/core" + "github.com/sdslabs/beastv4/utils" + log "github.com/sirupsen/logrus" +) + +var ( + CacheMutex *sync.Mutex + Cache *redis.Client + cacheError error +) + +var ( + BEAST_GLOBAL_DIR string = filepath.Join(os.Getenv("HOME"), ".beast") + cacheConfig Config +) + +type Config struct { + RedisConfig RedisConfig `toml:"redis_config"` +} +type RedisConfig struct { + User string `toml:"user"` + Password string `toml:"password"` + Host string `toml:"host"` + Port string `toml:"port"` + DB int `toml:"db"` +} + +// Db config is loaded separately here for temp use because init() function is +// called during initialization of package. +// It is also loaded during db backup/reset +func LoadCacheConfig() { + if _, err := toml.DecodeFile(filepath.Join(core.BEAST_GLOBAL_DIR, core.BEAST_CONFIG_FILE_NAME), &cacheConfig); err != nil { + log.Fatalf("Error loading TOML file: %v", err) + } +} + +// Connect redis +func ConnectCache() error { + LoadCacheConfig() + Cache = redis.NewClient(&redis.Options{ + Addr: fmt.Sprintf("%s:%s", cacheConfig.RedisConfig.Host, cacheConfig.RedisConfig.Port), + Username: cacheConfig.RedisConfig.User, + Password: cacheConfig.RedisConfig.Password, + DB: cacheConfig.RedisConfig.DB, + }) + + _, err := Cache.Ping(context.Background()).Result() + if err != nil { + return fmt.Errorf("failed to connected to redis: %s", err.Error()) + } + + log.Debug("Cache initialized") + return nil +} + +// Set up the initial bootstrapping for interacting with the +// Postgresql database for beast. The Db variable is the connection variable for the +// database, which is not closed after creating a connection here and can +// be used further after this. +func Init() { + CacheMutex = &sync.Mutex{} + if Cache == nil { + cacheError = ConnectCache() + if cacheError != nil { + log.Errorf("Error while initializing cache: %s", cacheError.Error()) + } + } +} + +func EnableKeyspaceExpiryNotifications() error { + if Cache == nil { + return fmt.Errorf("redis cache not initialized") + } + + ctx := context.Background() + result, err := Cache.Do(ctx, "CONFIG", "GET", "notify-keyspace-events").Result() + if err != nil { + return err + } + + current := "" + if values, ok := result.([]interface{}); ok && len(values) >= 2 { + current = fmt.Sprint(values[1]) + } + + next := current + if !strings.Contains(next, "E") { + next += "E" + } + if !strings.Contains(next, "x") { + next += "x" + } + + if next == current { + return nil + } + + return Cache.Do(ctx, "CONFIG", "SET", "notify-keyspace-events", next).Err() +} + +func SubscribeExpiredInstanceMarkers(ctx context.Context, handler func(instanceID string)) error { + if Cache == nil { + return fmt.Errorf("redis cache not initialized") + } + + pattern := fmt.Sprintf("__keyevent@%d__:expired", cacheConfig.RedisConfig.DB) + pubsub := Cache.PSubscribe(ctx, pattern) + defer pubsub.Close() + + if _, err := pubsub.Receive(ctx); err != nil { + return err + } + + ch := pubsub.Channel() + for { + select { + case <-ctx.Done(): + return ctx.Err() + case msg, ok := <-ch: + if !ok { + return nil + } + + instanceID, ok := utils.InstanceIDFromExpiryKey(msg.Payload) + if !ok { + continue + } + handler(instanceID) + } + } +} + +func Close() error { + if Cache == nil { + log.Warnln(fmt.Sprintf("Trying to close database connection when no connection is established...")) + return nil + } + + err := Cache.Close() + if err != nil { + log.Errorln(fmt.Sprintf("Error while closing cache connection gracefully: %s, attempting to terminate forcefully", err.Error())) + return TerminateCacheConnections() + } + + return nil +} + +func BackupAndReset() { + LoadCacheConfig() + + err := BackupCache() + if err != nil { + log.Errorf("Error while backing up cache: %s", err) + return + } + err = ResetCache() + if err != nil { + log.Errorf("Error while resetting up cache: %s", err) + return + } + + backupPath := filepath.Join(core.BEAST_GLOBAL_DIR, core.BEAST_BACKUP_DIR, core.BEAST_REMOTES_DIR) + err = utils.CreateIfNotExistDir(backupPath) + if err != nil { + log.Errorf("Error while creating backup directory: %s", err) + return + } + + backupPath = filepath.Join(backupPath, core.BEAST_REMOTES_DIR+time.Now().Format("20060102150405")+".bak") + oldPath := filepath.Join(core.BEAST_GLOBAL_DIR, core.BEAST_REMOTES_DIR) + err = os.Rename(oldPath, backupPath) + if err != nil { + log.Errorf("Error while backing up remote dir: %s", err) + return + } + + backupPath = filepath.Join(core.BEAST_GLOBAL_DIR, core.BEAST_BACKUP_DIR, core.BEAST_STAGING_DIR) + + err = utils.CreateIfNotExistDir(backupPath) + if err != nil { + log.Errorf("Error while creating backup directory: %s", err) + return + } + + oldPath = filepath.Join(core.BEAST_GLOBAL_DIR, core.BEAST_STAGING_DIR) + backupPath = filepath.Join(backupPath, core.BEAST_STAGING_DIR+time.Now().Format("20060102150405")+".bak") + err = os.Rename(oldPath, backupPath) + if err != nil { + log.Errorf("Error while backing up staging dir: %s", err) + return + } +} + +func BackupCache() error { + if cacheConfig == (Config{}) { + LoadCacheConfig() + } + + backupPath := filepath.Join(core.BEAST_GLOBAL_DIR, core.BEAST_BACKUP_DIR, core.BEAST_CACHE_DIR) + err := utils.CreateIfNotExistDir(backupPath) + if err != nil { + log.Errorf("Error while creating backup directory: %s", err) + return err + } + + backupFile := fmt.Sprintf("%d_%s.bak", cacheConfig.RedisConfig.DB, time.Now().Format("20060102150405")) + + args := []string{ + "-h", cacheConfig.RedisConfig.Host, + "-p", cacheConfig.RedisConfig.Port, + "-n", strconv.Itoa(cacheConfig.RedisConfig.DB), + "--rdb", filepath.Join(backupPath, backupFile), + } + if cacheConfig.RedisConfig.User != "" { + args = append(args, "--user", cacheConfig.RedisConfig.User) + } + if cacheConfig.RedisConfig.Password != "" { + args = append(args, "--pass", cacheConfig.RedisConfig.Password) + } + + cmd := exec.Command("redis-cli", args...) + + cmd.Env = append(os.Environ(), fmt.Sprintf("REDISCLI_AUTH=%s", cacheConfig.RedisConfig.Password)) + output, err := cmd.CombinedOutput() + if err != nil { + log.Printf("Backup error: %s\n", string(output)) + return err + } + log.Debug("Backup successful.") + return nil +} + +func ResetCache() error { + if cacheConfig == (Config{}) { + LoadCacheConfig() + } + err := TerminateCacheConnections() + if err != nil { + log.Errorf("Unable to terminate connections %s", err) + return err + } + + dropCmd := exec.Command( + "redis-cli", + "-h", cacheConfig.RedisConfig.Host, + "-p", cacheConfig.RedisConfig.Port, + "--user", cacheConfig.RedisConfig.User, + "-n", strconv.Itoa(cacheConfig.RedisConfig.DB), + "FLUSHDB", + ) + + dropCmd.Env = append(os.Environ(), fmt.Sprintf("REDISCLI_AUTH=%s", cacheConfig.RedisConfig.Password)) + + output, err := dropCmd.CombinedOutput() + if err != nil { + log.Printf("Drop Cache error: %s\n", string(output)) + return err + } + + log.Debug("Reset successful.") + return nil +} + +// Terminate all active connections before dropping +func TerminateCacheConnections() error { + if cacheConfig == (Config{}) { + LoadCacheConfig() + } + + cache := redis.NewClient(&redis.Options{ + Addr: fmt.Sprintf("%s:%s", cacheConfig.RedisConfig.Host, cacheConfig.RedisConfig.Port), + Username: core.REDIS_DEFAULT_USER, + Password: utils.PromptSecret("Enter default redis user password"), + }) + + _, err := cache.Ping(context.Background()).Result() + if err != nil { + log.Errorf("Terminate connections error: %s\n", err.Error()) + } + + defer cache.Close() + + _, err = cache.Do(context.Background(), + "CLIENT", "KILL", + "USER", cacheConfig.RedisConfig.User, + "SKIPME", "yes", + ).Result() + + if err != nil { + log.Errorf("Terminate connections error: %s\n", err.Error()) + } + + return nil +} + +func RestoreCache(backupFile string) error { + /* + The primary issue with restoring cache is that it needs to be written to /var/lib and redis needs to be restarted. + Redis will then pick up the changes and continue from there. + */ + return nil +} diff --git a/core/cache/instance.go b/core/cache/instance.go new file mode 100644 index 00000000..9f39c35c --- /dev/null +++ b/core/cache/instance.go @@ -0,0 +1,541 @@ +package cache + +import ( + "context" + "encoding/json" + "fmt" + "github.com/sdslabs/beastv4/utils" + "time" + + log "github.com/sirupsen/logrus" +) + +type Instance struct { + InstanceID string `json:"instance_id"` + ChallengeName string `json:"challenge_name"` + ContainerID string `json:"container_id"` + PortOwner string `json:"port_owner"` + Port uint32 `json:"port"` + UserID string `json:"user_id"` + Username string `json:"username"` + CreatedAt time.Time `json:"created_at"` + ExpiresAt time.Time `json:"expires_at"` + DeploymentType string `json:"deployment_type"` + ServerDeployed string `json:"server_deployed"` +} + +func (instance *Instance) PortOwnerID() string { + if instance.PortOwner != "" { + return instance.PortOwner + } + + return instance.ContainerID +} + +func SaveInstance(instance *Instance, ttl time.Duration) error { + if Cache == nil { + return fmt.Errorf("redis cache not initialized") + } + + ctx := context.Background() + CacheMutex.Lock() + defer CacheMutex.Unlock() + + data, err := json.Marshal(instance) + if err != nil { + return fmt.Errorf("failed to marshal instance: %w", err) + } + + key := utils.InstanceToKey(instance.InstanceID) + err = Cache.Set(ctx, key, data, 0).Err() + if err != nil { + return fmt.Errorf("failed to save instance: %w", err) + } + + expiryKey := utils.InstanceExpiryToKey(instance.InstanceID) + err = Cache.Set(ctx, expiryKey, instance.InstanceID, ttl).Err() + if err != nil { + return fmt.Errorf("failed to save instance expiry marker: %w", err) + } + + userKey := utils.UserChallengeToKey(instance.UserID, instance.ChallengeName) + err = Cache.Set(ctx, userKey, instance.InstanceID, ttl).Err() + if err != nil { + return fmt.Errorf("failed to save user instance mapping: %w", err) + } + + err = Cache.SAdd(ctx, utils.InstancesSetKey, instance.InstanceID).Err() + if err != nil { + log.Warnf("failed to add instance to set: %v", err) + } + + log.Debugf("Saved instance %s for user %s, challenge %s, port %d, expires in %v", + instance.InstanceID, instance.UserID, instance.ChallengeName, instance.Port, ttl) + + return nil +} + +func GetInstance(instanceID string) (*Instance, error) { + if Cache == nil { + return nil, fmt.Errorf("redis cache not initialized") + } + + ctx := context.Background() + CacheMutex.Lock() + defer CacheMutex.Unlock() + + key := utils.InstanceToKey(instanceID) + data, err := Cache.Get(ctx, key).Bytes() + if err != nil { + return nil, fmt.Errorf("instance not found: %w", err) + } + + var instance Instance + err = json.Unmarshal(data, &instance) + if err != nil { + return nil, fmt.Errorf("failed to unmarshal instance: %w", err) + } + + return &instance, nil +} + +func GetUserInstance(userID, challengeName string) (*Instance, error) { + if Cache == nil { + return nil, fmt.Errorf("redis cache not initialized") + } + + ctx := context.Background() + CacheMutex.Lock() + defer CacheMutex.Unlock() + + userKey := utils.UserChallengeToKey(userID, challengeName) + instanceID, err := Cache.Get(ctx, userKey).Result() + if err != nil { + return nil, fmt.Errorf("user instance not found: %w", err) + } + + key := utils.InstanceToKey(instanceID) + data, err := Cache.Get(ctx, key).Bytes() + if err != nil { + return nil, fmt.Errorf("instance not found: %w", err) + } + + var instance Instance + err = json.Unmarshal(data, &instance) + if err != nil { + return nil, fmt.Errorf("failed to unmarshal instance: %w", err) + } + + return &instance, nil +} + +func GetUserInstances(userID string) ([]*Instance, error) { + if Cache == nil { + return nil, fmt.Errorf("redis cache not initialized") + } + + ctx := context.Background() + CacheMutex.Lock() + defer CacheMutex.Unlock() + + pattern := utils.UserChallengesAllKey(userID) + var instances []*Instance + + iter := Cache.Scan(ctx, 0, pattern, 0).Iterator() + for iter.Next(ctx) { + userKey := iter.Val() + instanceID, err := Cache.Get(ctx, userKey).Result() + if err != nil { + continue + } + + key := utils.InstanceToKey(instanceID) + data, err := Cache.Get(ctx, key).Bytes() + if err != nil { + continue + } + + var instance Instance + err = json.Unmarshal(data, &instance) + if err != nil { + continue + } + + instances = append(instances, &instance) + } + + return instances, nil +} + +func GetAllInstances() ([]*Instance, error) { + if Cache == nil { + return nil, fmt.Errorf("redis cache not initialized") + } + + ctx := context.Background() + CacheMutex.Lock() + defer CacheMutex.Unlock() + + instanceIDs, err := Cache.SMembers(ctx, utils.InstancesSetKey).Result() + if err != nil { + return nil, fmt.Errorf("failed to get instance IDs: %w", err) + } + + var instances []*Instance + for _, id := range instanceIDs { + key := utils.InstanceToKey(id) + data, err := Cache.Get(ctx, key).Bytes() + if err != nil { + Cache.SRem(ctx, utils.InstancesSetKey, id) + continue + } + + var instance Instance + err = json.Unmarshal(data, &instance) + if err != nil { + continue + } + + instances = append(instances, &instance) + } + + return instances, nil +} + +func GetChallengeInstances(challengeName string) ([]*Instance, error) { + if Cache == nil { + return nil, fmt.Errorf("redis cache not initialized") + } + + ctx := context.Background() + CacheMutex.Lock() + defer CacheMutex.Unlock() + + instanceIDs, err := Cache.SMembers(ctx, utils.InstancesSetKey).Result() + if err != nil { + return nil, fmt.Errorf("failed to get instance IDs: %w", err) + } + + var instances []*Instance + for _, id := range instanceIDs { + key := utils.InstanceToKey(id) + data, err := Cache.Get(ctx, key).Bytes() + if err != nil { + Cache.SRem(ctx, utils.InstancesSetKey, id) + continue + } + + var instance Instance + err = json.Unmarshal(data, &instance) + if err != nil { + continue + } + + if instance.ChallengeName == challengeName { + instances = append(instances, &instance) + } + } + + return instances, nil +} + +func DeleteInstance(instanceID string) error { + if Cache == nil { + return fmt.Errorf("redis cache not initialized") + } + + ctx := context.Background() + CacheMutex.Lock() + defer CacheMutex.Unlock() + + key := utils.InstanceToKey(instanceID) + data, err := Cache.Get(ctx, key).Bytes() + if err != nil { + return fmt.Errorf("instance not found: %w", err) + } + + var instance Instance + err = json.Unmarshal(data, &instance) + if err != nil { + return fmt.Errorf("failed to unmarshal instance: %w", err) + } + + err = Cache.Del(ctx, key).Err() + if err != nil { + return fmt.Errorf("failed to delete instance: %w", err) + } + + userKey := utils.UserChallengeToKey(instance.UserID, instance.ChallengeName) + expiryKey := utils.InstanceExpiryToKey(instanceID) + Cache.Del(ctx, userKey) + Cache.Del(ctx, expiryKey) + Cache.SRem(ctx, utils.InstancesSetKey, instanceID) + + log.Debugf("Deleted instance %s for user %s, challenge %s", + instanceID, instance.UserID, instance.ChallengeName) + + return nil +} + +func ExtendInstance(instanceID string, additionalTime time.Duration) error { + if Cache == nil { + return fmt.Errorf("redis cache not initialized") + } + + ctx := context.Background() + CacheMutex.Lock() + defer CacheMutex.Unlock() + + key := utils.InstanceToKey(instanceID) + + data, err := Cache.Get(ctx, key).Bytes() + if err != nil { + return fmt.Errorf("instance not found: %w", err) + } + + var instance Instance + err = json.Unmarshal(data, &instance) + if err != nil { + return fmt.Errorf("failed to unmarshal instance: %w", err) + } + + newExpiresAt := instance.ExpiresAt.Add(additionalTime) + instance.ExpiresAt = newExpiresAt + + newTTL := time.Until(newExpiresAt) + if newTTL <= 0 { + return fmt.Errorf("instance has already expired") + } + + updatedData, err := json.Marshal(instance) + if err != nil { + return fmt.Errorf("failed to marshal instance: %w", err) + } + + err = Cache.Set(ctx, key, updatedData, 0).Err() + if err != nil { + return fmt.Errorf("failed to extend instance: %w", err) + } + + userKey := utils.UserChallengeToKey(instance.UserID, instance.ChallengeName) + expiryKey := utils.InstanceExpiryToKey(instanceID) + Cache.Set(ctx, expiryKey, instanceID, newTTL) + Cache.Expire(ctx, userKey, newTTL) + + log.Debugf("Extended instance %s by %v, new expiration: %v", instanceID, additionalTime, newExpiresAt) + + return nil +} + +func CountUserInstances(userID string) (int, error) { + if Cache == nil { + return 0, fmt.Errorf("redis cache not initialized") + } + + ctx := context.Background() + CacheMutex.Lock() + defer CacheMutex.Unlock() + + pattern := utils.UserChallengesAllKey(userID) + count := 0 + + iter := Cache.Scan(ctx, 0, pattern, 0).Iterator() + for iter.Next(ctx) { + count++ + } + + return count, nil +} + +func GetInstanceTTL(instanceID string) (time.Duration, error) { + if Cache == nil { + return 0, fmt.Errorf("redis cache not initialized") + } + + ctx := context.Background() + CacheMutex.Lock() + defer CacheMutex.Unlock() + + expiryKey := utils.InstanceExpiryToKey(instanceID) + ttl, err := Cache.TTL(ctx, expiryKey).Result() + if err != nil { + return 0, fmt.Errorf("failed to get TTL: %w", err) + } + + return ttl, nil +} + +func QueueInstanceForDeletion(instanceID string) error { + if Cache == nil { + return fmt.Errorf("redis cache not initialized") + } + + ctx := context.Background() + CacheMutex.Lock() + defer CacheMutex.Unlock() + + key := utils.InstanceToKey(instanceID) + + data, err := Cache.Get(ctx, key).Bytes() + if err != nil { + Cache.SRem(ctx, utils.InstancesSetKey, instanceID) + return fmt.Errorf("instance not found: %w", err) + } + + var instance Instance + err = json.Unmarshal(data, &instance) + if err != nil { + return fmt.Errorf("failed to unmarshal instance: %w", err) + } + + pipe := Cache.TxPipeline() + pipe.LPush(ctx, utils.InstanceDeletionQueue, data) + pipe.SRem(ctx, utils.InstancesSetKey, instanceID) + + userKey := utils.UserChallengeToKey(instance.UserID, instance.ChallengeName) + expiryKey := utils.InstanceExpiryToKey(instanceID) + pipe.Del(ctx, userKey) + pipe.Del(ctx, expiryKey) + + _, err = pipe.Exec(ctx) + if err != nil { + return fmt.Errorf("failed to queue instance for deletion: %w", err) + } + + log.Debugf("Queued instance %s for deletion (user: %s, challenge: %s)", + instanceID, instance.UserID, instance.ChallengeName) + + return nil +} + +func PopInstanceForDeletion() (*Instance, error) { + if Cache == nil { + return nil, fmt.Errorf("redis cache not initialized") + } + + ctx := context.Background() + CacheMutex.Lock() + defer CacheMutex.Unlock() + + data, err := Cache.RPop(ctx, utils.InstanceDeletionQueue).Bytes() + if err != nil { + if err.Error() == "redis: nil" { + return nil, nil + } + return nil, fmt.Errorf("failed to pop from deletion queue: %w", err) + } + + var instance Instance + err = json.Unmarshal(data, &instance) + if err != nil { + return nil, fmt.Errorf("failed to unmarshal instance from queue: %w", err) + } + + log.Debugf("Popped instance %s from deletion queue", instance.InstanceID) + return &instance, nil +} + +func DeleteInstanceMetadata(instanceID string) error { + if Cache == nil { + return fmt.Errorf("redis cache not initialized") + } + + ctx := context.Background() + CacheMutex.Lock() + defer CacheMutex.Unlock() + + key := utils.InstanceToKey(instanceID) + expiryKey := utils.InstanceExpiryToKey(instanceID) + + pipe := Cache.TxPipeline() + pipe.Del(ctx, key) + pipe.Del(ctx, expiryKey) + pipe.SRem(ctx, utils.InstancesSetKey, instanceID) + + _, err := pipe.Exec(ctx) + if err != nil { + return fmt.Errorf("failed to delete instance metadata: %w", err) + } + + return nil +} + +func RestoreQueuedInstance(instance *Instance) error { + if Cache == nil { + return fmt.Errorf("redis cache not initialized") + } + if instance == nil { + return fmt.Errorf("instance is nil") + } + + ctx := context.Background() + CacheMutex.Lock() + defer CacheMutex.Unlock() + + data, err := json.Marshal(instance) + if err != nil { + return fmt.Errorf("failed to marshal instance: %w", err) + } + + pipe := Cache.TxPipeline() + pipe.Set(ctx, utils.InstanceToKey(instance.InstanceID), data, 0) + pipe.SAdd(ctx, utils.InstancesSetKey, instance.InstanceID) + + _, err = pipe.Exec(ctx) + if err != nil { + return fmt.Errorf("failed to restore queued instance metadata: %w", err) + } + + return nil +} + +func GetDeletionQueueLength() (int64, error) { + if Cache == nil { + return 0, fmt.Errorf("redis cache not initialized") + } + + ctx := context.Background() + CacheMutex.Lock() + defer CacheMutex.Unlock() + + return Cache.LLen(ctx, utils.InstanceDeletionQueue).Result() +} + +func GetExpiredInstances() ([]*Instance, error) { + if Cache == nil { + return nil, fmt.Errorf("redis cache not initialized") + } + + ctx := context.Background() + CacheMutex.Lock() + defer CacheMutex.Unlock() + + instanceIDs, err := Cache.SMembers(ctx, utils.InstancesSetKey).Result() + if err != nil { + return nil, fmt.Errorf("failed to get instance IDs: %w", err) + } + + now := time.Now() + var expired []*Instance + + for _, id := range instanceIDs { + key := utils.InstanceToKey(id) + data, err := Cache.Get(ctx, key).Bytes() + if err != nil { + Cache.SRem(ctx, utils.InstancesSetKey, id) + continue + } + + var instance Instance + err = json.Unmarshal(data, &instance) + if err != nil { + continue + } + + if instance.ExpiresAt.Before(now) { + expired = append(expired, &instance) + } + } + + return expired, nil +} diff --git a/core/cache/ports.go b/core/cache/ports.go new file mode 100644 index 00000000..ab9e4b65 --- /dev/null +++ b/core/cache/ports.go @@ -0,0 +1,230 @@ +package cache + +import ( + "context" + "fmt" + "github.com/redis/go-redis/v9" + "github.com/sdslabs/beastv4/utils" + "strconv" +) + +const reservePortsScript = ` +local hostKey = KEYS[1] +local firstPort = tonumber(ARGV[1]) +local portRange = tonumber(ARGV[2]) +local count = tonumber(ARGV[3]) +local selected = {} + +for offset = 0, portRange - 1 do + local port = firstPort + offset + if redis.call("SISMEMBER", hostKey, port) == 0 then + table.insert(selected, port) + if #selected == count then + break + end + end +end + +if #selected < count then + return {} +end + +for _, port in ipairs(selected) do + redis.call("SADD", hostKey, port) +end + +return selected +` + +// GetFreePortOnHost gets the first available port in the specific range by checking its existance in the cache. +// algorithm can be imprived later on if it bottlenecks performance. +func GetFreePortOnHost(host string, firstPort uint32, portRange uint32) (uint32, error) { + ports, err := GetFreePortsOnHost(host, firstPort, portRange, 1) + if err != nil { + return 0, err + } + if len(ports) == 0 { + return 0, fmt.Errorf("no free port found on host: %s", host) + } + + return ports[0], nil +} + +func GetFreePortsOnHost(host string, firstPort uint32, portRange uint32, count int) ([]uint32, error) { + if Cache == nil { + Init() + } + + if count <= 0 { + return []uint32{}, nil + } + + CacheMutex.Lock() + defer CacheMutex.Unlock() + + ctx := context.Background() + hostKey := utils.HostToKey(host) + + result, err := Cache.Eval(ctx, reservePortsScript, []string{hostKey}, firstPort, portRange, count).Result() + if err != nil { + return nil, err + } + + values, ok := result.([]interface{}) + if !ok || len(values) != count { + return nil, fmt.Errorf("no free port found on host: %s", host) + } + + ports := make([]uint32, len(values)) + for i, value := range values { + port, err := redisValueToUint32(value) + if err != nil { + return nil, err + } + ports[i] = port + } + + return ports, nil +} + +// AssignFreePortOnHostToContainer allocates a port for a container on a given host machine +func AssignFreePortOnHostToContainer(host string, containerId string, port uint32) error { + return AssignPortsOnHostToContainer(host, containerId, []uint32{port}) +} + +func AssignPortsOnHostToContainer(host string, containerId string, ports []uint32) error { + CacheMutex.Lock() + defer CacheMutex.Unlock() + + ctx := context.Background() + instanceKey := utils.ContainerToKey(host, containerId) + + pipe := Cache.TxPipeline() + for _, port := range ports { + pipe.SAdd(ctx, instanceKey, port) + } + + results, err := pipe.Exec(ctx) + if err != nil { + return err + } + + for i, result := range results { + cmd, ok := result.(*redis.IntCmd) + if !ok { + continue + } + added, err := cmd.Result() + if err != nil { + return err + } + if added == 0 { + return fmt.Errorf("port: %v on host: %s is already registered to instance: %s", ports[i], host, containerId) + } + } + + return nil +} + +// GetContainerPortsOnHost gets all the assigned ports for a given container on a given host +func GetContainerPortsOnHost(host string, containerId string) ([]uint32, error) { + if Cache == nil { + Init() + } + + CacheMutex.Lock() + defer CacheMutex.Unlock() + + ctx := context.Background() + instanceKey := utils.ContainerToKey(host, containerId) + + result, err := Cache.SMembers(ctx, instanceKey).Result() + if err != nil { + return nil, err + } + + ports := make([]uint32, len(result)) + for i, s := range result { + port, err := strconv.ParseUint(s, 10, 32) + if err != nil { + return nil, err + } + + ports[i] = uint32(port) + } + + return ports, nil +} + +func redisValueToUint32(value interface{}) (uint32, error) { + switch v := value.(type) { + case int64: + return uint32(v), nil + case string: + port, err := strconv.ParseUint(v, 10, 32) + return uint32(port), err + case []byte: + port, err := strconv.ParseUint(string(v), 10, 32) + return uint32(port), err + default: + return 0, fmt.Errorf("unexpected Redis port value %T", value) + } +} + +func FreePortOnHost(host string, port uint32) error { + if Cache == nil { + Init() + } + + CacheMutex.Lock() + defer CacheMutex.Unlock() + + ctx := context.Background() + hostKey := utils.HostToKey(host) + + _, err := Cache.SRem(ctx, hostKey, port).Result() + if err != nil { + return err + } + + return nil +} + +// FreeContainerPortsOnHost frees all allocated host ports on a machine, at present occupied by a container +func FreeContainerPortsOnHost(host string, containerId string) error { + if Cache == nil { + Init() + } + + CacheMutex.Lock() + defer CacheMutex.Unlock() + + ctx := context.Background() + hostKey := utils.HostToKey(host) + instanceKey := utils.ContainerToKey(host, containerId) + + result, err := Cache.SMembers(ctx, instanceKey).Result() + if err != nil { + return err + } + + ports := make([]uint32, len(result)) + for i, portString := range result { + port, err := strconv.ParseUint(portString, 10, 32) + if err != nil { + return err + } + + ports[i] = uint32(port) + Cache.SRem(ctx, instanceKey, port) + } + + for _, port := range ports { + _, err = Cache.SRem(ctx, hostKey, port).Result() + if err != nil { + return err + } + } + + return nil +} diff --git a/core/cache/ports_instance_test.go b/core/cache/ports_instance_test.go new file mode 100644 index 00000000..baefcb66 --- /dev/null +++ b/core/cache/ports_instance_test.go @@ -0,0 +1,270 @@ +package cache + +import ( + "context" + "fmt" + "os" + "strconv" + "sync" + "testing" + "time" + + "github.com/redis/go-redis/v9" + "github.com/sdslabs/beastv4/utils" +) + +func setupRedisIntegrationTest(t *testing.T) func() { + t.Helper() + + addr := os.Getenv("BEAST_TEST_REDIS_ADDR") + if addr == "" { + t.Skip("set BEAST_TEST_REDIS_ADDR to run Redis cache integration tests") + } + + db := 0 + if rawDB := os.Getenv("BEAST_TEST_REDIS_DB"); rawDB != "" { + parsed, err := strconv.Atoi(rawDB) + if err != nil { + t.Fatalf("invalid BEAST_TEST_REDIS_DB: %v", err) + } + db = parsed + } + + previousCache := Cache + previousMutex := CacheMutex + previousConfig := cacheConfig + + CacheMutex = &sync.Mutex{} + Cache = redis.NewClient(&redis.Options{ + Addr: addr, + Username: os.Getenv("BEAST_TEST_REDIS_USER"), + Password: os.Getenv("BEAST_TEST_REDIS_PASSWORD"), + DB: db, + }) + cacheConfig.RedisConfig.DB = db + + ctx := context.Background() + if err := Cache.Ping(ctx).Err(); err != nil { + t.Fatalf("ping redis: %v", err) + } + if os.Getenv("BEAST_TEST_REDIS_FLUSH") == "1" { + if err := Cache.FlushDB(ctx).Err(); err != nil { + t.Fatalf("flush redis db: %v", err) + } + } + + return func() { + _ = Cache.Close() + Cache = previousCache + CacheMutex = previousMutex + cacheConfig = previousConfig + } +} + +func TestConcurrentMultiPortReservationDoesNotOverlap(t *testing.T) { + cleanup := setupRedisIntegrationTest(t) + defer cleanup() + + host := fmt.Sprintf("beast-test-host-%d", time.Now().UnixNano()) + defer func() { + Cache.Del(context.Background(), utils.HostToKey(host)) + }() + + const workers = 40 + const portsPerWorker = 3 + errCh := make(chan error, workers) + results := make(chan []uint32, workers) + start := make(chan struct{}) + var wg sync.WaitGroup + + for i := 0; i < workers; i++ { + wg.Add(1) + go func() { + defer wg.Done() + <-start + ports, err := GetFreePortsOnHost(host, 30000, 1000, portsPerWorker) + if err != nil { + errCh <- err + return + } + results <- ports + }() + } + + close(start) + wg.Wait() + close(errCh) + close(results) + + for err := range errCh { + t.Fatalf("reserve ports: %v", err) + } + + seen := map[uint32]bool{} + for ports := range results { + if len(ports) != portsPerWorker { + t.Fatalf("expected %d ports per reservation, got %d", portsPerWorker, len(ports)) + } + for _, port := range ports { + if seen[port] { + t.Fatalf("port %d was allocated more than once", port) + } + seen[port] = true + } + } + + expected := workers * portsPerWorker + if len(seen) != expected { + t.Fatalf("expected %d unique reserved ports, got %d", expected, len(seen)) + } +} + +func TestAssignAndFreeContainerPortsOnHost(t *testing.T) { + cleanup := setupRedisIntegrationTest(t) + defer cleanup() + + host := fmt.Sprintf("beast-test-free-host-%d", time.Now().UnixNano()) + owner := fmt.Sprintf("beast-test-owner-%d", time.Now().UnixNano()) + defer func() { + Cache.Del(context.Background(), utils.HostToKey(host), utils.ContainerToKey(host, owner)) + }() + + ports, err := GetFreePortsOnHost(host, 31000, 10, 2) + if err != nil { + t.Fatalf("reserve ports: %v", err) + } + if err := AssignPortsOnHostToContainer(host, owner, ports); err != nil { + t.Fatalf("assign ports: %v", err) + } + + assigned, err := GetContainerPortsOnHost(host, owner) + if err != nil { + t.Fatalf("get assigned ports: %v", err) + } + if len(assigned) != len(ports) { + t.Fatalf("expected %d assigned ports, got %d", len(ports), len(assigned)) + } + + if err := FreeContainerPortsOnHost(host, owner); err != nil { + t.Fatalf("free assigned ports: %v", err) + } + + assignedAfterFree, err := GetContainerPortsOnHost(host, owner) + if err != nil { + t.Fatalf("get assigned ports after free: %v", err) + } + if len(assignedAfterFree) != 0 { + t.Fatalf("expected no assigned ports after free, got %v", assignedAfterFree) + } + + reallocated, err := GetFreePortsOnHost(host, 31000, 10, 2) + if err != nil { + t.Fatalf("reserve ports after free: %v", err) + } + for i := range ports { + if reallocated[i] != ports[i] { + t.Fatalf("expected freed port %d to be reusable, got %d", ports[i], reallocated[i]) + } + } +} + +func TestInstanceMetadataOutlivesExpiryMarkerAndQueue(t *testing.T) { + cleanup := setupRedisIntegrationTest(t) + defer cleanup() + if os.Getenv("BEAST_TEST_REDIS_FLUSH") != "1" { + t.Skip("set BEAST_TEST_REDIS_FLUSH=1 for deletion queue tests") + } + + instanceID := fmt.Sprintf("inst-%d", time.Now().UnixNano()) + instance := &Instance{ + InstanceID: instanceID, + ChallengeName: "durable-instance", + ContainerID: "container-" + instanceID, + PortOwner: "owner-" + instanceID, + Port: 31337, + UserID: "user-1", + Username: "user-1", + CreatedAt: time.Now(), + ExpiresAt: time.Now().Add(50 * time.Millisecond), + DeploymentType: "standard_docker", + ServerDeployed: "localhost", + } + + if err := SaveInstance(instance, 50*time.Millisecond); err != nil { + t.Fatalf("save instance: %v", err) + } + time.Sleep(100 * time.Millisecond) + + if _, err := GetInstance(instanceID); err != nil { + t.Fatalf("durable instance metadata expired with marker: %v", err) + } + + expired, err := GetExpiredInstances() + if err != nil { + t.Fatalf("get expired instances: %v", err) + } + if len(expired) != 1 || expired[0].InstanceID != instanceID { + t.Fatalf("expected durable expired instance %s, got %#v", instanceID, expired) + } + + if err := QueueInstanceForDeletion(instanceID); err != nil { + t.Fatalf("queue instance for deletion: %v", err) + } + if _, err := GetInstance(instanceID); err != nil { + t.Fatalf("queueing deletion should keep durable metadata readable: %v", err) + } + + queued, err := PopInstanceForDeletion() + if err != nil { + t.Fatalf("pop queued instance: %v", err) + } + if queued == nil || queued.InstanceID != instanceID { + t.Fatalf("expected queued instance %s, got %#v", instanceID, queued) + } + + if err := DeleteInstanceMetadata(instanceID); err != nil { + t.Fatalf("delete instance metadata: %v", err) + } + if _, err := GetInstance(instanceID); err == nil { + t.Fatalf("expected instance metadata to be deleted after successful cleanup") + } +} + +func TestRedisExpiryMarkerSubscription(t *testing.T) { + cleanup := setupRedisIntegrationTest(t) + defer cleanup() + + if err := EnableKeyspaceExpiryNotifications(); err != nil { + t.Skipf("redis keyspace notifications are unavailable: %v", err) + } + + instanceID := fmt.Sprintf("event-%d", time.Now().UnixNano()) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + events := make(chan string, 1) + errCh := make(chan error, 1) + go func() { + if err := SubscribeExpiredInstanceMarkers(ctx, func(expiredInstanceID string) { + events <- expiredInstanceID + }); err != nil && ctx.Err() == nil { + errCh <- err + } + }() + + time.Sleep(50 * time.Millisecond) + if err := Cache.Set(context.Background(), utils.InstanceExpiryToKey(instanceID), instanceID, 50*time.Millisecond).Err(); err != nil { + t.Fatalf("set expiry marker: %v", err) + } + + select { + case got := <-events: + if got != instanceID { + t.Fatalf("expected expiry event for %s, got %s", instanceID, got) + } + case err := <-errCh: + t.Fatalf("expiry subscriber failed: %v", err) + case <-time.After(2 * time.Second): + t.Fatalf("timed out waiting for Redis expiry marker event") + } +} diff --git a/core/config/challenge.go b/core/config/challenge.go index 1452d31a..47da75cb 100644 --- a/core/config/challenge.go +++ b/core/config/challenge.go @@ -137,15 +137,31 @@ type ChallengeMetadata struct { Text string `toml:"text"` Points uint `toml:"points"` } `toml:"hints"` - MaxAttemptLimit int `toml:"maxAttemptLimit"` - PreReqs []string `toml:"preReqs"` - DynamicFlag bool `toml:"dynamicFlag"` - Points uint `toml:"points"` - MaxPoints uint `toml:"maxPoints"` - MinPoints uint `toml:"minPoints"` - Assets []string `toml:"assets"` - AdditionalLinks []string `toml:"additionalLinks"` - Difficulty string `toml:"difficulty"` + MaxAttemptLimit int `toml:"maxAttemptLimit"` + PreReqs []string `toml:"preReqs"` + DynamicFlag bool `toml:"dynamicFlag"` + Points uint `toml:"points"` + MaxPoints uint `toml:"maxPoints"` + MinPoints uint `toml:"minPoints"` + Assets []string `toml:"assets"` + AdditionalLinks []string `toml:"additionalLinks"` + Difficulty string `toml:"difficulty"` + Instanced bool `toml:"instanced"` + InstanceExpiration int64 `toml:"instance_expiration"` +} + +func (config *ChallengeMetadata) IsInstanced() bool { + return config.Instanced +} + +func (config *ChallengeMetadata) GetInstanceExpiration() int64 { + if config.InstanceExpiration > 0 { + return config.InstanceExpiration + } + if Cfg != nil && Cfg.InstanceConfig.DefaultExpiration > 0 { + return Cfg.InstanceConfig.DefaultExpiration + } + return 300 } // In this validation returned boolean value represents if the challenge type is @@ -252,7 +268,8 @@ type ChallengeEnv struct { AptDeps []string `toml:"apt_deps"` Ports []uint32 `toml:"ports"` DefaultPort uint32 `toml:"default_port"` - PortMappings []string `toml:"port_mappings"` + PortVariables []string `toml:"-"` + DefaultPortVar string `toml:"default_port_var"` SetupScripts []string `toml:"setup_scripts"` StaticContentDir string `toml:"static_dir"` RunCmd string `toml:"run_cmd"` @@ -274,107 +291,15 @@ func (config *ChallengeEnv) TrafficType() cr.TrafficType { return cr.TrafficType(config.Traffic) } -// NewPortMapping returns a new port mapping instance. -func NewPortMapping(hp, cp uint32) cr.PortMapping { - return cr.PortMapping{ - HostPort: hp, - ContainerPort: cp, - } -} - -// Given a port mapping array and a port the function checks whether the port exists in the mapping -// as a container port. -func checkIfPortExistInMapping(portMapping []cr.PortMapping, port uint32) bool { - for _, portMap := range portMapping { - if port == portMap.ContainerPort { - return true - } - } - - return false -} - -// GetPortMappings returns the entire port mapping for the challenge from the challenge -// environment configuration. -func (config *ChallengeEnv) GetPortMappings() ([]cr.PortMapping, error) { - var mapping []cr.PortMapping - - var containerPorts []uint32 - for _, portMap := range config.PortMappings { - hp, cp, err := utils.ParsePortMapping(portMap) - if err != nil { - return mapping, err - } - mapping = append(mapping, NewPortMapping(hp, cp)) - containerPorts = append(containerPorts, cp) - } - - for _, port := range config.Ports { - if !utils.UInt32InList(port, containerPorts) { - containerPorts = append(containerPorts, port) - mapping = append(mapping, NewPortMapping(port, port)) - } - } - - return mapping, nil -} - -// GetAllHostPorts is utility function for the ChallengeEnv configuration which returns -// the entire list of all the host ports which are being used by the challenge. -func (config *ChallengeEnv) GetAllHostPorts() ([]uint32, error) { - var hostPorts []uint32 - var containerPorts []uint32 - - for _, portMap := range config.PortMappings { - hp, cp, err := utils.ParsePortMapping(portMap) - if err != nil { - return hostPorts, err - } - hostPorts = append(hostPorts, hp) - containerPorts = append(containerPorts, cp) - } - - for _, port := range config.Ports { - if !utils.UInt32InList(port, containerPorts) { - hostPorts = append(hostPorts, port) - containerPorts = append(containerPorts, port) - } - } - - return hostPorts, nil -} - -// GetAllContainerPorts is utility function for the ChallengeEnv configuration which returns -// the entire list of all the container ports which are being used by the challenge. -func (config *ChallengeEnv) GetAllContainerPorts() ([]uint32, error) { - var containerPorts []uint32 - - for _, portMap := range config.PortMappings { - _, cp, err := utils.ParsePortMapping(portMap) - if err != nil { - return containerPorts, err - } - containerPorts = append(containerPorts, cp) - } - - for _, port := range config.Ports { - if !utils.UInt32InList(port, containerPorts) { - containerPorts = append(containerPorts, port) - } - } - - return containerPorts, nil -} - // GetDefaultPort returns the default port used by the challenge from the challenge environment // configuration. func (config *ChallengeEnv) GetDefaultPort() uint32 { - mappings, err := config.GetPortMappings() - if err != nil || len(mappings) == 0 { + ports := config.Ports + if len(ports) == 0 { return 0 } - return mappings[0].ContainerPort + return ports[0] } // ValidateRequiredFields validates required fields for the Challenge environment configuration. @@ -382,34 +307,6 @@ func (config *ChallengeEnv) GetDefaultPort() uint32 { // of the challenge. func (config *ChallengeEnv) ValidateRequiredFields(challType string, challdir string) error { // Validate port related stuff for the challenge environment configuration. - if len(config.Ports) == 0 && len(config.PortMappings) == 0 { - return errors.New("some port is required to be specified by the challenge") - } - - if len(config.Ports)+len(config.PortMappings) > int(core.MAX_PORT_PER_CHALL) { - return fmt.Errorf("max ports allowed for challenge : %d given : %d", core.MAX_PORT_PER_CHALL, len(config.Ports)) - } - - portMappings, err := config.GetPortMappings() - if err != nil { - return fmt.Errorf("error while parsing port mapping: %s", err) - } - - // By default if no port is specified to be default, the first port - // from the list is assumed to be default and the service is deployed accordingly. - if config.DefaultPort == 0 { - config.DefaultPort = portMappings[0].ContainerPort - } - - if !checkIfPortExistInMapping(portMappings, config.DefaultPort) { - return fmt.Errorf("`default_port` must be one of the Ports in the `ports` list") - } - - for _, portMap := range portMappings { - if portMap.HostPort < core.ALLOWED_MIN_PORT_VALUE || portMap.HostPort > core.ALLOWED_MAX_PORT_VALUE { - return fmt.Errorf("port value must be between %d and %d", core.ALLOWED_MIN_PORT_VALUE, core.ALLOWED_MAX_PORT_VALUE) - } - } if config.StaticContentDir != "" { if filepath.IsAbs(config.StaticContentDir) { @@ -457,9 +354,17 @@ func (config *ChallengeEnv) ValidateRequiredFields(challType string, challdir st if len(config.SetupScripts) > 0 { log.Warn("setup_scripts will be ignored when docker_compose is specified") } + + if err := config.ExtractPortsCompose(challdir); err != nil { + return err + } return nil } + if err := config.ExtractPorts(); err != nil { + return err + } + // Run command is only a required value in case of bare challenge types. if config.RunCmd == "" && config.Entrypoint == "" && config.DockerCtx == "" && config.DockerCompose == "" && challType == core.BARE_CHALLENGE_TYPE_NAME { return fmt.Errorf("a valid run_cmd should be provided for the challenge environment") @@ -528,6 +433,49 @@ func (config *ChallengeEnv) ValidateRequiredFields(challType string, challdir st return nil } +func (config *ChallengeEnv) ExtractPorts() error { + if config.DockerCompose == "" { + if len(config.Ports) == 0 && config.DefaultPort == 0 { + return errors.New("some port is required to be specified by the challenge") + } + if len(config.Ports) > int(core.MAX_PORT_PER_CHALL) { + return fmt.Errorf("max ports allowed for challenge : %d given : %d", core.MAX_PORT_PER_CHALL, len(config.Ports)) + } + + if config.DefaultPort == 0 { + config.DefaultPort = config.Ports[0] + log.Warnf("default port is 0 for challenge with default port : %d", config.Ports[0]) + } else if !utils.UInt32InList(config.DefaultPort, config.Ports) { + return fmt.Errorf("default port %d was not found in assigned ports", config.DefaultPort) + } + } + + return nil +} + +func (config *ChallengeEnv) ExtractPortsCompose(challdir string) error { + if config.DockerCompose != "" { + portVariables, err := utils.ExtractPortsFromCompose(filepath.Join(challdir, config.DockerCompose)) + if err != nil { + log.Warnf("failed to extract port variables from compose file with the following error : %s", err.Error()) + } + if len(portVariables) == 0 { + return errors.New("some port is required to be specified by the challenge") + } + + config.PortVariables = portVariables + if config.DefaultPortVar == "" { + config.DefaultPortVar = config.PortVariables[0] + log.Warnf("default port variable is empty, settting it to %s", config.PortVariables[0]) + } + if !utils.StringInSlice(config.DefaultPortVar, config.PortVariables) { + return fmt.Errorf("default port variable: %s was not found", config.DefaultPortVar) + } + } + + return nil +} + // Metadata related to author of the challenge, this structure includes // // - Name - Name of the author of the challenge @@ -567,9 +515,10 @@ type EnvironmentVar struct { } type Resources struct { - CPUShares int64 `toml:"cpu_shares"` - Memory int64 `toml:"memory_limit"` - PidsLimit int64 `toml:"pids_limit"` + CPUShares int64 `toml:"cpu_shares"` + Memory int64 `toml:"memory_limit"` + PidsLimit int64 `toml:"pids_limit"` + CPUsLimit float32 `toml:"cpuslimit"` } func (config *Resources) ValidateRequiredFields() { @@ -587,4 +536,9 @@ func (config *Resources) ValidateRequiredFields() { log.Debug("Pids Limit not provided in configuration, using default.") config.PidsLimit = Cfg.PidsLimit } + + if config.CPUsLimit <= 0 { + log.Debug("CPUsLimit not provided in configuration, using default.") + config.CPUsLimit = Cfg.CPUsLimit + } } diff --git a/core/config/config.go b/core/config/config.go index 167458f2..98255850 100644 --- a/core/config/config.go +++ b/core/config/config.go @@ -8,6 +8,7 @@ import ( "os" "path/filepath" "regexp" + "strings" "time" "github.com/sdslabs/beastv4/core" @@ -108,7 +109,7 @@ import ( // dbname = "beast" // host = "localhost" // port = "5432" -// sslmode = "prefer" +// sslmode = "prefer" // ``` type BeastConfig struct { AuthorizedKeysFile string `toml:"authorized_keys_file"` @@ -117,6 +118,7 @@ type BeastConfig struct { AvailableServers map[string]AvailableServer `toml:"available_servers"` GitRemotes []GitRemote `toml:"remote"` PsqlConf PsqlConfig `toml:"psql_config"` + RedisConf RedisConfig `toml:"redis_config"` JWTSecret string `toml:"jwt_secret"` NotificationWebhooks []NotificationWebhook `toml:"notification_webhooks"` CompetitionInfo CompetitionInfo `toml:"competition_info"` @@ -125,15 +127,60 @@ type BeastConfig struct { HealthProber bool `toml:"health_prober"` RemoteSyncPeriod time.Duration `toml:"-"` Rsp string `toml:"remote_sync_period"` + InstanceConfig InstanceConfig `toml:"instance_config"` - CPUShares int64 `toml:"default_cpu_shares"` - Memory int64 `toml:"default_memory_limit"` - PidsLimit int64 `toml:"default_pids_limit"` + CPUShares int64 `toml:"default_cpu_shares"` + Memory int64 `toml:"default_memory_limit"` + PidsLimit int64 `toml:"default_pids_limit"` + CPUsLimit float32 `toml:"default_cpus_limit"` - // For SMTP Configuration MailConfig MailConfig `toml:"mail_config"` } +type InstanceConfig struct { + DefaultExpiration int64 `toml:"default_expiration"` + MaxExtension int64 `toml:"max_extension"` + MaxInstancesPerUser int `toml:"max_instances_per_user"` +} + +func (config *InstanceConfig) Validate() { + if config.DefaultExpiration <= 0 { + config.DefaultExpiration = core.DEFAULT_MINIMUM_EXTEND_TIME + } + if config.MaxExtension <= config.DefaultExpiration { + config.DefaultExpiration = core.DEFAULT_MINIMUM_EXTEND_TIME + config.MaxExtension = core.DEFAULT_MAXIMUM_EXTEND_TIME + } + if config.MaxInstancesPerUser <= 0 { + config.MaxInstancesPerUser = core.DEFAULT_MAXIMUM_INSTANCES_PER_USER + } +} + +func ValidatePortRange(portRange string) error { + if portRange == "" { + return fmt.Errorf("port range is empty") + } + + firstPort, lastPort, err := utils.ParsePortMapping(portRange) + if err != nil { + return fmt.Errorf("error while parsing port range in global beast config: %s", err) + } + + if firstPort > lastPort { + return fmt.Errorf("invalid port range, %v cannot be greater than %v", firstPort, lastPort) + } + + if firstPort < core.ALLOWED_MIN_PORT_VALUE { + return fmt.Errorf("invalid port range, range cannot precede %v", core.ALLOWED_MIN_PORT_VALUE) + } + + if lastPort > core.ALLOWED_MAX_PORT_VALUE { + return fmt.Errorf("invalid port range, range cannot exceed %v", core.ALLOWED_MAX_PORT_VALUE) + } + + return nil +} + func (config *BeastConfig) ValidateConfig() error { log.Debug("Validating BeastConfig structure") @@ -170,19 +217,32 @@ func (config *BeastConfig) ValidateConfig() error { return fmt.Errorf("error while validating db config : %s", err) } + err = config.RedisConf.ValidateRedisConfig() + if err != nil { + return fmt.Errorf("error while validating redis config : %s", err) + } + if len(config.AvailableServers) == 0 { log.Warn("No available servers provided for challenges. Using default localhost") config.AvailableServers = map[string]AvailableServer{ core.LOCALHOST: { + Name: core.LOCALHOST, Host: core.LOCALHOST, Username: os.Getenv("USER"), SSHKeyPath: "", Active: true, + PortRange: fmt.Sprintf("%v%s%v", core.ALLOWED_MIN_PORT_VALUE, core.MappingDelimiter, core.ALLOWED_MAX_PORT_VALUE), }, } } - for _, server := range config.AvailableServers { + for name, server := range config.AvailableServers { + if strings.Contains(name, ":") { + return fmt.Errorf("server key %q contains invalid character ':'", name) + } + + server.Name = name + config.AvailableServers[name] = server if server.Active { err := server.ValidateServerConfig() if err != nil { @@ -248,32 +308,64 @@ func (config *BeastConfig) ValidateConfig() error { config.PidsLimit = core.DEFAULT_PIDS_LIMIT } + if config.CPUsLimit <= 0 { + log.Debug("Per container CPUsLimit Limit not provided using default value") + config.CPUsLimit = core.DEFAULT_CPU_LIMIT + } + if config.MailConfig.From == "" || config.MailConfig.Password == "" || config.MailConfig.SMTPHost == "" || config.MailConfig.SMTPPort == "" { log.Warn("Mail configuration not provided, email notifications will not work") } + config.InstanceConfig.Validate() + return nil } +func (config *BeastConfig) UseLocalDockerDaemon(serverName string) bool { + server, ok := config.AvailableServers[serverName] + if !ok { + return true + } + return server.Host == core.LOCALHOST || server.Host == core.LOCALHOST_IP +} + type AvailableServer struct { + Name string `toml:"-"` Host string `toml:"host"` Username string `toml:"username"` SSHKeyPath string `toml:"ssh_key_path"` Active bool `toml:"active"` + PortRange string `toml:"port_range"` } func (config *AvailableServer) ValidateServerConfig() error { - if config.Host == core.LOCALHOST { + if config.Host == "" { + return fmt.Errorf("host is empty") + } + config.Host = strings.TrimSpace(config.Host) + + err := ValidatePortRange(config.PortRange) + if err != nil { + return fmt.Errorf("error while validating port range for server %s: %s", config.Host, err) + } + + if config.Host == core.LOCALHOST || config.Host == core.LOCALHOST_IP { return nil } - if config.Host == "" || config.Username == "" || config.SSHKeyPath == "" { - log.Error("One of host, username or ssh_key_path is missing in the config") - return errors.New("server config not valid, config parameters missing") + + if config.Username == "" { + return fmt.Errorf("username is empty") } - err := utils.ValidateFileExists(config.SSHKeyPath) + if config.SSHKeyPath == "" { + return fmt.Errorf("ssh_key_path is empty") + } + + err = utils.ValidateFileExists(config.SSHKeyPath) if err != nil { return fmt.Errorf("provided ssh key file(%s) does not exists : %s", config.SSHKeyPath, err) } + return nil } @@ -324,6 +416,14 @@ type PsqlConfig struct { SslMode string `toml:"sslmode"` } +type RedisConfig struct { + User string `toml:"user"` + Password string `toml:"password"` + Host string `toml:"host"` + Port string `toml:"port"` + Db uint32 `toml:"db"` +} + func (config *PsqlConfig) ValidatePsqlConfig() error { if config.User == "" || config.Password == "" || config.Dbname == "" || config.Host == "" || config.Port == "" { log.Error("One of username, password, dbname, hostname, port is missing in the config") @@ -336,6 +436,14 @@ func (config *PsqlConfig) ValidatePsqlConfig() error { return nil } +func (config *RedisConfig) ValidateRedisConfig() error { + if config.Host == "" || config.Port == "" { + log.Error("One of hostname or port is missing in the config") + return errors.New("redis config not valid, config parameters missing") + } + return nil +} + type NotificationWebhook struct { URL string `toml:"url"` ServiceName string `toml:"service_name"` @@ -438,44 +546,9 @@ func LoadBeastConfig(configPath string) (BeastConfig, error) { return config, nil } -// Update the USED_PORT_LIST variable in config. -// Don't do this very often, we do this once during syncing the git repository -// then whenever you need updated used port list you need to sync the git remote -// by beast. -func UpdateUsedPortList() { - USED_PORTS_LIST = make([]uint32, 0) - - beastRemoteDir := filepath.Join(core.BEAST_GLOBAL_DIR, core.BEAST_REMOTES_DIR) - - for _, gitRemote := range Cfg.GitRemotes { - if !gitRemote.Active { - continue - } - - challengeDir := filepath.Join(beastRemoteDir, gitRemote.RemoteName, core.BEAST_REMOTE_CHALLENGE_DIR) - dirs := utils.GetAllDirectoriesName(challengeDir) - for _, dir := range dirs { - configFilePath := filepath.Join(dir, core.CHALLENGE_CONFIG_FILE_NAME) - var config BeastChallengeConfig - _, err := toml.DecodeFile(configFilePath, &config) - if err == nil { - hostPorts, err := config.Challenge.Env.GetAllHostPorts() - if err != nil { - log.Errorf("Error while parsing host ports for challenge %s", dir) - continue - } - - USED_PORTS_LIST = append(USED_PORTS_LIST, hostPorts...) - } - } - } - log.Debugf("Used port list updated: %v", USED_PORTS_LIST) -} - var Cfg *BeastConfig var SkipAuthorization bool var NoCache bool -var USED_PORTS_LIST []uint32 // InitConfig loads the config from the global config file and populate // the Cfg global variable used everywhere else. diff --git a/core/constants.go b/core/constants.go index 74eb1230..a46cf279 100644 --- a/core/constants.go +++ b/core/constants.go @@ -35,11 +35,13 @@ const ( //names ISSUER string = "beast-sds" DELIMITER string = "::::" LOCALHOST string = "localhost" + LOCALHOST_IP string = "127.0.0.1" BEAST_REMOTE_GLOBAL_DIR string = "~/.beast" // This should always be used for remote only. DOCKER_PID string = "/var/run/docker.pid" BEAST_GRAPH_CACHE string = "graph_cache.json" BEAST_LEADERBOARD_CACHE string = "leaderboard.json" POSTGRES_SUPER_USER string = "postgres" + REDIS_DEFAULT_USER string = "default" ) const ( //paths @@ -57,6 +59,9 @@ const ( //paths BEAST_SECRETS_DIR string = "secrets" BEAST_EXAMPLE_DIR string = "_examples" BEAST_CACHE_DIR string = "cache" + BEAST_BACKUP_DIR string = "backup" + DB_BACKUP_DIR string = "db" + CACHE_BACKUP_DIR string = "cache" ) const ( //chall types @@ -84,20 +89,21 @@ const ( // chall env ALLOWED_MAX_PORT_VALUE uint32 = 20000 ) const ( // default config - IMAGE_NA string = "IMAGE_NA" - CONTAINER_NA string = "CONTAINER_NA" - MAX_QUEUE_SIZE uint32 = 100 - DEFAULT_TICKER_FREQUENCY int = 1500 - DEFAULT_PROBE_TIMEOUT int = 10 - DEFAULT_USER_NAME string = "ghost" - DEFAULT_USER_EMAIL string = "ghost@ghost.com" - DEFAULT_CPU_SHARE int64 = (1 << 9) - DEFAULT_MEMORY_LIMIT int64 = (1 << 29) - DEFAULT_PIDS_LIMIT int64 = 100 - ITERATIONS int = 65536 - HASH_LENGTH int = 32 - TIMEPERIOD int64 = 6 * 60 * 60 - SSH_PORT int = 22 + IMAGE_NA string = "IMAGE_NA" + CONTAINER_NA string = "CONTAINER_NA" + MAX_QUEUE_SIZE uint32 = 100 + DEFAULT_TICKER_FREQUENCY int = 1500 + DEFAULT_PROBE_TIMEOUT int = 10 + DEFAULT_USER_NAME string = "ghost" + DEFAULT_USER_EMAIL string = "ghost@ghost.com" + DEFAULT_CPU_SHARE int64 = (1 << 9) + DEFAULT_MEMORY_LIMIT int64 = (1 << 29) + DEFAULT_PIDS_LIMIT int64 = 100 + DEFAULT_CPU_LIMIT float32 = .25 + ITERATIONS int = 65536 + HASH_LENGTH int = 32 + TIMEPERIOD int64 = 6 * 60 * 60 + SSH_PORT int = 22 ) const ( // roles @@ -106,8 +112,15 @@ const ( // roles USER int = 1 << 2 ) +const ( + DEFAULT_MINIMUM_EXTEND_TIME int64 = 300 + DEFAULT_MAXIMUM_EXTEND_TIME int64 = 600 + DEFAULT_MAXIMUM_INSTANCES_PER_USER int = 3 +) + var ( DEFAULT_REMOTE_PERIODIC_SYNC_TIME = time.Second * 120 + DEFAULT_HEALTH_CHECK_TIME = time.Second * 30 ) var DEPLOY_STATUS = map[string]string{ @@ -198,3 +211,5 @@ var NOTIFICATION_SERVICES = []string{ "slack", "discord", } + +const MappingDelimiter = ":" diff --git a/core/database/challenges.go b/core/database/challenges.go index 20e0a434..8c6bd547 100644 --- a/core/database/challenges.go +++ b/core/database/challenges.go @@ -8,6 +8,7 @@ import ( "html/template" "io/ioutil" "path/filepath" + "sort" "strings" "time" @@ -46,30 +47,32 @@ import ( type Challenge struct { gorm.Model - Name string `gorm:"not null;type:varchar(64);unique"` - DynamicFlag bool `gorm:"not null;default:false"` - Flag string `gorm:"type:text"` - Type string `gorm:"type:varchar(64)"` - Difficulty string `gorm:"not null;default:'medium'"` - MaxAttemptLimit int `gorm:"default:-1"` - PreReqs string `gorm:"type:text"` - Assets string `gorm:"type:text"` - AdditionalLinks string `gorm:"type:text"` - Description string `gorm:"type:text"` - Format string `gorm:"not null"` - ContainerId string `gorm:"size:64;unique"` - ImageId string `gorm:"size:64;unique"` - Status string `gorm:"not null;default:'Undeployed'"` - DeploymentType string `gorm:"not null;default:'standard_docker'"` - AuthorID uint `gorm:"not null"` - HealthCheck uint `gorm:"not null;default:1"` - Points uint `gorm:"default:0"` - MaxPoints uint `gorm:"default:0"` - MinPoints uint `gorm:"default:0"` - Ports []Port - Tags []*Tag `gorm:"many2many:tag_challenges;constraint:OnUpdate:CASCADE,OnDelete:CASCADE;"` - Users []*User `gorm:"many2many:user_challenges;"` - ServerDeployed string `gorm:"type:varchar(64)"` + Name string `gorm:"not null;type:varchar(64);unique"` + DynamicFlag bool `gorm:"not null;default:false"` + Flag string `gorm:"type:text"` + Type string `gorm:"type:varchar(64)"` + Difficulty string `gorm:"not null;default:'medium'"` + MaxAttemptLimit int `gorm:"default:-1"` + PreReqs string `gorm:"type:text"` + Assets string `gorm:"type:text"` + AdditionalLinks string `gorm:"type:text"` + Description string `gorm:"type:text"` + Format string `gorm:"not null"` + ContainerId string `gorm:"size:64;unique"` + ImageId string `gorm:"size:64;unique"` + Status string `gorm:"not null;default:'Undeployed'"` + DeploymentType string `gorm:"not null;default:'standard_docker'"` + AuthorID uint `gorm:"not null"` + HealthCheck uint `gorm:"not null;default:1"` + Points uint `gorm:"default:0"` + MaxPoints uint `gorm:"default:0"` + MinPoints uint `gorm:"default:0"` + Ports []Port + Tags []*Tag `gorm:"many2many:tag_challenges;constraint:OnUpdate:CASCADE,OnDelete:CASCADE;"` + Users []*User `gorm:"many2many:user_challenges;"` + ServerDeployed string `gorm:"type:varchar(64)"` + Instanced bool `gorm:"not null;default:false"` + InstanceExpiration int64 `gorm:"default:0"` } type UserChallenges struct { @@ -176,7 +179,7 @@ func QueryAllChallengesMetadata() ([]Challenge, error) { DBMux.Lock() defer DBMux.Unlock() - tx := Db.Select("id", "name", "created_at", "points", "difficulty"). + tx := Db.Select("id", "name", "created_at", "points", "difficulty", "instanced", "instance_expiration", "status"). Preload("Tags"). Find(&challenges) @@ -209,7 +212,7 @@ func QueryChallengeEntries(key string, value string) ([]Challenge, error) { return challenges, nil } -// QueryChallengeEntriesMetadata returns only selected columns: Name, ID, Tags, CreatedAt, Points, Difficulty +// QueryChallengeEntriesMetadata returns only selected columns: Name, ID, Tags, CreatedAt, Points, Difficulty, Instanced, InstanceExpiration, Status func QueryChallengeEntriesMetadata(key string, value string) ([]Challenge, error) { queryKey := fmt.Sprintf("%s = ?", key) @@ -219,7 +222,7 @@ func QueryChallengeEntriesMetadata(key string, value string) ([]Challenge, error defer DBMux.Unlock() // Only select the required columns, but preload Tags for tag names - tx := Db.Select("id", "name", "created_at", "points", "difficulty"). + tx := Db.Select("id", "name", "created_at", "points", "difficulty", "instanced", "instance_expiration", "status"). Preload("Tags"). Where(queryKey, value). Find(&challenges) @@ -332,7 +335,7 @@ func UpdateUserChallengeTries(userID uint, challengeID uint) error { } updates := map[string]interface{}{ - "tries": userChallenges.Tries + 1, + "tries": userChallenges.Tries + 1, } tx := Db.Model(&UserChallenges{}).Where("user_id = ? AND challenge_id = ?", userID, challengeID).Updates(updates) @@ -525,10 +528,29 @@ func SaveFlagSubmission(user_challenges *UserChallenges) error { return fmt.Errorf("error while saving record: %s", tx.Error) } - if err := tx.FirstOrCreate(user_challenges, *user_challenges).Error; err != nil { + // Check if a row already exists for this user+challenge pair (created by UpdateUserChallengeTries) + var existing UserChallenges + err := tx.Where("user_id = ? AND challenge_id = ?", user_challenges.UserID, user_challenges.ChallengeID).First(&existing).Error + if err == nil { + // Row exists: update it to mark as solved with the current timestamp + if updateErr := tx.Model(&existing).Updates(map[string]interface{}{ + "solved": user_challenges.Solved, + "created_at": user_challenges.CreatedAt, + }).Error; updateErr != nil { + tx.Rollback() + return updateErr + } + } else if errors.Is(err, gorm.ErrRecordNotFound) { + // No existing row: create a new one + if createErr := tx.Create(user_challenges).Error; createErr != nil { + tx.Rollback() + return createErr + } + } else { tx.Rollback() return err } + return tx.Commit().Error } @@ -641,6 +663,7 @@ func QueryDynamicFlagEntries(whereMap map[string]interface{}) ([]DynamicFlag, er DBMux.Lock() defer DBMux.Unlock() + whereMap = normalizeDynamicFlagWhereMap(whereMap) tx := Db.Where(whereMap).Find(&dynamicFlags) if errors.Is(tx.Error, gorm.ErrRecordNotFound) { return []DynamicFlag{}, nil @@ -649,6 +672,21 @@ func QueryDynamicFlagEntries(whereMap map[string]interface{}) ([]DynamicFlag, er return dynamicFlags, tx.Error } +func normalizeDynamicFlagWhereMap(whereMap map[string]interface{}) map[string]interface{} { + normalized := make(map[string]interface{}, len(whereMap)) + for key, value := range whereMap { + switch key { + case "Name": + normalized["name"] = value + case "Flag": + normalized["flag"] = value + default: + normalized[key] = value + } + } + return normalized +} + func DeleteDynamicFlagsByChallengeName(name string) error { DBMux.Lock() defer DBMux.Unlock() @@ -870,19 +908,40 @@ func QueryTimeSeriesForTopUsers(topUserId []uint) []UserLeaderboardResp { Username string CreatedAt time.Time Points uint + IsHint bool } var allRows []userChallengeRow if err := Db.Table("user_challenges"). - Select("user_challenges.user_id, users.username, user_challenges.created_at, challenges.points"). + Select("DISTINCT ON (user_challenges.user_id, user_challenges.challenge_id) user_challenges.user_id, users.username, user_challenges.created_at, challenges.points, false AS is_hint"). Joins("JOIN challenges ON user_challenges.challenge_id = challenges.id"). Joins("JOIN users ON user_challenges.user_id = users.id"). Where("user_challenges.user_id IN ? AND user_challenges.solved = ?", topUserId, true). - Order("user_challenges.user_id ASC, user_challenges.created_at ASC"). + Order("user_challenges.user_id, user_challenges.challenge_id, user_challenges.created_at ASC"). Scan(&allRows).Error; err != nil { return results } + var hintRows []userChallengeRow + if err := Db.Table("user_hints"). + Select("user_hints.user_id, users.username, COALESCE(user_hints.created_at, NOW()) AS created_at, hints.points, true AS is_hint"). + Joins("JOIN hints ON user_hints.hint_id = hints.hint_id"). + Joins("JOIN users ON user_hints.user_id = users.id"). + Where("user_hints.user_id IN ?", topUserId). + Scan(&hintRows).Error; err != nil { + return results + } + + allRows = append(allRows, hintRows...) + + // Re-sort by user_id then created_at for proper cumulative score calculation + sort.Slice(allRows, func(i, j int) bool { + if allRows[i].UserID != allRows[j].UserID { + return allRows[i].UserID < allRows[j].UserID + } + return allRows[i].CreatedAt.Before(allRows[j].CreatedAt) + }) + userRows := make(map[uint][]userChallengeRow) userMap := make(map[uint]string) for _, row := range allRows { @@ -900,7 +959,16 @@ func QueryTimeSeriesForTopUsers(topUserId []uint) []UserLeaderboardResp { var timeSeriesRaw []TimeSeries var cumulativeScore uint = 0 for _, r := range rows { - cumulativeScore += r.Points + if r.IsHint { + if cumulativeScore < r.Points { + cumulativeScore = 0 + } else { + cumulativeScore -= r.Points + } + } else { + cumulativeScore += r.Points + } + timeSeriesRaw = append(timeSeriesRaw, TimeSeries{ Timestamp: r.CreatedAt, Score: cumulativeScore, @@ -920,7 +988,7 @@ func QueryTimeSeriesForTopUsers(topUserId []uint) []UserLeaderboardResp { results = append(results, UserLeaderboardResp{ Id: userId, Username: username, - Score: cumulativeScore, + Score: uint(cumulativeScore), Rank: rank, TimeSeriesdata: timeSeries, }) diff --git a/core/database/database.go b/core/database/database.go index 3455ee29..b632f110 100644 --- a/core/database/database.go +++ b/core/database/database.go @@ -86,10 +86,16 @@ func Init() { log.Fatalf("Cannot create related models: %s", err) } - err := Db.AutoMigrate(&Challenge{}, &Transaction{}, &Port{}, &User{}, &Tag{}, &Notification{}, &Hint{}, &DynamicFlag{}, &OTP{}) + // UserHint must be explicitly migrated since GORM's AutoMigrate on User only handles + // the users table, not custom join table structs. Without this, the created_at and + // challenge_id columns on user_hints won't be added to existing databases. + err := Db.AutoMigrate(&Challenge{}, &Transaction{}, &Port{}, &User{}, &UserChallenges{}, &Tag{}, &Notification{}, &Hint{}, &DynamicFlag{}, &DynamicFlagClaim{}, &DynamicScoreDirty{}, &OTP{}, &UserHint{}) if err != nil { log.Fatalf("failed to migrate database with error: %s", err) } + if err := MigrateSubmissionGuards(); err != nil { + log.Fatalf("failed to migrate submission guards with error: %s", err) + } users, err := QueryUserEntries("email", core.DEFAULT_USER_EMAIL) if err != nil { @@ -131,7 +137,7 @@ func BackupAndReset() { return } - backupPath := filepath.Join(core.BEAST_GLOBAL_DIR, "backup", core.BEAST_REMOTES_DIR) + backupPath := filepath.Join(core.BEAST_GLOBAL_DIR, core.BEAST_BACKUP_DIR, core.BEAST_REMOTES_DIR) err = utils.CreateIfNotExistDir(backupPath) if err != nil { log.Errorf("Error while creating backup directory: %s", err) @@ -146,7 +152,7 @@ func BackupAndReset() { return } - backupPath = filepath.Join(core.BEAST_GLOBAL_DIR, "backup", core.BEAST_STAGING_DIR) + backupPath = filepath.Join(core.BEAST_GLOBAL_DIR, core.BEAST_BACKUP_DIR, core.BEAST_STAGING_DIR) err = utils.CreateIfNotExistDir(backupPath) if err != nil { @@ -168,7 +174,7 @@ func BackupDatabase() error { LoadDbConfig() } - backupPath := filepath.Join(core.BEAST_GLOBAL_DIR, "backup", "db") + backupPath := filepath.Join(core.BEAST_GLOBAL_DIR, core.BEAST_BACKUP_DIR, core.DB_BACKUP_DIR) err := utils.CreateIfNotExistDir(backupPath) if err != nil { log.Errorf("Error while creating backup directory: %s", err) diff --git a/core/database/hints.go b/core/database/hints.go index 1b6561d0..d45e087f 100644 --- a/core/database/hints.go +++ b/core/database/hints.go @@ -3,6 +3,7 @@ package database import ( "errors" "fmt" + "time" "gorm.io/gorm" ) @@ -23,6 +24,8 @@ type UserHint struct { ChallengeID uint Challenge Challenge `gorm:"foreignKey:ChallengeID"` + + CreatedAt time.Time } func CreateHintEntry(hint *Hint) error { diff --git a/core/database/submission.go b/core/database/submission.go new file mode 100644 index 00000000..3ab84bfa --- /dev/null +++ b/core/database/submission.go @@ -0,0 +1,313 @@ +package database + +import ( + "errors" + "fmt" + "time" + + "gorm.io/gorm" +) + +type DynamicFlagClaim struct { + gorm.Model + + ChallengeID uint `gorm:"not null"` + Flag string `gorm:"type:text;not null"` + UserID uint `gorm:"not null"` +} + +func (DynamicFlagClaim) TableName() string { + return "dynamic_flag_claims" +} + +type DynamicScoreDirty struct { + gorm.Model + + ChallengeID uint `gorm:"not null"` + LastSolveAt time.Time `gorm:"not null"` + LastSolverID uint `gorm:"not null"` +} + +func (DynamicScoreDirty) TableName() string { + return "dynamic_score_dirty" +} + +type SubmissionAttemptStatus uint8 + +const ( + SubmissionAttemptAccepted SubmissionAttemptStatus = iota + SubmissionAttemptAlreadySolved + SubmissionAttemptMaxAttempts +) + +type SubmissionAttemptResult struct { + Status SubmissionAttemptStatus + Tries uint +} + +type DynamicFlagClaimStatus uint8 + +const ( + DynamicFlagClaimCreated DynamicFlagClaimStatus = iota + DynamicFlagClaimedBySameUser + DynamicFlagClaimedByOtherUser +) + +type DynamicFlagClaimResult struct { + Status DynamicFlagClaimStatus + ClaimedByID uint +} + +func MigrateSubmissionGuards() error { + if err := dedupeUserChallengeRows(); err != nil { + return err + } + + statements := []string{ + `CREATE UNIQUE INDEX IF NOT EXISTS idx_user_challenges_user_challenge ON user_challenges (user_id, challenge_id)`, + `CREATE UNIQUE INDEX IF NOT EXISTS idx_dynamic_flag_claims_challenge_flag ON dynamic_flag_claims (challenge_id, flag)`, + `CREATE UNIQUE INDEX IF NOT EXISTS idx_dynamic_score_dirty_challenge ON dynamic_score_dirty (challenge_id)`, + } + + for _, statement := range statements { + if err := Db.Exec(statement).Error; err != nil { + return err + } + } + + return nil +} + +func dedupeUserChallengeRows() error { + return Db.Transaction(func(tx *gorm.DB) error { + merge := ` +WITH ranked AS ( + SELECT + id, + user_id, + challenge_id, + SUM(tries) OVER (PARTITION BY user_id, challenge_id) AS total_tries, + BOOL_OR(solved) OVER (PARTITION BY user_id, challenge_id) AS any_solved, + ROW_NUMBER() OVER ( + PARTITION BY user_id, challenge_id + ORDER BY solved DESC, created_at ASC, id ASC + ) AS rn + FROM user_challenges +), +keepers AS ( + SELECT id, total_tries, any_solved + FROM ranked + WHERE rn = 1 +) +UPDATE user_challenges uc +SET tries = keepers.total_tries, + solved = keepers.any_solved +FROM keepers +WHERE uc.id = keepers.id` + if err := tx.Exec(merge).Error; err != nil { + return fmt.Errorf("failed to merge duplicate user_challenges rows: %w", err) + } + + removeDuplicates := ` +WITH ranked AS ( + SELECT + id, + ROW_NUMBER() OVER ( + PARTITION BY user_id, challenge_id + ORDER BY solved DESC, created_at ASC, id ASC + ) AS rn + FROM user_challenges +) +DELETE FROM user_challenges uc +USING ranked +WHERE uc.id = ranked.id AND ranked.rn > 1` + if err := tx.Exec(removeDuplicates).Error; err != nil { + return fmt.Errorf("failed to delete duplicate user_challenges rows: %w", err) + } + + return nil + }) +} + +func ReserveSubmissionAttempt(userID, challengeID uint, maxAttemptLimit int, flag string, now time.Time) (SubmissionAttemptResult, error) { + var row struct { + ID uint + Tries uint + Solved bool + } + + tx := Db.Raw(` +INSERT INTO user_challenges (created_at, user_id, challenge_id, tries, solved, flag, cheating) +VALUES (?, ?, ?, 1, false, ?, false) +ON CONFLICT (user_id, challenge_id) DO UPDATE +SET tries = user_challenges.tries + 1, + flag = EXCLUDED.flag, + created_at = EXCLUDED.created_at +WHERE user_challenges.solved = false + AND (? <= 0 OR user_challenges.tries < ?) +RETURNING id, tries, solved`, + now, userID, challengeID, flag, maxAttemptLimit, maxAttemptLimit, + ).Scan(&row) + if tx.Error != nil { + return SubmissionAttemptResult{}, tx.Error + } + + if tx.RowsAffected > 0 { + return SubmissionAttemptResult{ + Status: SubmissionAttemptAccepted, + Tries: row.Tries, + }, nil + } + + var existing UserChallenges + err := Db.Where("user_id = ? AND challenge_id = ?", userID, challengeID).First(&existing).Error + if errors.Is(err, gorm.ErrRecordNotFound) { + return SubmissionAttemptResult{}, fmt.Errorf("submission attempt was not persisted") + } + if err != nil { + return SubmissionAttemptResult{}, err + } + + if existing.Solved { + return SubmissionAttemptResult{ + Status: SubmissionAttemptAlreadySolved, + Tries: existing.Tries, + }, nil + } + + return SubmissionAttemptResult{ + Status: SubmissionAttemptMaxAttempts, + Tries: existing.Tries, + }, nil +} + +func MarkSubmissionSolved(userID, challengeID uint, flag string, cheating bool, now time.Time) (bool, error) { + tx := Db.Model(&UserChallenges{}). + Where("user_id = ? AND challenge_id = ? AND solved = ?", userID, challengeID, false). + Updates(map[string]interface{}{ + "solved": true, + "flag": flag, + "cheating": cheating, + "created_at": now, + }) + if tx.Error != nil { + return false, tx.Error + } + + return tx.RowsAffected > 0, nil +} + +func MarkSubmissionCheating(userID, challengeID uint, flag string) error { + return Db.Model(&UserChallenges{}). + Where("user_id = ? AND challenge_id = ?", userID, challengeID). + Updates(map[string]interface{}{ + "flag": flag, + "cheating": true, + }).Error +} + +func ClaimDynamicFlag(challengeID, userID uint, flag string, now time.Time) (DynamicFlagClaimResult, error) { + var row struct { + UserID uint + } + + tx := Db.Raw(` +INSERT INTO dynamic_flag_claims (created_at, updated_at, challenge_id, flag, user_id) +VALUES (?, ?, ?, ?, ?) +ON CONFLICT (challenge_id, flag) DO NOTHING +RETURNING user_id`, + now, now, challengeID, flag, userID, + ).Scan(&row) + if tx.Error != nil { + return DynamicFlagClaimResult{}, tx.Error + } + + if tx.RowsAffected > 0 { + return DynamicFlagClaimResult{ + Status: DynamicFlagClaimCreated, + ClaimedByID: userID, + }, nil + } + + var claim DynamicFlagClaim + err := Db.Where("challenge_id = ? AND flag = ?", challengeID, flag).First(&claim).Error + if err != nil { + return DynamicFlagClaimResult{}, err + } + + if claim.UserID == userID { + return DynamicFlagClaimResult{ + Status: DynamicFlagClaimedBySameUser, + ClaimedByID: claim.UserID, + }, nil + } + + return DynamicFlagClaimResult{ + Status: DynamicFlagClaimedByOtherUser, + ClaimedByID: claim.UserID, + }, nil +} + +func AwardUserScore(userID uint, delta int64) error { + return Db.Model(&User{}). + Where("id = ?", userID). + UpdateColumn("score", gorm.Expr("CASE WHEN score + ? < 0 THEN 0 ELSE score + ? END", delta, delta)). + Error +} + +func MarkDynamicScoreDirty(challengeID, userID uint, solvedAt time.Time) error { + return Db.Exec(` +INSERT INTO dynamic_score_dirty (created_at, updated_at, challenge_id, last_solve_at, last_solver_id) +VALUES (?, ?, ?, ?, ?) +ON CONFLICT (challenge_id) DO UPDATE +SET updated_at = EXCLUDED.updated_at, + last_solve_at = EXCLUDED.last_solve_at, + last_solver_id = EXCLUDED.last_solver_id`, + solvedAt, solvedAt, challengeID, solvedAt, userID, + ).Error +} + +func QueryDirtyDynamicScores(limit int) ([]DynamicScoreDirty, error) { + var dirty []DynamicScoreDirty + err := Db.Order("updated_at ASC").Limit(limit).Find(&dirty).Error + return dirty, err +} + +func ClearDynamicScoreDirty(challengeID uint, seenUpdatedAt time.Time) error { + return Db.Unscoped(). + Where("challenge_id = ? AND updated_at <= ?", challengeID, seenUpdatedAt). + Delete(&DynamicScoreDirty{}). + Error +} + +func CountSolvedSubmissionsForChallenge(challengeID uint) (uint, error) { + var count int64 + err := Db.Table("user_challenges"). + Joins("JOIN users ON users.id = user_challenges.user_id"). + Where("user_challenges.challenge_id = ? AND user_challenges.solved = ? AND users.role = ?", challengeID, true, "contestant"). + Count(&count).Error + return uint(count), err +} + +func ApplyDynamicScoreDelta(challengeID uint, newPoints uint, delta int64) error { + return Db.Transaction(func(tx *gorm.DB) error { + if err := tx.Model(&Challenge{}).Where("id = ?", challengeID).Update("points", newPoints).Error; err != nil { + return err + } + + if delta == 0 { + return nil + } + + return tx.Exec(` +UPDATE users +SET score = CASE WHEN users.score + ? < 0 THEN 0 ELSE users.score + ? END +FROM user_challenges +WHERE users.id = user_challenges.user_id + AND user_challenges.challenge_id = ? + AND user_challenges.solved = true + AND users.role = ?`, + delta, delta, challengeID, "contestant", + ).Error + }) +} diff --git a/core/database/submission_test.go b/core/database/submission_test.go new file mode 100644 index 00000000..d5750a46 --- /dev/null +++ b/core/database/submission_test.go @@ -0,0 +1,355 @@ +package database + +import ( + "fmt" + "net/url" + "os" + "strings" + "sync" + "sync/atomic" + "testing" + "time" + + "github.com/sdslabs/beastv4/core" + "github.com/sdslabs/beastv4/pkg/auth" + "gorm.io/driver/postgres" + "gorm.io/gorm" +) + +func setupSubmissionTestDB(t *testing.T) func() { + t.Helper() + + dsn := os.Getenv("BEAST_TEST_PG_DSN") + if dsn == "" { + t.Skip("set BEAST_TEST_PG_DSN to run PostgreSQL submission race tests") + } + + adminDB, err := gorm.Open(postgres.Open(dsn), &gorm.Config{}) + if err != nil { + t.Fatalf("open admin postgres connection: %v", err) + } + + schema := fmt.Sprintf("beast_test_%d", time.Now().UnixNano()) + if err := adminDB.Exec(fmt.Sprintf(`CREATE SCHEMA "%s"`, schema)).Error; err != nil { + t.Fatalf("create test schema: %v", err) + } + + testDB, err := gorm.Open(postgres.Open(withSearchPath(dsn, schema)), &gorm.Config{}) + if err != nil { + _ = adminDB.Exec(fmt.Sprintf(`DROP SCHEMA "%s" CASCADE`, schema)).Error + t.Fatalf("open test postgres connection: %v", err) + } + + sqlDB, err := testDB.DB() + if err != nil { + t.Fatalf("get sql db: %v", err) + } + sqlDB.SetMaxOpenConns(32) + sqlDB.SetMaxIdleConns(32) + + previousDB := Db + previousMux := DBMux + Db = testDB + DBMux = &sync.Mutex{} + + if err := Db.AutoMigrate(&Challenge{}, &User{}, &UserChallenges{}, &DynamicFlag{}, &DynamicFlagClaim{}, &DynamicScoreDirty{}); err != nil { + t.Fatalf("auto migrate: %v", err) + } + if err := MigrateSubmissionGuards(); err != nil { + t.Fatalf("migrate submission guards: %v", err) + } + + return func() { + Db = previousDB + DBMux = previousMux + _ = sqlDB.Close() + _ = adminDB.Exec(fmt.Sprintf(`DROP SCHEMA "%s" CASCADE`, schema)).Error + if adminSQL, err := adminDB.DB(); err == nil { + _ = adminSQL.Close() + } + } +} + +func withSearchPath(dsn, schema string) string { + if strings.HasPrefix(dsn, "postgres://") || strings.HasPrefix(dsn, "postgresql://") { + parsed, err := url.Parse(dsn) + if err == nil { + query := parsed.Query() + query.Set("search_path", schema) + parsed.RawQuery = query.Encode() + return parsed.String() + } + } + + return dsn + " search_path=" + schema +} + +func createSubmissionTestUser(t *testing.T, username string) User { + t.Helper() + + user := User{ + Name: username, + Email: username + "@example.test", + AuthModel: auth.AuthModel{ + Username: username, + Role: core.USER_ROLES["contestant"], + }, + } + if err := Db.Create(&user).Error; err != nil { + t.Fatalf("create user %s: %v", username, err) + } + return user +} + +func createSubmissionTestChallenge(t *testing.T, name string, maxAttempts int, dynamic bool) Challenge { + t.Helper() + + challenge := Challenge{ + Name: name, + Type: "web", + Difficulty: "easy", + Flag: "flag{correct}", + DynamicFlag: dynamic, + Points: 500, + MaxPoints: 500, + MinPoints: 100, + MaxAttemptLimit: maxAttempts, + Status: core.DEPLOY_STATUS["deployed"], + } + if err := Db.Create(&challenge).Error; err != nil { + t.Fatalf("create challenge %s: %v", name, err) + } + return challenge +} + +func TestConcurrentCorrectSubmissionsAwardOnce(t *testing.T) { + cleanup := setupSubmissionTestDB(t) + defer cleanup() + + user := createSubmissionTestUser(t, "raceuser") + challenge := createSubmissionTestChallenge(t, "race-correct", -1, false) + + var awards int32 + errCh := make(chan error, 64) + start := make(chan struct{}) + var wg sync.WaitGroup + + for i := 0; i < 64; i++ { + wg.Add(1) + go func() { + defer wg.Done() + <-start + + now := time.Now() + attempt, err := ReserveSubmissionAttempt(user.ID, challenge.ID, challenge.MaxAttemptLimit, challenge.Flag, now) + if err != nil { + errCh <- err + return + } + if attempt.Status != SubmissionAttemptAccepted { + return + } + + won, err := MarkSubmissionSolved(user.ID, challenge.ID, challenge.Flag, false, now) + if err != nil { + errCh <- err + return + } + if !won { + return + } + + if err := AwardUserScore(user.ID, int64(challenge.Points)); err != nil { + errCh <- err + return + } + atomic.AddInt32(&awards, 1) + }() + } + + close(start) + wg.Wait() + close(errCh) + + for err := range errCh { + t.Fatalf("concurrent submission failed: %v", err) + } + if got := atomic.LoadInt32(&awards); got != 1 { + t.Fatalf("expected exactly one score award, got %d", got) + } + + var submissions []UserChallenges + if err := Db.Where("user_id = ? AND challenge_id = ?", user.ID, challenge.ID).Find(&submissions).Error; err != nil { + t.Fatalf("query submissions: %v", err) + } + if len(submissions) != 1 { + t.Fatalf("expected one user_challenges row, got %d", len(submissions)) + } + if !submissions[0].Solved { + t.Fatalf("expected submission row to be solved") + } + + var refreshed User + if err := Db.First(&refreshed, user.ID).Error; err != nil { + t.Fatalf("query user: %v", err) + } + if refreshed.Score != challenge.Points { + t.Fatalf("expected user score %d, got %d", challenge.Points, refreshed.Score) + } +} + +func TestConcurrentWrongSubmissionsRespectMaxAttempts(t *testing.T) { + cleanup := setupSubmissionTestDB(t) + defer cleanup() + + user := createSubmissionTestUser(t, "wronguser") + challenge := createSubmissionTestChallenge(t, "race-wrong", 3, false) + + var accepted int32 + errCh := make(chan error, 64) + start := make(chan struct{}) + var wg sync.WaitGroup + + for i := 0; i < 64; i++ { + wg.Add(1) + go func(i int) { + defer wg.Done() + <-start + + attempt, err := ReserveSubmissionAttempt(user.ID, challenge.ID, challenge.MaxAttemptLimit, fmt.Sprintf("wrong-%d", i), time.Now()) + if err != nil { + errCh <- err + return + } + if attempt.Status == SubmissionAttemptAccepted { + atomic.AddInt32(&accepted, 1) + } + }(i) + } + + close(start) + wg.Wait() + close(errCh) + + for err := range errCh { + t.Fatalf("concurrent wrong submission failed: %v", err) + } + if got := atomic.LoadInt32(&accepted); got != int32(challenge.MaxAttemptLimit) { + t.Fatalf("expected %d accepted attempts, got %d", challenge.MaxAttemptLimit, got) + } + + var submission UserChallenges + if err := Db.Where("user_id = ? AND challenge_id = ?", user.ID, challenge.ID).First(&submission).Error; err != nil { + t.Fatalf("query submission: %v", err) + } + if submission.Tries != uint(challenge.MaxAttemptLimit) { + t.Fatalf("expected tries %d, got %d", challenge.MaxAttemptLimit, submission.Tries) + } + if submission.Solved { + t.Fatalf("wrong submissions must not mark challenge solved") + } +} + +func TestDynamicFlagClaimFirstClaimWins(t *testing.T) { + cleanup := setupSubmissionTestDB(t) + defer cleanup() + + userA := createSubmissionTestUser(t, "claimusera") + userB := createSubmissionTestUser(t, "claimuserb") + challenge := createSubmissionTestChallenge(t, "dynamic-claim", -1, true) + + start := make(chan struct{}) + results := make(chan DynamicFlagClaimResult, 2) + errCh := make(chan error, 2) + for _, user := range []User{userA, userB} { + go func(user User) { + <-start + result, err := ClaimDynamicFlag(challenge.ID, user.ID, "flag{dynamic}", time.Now()) + if err != nil { + errCh <- err + return + } + results <- result + }(user) + } + + close(start) + + var got []DynamicFlagClaimResult + for len(got) < 2 { + select { + case err := <-errCh: + t.Fatalf("claim dynamic flag: %v", err) + case result := <-results: + got = append(got, result) + case <-time.After(5 * time.Second): + t.Fatalf("timed out waiting for dynamic flag claims") + } + } + + created := 0 + claimedByOther := 0 + for _, result := range got { + switch result.Status { + case DynamicFlagClaimCreated: + created++ + case DynamicFlagClaimedByOtherUser: + claimedByOther++ + default: + t.Fatalf("unexpected dynamic claim status %v", result.Status) + } + } + if created != 1 || claimedByOther != 1 { + t.Fatalf("expected one created and one duplicate claim, got created=%d duplicate=%d", created, claimedByOther) + } + + var claims []DynamicFlagClaim + if err := Db.Where("challenge_id = ? AND flag = ?", challenge.ID, "flag{dynamic}").Find(&claims).Error; err != nil { + t.Fatalf("query claims: %v", err) + } + if len(claims) != 1 { + t.Fatalf("expected one dynamic flag claim row, got %d", len(claims)) + } +} + +func TestDynamicScoreDirtyCoalescesConcurrentMarks(t *testing.T) { + cleanup := setupSubmissionTestDB(t) + defer cleanup() + + challenge := createSubmissionTestChallenge(t, "dynamic-score-dirty", -1, true) + user := createSubmissionTestUser(t, "dirtyuser") + + errCh := make(chan error, 64) + start := make(chan struct{}) + var wg sync.WaitGroup + + for i := 0; i < 64; i++ { + wg.Add(1) + go func() { + defer wg.Done() + <-start + if err := MarkDynamicScoreDirty(challenge.ID, user.ID, time.Now()); err != nil { + errCh <- err + } + }() + } + + close(start) + wg.Wait() + close(errCh) + + for err := range errCh { + t.Fatalf("mark dynamic score dirty: %v", err) + } + + dirty, err := QueryDirtyDynamicScores(10) + if err != nil { + t.Fatalf("query dirty dynamic scores: %v", err) + } + if len(dirty) != 1 { + t.Fatalf("expected one coalesced dirty-score marker, got %d", len(dirty)) + } + if dirty[0].ChallengeID != challenge.ID { + t.Fatalf("expected dirty marker for challenge %d, got %d", challenge.ID, dirty[0].ChallengeID) + } +} diff --git a/core/database/tag.go b/core/database/tag.go index 0408a1c3..376e9061 100644 --- a/core/database/tag.go +++ b/core/database/tag.go @@ -61,7 +61,7 @@ func QueryRelatedChallengesMetadata(tag *Tag) ([]Challenge, error) { Db.Where(&Tag{TagName: tag.TagName}).First(&tagName) if err := Db.Model(&tagName). - Select("id", "name", "created_at", "points", "difficulty"). + Select("id", "name", "created_at", "points", "difficulty", "instanced", "instance_expiration", "status"). Preload("Tags"). Association("Challenges"). Find(&challenges); err != nil { diff --git a/core/database/user.go b/core/database/user.go index 983bb93d..476a249f 100644 --- a/core/database/user.go +++ b/core/database/user.go @@ -164,6 +164,89 @@ func GetRelatedChallenges(user *User) ([]Challenge, error) { return challenges, nil } +// UserSolvedChallenge represents a challenge solved by a user with the actual solve timestamp +type UserSolvedChallenge struct { + ChallengeID uint + Name string + Type string + Points uint + Tags []*Tag + SolvedAt time.Time +} + +// GetUserSolvedChallenges returns distinct challenges solved by a user, with the actual solve timestamps. +// Unlike GetRelatedChallenges, this avoids duplicates from multiple user_challenges rows per challenge. +func GetUserSolvedChallenges(userID uint) ([]UserSolvedChallenge, error) { + type solveRow struct { + ChallengeID uint + Name string + Type string + Points uint + SolvedAt time.Time + } + var rows []solveRow + + DBMux.Lock() + defer DBMux.Unlock() + + err := Db.Table("user_challenges"). + Select("DISTINCT ON (user_challenges.challenge_id) user_challenges.challenge_id, challenges.name, challenges.type, challenges.points, user_challenges.created_at as solved_at"). + Joins("JOIN challenges ON challenges.id = user_challenges.challenge_id"). + Where("user_challenges.user_id = ? AND user_challenges.solved = ?", userID, true). + Order("user_challenges.challenge_id, user_challenges.created_at ASC"). + Scan(&rows).Error + if err != nil { + return nil, err + } + + if len(rows) == 0 { + return []UserSolvedChallenge{}, nil + } + + // Load tags for all solved challenges in a single query instead of one per challenge. + challengeIDs := make([]uint, len(rows)) + for i, r := range rows { + challengeIDs[i] = r.ChallengeID + } + + type tagRow struct { + ChallengeID uint + TagID uint + TagName string + } + var tagRows []tagRow + err = Db.Table("tags"). + Select("tag_challenges.challenge_id, tags.id as tag_id, tags.tag_name"). + Joins("JOIN tag_challenges ON tag_challenges.tag_id = tags.id"). + Where("tag_challenges.challenge_id IN ?", challengeIDs). + Scan(&tagRows).Error + if err != nil { + return nil, fmt.Errorf("failed to load tags for solved challenges: %w", err) + } + + tagMap := make(map[uint][]*Tag) + for _, tr := range tagRows { + tagMap[tr.ChallengeID] = append(tagMap[tr.ChallengeID], &Tag{ + Model: gorm.Model{ID: tr.TagID}, + TagName: tr.TagName, + }) + } + + results := make([]UserSolvedChallenge, 0, len(rows)) + for _, r := range rows { + results = append(results, UserSolvedChallenge{ + ChallengeID: r.ChallengeID, + Name: r.Name, + Type: r.Type, + Points: r.Points, + Tags: tagMap[r.ChallengeID], + SolvedAt: r.SolvedAt, + }) + } + + return results, nil +} + // Check whether challenge is submitted by the user func CheckPreviousSubmissions(userId uint, challId uint) (bool, error) { var userChallenges []UserChallenges @@ -419,12 +502,10 @@ func QueryAllUniqueTags() ([]string, error) { var tags []string DBMux.Lock() defer DBMux.Unlock() - + tx := Db.Model(&Challenge{}).Distinct().Pluck("tag", &tags) if tx.Error != nil { return nil, tx.Error } return tags, nil } - - diff --git a/core/manager/challenge.go b/core/manager/challenge.go index 72f63b9f..70680bab 100644 --- a/core/manager/challenge.go +++ b/core/manager/challenge.go @@ -3,6 +3,7 @@ package manager import ( "errors" "fmt" + "github.com/sdslabs/beastv4/core/cache" "path/filepath" "strings" @@ -57,11 +58,11 @@ func CommitChallengeContainer(challName string) error { return fmt.Errorf("challenge is not deployed") } var imageId string - if chall.ServerDeployed != core.LOCALHOST && chall.ServerDeployed != "" { + if config.Cfg.UseLocalDockerDaemon(chall.ServerDeployed) { + imageId, err = cr.CommitContainer(chall.ContainerId) + } else { server := config.Cfg.AvailableServers[chall.ServerDeployed] imageId, err = remoteManager.CommitContainerRemote(chall.ContainerId, server) - } else { - imageId, err = cr.CommitContainer(chall.ContainerId) } if err != nil { log.Errorf("Error while commiting the container : %s", err.Error()) @@ -181,17 +182,17 @@ func GetDeployWork(challengeName string) (*wpool.Task, error) { } } else if coreUtils.IsContainerIdValid(challenge.ContainerId) { var containers, remoteContainers []containerType.Container - if challenge.ServerDeployed != core.LOCALHOST && challenge.ServerDeployed != "" { - server := config.Cfg.AvailableServers[challenge.ServerDeployed] - remoteContainers, err = remoteManager.SearchRunningContainerByFilterRemote(map[string]string{"id": challenge.ContainerId}, server) + if config.Cfg.UseLocalDockerDaemon(challenge.ServerDeployed) { + containers, err = cr.SearchRunningContainerByFilter(map[string]string{"id": challenge.ContainerId}) if err != nil { - log.Errorf("error while searching for remote container with id %s", challenge.ContainerId) + log.Errorf("error while searching for container with id %s", challenge.ContainerId) return nil, errors.New("CONTAINER RUNTIME ERROR") } } else { - containers, err = cr.SearchRunningContainerByFilter(map[string]string{"id": challenge.ContainerId}) + server := config.Cfg.AvailableServers[challenge.ServerDeployed] + remoteContainers, err = remoteManager.SearchRunningContainerByFilterRemote(map[string]string{"id": challenge.ContainerId}, server) if err != nil { - log.Errorf("error while searching for container with id %s", challenge.ContainerId) + log.Errorf("error while searching for remote container with id %s", challenge.ContainerId) return nil, errors.New("CONTAINER RUNTIME ERROR") } } @@ -235,11 +236,12 @@ func GetDeployWork(challengeName string) (*wpool.Task, error) { if coreUtils.IsImageIdValid(challenge.ImageId) { var imageExist bool var err error - if challenge.ServerDeployed != core.LOCALHOST && challenge.ServerDeployed != "" { + log.Warnf("server: %s", challenge.ServerDeployed) + if config.Cfg.UseLocalDockerDaemon(challenge.ServerDeployed) { + imageExist, err = cr.CheckIfImageExists(challenge.ImageId) + } else { server := config.Cfg.AvailableServers[challenge.ServerDeployed] imageExist, err = remoteManager.CheckIfImageExistsOnRemote(challenge.ImageId, server) - } else { - imageExist, err = cr.CheckIfImageExists(challenge.ImageId) } if err != nil { log.Errorf("Error while searching for image with id %s: %s", challenge.ImageId, err) @@ -281,11 +283,11 @@ func GetDeployWork(challengeName string) (*wpool.Task, error) { // Check if the challenge is in staged state, it it is start the // pipeline from there on, else start deploy pipeline for the challenge // from remote - if challenge.ServerDeployed != core.LOCALHOST && challenge.ServerDeployed != "" { + if config.Cfg.UseLocalDockerDaemon(challenge.ServerDeployed) { + err = utils.ValidateFileExists(stagedFileName) + } else { server := config.Cfg.AvailableServers[challenge.ServerDeployed] err = remoteManager.ValidateFileRemoteExists(server, stagedFileName) - } else { - err = utils.ValidateFileExists(stagedFileName) } if err != nil { log.Infof("The requested challenge with Name %s is not already staged", challengeName) @@ -640,52 +642,71 @@ func undeployChallenge(challengeName string, purge bool) error { return fmt.Errorf("ChallengeName %s not valid", challengeName) } - if challenge.DeploymentType == core.DEPLOYMENT_TYPES["docker_compose"] { - log.Debugf("Detected Docker Compose deployment for challenge %s", challengeName) - - stagedDir := filepath.Join(core.BEAST_GLOBAL_DIR, core.BEAST_STAGING_DIR, challengeName) - if challenge.ServerDeployed != core.LOCALHOST && challenge.ServerDeployed != "" { - server := config.Cfg.AvailableServers[challenge.ServerDeployed] - - if !purge { - err = remoteManager.ComposeDownRemote(challengeName, stagedDir, server) + /* TODO: verify this cleanup */ + if challenge.Instanced { + // Kill all active instances of this challenge before undeploying + if err := KillChallengeInstances(challengeName); err != nil { + log.Warnf("Error killing instances for challenge %s: %v", challengeName, err) + // Continue with undeploy even if some instances failed to kill + } + } else { + if challenge.DeploymentType == core.DEPLOYMENT_TYPES["docker_compose"] { + log.Debugf("Detected Docker Compose deployment for challenge %s", challengeName) + + composeProjectName := utils.ProjectNameNotInstanced(challengeName) + if config.Cfg.UseLocalDockerDaemon(challenge.ServerDeployed) { + if !purge { + err = cr.ComposeDownProject(composeProjectName) + } else { + err = cr.ComposePurgeProject(composeProjectName) + } } else { - err = remoteManager.ComposePurgeRemote(challengeName, stagedDir, server) + server := config.Cfg.AvailableServers[challenge.ServerDeployed] + + if !purge { + err = remoteManager.ComposeDownProjectRemote(composeProjectName, server) + } else { + err = remoteManager.ComposePurgeProjectRemote(composeProjectName, server) + } + } + if err != nil { + log.Errorf("Error while removing challenge instance : %s", err) + return fmt.Errorf("error while removing challenge instance : %s", err) } } else { - if !purge { - err = cr.ComposeDown(challengeName, stagedDir) + // If a existing container ID is not found make sure that you atleast + // set the deploy status to undeployed. This earlier caused problem since if a challenge + // was in staging state(and deployed is cancled) then we can neither deploy new + // version nor we can undeploy the existing version(since it does not exist) + // So this.... + if challenge.ContainerId == coreUtils.GetTempContainerId(challengeName) { + log.Warnf("No instance of challenge(%s) deployed", challengeName) } else { - err = cr.ComposePurge(challengeName, stagedDir) + log.Debug("Removing challenge instance for ", challengeName) + if config.Cfg.UseLocalDockerDaemon(challenge.ServerDeployed) { + err = cr.StopAndRemoveContainer(challenge.ContainerId) + } else { + server := config.Cfg.AvailableServers[challenge.ServerDeployed] + err = remoteManager.StopAndRemoveContainerRemote(challenge.ContainerId, server) + } + if err != nil { + // This should not return from here, this should assume that + // the container instance does not exist and hence should update the database + // with the container ID. + p := fmt.Errorf("error while removing challenge instance : %s", err) + log.Error(p.Error()) + } } } - if err != nil { - log.Errorf("Error while removing challenge instance : %s", err) - return fmt.Errorf("error while removing challenge instance : %s", err) + + portOwner := challenge.ContainerId + if challenge.DeploymentType == core.DEPLOYMENT_TYPES["docker_compose"] { + portOwner = utils.ProjectNameNotInstanced(challengeName) } - } else { - // If a existing container ID is not found make sure that you atleast - // set the deploy status to undeployed. This earlier caused problem since if a challenge - // was in staging state(and deployed is cancled) then we can neither deploy new - // version nor we can undeploy the existing version(since it does not exist) - // So this.... - if challenge.ContainerId == coreUtils.GetTempContainerId(challengeName) { - log.Warnf("No instance of challenge(%s) deployed", challengeName) - } else { - log.Debug("Removing challenge instance for ", challengeName) - if challenge.ServerDeployed != core.LOCALHOST && challenge.ServerDeployed != "" { - server := config.Cfg.AvailableServers[challenge.ServerDeployed] - err = remoteManager.StopAndRemoveContainerRemote(challenge.ContainerId, server) - } else { - err = cr.StopAndRemoveContainer(challenge.ContainerId) - } - if err != nil { - // This should not return from here, this should assume that - // the container instance does not exist and hence should update the database - // with the container ID. - p := fmt.Errorf("error while removing challenge instance : %s", err) - log.Error(p.Error()) - } + + err = cache.FreeContainerPortsOnHost(challenge.ServerDeployed, portOwner) + if err != nil { + return fmt.Errorf("error while freeing ports for container %s on host %s: %s", portOwner, challenge.ServerDeployed, err) } } diff --git a/core/manager/health_check.go b/core/manager/health_check.go index 767dcd10..37cfeac2 100644 --- a/core/manager/health_check.go +++ b/core/manager/health_check.go @@ -1,12 +1,15 @@ package manager import ( + "context" "fmt" "path/filepath" "strings" + "sync" "time" "github.com/sdslabs/beastv4/core" + "github.com/sdslabs/beastv4/core/cache" "github.com/sdslabs/beastv4/core/config" "github.com/sdslabs/beastv4/core/database" "github.com/sdslabs/beastv4/pkg/cr" @@ -18,6 +21,10 @@ import ( ) var HEALTH_CHECKER = false +var ( + instanceCleanupOnce sync.Once + instanceExpirySubscriberOnce sync.Once +) // Check for static challenegs' assets to be present on staging server. // At the time of writing, Beast deploys assets to localhost only. @@ -42,8 +49,8 @@ func CheckStaticChallenge(chall database.Challenge) error { // Check for container running or not. func containerProber(chall database.Challenge) error { - challHost := chall.ServerDeployed - if challHost == core.LOCALHOST || challHost == "" { + serverDeployed := chall.ServerDeployed + if config.Cfg.UseLocalDockerDaemon(serverDeployed) { containers, err := cr.SearchRunningContainerByFilter(map[string]string{"id": chall.ContainerId}) if err != nil || len(containers) <= 0 { err = fmt.Errorf("error while searching for container with id %s on server: %s", chall.ContainerId, chall.ServerDeployed) @@ -86,8 +93,14 @@ func ChallengesHealthProber(waitTime int) { // Do a better job at health probing mechanism. if len(allocatedPorts) > 0 { port := int(allocatedPorts[0].PortNo) + serverDeployed := chall.ServerDeployed + if config.Cfg.UseLocalDockerDaemon(serverDeployed) { + serverDeployed = core.LOCALHOST + } else if s, ok := config.Cfg.AvailableServers[serverDeployed]; ok { + serverDeployed = s.Host + } prober := probes.NewTcpProber() - result, err := prober.Probe(chall.ServerDeployed, port, time.Duration(core.DEFAULT_PROBE_TIMEOUT)*time.Second) + result, err := prober.Probe(serverDeployed, port, time.Duration(core.DEFAULT_PROBE_TIMEOUT)*time.Second) if err != nil { msg := fmt.Sprintf("NETWORK HEALTH CHECK %s: %s : %s", result, chall.Name, err) log.WithFields(log.Fields{ @@ -127,8 +140,8 @@ func ChallengesHealthProber(waitTime int) { // Check for Remote Server running or not func ServerHealthProber(waitTime int) { - for _, server := range config.Cfg.AvailableServers { - if server.Active && server.Host != core.LOCALHOST { + for serverDeployed, server := range config.Cfg.AvailableServers { + if server.Active && !config.Cfg.UseLocalDockerDaemon(serverDeployed) { err := remoteManager.PingServer(server) if err != nil { msg := fmt.Sprintf("SERVER HEALTH CHECK Faliure: %s : %s", server.Host, err) @@ -145,19 +158,149 @@ func ServerHealthProber(waitTime int) { } } -// Check for beast services running or not func BeastHeathCheckProber(waitTime int) { if !HEALTH_CHECKER { log.Info("Starting Health Check prober.") HEALTH_CHECKER = true + + go InstanceCleanupProber() + for { go ChallengesHealthProber(waitTime) go ServerHealthProber(waitTime) go database.BackupDatabase() - // Wait for some time before next probing. + go cache.BackupCache() time.Sleep(time.Duration(waitTime) * time.Second) } } else { log.Warn("Health Checker Already Running. Not Starting Again") } } + +func InstanceCleanupProber() { + started := false + instanceCleanupOnce.Do(func() { + started = true + }) + if !started { + log.Warn("Instance cleanup prober already running. Not starting again") + return + } + + log.Info("Starting Instance Cleanup prober with event-driven expiry and reconciliation interval: ", core.DEFAULT_HEALTH_CHECK_TIME) + startInstanceExpirySubscriber() + + for { + ProcessInstanceDeletionQueue() + CleanupOrphanedInstanceContainers() + QueueExpiredInstances() + time.Sleep(core.DEFAULT_HEALTH_CHECK_TIME) + } +} + +func startInstanceExpirySubscriber() { + instanceExpirySubscriberOnce.Do(func() { + if err := cache.EnableKeyspaceExpiryNotifications(); err != nil { + log.Warnf("Redis keyspace expiry notifications unavailable, relying on reconciliation: %v", err) + } + + go func() { + err := cache.SubscribeExpiredInstanceMarkers(context.Background(), func(instanceID string) { + log.Infof("Instance expiry marker fired for %s, queueing cleanup", instanceID) + if err := cache.QueueInstanceForDeletion(instanceID); err != nil { + log.Warnf("Failed to queue expired instance %s from Redis event: %v", instanceID, err) + return + } + ProcessInstanceDeletionQueue() + }) + if err != nil { + log.Warnf("Redis expiry subscriber stopped, reconciliation will continue cleanup: %v", err) + } + }() + }) +} + +func QueueExpiredInstances() { + log.Debug("Checking for expired instances") + + expired, err := cache.GetExpiredInstances() + if err != nil { + log.Warnf("Failed to get expired instances: %v", err) + return + } + + for _, instance := range expired { + log.Infof("Instance %s expired (challenge: %s, user: %s), queueing for deletion", + instance.InstanceID, instance.ChallengeName, instance.UserID) + + err := cache.QueueInstanceForDeletion(instance.InstanceID) + if err != nil { + log.Warnf("Failed to queue instance %s for deletion: %v", instance.InstanceID, err) + } + } +} + +func ProcessInstanceDeletionQueue() { + log.Debug("Processing instance deletion queue") + + for i := 0; i < 10; i++ { + instance, err := cache.PopInstanceForDeletion() + if err != nil { + log.Warnf("Error popping from deletion queue: %v", err) + return + } + + if instance == nil { + return + } + + log.Infof("Processing deletion for instance %s (challenge: %s, container: %s, server: %s)", + instance.InstanceID, instance.ChallengeName, instance.ContainerID, instance.ServerDeployed) + + if _, err := cache.GetInstance(instance.InstanceID); err != nil { + log.Debugf("Skipping stale deletion queue item for instance %s: %v", instance.InstanceID, err) + continue + } + + err = killInstanceContainer(instance.ContainerID, instance.DeploymentType, instance.InstanceID, instance.ChallengeName, instance.ServerDeployed) + if err != nil { + log.Warnf("Failed to kill container for instance %s: %v", instance.InstanceID, err) + if restoreErr := cache.RestoreQueuedInstance(instance); restoreErr != nil { + log.Warnf("Failed to restore metadata for instance %s after cleanup failure: %v", instance.InstanceID, restoreErr) + } + continue + } + + log.Infof("Successfully killed container for instance %s", instance.InstanceID) + + if err := cache.FreeContainerPortsOnHost(instance.ServerDeployed, instance.PortOwnerID()); err != nil { + log.Warnf("Failed to free ports for instance %s: %v", instance.InstanceID, err) + if restoreErr := cache.RestoreQueuedInstance(instance); restoreErr != nil { + log.Warnf("Failed to restore metadata for instance %s after port cleanup failure: %v", instance.InstanceID, restoreErr) + } + continue + } + if err := cache.DeleteInstanceMetadata(instance.InstanceID); err != nil { + log.Warnf("Failed to delete metadata for instance %s: %v", instance.InstanceID, err) + } + } + + queueLen, _ := cache.GetDeletionQueueLength() + if queueLen > 0 { + log.Debugf("Deletion queue still has %d items, will process in next cycle", queueLen) + } +} + +func CleanupOrphanedInstanceContainers() { + log.Debug("Checking for orphaned instance containers") + + cr.CleanupOrphans() + cr.CleanupOrphanedComposeInstances() + + for serverDeployed, server := range config.Cfg.AvailableServers { + if server.Active && !config.Cfg.UseLocalDockerDaemon(serverDeployed) { + remoteManager.CleanupOrphanedOnServer(serverDeployed) + remoteManager.CleanupOrphanedComposeInstancesOnServer(serverDeployed) + } + } +} diff --git a/core/manager/instance.go b/core/manager/instance.go new file mode 100644 index 00000000..9e0cfdfe --- /dev/null +++ b/core/manager/instance.go @@ -0,0 +1,445 @@ +package manager + +import ( + "fmt" + "path/filepath" + "time" + + "github.com/BurntSushi/toml" + "github.com/google/uuid" + "github.com/sdslabs/beastv4/core" + "github.com/sdslabs/beastv4/core/cache" + cfg "github.com/sdslabs/beastv4/core/config" + "github.com/sdslabs/beastv4/core/database" + coreUtils "github.com/sdslabs/beastv4/core/utils" + "github.com/sdslabs/beastv4/pkg/cr" + "github.com/sdslabs/beastv4/pkg/remoteManager" + "github.com/sdslabs/beastv4/utils" + + log "github.com/sirupsen/logrus" +) + +func SpawnInstance(challengeName, userID, username string) (*cache.Instance, error) { + log.Infof("Spawning instance of challenge %s for user %s", challengeName, userID) + + existingInstance, err := cache.GetUserInstance(userID, challengeName) + if err == nil && existingInstance != nil { + return existingInstance, fmt.Errorf("user already has an active instance of this challenge") + } + + instanceCount, err := cache.CountUserInstances(userID) + if err != nil { + log.Warnf("Failed to count user instances: %v", err) + } else if instanceCount >= cfg.Cfg.InstanceConfig.MaxInstancesPerUser { + return nil, fmt.Errorf("maximum instances limit reached (%d)", cfg.Cfg.InstanceConfig.MaxInstancesPerUser) + } + + challenge, err := database.QueryFirstChallengeEntry("name", challengeName) + if err != nil { + return nil, fmt.Errorf("failed to query challenge: %w", err) + } + + challengeStagingDir := filepath.Join(core.BEAST_GLOBAL_DIR, core.BEAST_STAGING_DIR, challengeName) + configFile := filepath.Join(challengeStagingDir, core.CHALLENGE_CONFIG_FILE_NAME) + + var config cfg.BeastChallengeConfig + _, err = toml.DecodeFile(configFile, &config) + if err != nil { + return nil, fmt.Errorf("failed to load challenge config: %w", err) + } + + config.Resources.ValidateRequiredFields() + + if !config.Challenge.Metadata.IsInstanced() { + return nil, fmt.Errorf("challenge %s is not configured for instancing", challengeName) + } + + if challenge.ImageId == "" && config.Challenge.Env.DockerCompose == "" { + return nil, fmt.Errorf("challenge %s has not been committed (no image available)", challengeName) + } + + serverDeployed := selectServerForInstance() + + instanceID := uuid.New().String()[:12] + + expirationSeconds := config.Challenge.Metadata.GetInstanceExpiration() + ttl := time.Duration(expirationSeconds) * time.Second + expiresAt := time.Now().Add(ttl) + + var port uint32 + var containerID string + var portOwner string + var deploymentType string + + if config.Challenge.Env.DockerCompose != "" { + err = config.Challenge.Env.ExtractPortsCompose(challengeStagingDir) + if err != nil { + return nil, fmt.Errorf("failed to extract port variables from compose file: %s", err.Error()) + } + + ports, err := allocateInstancePortsCompose(serverDeployed, config.Challenge.Env) + if err != nil { + return nil, fmt.Errorf("failed to allocate instance ports: %s", err.Error()) + } + + port = ports[config.Challenge.Env.DefaultPortVar] + + containerID, err = deployInstanceFromCompose(instanceID, challengeName, &config, challengeStagingDir, serverDeployed, ports) + portOwner = utils.ComposeDockerProjectNameInstanced(challengeName, instanceID) + deploymentType = core.DEPLOYMENT_TYPES["docker_compose"] + + if err != nil { + coreUtils.FreePortsOnHostCompose(serverDeployed, ports) + return nil, err + } + + if err := coreUtils.AssignPortsOnContainerToHostCompose(serverDeployed, portOwner, ports); err != nil { + if cleanupErr := killInstanceContainer(containerID, deploymentType, instanceID, challengeName, serverDeployed); cleanupErr != nil { + log.Warnf("failed to cleanup instance %s after port registration failure: %v", instanceID, cleanupErr) + } + coreUtils.FreePortsOnHostCompose(serverDeployed, ports) + return nil, fmt.Errorf("failed to register instance ports: %w", err) + } + } else { + err = config.Challenge.Env.ExtractPorts() + if err != nil { + return nil, fmt.Errorf("failed to extract port variables from compose file: %s", err.Error()) + } + + ports, err := allocateInstancePorts(serverDeployed, config.Challenge.Env) + if err != nil { + return nil, fmt.Errorf("failed to allocate instance ports: %s", err.Error()) + } + + var found bool + found, port = utils.Uint32InIndexList(config.Challenge.Env.DefaultPort, config.Challenge.Env.Ports, ports) + if !found { + coreUtils.FreePortsOnHost(serverDeployed, ports) + return nil, fmt.Errorf("failed to allocate instance port for challenge %s", challengeName) + } + + containerID, err = deployInstanceContainer(instanceID, challengeName, challenge.ImageId, &config, serverDeployed, ports) + portOwner = containerID + deploymentType = core.DEPLOYMENT_TYPES["standard_docker"] + + if err != nil { + coreUtils.FreePortsOnHost(serverDeployed, ports) + return nil, fmt.Errorf("error while creating container for challenge %s: %s", challenge.Name, err.Error()) + } + + if err := coreUtils.AssignPortsOnContainerToHost(serverDeployed, containerID, ports); err != nil { + if cleanupErr := killInstanceContainer(containerID, deploymentType, instanceID, challengeName, serverDeployed); cleanupErr != nil { + log.Warnf("failed to cleanup instance %s after port registration failure: %v", instanceID, cleanupErr) + } + coreUtils.FreePortsOnHost(serverDeployed, ports) + return nil, fmt.Errorf("failed to register instance ports: %w", err) + } + } + + instance := &cache.Instance{ + InstanceID: instanceID, + ChallengeName: challengeName, + ContainerID: containerID, + PortOwner: portOwner, + Port: port, + UserID: userID, + Username: username, + CreatedAt: time.Now(), + ExpiresAt: expiresAt, + DeploymentType: deploymentType, + ServerDeployed: serverDeployed, + } + + err = cache.SaveInstance(instance, ttl) + if err != nil { + if err := killInstanceContainer(containerID, deploymentType, instanceID, challengeName, serverDeployed); err != nil { + return nil, fmt.Errorf("failed to kill instance container: %w", err) + } + if err := cache.FreeContainerPortsOnHost(serverDeployed, portOwner); err != nil { + return nil, fmt.Errorf("failed to free container ports: %w", err) + } + + return nil, fmt.Errorf("failed to save instance: %w", err) + } + + log.Infof("Successfully spawned instance %s for user %s, challenge %s on port %d (server: %s)", + instanceID, userID, challengeName, port, serverDeployed) + + return instance, nil +} + +func KillInstance(instanceID string) error { + log.Infof("Killing instance %s", instanceID) + + instance, err := cache.GetInstance(instanceID) + if err != nil { + return fmt.Errorf("instance not found: %w", err) + } + + err = killInstanceContainer(instance.ContainerID, instance.DeploymentType, instanceID, instance.ChallengeName, instance.ServerDeployed) + if err != nil { + log.Warnf("Error killing container for instance %s: %v", instanceID, err) + } + + err = cache.FreeContainerPortsOnHost(instance.ServerDeployed, instance.PortOwnerID()) + if err != nil { + return fmt.Errorf("failed to free container ports: %w", err) + } + + err = cache.DeleteInstance(instanceID) + if err != nil { + return fmt.Errorf("failed to delete instance from cache: %w", err) + } + + log.Infof("Successfully killed instance %s", instanceID) + return nil +} + +// KillUserInstance kills a user's instance of a specific challenge +func KillUserInstance(userID, challengeName string) error { + instance, err := cache.GetUserInstance(userID, challengeName) + if err != nil { + return fmt.Errorf("instance not found: %w", err) + } + + return KillInstance(instance.InstanceID) +} + +// ExtendInstance extends the lifetime of an instance +func ExtendInstance(instanceID string, additionalSeconds int64) error { + // Check max extension limit + maxExtension := cfg.Cfg.InstanceConfig.MaxExtension + if additionalSeconds > maxExtension { + additionalSeconds = maxExtension + } + + additionalTime := time.Duration(additionalSeconds) * time.Second + return cache.ExtendInstance(instanceID, additionalTime) +} + +// GetInstance retrieves an instance by ID +func GetInstance(instanceID string) (*cache.Instance, error) { + return cache.GetInstance(instanceID) +} + +// GetUserInstance retrieves a user's instance of a challenge +func GetUserInstance(userID, challengeName string) (*cache.Instance, error) { + return cache.GetUserInstance(userID, challengeName) +} + +// GetUserInstances retrieves all instances for a user +func GetUserInstances(userID string) ([]*cache.Instance, error) { + return cache.GetUserInstances(userID) +} + +// GetAllInstances retrieves all active instances (admin only) +func GetAllInstances() ([]*cache.Instance, error) { + return cache.GetAllInstances() +} + +// GetChallengeInstances retrieves all active instances for a specific challenge +func GetChallengeInstances(challengeName string) ([]*cache.Instance, error) { + return cache.GetChallengeInstances(challengeName) +} + +// KillChallengeInstances kills all active instances of a challenge. +// This should be called when undeploying or purging a challenge. +func KillChallengeInstances(challengeName string) error { + instances, err := cache.GetChallengeInstances(challengeName) + if err != nil { + return fmt.Errorf("failed to get instances for challenge %s: %w", challengeName, err) + } + + if len(instances) == 0 { + log.Debugf("No active instances found for challenge %s", challengeName) + return nil + } + + log.Infof("Killing %d active instance(s) for challenge %s", len(instances), challengeName) + + var lastErr error + for _, instance := range instances { + log.Infof("Killing instance %s for user %s (challenge: %s)", + instance.InstanceID, instance.UserID, challengeName) + + if err := KillInstance(instance.InstanceID); err != nil { + log.Warnf("Failed to kill instance %s: %v", instance.InstanceID, err) + lastErr = err + } + } + + return lastErr +} + +func allocateInstancePorts(host string, env cfg.ChallengeEnv) ([]uint32, error) { + var firstPort, lastPort uint32 + var err error + + server := cfg.Cfg.AvailableServers[host] + firstPort, lastPort, err = utils.ParsePortMapping(server.PortRange) + + if err != nil { + return nil, fmt.Errorf("failed to parse port range: %w", err) + } + + portRange := lastPort - firstPort + 1 + + ports, err := cache.GetFreePortsOnHost(host, firstPort, portRange, len(env.Ports)) + if err != nil { + return nil, fmt.Errorf("failed to allocate ports: %w", err) + } + + return ports, nil +} + +func allocateInstancePortsCompose(host string, env cfg.ChallengeEnv) (map[string]uint32, error) { + var err error + var firstPort, lastPort uint32 + + server := cfg.Cfg.AvailableServers[host] + firstPort, lastPort, err = utils.ParsePortMapping(server.PortRange) + + if err != nil { + return nil, fmt.Errorf("failed to parse port range: %w", err) + } + + portRange := lastPort - firstPort + 1 + + allocatedPorts, err := cache.GetFreePortsOnHost(host, firstPort, portRange, len(env.PortVariables)) + if err != nil { + return nil, fmt.Errorf("failed to allocate instance ports: %w", err) + } + + ports := make(map[string]uint32, len(env.PortVariables)) + for i, portVariable := range env.PortVariables { + ports[portVariable] = allocatedPorts[i] + } + + return ports, nil +} + +func selectServerForInstance() string { + availableServer, err := remoteManager.ServerQueue.GetNextAvailableInstance() + if err == nil && availableServer.Name != "" { + return availableServer.Name + } + return core.LOCALHOST +} + +func deployInstanceContainer(instanceID, challengeName string, imageID string, config *cfg.BeastChallengeConfig, serverDeployed string, ports []uint32) (string, error) { + // Instanced non compose challenges are managed by the container ID + containerName := utils.ComposeDockerProjectNameInstanced(challengeName, instanceID) + + containerPort := config.Challenge.Env.DefaultPort + if containerPort == 0 { + containerPort = 8080 + } + + portMapping := make([]cr.PortMapping, len(ports)) + for i, port := range ports { + portMapping[i] = cr.PortMapping{ + HostPort: port, + ContainerPort: config.Challenge.Env.Ports[i], + } + } + + var containerEnv []string + for _, env := range config.Challenge.Env.EnvironmentVars { + containerEnv = append(containerEnv, fmt.Sprintf("%s=%s", env.Key, filepath.Join(core.BEAST_DOCKER_CHALLENGE_DIR, env.Value))) + } + + containerConfig := cr.CreateContainerConfig{ + PortMapping: portMapping, + MountsMap: make(map[string]string), + ImageId: imageID, + ContainerName: containerName, + ChallengeName: challengeName, + ContainerEnv: containerEnv, + Traffic: config.Challenge.Env.TrafficType(), + CPUsLimit: config.Resources.CPUsLimit, + CPUShares: config.Resources.CPUShares, + Memory: config.Resources.Memory, + PidsLimit: config.Resources.PidsLimit, + Labels: map[string]string{ + "beast.instance": "true", + "beast.instance.id": instanceID, + }, + } + + var containerId string + var err error + + if cfg.Cfg.UseLocalDockerDaemon(serverDeployed) { + containerId, err = cr.CreateContainerFromImage(&containerConfig) + } else { + server := cfg.Cfg.AvailableServers[serverDeployed] + containerId, err = remoteManager.CreateContainerFromImageRemote(containerConfig, server) + } + + if err != nil { + return "", fmt.Errorf("failed to create container: %w", err) + } + + return containerId, nil +} + +func deployInstanceFromCompose(instanceID, challengeName string, config *cfg.BeastChallengeConfig, stagingDir string, serverDeployed string, ports map[string]uint32) (string, error) { + // Instanced compose challenges are managed by the projectName + projectName := utils.ComposeDockerProjectNameInstanced(challengeName, instanceID) + + if cfg.Cfg.UseLocalDockerDaemon(serverDeployed) { + primaryContainer, err := cr.DeployContainerFromCompose(challengeName, projectName, stagingDir, config.Challenge.Env.DockerCompose, ports) + if err != nil { + return "", fmt.Errorf("failed to deploy instance %s: %w", instanceID, err) + } + + return primaryContainer, nil + } else { + server := cfg.Cfg.AvailableServers[serverDeployed] + containerId, err := remoteManager.DeployContainerFromComposeRemote(challengeName, projectName, stagingDir, config.Challenge.Env.DockerCompose, server, ports) + if err != nil { + return "", fmt.Errorf("failed to deploy compose on remote: %w", err) + } + + return containerId, nil + } +} + +func killInstanceContainer(containerID, deploymentType, instanceID, challengeName, serverDeployed string) error { + if cfg.Cfg.UseLocalDockerDaemon(serverDeployed) { + if deploymentType == core.DEPLOYMENT_TYPES["docker_compose"] { + projectName := utils.ComposeDockerProjectNameInstanced(challengeName, instanceID) + + // managed by the projectName + err := cr.ComposePurgeProject(projectName) + if err != nil { + return fmt.Errorf("docker compose down failed: %s", err.Error()) + } + } else { + // managed by the containerID + err := cr.StopAndRemoveContainer(containerID) + if err != nil { + return fmt.Errorf("failed to stop container: %w", err) + } + } + } else { + server := cfg.Cfg.AvailableServers[serverDeployed] + if deploymentType == core.DEPLOYMENT_TYPES["docker_compose"] { + projectName := utils.ComposeDockerProjectNameInstanced(challengeName, instanceID) + + // managed by the projectName + err := remoteManager.ComposePurgeProjectRemote(projectName, server) + if err != nil { + return fmt.Errorf("failed to stop compose on remote: %w", err) + } + } else { + // managed by the containerID + err := remoteManager.StopAndRemoveContainerRemote(containerID, server) + if err != nil { + return fmt.Errorf("failed to stop container on remote: %w", err) + } + } + } + + return nil +} diff --git a/core/manager/pipeline.go b/core/manager/pipeline.go index f13eb40d..08a257f8 100644 --- a/core/manager/pipeline.go +++ b/core/manager/pipeline.go @@ -157,14 +157,21 @@ func commitChallenge(challenge *database.Challenge, config cfg.BeastChallengeCon ) imageId = "" - challengeTag := coreUtils.EncodeID(challengeName) + challengeTag := utils.EncodeID(challengeName) log.Printf("== Server for challenge %s : %s", challengeName, challenge.ServerDeployed) if config.Challenge.Env.DockerCompose != "" { // Should add some validation for the compose file + if cfg.Cfg.UseLocalDockerDaemon(challenge.ServerDeployed) { + var buff *bytes.Buffer - if challenge.ServerDeployed != core.LOCALHOST && challenge.ServerDeployed != "" { - + buff, buildErr = cr.BuildImagesFromCompose(challengeName, challengeTag, stagedPath, config.Challenge.Env.DockerCompose, noCache) + if buff != nil { + logBytes = buff.Bytes() + } else { + logBytes = []byte("BuildImagesFromCompose returned nil buffer") + } + } else { server := cfg.Cfg.AvailableServers[challenge.ServerDeployed] logBytes, buildErr = remoteManager.BuildImagesFromComposeRemote( challengeName, @@ -173,22 +180,21 @@ func commitChallenge(challenge *database.Challenge, config cfg.BeastChallengeCon server, noCache, ) - } else { - var buff *bytes.Buffer - - buff, buildErr = cr.BuildImagesFromCompose(challengeName, challengeTag, stagedPath, config.Challenge.Env.DockerCompose, noCache) - if buff != nil { - logBytes = buff.Bytes() - } else { - logBytes = []byte("BuildImagesFromCompose returned nil buffer") - } } // For Docker Compose challenges, ensure ImageId is empty in the database if err := database.UpdateChallenge(challenge, map[string]any{"ImageId": ""}); err != nil { return fmt.Errorf("error while setting empty ImageId for Docker Compose challenge: %s", err) } } else { - if challenge.ServerDeployed != core.LOCALHOST && challenge.ServerDeployed != "" { + if cfg.Cfg.UseLocalDockerDaemon(challenge.ServerDeployed) { + var buff *bytes.Buffer + buff, imageId, buildErr = cr.BuildImageFromTarContext(challengeName, challengeTag, stagedPath, config.Challenge.Env.DockerCtx, noCache) + if buff != nil { + logBytes = buff.Bytes() + } else { + logBytes = []byte("BuildImageFromTarContext returned nil buffer") + } + } else { server := cfg.Cfg.AvailableServers[challenge.ServerDeployed] stagedRemoteChallengePath := filepath.Join(core.BEAST_REMOTE_GLOBAL_DIR, core.BEAST_STAGING_DIR, challengeName) remoteStagedPath := filepath.Join(stagedRemoteChallengePath, fmt.Sprintf("%s.tar.gz", challengeName)) @@ -197,14 +203,6 @@ func commitChallenge(challenge *database.Challenge, config cfg.BeastChallengeCon return fmt.Errorf("error while checking if the challenge is staged on the remote server") } logBytes, imageId, buildErr = remoteManager.BuildImageFromTarContextRemote(challengeName, challengeTag, remoteStagedPath, server) - } else { - var buff *bytes.Buffer - buff, imageId, buildErr = cr.BuildImageFromTarContext(challengeName, challengeTag, stagedPath, config.Challenge.Env.DockerCtx, noCache) - if buff != nil { - logBytes = buff.Bytes() - } else { - logBytes = []byte("BuildImageFromTarContext returned nil buffer") - } } } // Create logs directory for the challenge in staging directory. @@ -259,6 +257,11 @@ func deployChallenge(challenge *database.Challenge, config cfg.BeastChallengeCon challengeName := config.Challenge.Metadata.Name stagingDir := filepath.Join(core.BEAST_GLOBAL_DIR, core.BEAST_STAGING_DIR, challengeName) + host := challenge.ServerDeployed + if host == "" { + host = core.LOCALHOST + } + if config.Challenge.Env.DockerCompose != "" { // currently the first container id returned var primaryContainerId string @@ -266,17 +269,38 @@ func deployChallenge(challenge *database.Challenge, config cfg.BeastChallengeCon composeFileName := config.Challenge.Env.DockerCompose - if challenge.ServerDeployed != core.LOCALHOST && challenge.ServerDeployed != "" { - server := cfg.Cfg.AvailableServers[challenge.ServerDeployed] - primaryContainerId, err = remoteManager.DeployContainerFromComposeRemote(challengeName, stagingDir, composeFileName, server) - if err != nil { - return fmt.Errorf("error while deploying challenge with docker-compose on remote: %v", err) - } + ports, err := allocateInstancePortsCompose(host, config.Challenge.Env) + if err != nil { + return fmt.Errorf("failed to allocate instance ports: %w", err) + } + + // Non instanced compose challenges are identified by the challenge name (without encoding) + composeProjectName := utils.ProjectNameNotInstanced(challengeName) + if cfg.Cfg.UseLocalDockerDaemon(host) { + primaryContainerId, err = cr.DeployContainerFromCompose(challengeName, composeProjectName, stagingDir, composeFileName, ports) } else { - primaryContainerId, err = cr.DeployContainerFromCompose(challengeName, stagingDir, composeFileName) - if err != nil { - return fmt.Errorf("error while deploying challenge with docker-compose: %v", err) + server := cfg.Cfg.AvailableServers[challenge.ServerDeployed] + primaryContainerId, err = remoteManager.DeployContainerFromComposeRemote(challengeName, composeProjectName, stagingDir, composeFileName, server, ports) + } + + if err != nil { + coreUtils.FreePortsOnHostCompose(host, ports) + return err + } + + if err := coreUtils.AssignPortsOnContainerToHostCompose(host, composeProjectName, ports); err != nil { + if cfg.Cfg.UseLocalDockerDaemon(host) { + if cleanupErr := cr.ComposePurgeProject(composeProjectName); cleanupErr != nil { + log.Warnf("failed to cleanup compose project %s after port registration failure: %v", composeProjectName, cleanupErr) + } + } else { + server := cfg.Cfg.AvailableServers[challenge.ServerDeployed] + if cleanupErr := remoteManager.ComposePurgeProjectRemote(composeProjectName, server); cleanupErr != nil { + log.Warnf("failed to cleanup remote compose project %s after port registration failure: %v", composeProjectName, cleanupErr) + } } + coreUtils.FreePortsOnHostCompose(host, ports) + return fmt.Errorf("error while registering ports for challenge %s: %s", challenge.Name, err) } // only for backward compatibility @@ -292,16 +316,16 @@ func deployChallenge(challenge *database.Challenge, config cfg.BeastChallengeCon staticMount := make(map[string]string) var staticMountDir string - if challenge.ServerDeployed == core.LOCALHOST || challenge.ServerDeployed == "" { + if cfg.Cfg.UseLocalDockerDaemon(challenge.ServerDeployed) { staticMountDir = filepath.Join(core.BEAST_GLOBAL_DIR, core.BEAST_STAGING_DIR, config.Challenge.Metadata.Name, core.BEAST_STATIC_FOLDER) } else { - staticMountDir = filepath.Join("$HOME/.beast", core.BEAST_STAGING_DIR, config.Challenge.Metadata.Name, core.BEAST_STATIC_FOLDER) + staticMountDir = filepath.Join(core.BEAST_REMOTE_GLOBAL_DIR, core.BEAST_STAGING_DIR, config.Challenge.Metadata.Name, core.BEAST_STATIC_FOLDER) } relativeStaticContentDir := config.Challenge.Env.StaticContentDir if relativeStaticContentDir == "" { relativeStaticContentDir = core.PUBLIC } - staticMount[staticMountDir] = filepath.Join("/challenge", relativeStaticContentDir) + staticMount[staticMountDir] = filepath.Join(core.BEAST_DOCKER_CHALLENGE_DIR, relativeStaticContentDir) log.Debugf("Static mount config for deploy : %s", staticMount) var containerEnv []string @@ -318,37 +342,58 @@ func deployChallenge(challenge *database.Challenge, config cfg.BeastChallengeCon config.Resources.PidsLimit, ) - portMapping, err := config.Challenge.Env.GetPortMappings() + ports, err := allocateInstancePorts(host, config.Challenge.Env) if err != nil { - return fmt.Errorf("error while parsing port mapping for the challenge %s: %s", config.Challenge.Metadata.Name, err) + return fmt.Errorf("failed to allocate instance ports: %s", err.Error()) + } + + portMapping := make([]cr.PortMapping, len(config.Challenge.Env.Ports)) + for i, hostPort := range ports { + portMapping[i] = cr.PortMapping{ + HostPort: hostPort, + ContainerPort: config.Challenge.Env.Ports[i], + } } + // Non instanced non compose challenges are managed by the containerID containerConfig := cr.CreateContainerConfig{ PortMapping: portMapping, MountsMap: staticMount, ImageId: challenge.ImageId, - ContainerName: coreUtils.EncodeID(config.Challenge.Metadata.Name), + ContainerName: utils.ProjectNameNotInstanced(config.Challenge.Metadata.Name), ContainerEnv: containerEnv, ContainerNetwork: containerNetwork, Traffic: config.Challenge.Env.TrafficType(), CPUShares: config.Resources.CPUShares, + CPUsLimit: config.Resources.CPUsLimit, Memory: config.Resources.Memory, PidsLimit: config.Resources.PidsLimit, } log.Debugf("create container config for challenge(%s): %v", config.Challenge.Metadata.Name, containerConfig) var containerId string - if challenge.ServerDeployed == core.LOCALHOST || challenge.ServerDeployed == "" { + if cfg.Cfg.UseLocalDockerDaemon(host) { containerId, err = cr.CreateContainerFromImage(&containerConfig) } else { - server := cfg.Cfg.AvailableServers[challenge.ServerDeployed] - containerId, err = remoteManager.CreateContainerFromImageRemote(containerConfig, server) + containerId, err = remoteManager.CreateContainerFromImageRemote(containerConfig, cfg.Cfg.AvailableServers[host]) } if err != nil { - if containerId != "" { - return fmt.Errorf("error while starting the container : %s", err) + coreUtils.FreePortsOnHost(host, ports) + return fmt.Errorf("error while creating container for challenge %s: %s", challenge.Name, err.Error()) + } + + if err := coreUtils.AssignPortsOnContainerToHost(host, containerId, ports); err != nil { + if cfg.Cfg.UseLocalDockerDaemon(host) { + if cleanupErr := cr.StopAndRemoveContainer(containerId); cleanupErr != nil { + log.Warnf("failed to cleanup container %s after port registration failure: %v", containerId, cleanupErr) + } + } else { + if cleanupErr := remoteManager.StopAndRemoveContainerRemote(containerId, cfg.Cfg.AvailableServers[host]); cleanupErr != nil { + log.Warnf("failed to cleanup remote container %s after port registration failure: %v", containerId, cleanupErr) + } } - return fmt.Errorf("error while trying to create a container for the challenge: %s", err) + coreUtils.FreePortsOnHost(host, ports) + return fmt.Errorf("error while registering ports for challenge %s: %s", challenge.Name, err) } if err = database.UpdateChallenge(challenge, map[string]any{ @@ -472,13 +517,13 @@ func bootstrapDeployPipeline(challengeDir string, skipStage bool, skipCommit boo database.UpdateChallenge(&challenge, map[string]interface{}{"status": core.DEPLOY_STATUS["undeployed"]}) return fmt.Errorf("STAGING ERROR: %s : %s", challengeName, err) } - if challenge.ServerDeployed != core.LOCALHOST && challenge.ServerDeployed != "" { + if !cfg.Cfg.UseLocalDockerDaemon(challenge.ServerDeployed) { remoteManager.StageChallRemote(cfg.Cfg.AvailableServers[challenge.ServerDeployed], challenge) } } else { log.Debugf("Checking if challenge already staged") - if challenge.ServerDeployed != core.LOCALHOST && challenge.ServerDeployed != "" { - err := remoteManager.ValidateFileRemoteExists(cfg.Cfg.AvailableServers[challenge.ServerDeployed], stagedRemoteChallengePath) + if cfg.Cfg.UseLocalDockerDaemon(challenge.ServerDeployed) { + err = utils.ValidateFileExists(stagedChallengePath) if err != nil { msg := "Challenge not already in staged(but skipping asked), could not proceed further" log.WithFields(log.Fields{ @@ -488,7 +533,7 @@ func bootstrapDeployPipeline(challengeDir string, skipStage bool, skipCommit boo return fmt.Errorf("STAGING ERROR: %s : %s", challengeName, msg) } } else { - err = utils.ValidateFileExists(stagedChallengePath) + err = remoteManager.ValidateFileRemoteExists(cfg.Cfg.AvailableServers[challenge.ServerDeployed], stagedRemoteChallengePath) if err != nil { msg := "Challenge not already in staged(but skipping asked), could not proceed further" log.WithFields(log.Fields{ @@ -522,6 +567,12 @@ func bootstrapDeployPipeline(challengeDir string, skipStage bool, skipCommit boo log.Debugf("Skipping commit phase") } + if challenge.Instanced { + database.UpdateChallenge(&challenge, map[string]interface{}{"status": core.DEPLOY_STATUS["deployed"]}) + log.Infof("Challenge %s is instanced, skipping deploy stage", challengeName) + return nil + } + database.UpdateChallenge(&challenge, map[string]interface{}{"status": core.DEPLOY_STATUS["deploying"]}) err = deployChallenge(&challenge, config) diff --git a/core/manager/sync.go b/core/manager/sync.go index ef47b72b..12c7c185 100644 --- a/core/manager/sync.go +++ b/core/manager/sync.go @@ -67,7 +67,6 @@ func SyncBeastRemote(defaultauthorpassword string) error { } } log.Info("Beast git base synced with remote") - go config.UpdateUsedPortList() UpdateChallenges(defaultauthorpassword) return fmt.Errorf("%s", strings.Join(errStrings, "\n")) } @@ -187,7 +186,6 @@ func SyncAndGetChangesFromRemote(defaultauthorpassword string) []string { } } log.Info("Beast git base synced with remote") - go config.UpdateUsedPortList() UpdateChallenges(defaultauthorpassword) return modifiedChallsNameList diff --git a/core/manager/utils.go b/core/manager/utils.go index 9861d251..497b8a14 100644 --- a/core/manager/utils.go +++ b/core/manager/utils.go @@ -4,6 +4,7 @@ import ( "archive/zip" "bytes" "fmt" + "github.com/sdslabs/beastv4/core/cache" "io" "io/ioutil" "os" @@ -491,7 +492,7 @@ func UpdateOrCreateChallengeDbEntry(challEntry *database.Challenge, config cfg.B availableServerHostname := core.LOCALHOST if config.Challenge.Metadata.Type != core.STATIC_CHALLENGE_TYPE_NAME { availableServer, _ := remoteManager.ServerQueue.GetNextAvailableInstance() - availableServerHostname = availableServer.Host + availableServerHostname = availableServer.Name } if config.Challenge.Metadata.Difficulty == "" { log.Debug("Setting difficulty to default(medium)") @@ -504,26 +505,28 @@ func UpdateOrCreateChallengeDbEntry(challEntry *database.Challenge, config cfg.B } *challEntry = database.Challenge{ - Name: config.Challenge.Metadata.Name, - AuthorID: userEntry.ID, - Format: config.Challenge.Metadata.Type, - Status: core.DEPLOY_STATUS["undeployed"], - ContainerId: coreUtils.GetTempContainerId(config.Challenge.Metadata.Name), - ImageId: coreUtils.GetTempImageId(config.Challenge.Metadata.Name), - MaxAttemptLimit: config.Challenge.Metadata.MaxAttemptLimit, - PreReqs: strings.Join(config.Challenge.Metadata.PreReqs, core.DELIMITER), - DynamicFlag: config.Challenge.Metadata.DynamicFlag, - Flag: config.Challenge.Metadata.Flag, - Type: config.Challenge.Metadata.Type, - Description: config.Challenge.Metadata.Description, - Assets: strings.Join(assetsURL, core.DELIMITER), - AdditionalLinks: strings.Join(config.Challenge.Metadata.AdditionalLinks, core.DELIMITER), - Points: config.Challenge.Metadata.Points, - MinPoints: config.Challenge.Metadata.MinPoints, - MaxPoints: config.Challenge.Metadata.MaxPoints, - Difficulty: config.Challenge.Metadata.Difficulty, - ServerDeployed: availableServerHostname, - DeploymentType: deploymentType, + Name: config.Challenge.Metadata.Name, + AuthorID: userEntry.ID, + Format: config.Challenge.Metadata.Type, + Status: core.DEPLOY_STATUS["undeployed"], + ContainerId: coreUtils.GetTempContainerId(config.Challenge.Metadata.Name), + ImageId: coreUtils.GetTempImageId(config.Challenge.Metadata.Name), + MaxAttemptLimit: config.Challenge.Metadata.MaxAttemptLimit, + PreReqs: strings.Join(config.Challenge.Metadata.PreReqs, core.DELIMITER), + DynamicFlag: config.Challenge.Metadata.DynamicFlag, + Flag: config.Challenge.Metadata.Flag, + Type: config.Challenge.Metadata.Type, + Description: config.Challenge.Metadata.Description, + Assets: strings.Join(assetsURL, core.DELIMITER), + AdditionalLinks: strings.Join(config.Challenge.Metadata.AdditionalLinks, core.DELIMITER), + Points: config.Challenge.Metadata.Points, + MinPoints: config.Challenge.Metadata.MinPoints, + MaxPoints: config.Challenge.Metadata.MaxPoints, + Difficulty: config.Challenge.Metadata.Difficulty, + ServerDeployed: availableServerHostname, + DeploymentType: deploymentType, + Instanced: config.Challenge.Metadata.Instanced, + InstanceExpiration: config.Challenge.Metadata.InstanceExpiration, } err = database.CreateChallengeEntry(challEntry) @@ -566,37 +569,44 @@ func UpdateOrCreateChallengeDbEntry(challEntry *database.Challenge, config cfg.B return false } - hostPorts, err := config.Challenge.Env.GetAllHostPorts() - if err != nil { - return fmt.Errorf("error while parsing host port for challenge %s : %s", challEntry.Name, err) - } - // Once the challenge entry has been created, add entries to the ports - // table in the database with the ports to expose - // for the challenge. - // TODO: Do all this under a database transaction so that if any port - // request is not available - for _, port := range hostPorts { - if isAllocated(port) { - // The port has already been allocated to the challenge - // Do nothing for this. - continue - } - - portEntry := database.Port{ - ChallengeID: challEntry.ID, - PortNo: port, + if challEntry.ContainerId != "" { + portOwner := challEntry.ContainerId + if challEntry.DeploymentType == core.DEPLOYMENT_TYPES["docker_compose"] { + portOwner = utils.ProjectNameNotInstanced(challEntry.Name) } - gotPort, err := database.PortEntryGetOrCreate(&portEntry) + hostPorts, err := cache.GetContainerPortsOnHost(challEntry.ServerDeployed, portOwner) if err != nil { - return err - } + return fmt.Errorf("error while parsing host port for challenge %s : %s", challEntry.Name, err) + } + // Once the challenge entry has been created, add entries to the ports + // table in the database with the ports to expose + // for the challenge. + // TODO: Do all this under a database transaction so that if any port + // request is not available + for _, port := range hostPorts { + if isAllocated(port) { + // The port has already been allocated to the challenge + // Do nothing for this. + continue + } - // var gotChall database.Challenge - // database.Db.Model(&gotPort).Related(&gotChall) + portEntry := database.Port{ + ChallengeID: challEntry.ID, + PortNo: port, + } + + gotPort, err := database.PortEntryGetOrCreate(&portEntry) + if err != nil { + return err + } - if gotPort.ChallengeID != challEntry.ID { - return fmt.Errorf("the port %d requested is already in use by another challenge", gotPort.PortNo) + // var gotChall database.Challenge + // database.Db.Model(&gotPort).Related(&gotChall) + + if gotPort.ChallengeID != challEntry.ID { + return fmt.Errorf("the port %d requested is already in use by another challenge", gotPort.PortNo) + } } } diff --git a/core/utils/cleanup.go b/core/utils/cleanup.go index 97388635..b30afd3f 100644 --- a/core/utils/cleanup.go +++ b/core/utils/cleanup.go @@ -4,6 +4,7 @@ import ( "fmt" container_types "github.com/docker/docker/api/types" "github.com/sdslabs/beastv4/core" + "github.com/sdslabs/beastv4/core/cache" "github.com/sdslabs/beastv4/core/config" cfg "github.com/sdslabs/beastv4/core/config" "github.com/sdslabs/beastv4/core/database" @@ -64,9 +65,10 @@ func CleanupContainerByFilter(filter, filterVal string) error { func CleanupChallengeContainers(chall *database.Challenge, config cfg.BeastChallengeConfig) error { if chall.DeploymentType == core.DEPLOYMENT_TYPES["docker_compose"] { log.Debugf("Cleaning up Docker Compose challenge: %s", chall.Name) - projectName := utils.GetProjectName(chall.Name) + // Same -p as deployPipeline / ComposeDown for non-instanced compose. + projectName := utils.ProjectNameNotInstanced(chall.Name) - if chall.ServerDeployed != core.LOCALHOST && chall.ServerDeployed != "" { + if !cfg.Cfg.UseLocalDockerDaemon(chall.ServerDeployed) { server := cfg.Cfg.AvailableServers[chall.ServerDeployed] downCommand := fmt.Sprintf("docker compose -p %s down", projectName) _, err := remoteManager.RunCommandOnServer(server, downCommand) @@ -74,9 +76,15 @@ func CleanupChallengeContainers(chall *database.Challenge, config cfg.BeastChall log.Errorf("Error running docker compose down on remote: %v", err) return err } + } else if err := cr.ComposeDownProject(projectName); err != nil { + log.Errorf("Error running docker compose down locally: %v", err) + return err } database.UpdateChallenge(chall, map[string]any{"ContainerId": GetTempContainerId(chall.Name)}) + if err := cache.FreeContainerPortsOnHost(chall.ServerDeployed, projectName); err != nil { + log.Warnf("Failed to free ports for compose challenge %s: %v", chall.Name, err) + } return nil } @@ -87,24 +95,27 @@ func CleanupChallengeContainers(chall *database.Challenge, config cfg.BeastChall } database.UpdateChallenge(chall, map[string]any{"ContainerId": GetTempContainerId(chall.Name)}) + if err := cache.FreeContainerPortsOnHost(chall.ServerDeployed, chall.ContainerId); err != nil { + log.Warnf("Failed to free ports for challenge %s: %v", chall.Name, err) + } } - err := CleanupContainerByFilter("name", EncodeID(config.Challenge.Metadata.Name)) + err := CleanupContainerByFilter("name", utils.EncodeID(config.Challenge.Metadata.Name)) return err } func CleanupChallengeImage(chall *database.Challenge) error { - if chall.ServerDeployed != core.LOCALHOST && chall.ServerDeployed != "" { - server := config.Cfg.AvailableServers[chall.ServerDeployed] - err := remoteManager.RemoveImageRemote(chall.ImageId, server) + if cfg.Cfg.UseLocalDockerDaemon(chall.ServerDeployed) { + err := cr.RemoveImage(chall.ImageId) if err != nil { - log.Errorf("Error while cleaning up image on remote %s with id %s", chall.ServerDeployed, chall.ImageId) + log.Errorf("Error while cleaning up image with id %s", chall.ImageId) return err } } else { - err := cr.RemoveImage(chall.ImageId) + server := config.Cfg.AvailableServers[chall.ServerDeployed] + err := remoteManager.RemoveImageRemote(chall.ImageId, server) if err != nil { - log.Errorf("Error while cleaning up image with id %s", chall.ImageId) + log.Errorf("Error while cleaning up image on remote %s with id %s", chall.ServerDeployed, chall.ImageId) return err } } diff --git a/core/utils/id.go b/core/utils/id.go index 68a726f0..7053543a 100644 --- a/core/utils/id.go +++ b/core/utils/id.go @@ -1,7 +1,6 @@ package utils import ( - "crypto/sha256" "fmt" "strings" @@ -24,10 +23,6 @@ func GetTempContainerId(a string) string { return b } -func EncodeID(a string) string { - return fmt.Sprintf("%x", sha256.Sum256([]byte(a)))[:30] -} - func IsImageIdValid(a string) bool { return (!strings.HasPrefix(a, core.IMAGE_NA) && a != "") } diff --git a/core/utils/logs.go b/core/utils/logs.go index af7d1a45..956358e7 100644 --- a/core/utils/logs.go +++ b/core/utils/logs.go @@ -47,21 +47,22 @@ func GetLogs(challname string, live bool) (*cr.Log, error) { } if live { - if chall.ServerDeployed != core.LOCALHOST && chall.ServerDeployed != "" { + if config.Cfg.UseLocalDockerDaemon(chall.ServerDeployed) { + cr.ShowLiveContainerLogs(chall.ContainerId) + } else { server := config.Cfg.AvailableServers[chall.ServerDeployed] remoteManager.ShowLiveContainerLogsRemote(chall.ContainerId, server) - } else { - cr.ShowLiveContainerLogs(chall.ContainerId) } return nil, nil } - if chall.ServerDeployed != core.LOCALHOST && chall.ServerDeployed != "" { - server := config.Cfg.AvailableServers[chall.ServerDeployed] - return remoteManager.GetContainerStdLogsRemote(chall.ContainerId, server) + + if config.Cfg.UseLocalDockerDaemon(chall.ServerDeployed) { + return cr.GetContainerStdLogs(chall.ContainerId) } - return cr.GetContainerStdLogs(chall.ContainerId) -} + server = config.Cfg.AvailableServers[chall.ServerDeployed] + return remoteManager.GetContainerStdLogsRemote(chall.ContainerId, server) +} func LogFlag(msg string, challName string) error { // log the cheating attempt in a file in cheat.log diff --git a/core/utils/ports.go b/core/utils/ports.go new file mode 100644 index 00000000..c7953375 --- /dev/null +++ b/core/utils/ports.go @@ -0,0 +1,47 @@ +package utils + +import ( + "github.com/sdslabs/beastv4/core/cache" + log "github.com/sirupsen/logrus" +) + +func AssignPortsOnContainerToHost(serverDeployed string, containerID string, ports []uint32) error { + /* Failure should be treated as fatal since this can lead to a leak... */ + if err := cache.AssignPortsOnHostToContainer(serverDeployed, containerID, ports); err != nil { + log.Warnf("Failed to register ports %v for container %s: %v", ports, containerID, err) + return err + } + + return nil +} + +func FreePortsOnHost(serverDeployed string, ports []uint32) { + for _, port := range ports { + if err := cache.FreePortOnHost(serverDeployed, port); err != nil { + log.Warnf("Failed to free port %d for host %s: %v", port, serverDeployed, err) + } + } +} + +func AssignPortsOnContainerToHostCompose(serverDeployed string, containerID string, ports map[string]uint32) error { + portList := make([]uint32, 0, len(ports)) + for _, port := range ports { + portList = append(portList, port) + } + + /* Failure should be treated as fatal since this can lead to a leak... */ + if err := cache.AssignPortsOnHostToContainer(serverDeployed, containerID, portList); err != nil { + log.Warnf("Failed to register ports %v for container %s: %v", portList, containerID, err) + return err + } + + return nil +} + +func FreePortsOnHostCompose(serverDeployed string, ports map[string]uint32) { + for _, port := range ports { + if err := cache.FreePortOnHost(serverDeployed, port); err != nil { + log.Warnf("Failed to free port %d for host %s: %v", port, serverDeployed, err) + } + } +} diff --git a/go.mod b/go.mod index cb7ee2f7..21e0d0cd 100644 --- a/go.mod +++ b/go.mod @@ -12,7 +12,6 @@ require ( github.com/gin-contrib/cors v1.3.1 github.com/gin-contrib/static v0.0.0-20200916080430-d45d9a37d28e github.com/gin-gonic/gin v1.7.0 - github.com/golang/protobuf v1.3.3 github.com/google/uuid v1.6.0 github.com/jackc/pgx/v5 v5.5.5 github.com/jinzhu/gorm v1.9.1 @@ -20,6 +19,7 @@ require ( github.com/manifoldco/promptui v0.9.0 github.com/mohae/struct2csv v0.0.0-20151122200941-e72239694eae github.com/olekukonko/tablewriter v0.0.5 + github.com/redis/go-redis/v9 v9.17.3 github.com/sirupsen/logrus v1.9.0 github.com/spf13/cobra v0.0.3 github.com/swaggo/gin-swagger v1.0.0 @@ -27,7 +27,6 @@ require ( golang.org/x/crypto v0.29.0 golang.org/x/net v0.31.0 golang.org/x/term v0.26.0 - google.golang.org/grpc v1.19.0 gopkg.in/src-d/go-git.v4 v4.7.0 gorm.io/driver/postgres v1.5.11 gorm.io/gorm v1.25.10 @@ -40,6 +39,7 @@ require ( github.com/alcortesm/tgz v0.0.0-20161220082320-9c5fe88206d7 // indirect github.com/anmitsu/go-shlex v0.0.0-20161002113705-648efa622239 // indirect github.com/aymanbagabas/go-osc52/v2 v2.0.1 // indirect + github.com/cespare/xxhash/v2 v2.3.0 // indirect github.com/charmbracelet/colorprofile v0.2.3-0.20250311203215-f60798e515dc // indirect github.com/charmbracelet/lipgloss v1.1.0 // indirect github.com/charmbracelet/x/ansi v0.8.0 // indirect @@ -48,6 +48,7 @@ require ( github.com/chzyer/readline v0.0.0-20180603132655-2972be24d48e // indirect github.com/cpuguy83/go-md2man v1.0.10 // indirect github.com/denisenkom/go-mssqldb v0.0.0-20180901172138-1eb28afdf9b6 // indirect + github.com/dgryski/go-rendezvous v0.0.0-20200823014737-9f7001d12a5f // indirect github.com/docker/distribution v2.8.1+incompatible // indirect github.com/docker/go-units v0.5.0 // indirect github.com/dustin/go-humanize v1.0.1 // indirect @@ -65,6 +66,7 @@ require ( github.com/go-playground/validator/v10 v10.4.1 // indirect github.com/go-sql-driver/mysql v1.4.0 // indirect github.com/gogo/protobuf v1.3.2 // indirect + github.com/golang/protobuf v1.3.3 // indirect github.com/inconshreveable/mousetrap v1.0.0 // indirect github.com/jackc/pgpassfile v1.0.0 // indirect github.com/jackc/pgservicefile v0.0.0-20221227161230-091c0ba34f0a // indirect @@ -109,7 +111,6 @@ require ( golang.org/x/time v0.3.0 // indirect golang.org/x/tools v0.27.0 // indirect google.golang.org/appengine v1.2.0 // indirect - google.golang.org/genproto v0.0.0-20180817151627-c66870c02cf8 // indirect gopkg.in/src-d/go-billy.v4 v4.3.0 // indirect gopkg.in/src-d/go-git-fixtures.v3 v3.3.0 // indirect gopkg.in/warnings.v0 v0.1.2 // indirect diff --git a/go.sum b/go.sum index 410c894b..cf05334d 100644 --- a/go.sum +++ b/go.sum @@ -1,9 +1,7 @@ -cloud.google.com/go v0.26.0/go.mod h1:aQUYkXzVsufM+DwF1aE+0xfcU+56JwCaLick0ClmMTw= cloud.google.com/go v0.28.0 h1:KZ/88LWSw8NxMkjdQyX7LQSGR9PkHr4PaVuNm8zgFq0= cloud.google.com/go v0.28.0/go.mod h1:aQUYkXzVsufM+DwF1aE+0xfcU+56JwCaLick0ClmMTw= github.com/Azure/go-ansiterm v0.0.0-20210617225240-d185dfc1b5a1 h1:UQHMgLO+TxOElx5B5HZ4hJQsoJ/PvUvKRhJHDQXO8P8= github.com/Azure/go-ansiterm v0.0.0-20210617225240-d185dfc1b5a1/go.mod h1:xomTg63KZ2rFqZQzSB4Vz2SUXa1BpHTVz9L5PTmPC4E= -github.com/BurntSushi/toml v0.3.1/go.mod h1:xHWCNGjB5oqiDr8zfno3MHue2Ht5sIBksp03qcyfWMU= github.com/BurntSushi/toml v1.4.0 h1:kuoIxZQy2WRRk1pttg9asf+WVv6tWQuBNVmK8+nqPr0= github.com/BurntSushi/toml v1.4.0/go.mod h1:ukJfTF/6rtPPRCnwkur4qwRxa8vTRFBF0uk2lLoLwho= github.com/KyleBanks/depth v1.2.1 h1:5h8fQADFrWtarTdtDudMmGsC7GPbOAu6RVB3ffsVFHc= @@ -20,6 +18,12 @@ github.com/araddon/dateparse v0.0.0-20210429162001-6b43995a97de h1:FxWPpzIjnTlhP github.com/araddon/dateparse v0.0.0-20210429162001-6b43995a97de/go.mod h1:DCaWoUhZrYW9p1lxo/cm8EmUOOzAPSEZNGF2DK1dJgw= github.com/aymanbagabas/go-osc52/v2 v2.0.1 h1:HwpRHbFMcZLEVr42D4p7XBqjyuxQH5SMiErDT4WkJ2k= github.com/aymanbagabas/go-osc52/v2 v2.0.1/go.mod h1:uYgXzlJ7ZpABp8OJ+exZzJJhRNQ2ASbcXHWsFqH8hp8= +github.com/bsm/ginkgo/v2 v2.12.0 h1:Ny8MWAHyOepLGlLKYmXG4IEkioBysk6GpaRTLC8zwWs= +github.com/bsm/ginkgo/v2 v2.12.0/go.mod h1:SwYbGRRDovPVboqFv0tPTcG1sN61LM1Z4ARdbAV9g4c= +github.com/bsm/gomega v1.27.10 h1:yeMWxP2pV2fG3FgAODIY8EiRE3dy0aeFYt4l7wh6yKA= +github.com/bsm/gomega v1.27.10/go.mod h1:JyEr/xRbxbtgWNi8tIEVPUYZ5Dzef52k01W3YH0H+O0= +github.com/cespare/xxhash/v2 v2.3.0 h1:UL815xU9SqsFlibzuggzjXhog7bL6oX9BbNZnL2UFvs= +github.com/cespare/xxhash/v2 v2.3.0/go.mod h1:VGX0DQ3Q6kWi7AoAeZDth3/j3BFtOZR5XLFGgcrjCOs= github.com/charmbracelet/bubbles v0.21.0 h1:9TdC97SdRVg/1aaXNVWfFH3nnLAwOXr8Fn6u6mfQdFs= github.com/charmbracelet/bubbles v0.21.0/go.mod h1:HF+v6QUR4HkEpz62dx7ym2xc71/KBHg+zKwJtMw+qtg= github.com/charmbracelet/bubbletea v1.3.4 h1:kCg7B+jSCFPLYRA52SDZjr51kG/fMUEoPoZrkaDHyoI= @@ -40,7 +44,6 @@ github.com/chzyer/readline v0.0.0-20180603132655-2972be24d48e h1:fY5BOSpyZCqRo5O github.com/chzyer/readline v0.0.0-20180603132655-2972be24d48e/go.mod h1:nSuG5e5PlCu98SY8svDHJxuZscDgtXS6KTTbou5AhLI= github.com/chzyer/test v0.0.0-20180213035817-a1ea475d72b1 h1:q763qf9huN11kDQavWsoZXJNW3xEE4JJyHa5Q25/sd8= github.com/chzyer/test v0.0.0-20180213035817-a1ea475d72b1/go.mod h1:Q3SI9o4m/ZMnBNeIyt5eFwwo7qiLfzFZmjNmxjkiQlU= -github.com/client9/misspell v0.3.4/go.mod h1:qj6jICC3Q7zFZvVWo7KLAzC3yx5G7kyvSDkc90ppPyw= github.com/cpuguy83/go-md2man v1.0.10 h1:BSKMNlYxDvnunlTymqtgONjNnaRV1sTpcovwwjF22jk= github.com/cpuguy83/go-md2man v1.0.10/go.mod h1:SmD6nW6nTyfqj6ABTjUi3V3JVMnlJmwcJI5acqYI6dE= github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= @@ -50,6 +53,8 @@ github.com/denisenkom/go-mssqldb v0.0.0-20180901172138-1eb28afdf9b6 h1:BZGp1dbKF github.com/denisenkom/go-mssqldb v0.0.0-20180901172138-1eb28afdf9b6/go.mod h1:xN/JuLBIz4bjkxNmByTiV1IbhfnYb6oo99phBn4Eqhc= github.com/dgrijalva/jwt-go v3.2.0+incompatible h1:7qlOGliEKZXTDg6OTjfoBKDXWrumCAMpl/TFQ4/5kLM= github.com/dgrijalva/jwt-go v3.2.0+incompatible/go.mod h1:E3ru+11k8xSBh+hMPgOLZmtrrCbhqsmaPHjLKYnJCaQ= +github.com/dgryski/go-rendezvous v0.0.0-20200823014737-9f7001d12a5f h1:lO4WD4F/rVNCu3HqELle0jiPLLBs70cWOduZpkS1E78= +github.com/dgryski/go-rendezvous v0.0.0-20200823014737-9f7001d12a5f/go.mod h1:cuUVRXasLTGF7a8hSLbxyZXjz+1KgoB3wDUb6vlszIc= github.com/docker/distribution v2.8.1+incompatible h1:Q50tZOPR6T/hjNsyc9g8/syEs6bk8XXApsHjKukMl68= github.com/docker/distribution v2.8.1+incompatible/go.mod h1:J2gT2udsDAN96Uj4KfcMRqY0/ypR+oyYUYmja8H+y+w= github.com/docker/docker v20.10.22+incompatible h1:6jX4yB+NtcbldT90k7vBSaWJDB3i+zkVJT9BEK8kQkk= @@ -100,9 +105,6 @@ github.com/go-sql-driver/mysql v1.4.0 h1:7LxgVwFb2hIQtMm87NdgAVfXjnt4OePseqT1tKx github.com/go-sql-driver/mysql v1.4.0/go.mod h1:zAC/RDZ24gD3HViQzih4MyKcchzm+sOG5ZlKdlhCg5w= github.com/gogo/protobuf v1.3.2 h1:Ov1cvc58UF3b5XjBnZv7+opcTcQFZebYjWzi34vdm4Q= github.com/gogo/protobuf v1.3.2/go.mod h1:P1XiOD3dCwIKUDQYPy72D8LYyHL2YPYrpS2s69NZV8Q= -github.com/golang/glog v0.0.0-20160126235308-23def4e6c14b h1:VKtxabqXZkF25pY9ekfRL6a582T4P37/31XEstQ5p58= -github.com/golang/glog v0.0.0-20160126235308-23def4e6c14b/go.mod h1:SBH7ygxi8pfUlaOkMMuAQtPIUF8ecWP5IEl/CR7VP2Q= -github.com/golang/mock v1.1.1/go.mod h1:oTYuIxOrZwtPieC+H1uAHpcLFnEyAGVDL/k47Jfbm0A= github.com/golang/protobuf v1.2.0/go.mod h1:6lQm79b+lXiMfvg/cZm0SGofjICqVBUtrP5yJMmIC1U= github.com/golang/protobuf v1.3.2/go.mod h1:6lQm79b+lXiMfvg/cZm0SGofjICqVBUtrP5yJMmIC1U= github.com/golang/protobuf v1.3.3 h1:gyjaxf+svBWX08ZjK86iN9geUJF0H6gp2IRKX6Nf6/I= @@ -202,6 +204,8 @@ github.com/pkg/errors v0.9.1 h1:FEBLx1zS214owpjy7qsBeixbURkuhQAwrK5UwLGTwt4= github.com/pkg/errors v0.9.1/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0= github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM= github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4= +github.com/redis/go-redis/v9 v9.17.3 h1:fN29NdNrE17KttK5Ndf20buqfDZwGNgoUr9qjl1DQx4= +github.com/redis/go-redis/v9 v9.17.3/go.mod h1:u410H11HMLoB+TP67dz8rL9s6QW2j76l0//kSOd3370= github.com/rivo/uniseg v0.1.0/go.mod h1:J6wj4VEh+S6ZtnVlnTBMWIodfgj8LQOQFoIToxlJtxc= github.com/rivo/uniseg v0.2.0/go.mod h1:J6wj4VEh+S6ZtnVlnTBMWIodfgj8LQOQFoIToxlJtxc= github.com/rivo/uniseg v0.4.7 h1:WUdvkW8uEhrYfLC4ZzdpI2ztxP1I582+49Oc5Mq64VQ= @@ -247,27 +251,22 @@ golang.org/x/crypto v0.29.0 h1:L5SG1JTTXupVV3n6sUqMTeWbjAyfPwoda2DLX8J8FrQ= golang.org/x/crypto v0.29.0/go.mod h1:+F4F4N5hv6v38hfeYwTdx20oUvLLc+QfrE9Ax9HtgRg= golang.org/x/exp v0.0.0-20220909182711-5c715a9e8561 h1:MDc5xs78ZrZr3HMQugiXOAkSZtfTpbJLDr/lwfgO53E= golang.org/x/exp v0.0.0-20220909182711-5c715a9e8561/go.mod h1:cyybsKvd6eL0RnXn6p/Grxp8F5bW7iYuBgsNCOHpMYE= -golang.org/x/lint v0.0.0-20181026193005-c67002cb31c3/go.mod h1:UVdnD1Gm6xHRNCYTkRU2/jEulfH38KcIWyp/GAMgvoE= golang.org/x/mod v0.2.0/go.mod h1:s0Qsj1ACt9ePp/hMypM3fl4fZqREWJwdYDEqhRiZZUA= golang.org/x/mod v0.3.0/go.mod h1:s0Qsj1ACt9ePp/hMypM3fl4fZqREWJwdYDEqhRiZZUA= golang.org/x/mod v0.22.0 h1:D4nJWe9zXqHOmWqj4VMOJhvzj7bEZg4wEYa759z1pH4= golang.org/x/mod v0.22.0/go.mod h1:6SkKJ3Xj0I0BrPOZoBy3bdMptDDU9oJrpohJ3eWZ1fY= golang.org/x/net v0.0.0-20180724234803-3673e40ba225/go.mod h1:mL1N/T3taQHkDXs73rZJwtUhF3w3ftmwwsq0BUmARs4= -golang.org/x/net v0.0.0-20180826012351-8a410e7b638d/go.mod h1:mL1N/T3taQHkDXs73rZJwtUhF3w3ftmwwsq0BUmARs4= golang.org/x/net v0.0.0-20190404232315-eb5bcb51f2a3/go.mod h1:t9HGtf8HONx5eT2rtn7q6eTqICYqUVnKs3thJo3Qplg= golang.org/x/net v0.0.0-20190620200207-3b0461eec859/go.mod h1:z5CRVTTTmAJ677TzLLGU+0bjPO0LkuOLi4/5GtJWs/s= golang.org/x/net v0.0.0-20200226121028-0de0cce0169b/go.mod h1:z5CRVTTTmAJ677TzLLGU+0bjPO0LkuOLi4/5GtJWs/s= golang.org/x/net v0.0.0-20201021035429-f5854403a974/go.mod h1:sp8m0HH+o8qH0wwXwYZr8TS3Oi6o0r6Gce1SSxlDquU= golang.org/x/net v0.31.0 h1:68CPQngjLL0r2AlUKiSxtQFKvzRVbnzLwMUn5SzcLHo= golang.org/x/net v0.31.0/go.mod h1:P4fl1q7dY2hnZFxEk4pPSkDHF+QqjitcnDjUQyMM+pM= -golang.org/x/oauth2 v0.0.0-20180821212333-d2e6202438be/go.mod h1:N/0e6XlmueqKjAGxoOufVs8QHGRruUQn6yWY3a++T0U= -golang.org/x/sync v0.0.0-20180314180146-1d60e4601c6f/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM= golang.org/x/sync v0.0.0-20190423024810-112230192c58/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM= golang.org/x/sync v0.0.0-20190911185100-cd5d95a43a6e/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM= golang.org/x/sync v0.0.0-20201020160332-67f06af15bc9/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM= golang.org/x/sync v0.11.0 h1:GGz8+XQP4FvTTrjZPzNKTMFtSXH80RAzG+5ghFPgK9w= golang.org/x/sync v0.11.0/go.mod h1:Czt+wKu1gCyEFDUtn0jG5QVvpJ6rzVqr5aXyt9drQfk= -golang.org/x/sys v0.0.0-20180830151530-49385e6e1522/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY= golang.org/x/sys v0.0.0-20180903190138-2b024373dcd9/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY= golang.org/x/sys v0.0.0-20181122145206-62eef0e2fa9b/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY= golang.org/x/sys v0.0.0-20190215142949-d0b11bdaac8a/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY= @@ -291,7 +290,6 @@ golang.org/x/text v0.20.0/go.mod h1:D4IsuqiFMhST5bX19pQ9ikHC2GsaKyk/oF+pn3ducp4= golang.org/x/time v0.3.0 h1:rg5rLMjNzMS1RkNLzCG38eapWhnYLFYXDXj2gOlr8j4= golang.org/x/time v0.3.0/go.mod h1:tRJNPiyCQ0inRvYxbN9jk5I+vvW/OXSQhTDSoE431IQ= golang.org/x/tools v0.0.0-20180917221912-90fa682c2a6e/go.mod h1:n7NCudcB/nEzxVGmLbDWY5pfWTLqBcC2KZ6jyYvM4mQ= -golang.org/x/tools v0.0.0-20190114222345-bf090417da8b/go.mod h1:n7NCudcB/nEzxVGmLbDWY5pfWTLqBcC2KZ6jyYvM4mQ= golang.org/x/tools v0.0.0-20191119224855-298f0cb1881e/go.mod h1:b+2E5dAYhXwXZwtnZ6UAqBI28+e2cm9otk0dWdXHAEo= golang.org/x/tools v0.0.0-20200619180055-7c47624df98f/go.mod h1:EkVYQZoAsY45+roYkvgYkIh4xh/qjgUK9TdY2XT94GE= golang.org/x/tools v0.0.0-20210106214847-113979e3529a/go.mod h1:emZCQorbCU4vsT4fOWvOPXz4eW1wZW4PmDk9uLelYpA= @@ -302,13 +300,8 @@ golang.org/x/xerrors v0.0.0-20190717185122-a985d3407aa7/go.mod h1:I/5z698sn9Ka8T golang.org/x/xerrors v0.0.0-20191011141410-1b5146add898/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0= golang.org/x/xerrors v0.0.0-20191204190536-9bdfabe68543/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0= golang.org/x/xerrors v0.0.0-20200804184101-5ec99f83aff1/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0= -google.golang.org/appengine v1.1.0/go.mod h1:EbEs0AVv82hx2wNQdGPgUI5lhzA/G0D9YwlJXL52JkM= google.golang.org/appengine v1.2.0 h1:S0iUepdCWODXRvtE+gcRDd15L+k+k1AiHlMiMjefH24= google.golang.org/appengine v1.2.0/go.mod h1:xpcJRLb0r/rnEns0DIKYYv+WjYCduHsrkT7/EB5XEv4= -google.golang.org/genproto v0.0.0-20180817151627-c66870c02cf8 h1:Nw54tB0rB7hY/N0NQvRW8DG4Yk3Q6T9cu9RcFQDu1tc= -google.golang.org/genproto v0.0.0-20180817151627-c66870c02cf8/go.mod h1:JiN7NxoALGmiZfu7CAH4rXhgtRTLTxftemlI0sWmxmc= -google.golang.org/grpc v1.19.0 h1:cfg4PD8YEdSFnm7qLV4++93WcmhH2nIUhMjhdCvl3j8= -google.golang.org/grpc v1.19.0/go.mod h1:mqu4LbDTu4XGKhr4mRzUsmM4RtVoemTSY81AxZiDr8c= gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0= gopkg.in/check.v1 v1.0.0-20180628173108-788fd7840127/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0= gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c h1:Hei/4ADfdWqJk1ZMxUNpqntNwaWcugrBjAiHlqqRiVk= @@ -336,4 +329,3 @@ gorm.io/gorm v1.25.10 h1:dQpO+33KalOA+aFYGlK+EfxcI5MbO7EP2yYygwh9h+s= gorm.io/gorm v1.25.10/go.mod h1:hbnx/Oo0ChWMn1BIhpy1oYozzpM15i4YPuHDmfYtwg8= gotest.tools/v3 v3.4.0 h1:ZazjZUfuVeZGLAmlKKuyv3IKP5orXcwtOwDQH6YVr6o= gotest.tools/v3 v3.4.0/go.mod h1:CtbdzLSsqVhDgMtKsx03ird5YTGB3ar27v0u/yKBW5g= -honnef.co/go/tools v0.0.0-20190102054323-c2f93a96b099/go.mod h1:rf3lG4BRIbNafJWhAfAdb/ePZxsR/4RtNHQocxwk9r4= diff --git a/pkg/cr/client.go b/pkg/cr/client.go new file mode 100644 index 00000000..61368402 --- /dev/null +++ b/pkg/cr/client.go @@ -0,0 +1,7 @@ +package cr + +import "github.com/docker/docker/client" + +func newDockerClient() (*client.Client, error) { + return client.NewClientWithOpts(client.FromEnv, client.WithAPIVersionNegotiation()) +} diff --git a/pkg/cr/containers.go b/pkg/cr/containers.go index ea102da8..a36e9cb3 100644 --- a/pkg/cr/containers.go +++ b/pkg/cr/containers.go @@ -5,6 +5,7 @@ import ( "encoding/json" "fmt" "io/ioutil" + "os" "os/exec" "path/filepath" "strconv" @@ -14,7 +15,6 @@ import ( "github.com/docker/docker/api/types/container" "github.com/docker/docker/api/types/filters" "github.com/docker/docker/api/types/mount" - "github.com/docker/docker/client" "github.com/docker/go-connections/nat" "github.com/sdslabs/beastv4/pkg/defaults" utils "github.com/sdslabs/beastv4/utils" @@ -66,8 +66,10 @@ type CreateContainerConfig struct { ContainerEnv []string ContainerNetwork string Traffic TrafficType + Labels map[string]string CPUShares int64 + CPUsLimit float32 Memory int64 PidsLimit int64 } @@ -87,7 +89,7 @@ type Log struct { // Function is equivalent to docker ps -a func SearchContainerByFilter(filterMap map[string]string) ([]types.Container, error) { - cli, err := client.NewEnvClient() + cli, err := newDockerClient() if err != nil { return []types.Container{}, err } @@ -107,7 +109,7 @@ func SearchContainerByFilter(filterMap map[string]string) ([]types.Container, er // Function is equivalent to docker ps func SearchRunningContainerByFilter(filterMap map[string]string) ([]types.Container, error) { - cli, err := client.NewEnvClient() + cli, err := newDockerClient() if err != nil { return []types.Container{}, err } @@ -125,7 +127,7 @@ func SearchRunningContainerByFilter(filterMap map[string]string) ([]types.Contai } func StopAndRemoveContainer(containerId string) error { - cli, err := client.NewEnvClient() + cli, err := newDockerClient() if err != nil { return err } @@ -148,9 +150,9 @@ func StopAndRemoveContainer(containerId string) error { } func CreateContainerFromImage(containerConfig *CreateContainerConfig) (string, error) { - containerName := fmt.Sprintf("beast_%s_%s", containerConfig.ChallengeName, containerConfig.ContainerName[:3]) + containerName := containerConfig.ContainerName ctx := context.Background() - cli, err := client.NewEnvClient() + cli, err := newDockerClient() if err != nil { return "", err } @@ -172,16 +174,21 @@ func CreateContainerFromImage(containerConfig *CreateContainerConfig) (string, e }} } + labels := map[string]string{ + "beast.challenge": containerConfig.ChallengeName, + "com.sdslabs.beast.project": utils.ProjectNameNotInstanced(containerConfig.ChallengeName), + "com.docker.compose.project": utils.ProjectNameNotInstanced(containerConfig.ChallengeName), + "com.sdslabs.beast.challenge": containerConfig.ChallengeName, + } + for k, v := range containerConfig.Labels { + labels[k] = v + } + config := &container.Config{ Image: containerConfig.ImageId, ExposedPorts: portSet, Env: containerConfig.ContainerEnv, - Labels: map[string]string{ - "beast.challenge": containerConfig.ChallengeName, - "com.sdslabs.beast.project": utils.GetProjectName(containerConfig.ChallengeName), - "com.docker.compose.project": utils.GetProjectName(containerConfig.ChallengeName), - "com.sdslabs.beast.challenge": containerConfig.ChallengeName, - }, + Labels: labels, } var mountBindings []mount.Mount @@ -196,6 +203,7 @@ func CreateContainerFromImage(containerConfig *CreateContainerConfig) (string, e } resources := container.Resources{ + NanoCPUs: int64(containerConfig.CPUsLimit * 1e9), CPUShares: containerConfig.CPUShares, Memory: containerConfig.Memory, PidsLimit: &containerConfig.PidsLimit, @@ -228,7 +236,7 @@ func CreateContainerFromImage(containerConfig *CreateContainerConfig) (string, e } func GetContainerStdLogs(containerID string) (*Log, error) { - cli, err := client.NewEnvClient() + cli, err := newDockerClient() if err != nil { return nil, err } @@ -259,7 +267,7 @@ func GetContainerStdLogs(containerID string) (*Log, error) { } func ShowLiveContainerLogs(containerID string) { - cli, err := client.NewEnvClient() + cli, err := newDockerClient() if err != nil { log.Error(err) } @@ -280,7 +288,7 @@ func ShowLiveContainerLogs(containerID string) { func CommitContainer(containerId string) (string, error) { ctx := context.Background() - cli, err := client.NewEnvClient() + cli, err := newDockerClient() if err != nil { return "", err } @@ -293,9 +301,8 @@ func CommitContainer(containerId string) (string, error) { return commitResp.ID, nil } -func DeployContainerFromCompose(challengeName, stagedPath, composeFileName string) (string, error) { +func DeployContainerFromCompose(challengeName string, projectName string, stagedPath string, composeFileName string, ports map[string]uint32) (string, error) { extractDir := filepath.Join(stagedPath, challengeName) - projectName := utils.GetProjectName(challengeName) composeFile := filepath.Join(extractDir, composeFileName) log.Debugf("Deploying challenge %s using docker compose with project name %s and file %s", challengeName, projectName, composeFileName) @@ -307,6 +314,13 @@ func DeployContainerFromCompose(challengeName, stagedPath, composeFileName strin "-p", projectName, "up", "-d") + environment := os.Environ() + for variable, port := range ports { + environment = append(environment, fmt.Sprintf("%s=%s", variable, strconv.FormatUint(uint64(port), 10))) + } + + upCmd.Env = environment + var upOutput bytes.Buffer upCmd.Stdout = &upOutput upCmd.Stderr = &upOutput @@ -411,9 +425,9 @@ func getPrimaryComposeContainerId(projectName string) (string, error) { return containerIds[0], nil } -func ComposeDown(challengeName, stagedDir string) error { - log.Debugf("Stopping challenge %s using docker compose", challengeName) - projectName := utils.GetProjectName(challengeName) +// ComposeDownProject runs docker compose down for an explicit -p project name (shared or instanced). +func ComposeDownProject(projectName string) error { + log.Debugf("Stopping docker compose project %s", projectName) downCmd := exec.Command("docker", "compose", "-p", projectName, "down") var downOutput bytes.Buffer @@ -421,16 +435,16 @@ func ComposeDown(challengeName, stagedDir string) error { downCmd.Stderr = &downOutput if err := downCmd.Run(); err != nil { - return fmt.Errorf("docker compose down failed for challenge %s: %v. Output: %s", challengeName, err, downOutput.String()) + return fmt.Errorf("docker compose down failed for project %s: %v. Output: %s", projectName, err, downOutput.String()) } - log.Debugf("Successfully stopped challenge %s", challengeName) + log.Debugf("Successfully stopped compose project %s", projectName) return nil } -func ComposePurge(challengeName, stagedDir string) error { - log.Debugf("Purging challenge %s using docker compose", challengeName) - projectName := utils.GetProjectName(challengeName) +// ComposePurgeProject runs compose down with volumes/images removal for an explicit -p name. +func ComposePurgeProject(projectName string) error { + log.Debugf("Purging docker compose project %s", projectName) purgeCmd := exec.Command("docker", "compose", "-p", projectName, "down", "--remove-orphans", "--volumes", "--rmi", "all") @@ -440,9 +454,9 @@ func ComposePurge(challengeName, stagedDir string) error { purgeCmd.Stderr = &purgeOutput if err := purgeCmd.Run(); err != nil { - return fmt.Errorf("docker compose purge failed for challenge %s: %v. Output: %s", challengeName, err, purgeOutput.String()) + return fmt.Errorf("docker compose purge failed for project %s: %v. Output: %s", projectName, err, purgeOutput.String()) } - log.Debugf("Successfully purged challenge %s", challengeName) + log.Debugf("Successfully purged compose project %s", projectName) return nil } diff --git a/pkg/cr/containers_integration_test.go b/pkg/cr/containers_integration_test.go new file mode 100644 index 00000000..f48fa24d --- /dev/null +++ b/pkg/cr/containers_integration_test.go @@ -0,0 +1,67 @@ +package cr + +import ( + "fmt" + "os" + "testing" + "time" + + "github.com/sdslabs/beastv4/utils" +) + +func TestCreateSearchAndRemoveContainerIntegration(t *testing.T) { + if os.Getenv("BEAST_TEST_DOCKER") != "1" { + t.Skip("set BEAST_TEST_DOCKER=1 to run Docker container integration tests") + } + + image := os.Getenv("BEAST_TEST_DOCKER_IMAGE") + if image == "" { + image = "redis:7-alpine" + } + + challengeName := fmt.Sprintf("beast-cr-integration-%d", time.Now().UnixNano()) + containerName := utils.ProjectNameNotInstanced(challengeName) + + containerID, err := CreateContainerFromImage(&CreateContainerConfig{ + ImageId: image, + ContainerName: containerName, + ChallengeName: challengeName, + MountsMap: map[string]string{}, + Labels: map[string]string{ + "beast.integration_test": "true", + }, + }) + if err != nil { + t.Fatalf("create container from image %s: %v", image, err) + } + defer func() { + _ = StopAndRemoveContainer(containerID) + }() + + containers, err := SearchRunningContainerByFilter(map[string]string{"id": containerID}) + if err != nil { + t.Fatalf("search running container by id: %v", err) + } + if len(containers) != 1 { + t.Fatalf("expected one running container, got %d", len(containers)) + } + + container := containers[0] + if container.Labels["beast.challenge"] != challengeName { + t.Fatalf("expected beast.challenge label %q, got %q", challengeName, container.Labels["beast.challenge"]) + } + if container.Labels["com.sdslabs.beast.project"] != utils.ProjectNameNotInstanced(challengeName) { + t.Fatalf("unexpected Beast project label %q", container.Labels["com.sdslabs.beast.project"]) + } + + if err := StopAndRemoveContainer(containerID); err != nil { + t.Fatalf("stop and remove container: %v", err) + } + containers, err = SearchContainerByFilter(map[string]string{"id": containerID}) + if err != nil { + t.Fatalf("search removed container by id: %v", err) + } + if len(containers) != 0 { + t.Fatalf("expected removed container to be absent, got %d matches", len(containers)) + } +} diff --git a/pkg/cr/health_check.go b/pkg/cr/health_check.go new file mode 100644 index 00000000..0291e2fb --- /dev/null +++ b/pkg/cr/health_check.go @@ -0,0 +1,166 @@ +package cr + +import ( + "bytes" + "encoding/json" + "fmt" + "github.com/docker/docker/api/types" + "github.com/sdslabs/beastv4/core" + "github.com/sdslabs/beastv4/core/cache" + log "github.com/sirupsen/logrus" + "os/exec" + "strings" +) + +func CleanupOrphans() { + var containers []types.Container + var err error + + containers, err = SearchContainerByFilter(map[string]string{ + "label": "beast.instance=true", + }) + + if err != nil { + log.Warnf("Failed to search for instance containers on %s: %v", core.LOCALHOST, err) + return + } + + for _, container := range containers { + instanceID := container.Labels["beast.instance.id"] + if instanceID == "" { + for _, name := range container.Names { + name = strings.TrimPrefix(name, "/") + if strings.HasPrefix(name, "beast_instance_") { + parts := strings.Split(name, "_") + if len(parts) >= 4 { + instanceID = parts[len(parts)-1] + break + } + } + } + } + + if instanceID == "" { + continue + } + + _, err = cache.GetInstance(instanceID) + if err != nil { + containerName := "" + if len(container.Names) > 0 { + containerName = strings.TrimPrefix(container.Names[0], "/") + } + log.Infof("Removing orphaned instance container: %s (instance %s) on %s", containerName, instanceID, core.LOCALHOST) + + err = StopAndRemoveContainer(container.ID) + if err != nil { + log.Warnf("Failed to remove orphaned container %s: %s", container.ID[:12], err.Error()) + } + + err = cache.FreeContainerPortsOnHost(core.LOCALHOST, container.ID) + if err != nil { + log.Warnf("Failed to free port for orphan container %s: %s", container.ID[:12], err.Error()) + } + } + } +} + +// CleanupOrphanedComposeInstances finds and removes orphaned docker compose instance projects. +// Instanced compose uses ComposeDockerProjectNameInstanced (-p = beast-instance--); +// compose ls project names are matched by prefix "beast-instance-". +func CleanupOrphanedComposeInstances() { + var projectNames []string + var err error + + projectNames, err = getOrphanedComposeInstanceProjects() + + if err != nil { + log.Warnf("Failed to get compose instance projects on %s: %v", core.LOCALHOST, err) + return + } + + for _, projectName := range projectNames { + // Extract instance ID from project name: beast-instance-{encoded_challenge}-{instanceID} + parts := strings.Split(projectName, "-") + if len(parts) < 4 { + continue + } + instanceID := parts[len(parts)-1] + + // Check if instance still exists in cache + _, err = cache.GetInstance(instanceID) + if err != nil { + log.Infof("Removing orphaned compose instance project: %s (instance %s) on %s", projectName, instanceID, core.LOCALHOST) + + err = composeDownProject(projectName) + if err != nil { + log.Warnf("Failed to remove orphaned compose project %s: %v", projectName, err) + } + if err := cache.FreeContainerPortsOnHost(core.LOCALHOST, projectName); err != nil { + log.Warnf("Failed to free ports for orphaned compose project %s: %v", projectName, err) + } + } + } +} + +// getOrphanedComposeInstanceProjects returns a list of docker compose project names +// that match the instance naming pattern (beast-instance-*) +func getOrphanedComposeInstanceProjects() ([]string, error) { + cmd := exec.Command("docker", "compose", "ls", "--format", "json") + var output bytes.Buffer + cmd.Stdout = &output + cmd.Stderr = &output + + if err := cmd.Run(); err != nil { + return nil, fmt.Errorf("docker compose ls failed: %v, output: %s", err, output.String()) + } + + type ComposeProject struct { + Name string `json:"Name"` + Status string `json:"Status"` + } + + var projects []ComposeProject + outputStr := strings.TrimSpace(output.String()) + if outputStr == "" { + return nil, nil + } + + if err := json.Unmarshal([]byte(outputStr), &projects); err != nil { + // Try parsing line by line (older docker compose versions) + for _, line := range strings.Split(outputStr, "\n") { + if strings.TrimSpace(line) == "" { + continue + } + var project ComposeProject + if err := json.Unmarshal([]byte(line), &project); err != nil { + continue + } + projects = append(projects, project) + } + } + + var instanceProjects []string + for _, project := range projects { + if strings.HasPrefix(project.Name, "beast-instance-") { + instanceProjects = append(instanceProjects, project.Name) + } + } + + return instanceProjects, nil +} + +// composeDownProject removes a docker compose project by name +func composeDownProject(projectName string) error { + cmd := exec.Command("docker", "compose", "-p", projectName, "down", "--remove-orphans", "-v") + var output bytes.Buffer + cmd.Stdout = &output + cmd.Stderr = &output + + if err := cmd.Run(); err != nil { + return fmt.Errorf("docker compose down failed: %v, output: %s", err, output.String()) + } + + log.Debugf("Successfully removed compose project %s", projectName) + return nil +} diff --git a/pkg/cr/images.go b/pkg/cr/images.go index 36836bf8..fe4f6a39 100644 --- a/pkg/cr/images.go +++ b/pkg/cr/images.go @@ -10,7 +10,6 @@ import ( "github.com/docker/docker/api/types" "github.com/docker/docker/api/types/filters" - "github.com/docker/docker/client" "github.com/sdslabs/beastv4/core" "github.com/sdslabs/beastv4/utils" @@ -19,7 +18,7 @@ import ( ) func RemoveImage(imageId string) error { - cli, err := client.NewEnvClient() + cli, err := newDockerClient() if err != nil { return err } @@ -34,7 +33,7 @@ func RemoveImage(imageId string) error { func CheckIfImageExists(imageId string) (bool, error) { ctx := context.Background() - cli, err := client.NewEnvClient() + cli, err := newDockerClient() if err != nil { return false, err } @@ -52,7 +51,7 @@ func CheckIfImageExists(imageId string) (bool, error) { } func SearchImageByFilter(filterMap map[string]string) ([]types.ImageSummary, error) { - cli, err := client.NewEnvClient() + cli, err := newDockerClient() if err != nil { return []types.ImageSummary{}, err } @@ -85,12 +84,12 @@ func BuildImageFromTarContext(challengeName, challengeTag, tarContextPath, docke NoCache: noCache, Labels: map[string]string{ "beast.challenge": challengeName, - "com.sdslabs.beast.project": utils.GetProjectName(challengeName), - "com.docker.compose.project": utils.GetProjectName(challengeName), + "com.sdslabs.beast.project": utils.ProjectNameNotInstanced(challengeName), + "com.docker.compose.project": utils.ProjectNameNotInstanced(challengeName), }, } - dockerClient, err := client.NewEnvClient() + dockerClient, err := newDockerClient() if err != nil { return nil, "", fmt.Errorf("error while creating a docker client for beast: %s", err) } diff --git a/pkg/probes/tcp.go b/pkg/probes/tcp.go index 27e57fbb..8df0c822 100644 --- a/pkg/probes/tcp.go +++ b/pkg/probes/tcp.go @@ -2,11 +2,11 @@ package probes import ( "fmt" + "github.com/sdslabs/beastv4/core" "net" "strconv" "time" - "github.com/sdslabs/beastv4/core" log "github.com/sirupsen/logrus" ) @@ -22,8 +22,8 @@ type TcpProber struct{} // If the socket fails to open, it returns Failure. func (pr TcpProber) Probe(host string, port int, timeout time.Duration) (ProbeResult, error) { var hostAddress string - if host == core.LOCALHOST || host == "" { - hostAddress = "127.0.0.1" + if host == core.LOCALHOST { + hostAddress = core.LOCALHOST_IP } else { ips, err := net.LookupIP(host) if err != nil { diff --git a/pkg/remoteManager/container.go b/pkg/remoteManager/container.go index a13d9935..25ec6360 100644 --- a/pkg/remoteManager/container.go +++ b/pkg/remoteManager/container.go @@ -8,7 +8,7 @@ import ( "strings" "github.com/docker/docker/api/types" - "github.com/sdslabs/beastv4/core" + _ "github.com/sdslabs/beastv4/core" "github.com/sdslabs/beastv4/core/config" "github.com/sdslabs/beastv4/core/database" "github.com/sdslabs/beastv4/pkg/cr" @@ -18,7 +18,7 @@ import ( ) func CreateContainerFromImageRemote(containerConfig cr.CreateContainerConfig, server config.AvailableServer) (string, error) { - var containerName, containerEnv, exposedPorts, portMap, cpuLimit, memoryLimit, pidLimit, imageID, mountBindings string + var containerName, containerEnv, exposedPorts, portMap, cpuShareLimit, cpuLimit, memoryLimit, pidLimit, imageID, mountBindings string if containerConfig.ContainerName != "" { containerName = fmt.Sprintf("--name %s ", containerConfig.ContainerName) } @@ -26,11 +26,14 @@ func CreateContainerFromImageRemote(containerConfig cr.CreateContainerConfig, se containerEnv += fmt.Sprintf("--env %s ", envVar) } for _, portMapping := range containerConfig.PortMapping { - portMap += fmt.Sprintf("-p 0.0.0.0:%d:%d/%s ", portMapping.ContainerPort, portMapping.HostPort, containerConfig.TrafficType()) + portMap += fmt.Sprintf("-p 0.0.0.0:%d:%d/%s ", portMapping.HostPort, portMapping.ContainerPort, containerConfig.TrafficType()) exposedPorts += fmt.Sprintf("--expose %d ", portMapping.ContainerPort) } if containerConfig.CPUShares != 0 { - cpuLimit = fmt.Sprintf("--cpu-shares %d ", containerConfig.CPUShares) + cpuShareLimit = fmt.Sprintf("--cpu-shares %d ", containerConfig.CPUShares) + } + if containerConfig.CPUsLimit != 0 { + cpuLimit = fmt.Sprintf("--cpus %f ", containerConfig.CPUsLimit) } if containerConfig.Memory != 0 { memoryLimit = fmt.Sprintf("--memory %d ", containerConfig.Memory) @@ -44,7 +47,7 @@ func CreateContainerFromImageRemote(containerConfig cr.CreateContainerConfig, se for src, dest := range containerConfig.MountsMap { mountBindings += fmt.Sprintf("--mount type=bind,source=%s,target=%s ", src, dest) } - dockerCommand := fmt.Sprintf("docker run -d %s %s %s %s %s %s %s %s %s", containerName, containerEnv, exposedPorts, mountBindings, cpuLimit, memoryLimit, pidLimit, portMap, imageID) + dockerCommand := fmt.Sprintf("docker run -d %s %s %s %s %s %s %s %s %s %s", containerName, containerEnv, exposedPorts, mountBindings, cpuShareLimit, cpuLimit, memoryLimit, pidLimit, portMap, imageID) // fmt.Printf("%s, %s, %s, %s\n", containerName, containerEnv, exposedPorts, portMap) // dockerCommand := fmt.Sprintf("docker run \\ // --name \\ @@ -108,9 +111,9 @@ func SearchContainerByFilterRemote(filterMap map[string]string, server config.Av for key, val := range filterMap { filterArgs += fmt.Sprintf("--filter='%s=%s' ", key, val) } - for _, server := range config.Cfg.AvailableServers { + for serverDeployed, server := range config.Cfg.AvailableServers { if server.Active { - if server.Host != core.LOCALHOST { + if !config.Cfg.UseLocalDockerDaemon(serverDeployed) { output, err = RunCommandOnServer(server, fmt.Sprintf("docker ps -a %s --format '{{.ID}}'", filterArgs)) if err != nil { return []types.Container{}, err @@ -136,9 +139,9 @@ func SearchRunningContainerByFilterRemote(filterMap map[string]string, server co for key, val := range filterMap { filterArgs += fmt.Sprintf("--filter='%s=%s' ", key, val) } - for _, server := range config.Cfg.AvailableServers { + for serverDeployed, server := range config.Cfg.AvailableServers { if server.Active { - if server.Host != core.LOCALHOST { + if !config.Cfg.UseLocalDockerDaemon(serverDeployed) { output, err = RunCommandOnServer(server, fmt.Sprintf("docker ps %s --format '{{.ID}}'", filterArgs)) if err != nil { return []types.Container{}, err @@ -198,12 +201,11 @@ func CommitContainerRemote(containerID string, server config.AvailableServer) (s return imageID, nil } -func DeployContainerFromComposeRemote(challengeName, stagedDir, composeFileName string, server config.AvailableServer) (string, error) { +func DeployContainerFromComposeRemote(challengeName string, projectName string, stagedDir string, composeFileName string, server config.AvailableServer, ports map[string]uint32) (string, error) { extractDir := filepath.Join(stagedDir, challengeName) - projectName := utils.GetProjectName(challengeName) composeFile := filepath.Join(extractDir, composeFileName) - upCommand := fmt.Sprintf("docker compose -f %s -p %s up -d", composeFile, projectName) + upCommand := fmt.Sprintf("%s docker compose -f %s -p %s up -d", utils.PortMappingToEnvironmentVariable(ports), composeFile, projectName) log.Debugf("Deploying challenge %s using docker compose remotely with project %s and file %s", challengeName, projectName, composeFileName) upOutput, err := RunCommandOnServer(server, upCommand) if err != nil { @@ -303,32 +305,26 @@ func getPrimaryComposeContainerIdRemote(projectName string, server config.Availa return containerId, nil } -func ComposeDownRemote(challengeName, stagedDir string, server config.AvailableServer) error { - log.Debugf("Stopping challenge %s using docker compose on remote", challengeName) - projectName := utils.GetProjectName(challengeName) - +// ComposeDownProjectRemote runs docker compose down for an explicit -p project name. +func ComposeDownProjectRemote(projectName string, server config.AvailableServer) error { + log.Debugf("Stopping docker compose project %s on remote", projectName) downCommand := fmt.Sprintf("docker compose -p %s down", projectName) - log.Debugf("Stopping challenge %s using docker compose remotely: %s", challengeName, downCommand) downOutput, err := RunCommandOnServer(server, downCommand) if err != nil { - return fmt.Errorf("docker compose down failed for challenge %s on remote: %v. Output: %s", challengeName, err, downOutput) + return fmt.Errorf("docker compose down failed for project %s on remote: %v. Output: %s", projectName, err, downOutput) } - - log.Debugf("Successfully stopped challenge %s on remote. Output: %s", challengeName, downOutput) + log.Debugf("Successfully stopped compose project %s on remote. Output: %s", projectName, downOutput) return nil } -func ComposePurgeRemote(challengeName, stagedDir string, server config.AvailableServer) error { - log.Debugf("Purging challenge %s using docker compose on remote", challengeName) - projectName := utils.GetProjectName(challengeName) +// ComposePurgeProjectRemote purges a compose project by explicit -p name (shared or instanced). +func ComposePurgeProjectRemote(projectName string, server config.AvailableServer) error { + log.Debugf("Purging docker compose project %s on remote", projectName) purgeCommand := fmt.Sprintf("docker compose -p %s down --remove-orphans --volumes --rmi all", projectName) - log.Debugf("Purge challenge %s using docker compose remotely: %s", challengeName, purgeCommand) purgeOutput, err := RunCommandOnServer(server, purgeCommand) if err != nil { - return fmt.Errorf("docker compose purge failed for challenge %s on remote: %v. Output: %s", challengeName, err, purgeOutput) - + return fmt.Errorf("docker compose purge failed for project %s on remote: %v. Output: %s", projectName, err, purgeOutput) } - - log.Debugf("Successfully purged challenge %s on remote. Output: %s", challengeName, purgeOutput) + log.Debugf("Successfully purged compose project %s on remote. Output: %s", projectName, purgeOutput) return nil } diff --git a/pkg/remoteManager/file.go b/pkg/remoteManager/file.go index 36d98655..8ffe4cd7 100644 --- a/pkg/remoteManager/file.go +++ b/pkg/remoteManager/file.go @@ -83,7 +83,7 @@ func BuildImageFromTarContextRemote(challengeName string, imageTag string, stage if err != nil { return []byte{}, "", fmt.Errorf("failed to extract tar: %s", err) } - projectName := utils.GetProjectName(challengeName) + projectName := utils.ProjectNameNotInstanced(challengeName) dockerBuildCmd := fmt.Sprintf("cd %s && docker build -t %s "+ "--label beast.challenge=%s "+ "--label com.sdslabs.beast.project=%s "+ diff --git a/pkg/remoteManager/health_check.go b/pkg/remoteManager/health_check.go new file mode 100644 index 00000000..659dd7ea --- /dev/null +++ b/pkg/remoteManager/health_check.go @@ -0,0 +1,156 @@ +package remoteManager + +import ( + "encoding/json" + "fmt" + "github.com/docker/docker/api/types" + "github.com/sdslabs/beastv4/core/cache" + "github.com/sdslabs/beastv4/core/config" + log "github.com/sirupsen/logrus" + "strings" +) + +func CleanupOrphanedOnServer(serverDeployed string) { + var containers []types.Container + var err error + + server := config.Cfg.AvailableServers[serverDeployed] + containers, err = SearchContainerByFilterRemote(map[string]string{ + "label": "beast.instance=true", + }, server) + + if err != nil { + log.Warnf("Failed to search for instance containers on %s: %v", serverDeployed, err) + return + } + + for _, container := range containers { + instanceID := container.Labels["beast.instance.id"] + if instanceID == "" { + for _, name := range container.Names { + name = strings.TrimPrefix(name, "/") + if strings.HasPrefix(name, "beast_instance_") { + parts := strings.Split(name, "_") + if len(parts) >= 4 { + instanceID = parts[len(parts)-1] + break + } + } + } + } + + if instanceID == "" { + continue + } + + _, err = cache.GetInstance(instanceID) + if err != nil { + containerName := "" + if len(container.Names) > 0 { + containerName = strings.TrimPrefix(container.Names[0], "/") + } + log.Infof("Removing orphaned instance container: %s (instance %s) on %s", containerName, instanceID, serverDeployed) + + err = StopAndRemoveContainerRemote(container.ID, server) + if err != nil { + log.Warnf("Failed to remove orphaned container %s on %s: %v", container.ID[:12], serverDeployed, err) + } + + err = cache.FreeContainerPortsOnHost(serverDeployed, container.ID) + if err != nil { + log.Warnf("Failed to free port for orphan container %s on server %s: %s", container.ID[:12], serverDeployed, err.Error()) + } + } + } +} + +// cleanupOrphanedComposeInstancesOnServer finds and removes orphaned docker compose instance projects. +// Instanced stacks use ComposeDockerProjectNameInstanced; compose ls names use prefix "beast-instance-". +func CleanupOrphanedComposeInstancesOnServer(serverDeployed string) { + var projectNames []string + var err error + + server := config.Cfg.AvailableServers[serverDeployed] + projectNames, err = getOrphanedComposeInstanceProjectsRemote(server) + + if err != nil { + log.Warnf("Failed to get compose instance projects on %s: %v", serverDeployed, err) + return + } + + for _, projectName := range projectNames { + // Extract instance ID from project name: beast-instance-{encoded_challenge}-{instanceID} + parts := strings.Split(projectName, "-") + if len(parts) < 4 { + continue + } + instanceID := parts[len(parts)-1] + + // Check if instance still exists in cache + _, err := cache.GetInstance(instanceID) + if err != nil { + log.Infof("Removing orphaned compose instance project: %s (instance %s) on %s", projectName, instanceID, serverDeployed) + + if err := composeDownProjectRemote(projectName, server); err != nil { + log.Warnf("Failed to remove orphaned compose project %s on %s: %v", projectName, serverDeployed, err) + } + if err := cache.FreeContainerPortsOnHost(serverDeployed, projectName); err != nil { + log.Warnf("Failed to free ports for orphaned compose project %s on %s: %v", projectName, serverDeployed, err) + } + } + } +} + +// getOrphanedComposeInstanceProjectsRemote returns compose instance projects on a remote server +func getOrphanedComposeInstanceProjectsRemote(server config.AvailableServer) ([]string, error) { + output, err := RunCommandOnServer(server, "docker compose ls --format json") + if err != nil { + return nil, fmt.Errorf("docker compose ls failed on remote: %v", err) + } + + type ComposeProject struct { + Name string `json:"Name"` + Status string `json:"Status"` + } + + var projects []ComposeProject + outputStr := strings.TrimSpace(output) + if outputStr == "" { + return nil, nil + } + + if err := json.Unmarshal([]byte(outputStr), &projects); err != nil { + // Try parsing line by line + for _, line := range strings.Split(outputStr, "\n") { + if strings.TrimSpace(line) == "" { + continue + } + var project ComposeProject + if err := json.Unmarshal([]byte(line), &project); err != nil { + continue + } + projects = append(projects, project) + } + } + + var instanceProjects []string + for _, project := range projects { + if strings.HasPrefix(project.Name, "beast-instance-") { + instanceProjects = append(instanceProjects, project.Name) + } + } + + return instanceProjects, nil +} + +// composeDownProjectRemote removes a docker compose project on a remote server +func composeDownProjectRemote(projectName string, server config.AvailableServer) error { + cmd := fmt.Sprintf("docker compose -p %s down --remove-orphans -v", projectName) + output, err := RunCommandOnServer(server, cmd) + if err != nil { + return fmt.Errorf("docker compose down failed on remote: %v, output: %s", err, output) + } + + log.Debugf("Successfully removed compose project %s on %s", projectName, server.Host) + return nil +} diff --git a/pkg/remoteManager/init.go b/pkg/remoteManager/init.go index 2a9aa893..e4c9a43f 100644 --- a/pkg/remoteManager/init.go +++ b/pkg/remoteManager/init.go @@ -1,6 +1,9 @@ package remoteManager import ( + "fmt" + "path/filepath" + "github.com/sdslabs/beastv4/core" "github.com/sdslabs/beastv4/core/config" log "github.com/sirupsen/logrus" @@ -8,11 +11,11 @@ import ( func Init() { ServerQueue = NewLoadBalancerQueue() - for _, server := range config.Cfg.AvailableServers { + for serverDeployed, server := range config.Cfg.AvailableServers { if server.Active { - if server.Host == core.LOCALHOST { + // Skip SSH bootstrap for loopback workers; they use the local Docker socket from Beast. + if config.Cfg.UseLocalDockerDaemon(serverDeployed) { continue - ServerQueue.Push(server) } client, err := CreateSSHClient(server) if err != nil { @@ -21,7 +24,10 @@ func Init() { } defer client.Close() ServerQueue.Push(server) - RunCommandOnServer(server, "mkdir -p $HOME/.beast/staging/") + _, err = RunCommandOnServer(server, fmt.Sprintf("mkdir -p %s", filepath.Join(core.BEAST_REMOTE_GLOBAL_DIR, core.BEAST_STAGING_DIR))) + if err != nil { + log.Errorf("failed to run command on server %s: %s", server.Host, err.Error()) + } } } } diff --git a/scripts/test/backend_submit_race.sh b/scripts/test/backend_submit_race.sh new file mode 100755 index 00000000..b17a1b0b --- /dev/null +++ b/scripts/test/backend_submit_race.sh @@ -0,0 +1,345 @@ +#!/usr/bin/env bash +set -euo pipefail + +ROOT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd)" + +PGHOST="${BEAST_TEST_PGHOST:-localhost}" +PGPORT="${BEAST_TEST_PGPORT:-55543}" +PGUSER="${BEAST_TEST_PGUSER:-beasttest}" +PGPASSWORD="${BEAST_TEST_PGPASSWORD:-beasttest}" +PGDATABASE="${BEAST_TEST_PGDATABASE:-beast_backend_test}" +REDIS_HOST="${BEAST_TEST_REDIS_HOST:-localhost}" +REDIS_PORT="${BEAST_TEST_REDIS_PORT:-56380}" +SERVER_PORT="${BEAST_TEST_SERVER_PORT:-5505}" +TEST_HOME="${BEAST_TEST_HOME:-/tmp/beast-backend-submit-race}" +LOG_FILE="$TEST_HOME/beast-api.log" +BASE_URL="http://localhost:$SERVER_PORT" + +export PGPASSWORD + +cleanup() { + local status=$? + if [[ -n "${SERVER_PID:-}" ]] && kill -0 "$SERVER_PID" 2>/dev/null; then + kill "$SERVER_PID" 2>/dev/null || true + wait "$SERVER_PID" 2>/dev/null || true + fi + exit "$status" +} +trap cleanup EXIT + +psql_root() { + PGPASSWORD="$PGPASSWORD" psql -q -h "$PGHOST" -p "$PGPORT" -U "$PGUSER" -d postgres "$@" +} + +psql_test() { + PGPASSWORD="$PGPASSWORD" psql -q -h "$PGHOST" -p "$PGPORT" -U "$PGUSER" -d "$PGDATABASE" "$@" +} + +wait_for_http() { + for _ in $(seq 1 60); do + if curl -fsS "$BASE_URL/" >/dev/null 2>&1; then + return 0 + fi + sleep 1 + done + + echo "backend did not become ready; last log lines:" >&2 + tail -120 "$LOG_FILE" >&2 || true + return 1 +} + +json_field() { + jq -r "$1" +} + +register_user() { + local username="$1" + local password="$2" + curl -fsS -X POST "$BASE_URL/auth/register" \ + -F "name=$username" \ + -F "username=$username" \ + -F "password=$password" \ + -F "email=$username@example.test" >/dev/null +} + +login_user() { + local username="$1" + local password="$2" + curl -fsS -X POST "$BASE_URL/auth/login" \ + -F "username=$username" \ + -F "password=$password" | json_field '.token' +} + +seed_challenge() { + local name="$1" + local flag="$2" + local max_attempts="$3" + local dynamic="$4" + local author_id + author_id="$(psql_test -Atc "SELECT id FROM users ORDER BY id LIMIT 1")" + + psql_test -Atc " + INSERT INTO challenges ( + created_at, + updated_at, + name, + dynamic_flag, + flag, + type, + difficulty, + max_attempt_limit, + format, + container_id, + image_id, + status, + deployment_type, + author_id, + health_check, + points, + max_points, + min_points, + server_deployed + ) + VALUES ( + now(), + now(), + '$name', + $dynamic, + '$flag', + 'web', + 'easy', + $max_attempts, + 'web', + 'container-$name', + 'image-$name', + 'Deployed', + 'standard_docker', + $author_id, + 0, + 500, + 500, + 100, + 'localhost' + ) + RETURNING id" +} + +submit_concurrently() { + local token="$1" + local challenge_id="$2" + local flag="$3" + local requests="$4" + + python3 - "$BASE_URL" "$token" "$challenge_id" "$flag" "$requests" <<'PY' +import concurrent.futures +import json +import sys +import urllib.error +import urllib.parse +import urllib.request + +base_url, token, challenge_id, flag, requests = sys.argv[1], sys.argv[2], sys.argv[3], sys.argv[4], int(sys.argv[5]) + +def submit(_): + data = urllib.parse.urlencode({"chall_id": challenge_id, "flag": flag}).encode() + request = urllib.request.Request( + base_url + "/api/submit/challenge", + data=data, + headers={"Authorization": "Bearer " + token}, + method="POST", + ) + try: + with urllib.request.urlopen(request, timeout=10) as response: + body = response.read().decode() + return response.status, json.loads(body) + except urllib.error.HTTPError as error: + body = error.read().decode() + try: + parsed = json.loads(body) + except json.JSONDecodeError: + parsed = {"raw": body} + return error.code, parsed + +with concurrent.futures.ThreadPoolExecutor(max_workers=min(32, requests)) as executor: + results = list(executor.map(submit, range(requests))) + +print(json.dumps(results)) +PY +} + +assert_one_success() { + local results_json="$1" + local successes + successes="$(jq '[.[] | select(.[1].success == true)] | length' <<<"$results_json")" + if [[ "$successes" != "1" ]]; then + echo "expected exactly one successful submit response, got $successes" >&2 + jq . <<<"$results_json" >&2 + return 1 + fi +} + +assert_zero_successes() { + local results_json="$1" + local successes + successes="$(jq '[.[] | select(.[1].success == true)] | length' <<<"$results_json")" + if [[ "$successes" != "0" ]]; then + echo "expected zero successful submit responses, got $successes" >&2 + jq . <<<"$results_json" >&2 + return 1 + fi +} + +rm -rf "$TEST_HOME" +mkdir -p "$TEST_HOME/.beast/scripts" "$TEST_HOME/.beast/cache" "$TEST_HOME/.beast/remotes" "$TEST_HOME/.beast/uploads" "$TEST_HOME/.beast/secrets" "$TEST_HOME/.beast/staging" "$TEST_HOME/.beast/assets/logo" + +psql_root -v ON_ERROR_STOP=1 -c "DROP DATABASE IF EXISTS $PGDATABASE WITH (FORCE)" >/dev/null +psql_root -v ON_ERROR_STOP=1 -c "CREATE DATABASE $PGDATABASE" >/dev/null + +cat >"$TEST_HOME/.beast/config.toml" <"$LOG_FILE" 2>&1 +) & +SERVER_PID=$! + +wait_for_http + +register_user "apiwinner" "pw" +TOKEN_WINNER="$(login_user "apiwinner" "pw")" +CHALLENGE_CORRECT_ID="$(seed_challenge "api-race-correct" "flag{api-correct}" -1 false)" +CORRECT_RESULTS="$(submit_concurrently "$TOKEN_WINNER" "$CHALLENGE_CORRECT_ID" "flag{api-correct}" 64)" +assert_one_success "$CORRECT_RESULTS" + +WINNER_SCORE="$(psql_test -Atc "SELECT score FROM users WHERE username = 'apiwinner'")" +if [[ "$WINNER_SCORE" != "500" ]]; then + echo "expected apiwinner score 500, got $WINNER_SCORE" >&2 + exit 1 +fi +SOLVED_ROWS="$(psql_test -Atc "SELECT COUNT(*) FROM user_challenges WHERE user_id = (SELECT id FROM users WHERE username = 'apiwinner') AND challenge_id = $CHALLENGE_CORRECT_ID AND solved = true")" +if [[ "$SOLVED_ROWS" != "1" ]]; then + echo "expected one solved user_challenges row, got $SOLVED_ROWS" >&2 + exit 1 +fi + +register_user "apiwrong" "pw" +TOKEN_WRONG="$(login_user "apiwrong" "pw")" +CHALLENGE_WRONG_ID="$(seed_challenge "api-race-wrong" "flag{api-wrong}" 3 false)" +WRONG_RESULTS="$(submit_concurrently "$TOKEN_WRONG" "$CHALLENGE_WRONG_ID" "not-the-flag" 64)" +assert_zero_successes "$WRONG_RESULTS" + +WRONG_TRIES="$(psql_test -Atc "SELECT tries FROM user_challenges WHERE user_id = (SELECT id FROM users WHERE username = 'apiwrong') AND challenge_id = $CHALLENGE_WRONG_ID")" +if [[ "$WRONG_TRIES" != "3" ]]; then + echo "expected apiwrong tries 3, got $WRONG_TRIES" >&2 + exit 1 +fi +WRONG_SCORE="$(psql_test -Atc "SELECT score FROM users WHERE username = 'apiwrong'")" +if [[ "$WRONG_SCORE" != "0" ]]; then + echo "expected apiwrong score 0, got $WRONG_SCORE" >&2 + exit 1 +fi + +register_user "apidynone" "pw" +register_user "apidyntwo" "pw" +TOKEN_DYN_ONE="$(login_user "apidynone" "pw")" +TOKEN_DYN_TWO="$(login_user "apidyntwo" "pw")" +CHALLENGE_DYNAMIC_ID="$(seed_challenge "api-race-dynamic" "unused-static-flag" -1 true)" +psql_test -v ON_ERROR_STOP=1 -c "INSERT INTO dynamic_flags (created_at, updated_at, name, flag) VALUES (now(), now(), 'api-race-dynamic', 'flag{dynamic-shared}')" >/dev/null + +DYNAMIC_RESULTS="$( + python3 - "$BASE_URL" "$TOKEN_DYN_ONE" "$TOKEN_DYN_TWO" "$CHALLENGE_DYNAMIC_ID" <<'PY' +import concurrent.futures +import json +import sys +import urllib.parse +import urllib.request + +base_url, token_one, token_two, challenge_id = sys.argv[1:5] + +def submit(token): + data = urllib.parse.urlencode({"chall_id": challenge_id, "flag": "flag{dynamic-shared}"}).encode() + request = urllib.request.Request( + base_url + "/api/submit/challenge", + data=data, + headers={"Authorization": "Bearer " + token}, + method="POST", + ) + with urllib.request.urlopen(request, timeout=10) as response: + return response.status, json.loads(response.read().decode()) + +with concurrent.futures.ThreadPoolExecutor(max_workers=2) as executor: + results = list(executor.map(submit, [token_one, token_two])) +print(json.dumps(results)) +PY +)" +assert_one_success "$DYNAMIC_RESULTS" + +DYNAMIC_CLAIMS="$(psql_test -Atc "SELECT COUNT(*) FROM dynamic_flag_claims WHERE challenge_id = $CHALLENGE_DYNAMIC_ID AND flag = 'flag{dynamic-shared}'")" +if [[ "$DYNAMIC_CLAIMS" != "1" ]]; then + echo "expected one dynamic flag claim, got $DYNAMIC_CLAIMS" >&2 + exit 1 +fi + +LEADERBOARD="$(curl -fsS -H "Authorization: Bearer $TOKEN_WINNER" "$BASE_URL/api/info/leaderboard?page=1")" +if ! jq -e '.[] | select(.username == "apiwinner" and .score == 500)' <<<"$LEADERBOARD" >/dev/null; then + echo "leaderboard did not include apiwinner score 500" >&2 + jq . <<<"$LEADERBOARD" >&2 + exit 1 +fi + +echo "backend submit race verification passed" diff --git a/utils/cache.go b/utils/cache.go new file mode 100644 index 00000000..74e5b0a9 --- /dev/null +++ b/utils/cache.go @@ -0,0 +1,51 @@ +package utils + +import ( + "fmt" + "strings" +) + +const ( + instanceKeyPrefix = "beast:instance" + instanceExpiryPrefix = "beast:instance_expiry" + userInstanceKeyPrefix = "beast:user_instance" + InstancesSetKey = "beast:instances" + InstanceDeletionQueue = "beast:instances:to_delete" + + hostPrefixKey = "beast:host" + containerPrefixKey = "container" +) + +func HostToKey(host string) string { + return fmt.Sprintf("%s:%s", hostPrefixKey, host) +} + +func ContainerToKey(host string, containerId string) string { + return fmt.Sprintf("%s:%s:%s:%s", hostPrefixKey, host, containerPrefixKey, containerId) +} + +func InstanceToKey(instanceID string) string { + return fmt.Sprintf("%s:%s", instanceKeyPrefix, instanceID) +} + +func InstanceExpiryToKey(instanceID string) string { + return fmt.Sprintf("%s:%s", instanceExpiryPrefix, instanceID) +} + +func InstanceIDFromExpiryKey(key string) (string, bool) { + prefix := instanceExpiryPrefix + ":" + if !strings.HasPrefix(key, prefix) { + return "", false + } + + instanceID := strings.TrimPrefix(key, prefix) + return instanceID, instanceID != "" +} + +func UserChallengeToKey(userID, challengeName string) string { + return fmt.Sprintf("%s:%s:%s", userInstanceKeyPrefix, userID, challengeName) +} + +func UserChallengesAllKey(userID string) string { + return fmt.Sprintf("%s:%s:*", userInstanceKeyPrefix, userID) +} diff --git a/utils/compose.go b/utils/compose.go new file mode 100644 index 00000000..650188eb --- /dev/null +++ b/utils/compose.go @@ -0,0 +1,58 @@ +package utils + +import ( + "fmt" + "gopkg.in/yaml.v2" + "os" + "regexp" + "strings" +) + +type Compose struct { + Services map[string]struct { + Ports []string `yaml:"ports"` + } `yaml:"services"` +} + +var portRegex = regexp.MustCompile(`\$\{([^}]+)}`) + +func ExtractPortsFromCompose(composeFile string) ([]string, error) { + data, err := os.ReadFile(composeFile) + if err != nil { + return nil, fmt.Errorf("error while reading compose file: %w", err) + } + var raw Compose + err = yaml.Unmarshal(data, &raw) + if err != nil { + return nil, fmt.Errorf("error while parsing compose file: %s", err.Error()) + } + + portVariables := make([]string, 0) + seen := make(map[string]bool) + for _, service := range raw.Services { + for _, port := range service.Ports { + matches := portRegex.FindAllStringSubmatch(port, -1) + + if len(matches) == 0 { + return nil, fmt.Errorf("port %s is not mapped using an env variable", port) + } + + for _, match := range matches { + varName := match[1] + + /* Only ${PORT} is valid, ${PORT:-DEFAULT} should fail */ + if strings.Contains(varName, ":-") { + return nil, fmt.Errorf("port variable ${%s} uses default value syntax (:-) which is not supported; use ${%s} instead", + varName, strings.SplitN(varName, ":-", 2)[0]) + } + + if !seen[varName] { + seen[varName] = true + portVariables = append(portVariables, varName) + } + } + } + } + + return portVariables, nil +} diff --git a/utils/datatypes.go b/utils/datatypes.go index df2d48d7..f21f1e6f 100644 --- a/utils/datatypes.go +++ b/utils/datatypes.go @@ -3,12 +3,11 @@ package utils import ( "errors" "fmt" + "github.com/sdslabs/beastv4/core" "strconv" "strings" ) -const mappingDelimeter = ":" - // From a list of strings generate a list containing only unique strings // from the list. func GetUniqueStrings(list []string) []string { @@ -44,25 +43,51 @@ func UInt32InList(a uint32, list []uint32) bool { return false } +func Uint32InIndexList(a uint32, la []uint32, lb []uint32) (bool, uint32) { + for i, a_ := range la { + if a == a_ { + return true, lb[i] + } + } + + return false, 0 +} + // ParsePortMapping parses the port mapping string and return the required ports // If the portMapping string is not valid, this returns an error. -// The format of the port mapping is `HOST_PORT:CONTAINER_PORT` +// The format of the port mapping is `PORT_FIRST:PORT_LAST` func ParsePortMapping(portMap string) (uint32, uint32, error) { - ports := strings.Split(portMap, mappingDelimeter) + ports := strings.Split(portMap, core.MappingDelimiter) if len(ports) != 2 { return 0, 0, errors.New("port mapping string is not valid") } - hostPort, err := strconv.ParseUint(ports[0], 10, 32) + firstPort, err := strconv.ParseUint(ports[0], 10, 32) if err != nil { - return 0, 0, fmt.Errorf("host port is not a valid port in: %s", portMap) + return 0, 0, fmt.Errorf("first port is not a valid port in: %s", portMap) } - containerPort, err := strconv.ParseUint(ports[1], 10, 32) + lastPort, err := strconv.ParseUint(ports[1], 10, 32) if err != nil { - return 0, 0, fmt.Errorf("container port is not a valid port in: %s", portMap) + return 0, 0, fmt.Errorf("second port is not a valid port in: %s", portMap) + } + + if firstPort > lastPort { + return 0, 0, fmt.Errorf("first port is greater than last port") + } + + return uint32(firstPort), uint32(lastPort), nil +} + +func PortMappingToEnvironmentVariable(ports map[string]uint32) string { + env := make([]string, len(ports)) + + i := 0 + for variable, port := range ports { + env[i] = fmt.Sprintf("%s=%s", variable, strconv.FormatUint(uint64(port), 10)) + i++ } - return uint32(hostPort), uint32(containerPort), nil + return strings.Join(env, " ") } diff --git a/utils/id.go b/utils/id.go index a8f31352..ab177bb0 100644 --- a/utils/id.go +++ b/utils/id.go @@ -6,6 +6,7 @@ package utils import ( cryptorand "crypto/rand" + "crypto/sha256" "encoding/hex" "fmt" "io" @@ -109,11 +110,37 @@ func (fn readerFunc) Read(p []byte) (int, error) { return fn(p) } -// GetProjectName generates the standard project name for both docker and docker-compose deployments. -// This name is used for: -// - Docker Compose project name (-p flag) -// - Container labels (com.sdslabs.beast.project, com.docker.compose.project) -// - Container naming conventions -func GetProjectName(challengeName string) string { - return fmt.Sprintf("beast-%s", challengeName) +func EncodeID(a string) string { + return fmt.Sprintf("%x", sha256.Sum256([]byte(a)))[:30] +} + +// GetInstanceIdentifier returns the inner segment for an instanced workload (compose project key +// or container name body) before the beast- prefix is applied. Format: +// +// instance-- +func GetInstanceIdentifier(challengeName string, instanceId string) string { + return fmt.Sprintf("instance-%s-%s", EncodeID(challengeName), instanceId) +} + +// GetChallengeIdentifier prefixes a logical key with "beast-" for Docker names and labels. +// For compose, the full docker compose -p value is often this prefix applied to either the +// challenge name (non-instanced) or GetInstanceIdentifier (instanced). Prefer the helpers +// ProjectNameNotInstanced / ComposeDockerProjectNameInstanced for -p so deploy +// and teardown stay aligned. +func GetChallengeIdentifier(challengeIdentifier string) string { + return fmt.Sprintf("beast-%s", challengeIdentifier) +} + +// ProjectNameNotInstanced is the exact docker compose -p project name for a +// non-instanced (shared) compose challenge. Use this in deployPipeline, ComposeDown, ComposePurge, +// and any cleanup that must target the same stack. +func ProjectNameNotInstanced(challengeName string) string { + return GetChallengeIdentifier(EncodeID(challengeName)) +} + +// ComposeDockerProjectNameInstanced is the exact docker compose -p project name for an instanced +// compose challenge. Format: beast-instance--. +// Must match deployInstanceFromCompose and instanced ComposePurge/teardown. +func ComposeDockerProjectNameInstanced(challengeName, instanceID string) string { + return GetChallengeIdentifier(GetInstanceIdentifier(challengeName, instanceID)) } diff --git a/utils/id_test.go b/utils/id_test.go new file mode 100644 index 00000000..c23abc15 --- /dev/null +++ b/utils/id_test.go @@ -0,0 +1,16 @@ +package utils + +import "testing" + +func TestDockerProjectNameHelpers(t *testing.T) { + challengeName := "Web Challenge 01" + encoded := EncodeID(challengeName) + + if got, want := ProjectNameNotInstanced(challengeName), "beast-"+encoded; got != want { + t.Fatalf("ProjectNameNotInstanced() = %q, want %q", got, want) + } + + if got, want := ComposeDockerProjectNameInstanced(challengeName, "abc123"), "beast-instance-"+encoded+"-abc123"; got != want { + t.Fatalf("ComposeDockerProjectNameInstanced() = %q, want %q", got, want) + } +} diff --git a/utils/prompt.go b/utils/prompt.go index cc0178c3..dda4e63b 100644 --- a/utils/prompt.go +++ b/utils/prompt.go @@ -67,6 +67,31 @@ func PromptInt64(prompt string, defaultValue int64) int64 { return tempInt } +func PromptFloat32(prompt string, defaultValue float32) float32 { + log.Println(fmt.Sprintf("%s (defaults to %v)", prompt, defaultValue)) + + scanner := bufio.NewScanner(os.Stdin) + scanner.Scan() + + if err := scanner.Err(); err != nil { + log.Errorln(fmt.Sprintf("Failed to read input... defaulting to %v...", defaultValue)) + return defaultValue + } + + temp := scanner.Text() + tempFloat, err := strconv.ParseFloat(temp, 32) + + if temp == "" { + log.Warnln(fmt.Sprintf("Input empty.. defaulting to %v...", defaultValue)) + return defaultValue + } else if err != nil { + log.Errorln(fmt.Sprintf("Failed to read input... defaulting to %v...", defaultValue)) + return defaultValue + } + + return float32(tempFloat) +} + func PromptSelection(prompt string, items []string) string { selection := promptui.Select{ Label: prompt,