feat: add WebSocket security vulnerability module (Issue #527)

build.gradle

@@ -111,6 +111,7 @@ tasks.register('GenerateSampleVulnerability'){
 
 dependencies {
     implementation group: 'org.springframework.boot', name: 'spring-boot-starter-web'
+    implementation group: 'org.springframework.boot', name: 'spring-boot-starter-websocket'
 
     // https://mvnrepository.com/artifact/org.mockito/mockito-core
     testImplementation group: 'org.mockito', name: 'mockito-core', version: '3.5.13'

src/main/java/org/sasanlabs/service/vulnerability/websocket/InjectionWebSocketHandler.java

@@ -0,0 +1,42 @@
+package org.sasanlabs.service.vulnerability.websocket;
+
+import org.apache.logging.log4j.LogManager;
+import org.apache.logging.log4j.Logger;
+import org.springframework.stereotype.Component;
+import org.springframework.web.socket.CloseStatus;
+import org.springframework.web.socket.TextMessage;
+import org.springframework.web.socket.WebSocketSession;
+import org.springframework.web.socket.handler.TextWebSocketHandler;
+
+/**
+ * Level 2: WebSocket handler that reflects user messages without input validation. Demonstrates
+ * message injection risk when payloads are reflected without sanitization.
+ */
+@Component
+public class InjectionWebSocketHandler extends TextWebSocketHandler {
+
+    private static final Logger LOGGER = LogManager.getLogger(InjectionWebSocketHandler.class);
+
+    @Override
+    public void afterConnectionEstablished(WebSocketSession session) throws Exception {
+        LOGGER.info("WebSocket connection established (injection level): {}", session.getId());
+        session.sendMessage(
+                new TextMessage(
+                        "Connected! Messages are processed and reflected without input"
+                                + " validation."));
+    }
+
+    @Override
+    protected void handleTextMessage(WebSocketSession session, TextMessage message)
+            throws Exception {
+        String payload = message.getPayload();
+        LOGGER.info("Received message without validation: {}", payload);
+        // Vulnerability: user input reflected directly without sanitization
+        session.sendMessage(new TextMessage("Server processed: " + payload));
+    }
+
+    @Override
+    public void afterConnectionClosed(WebSocketSession session, CloseStatus status) {
+        LOGGER.info("WebSocket connection closed: {}", session.getId());
+    }
+}

src/main/java/org/sasanlabs/service/vulnerability/websocket/NoOriginValidationWebSocketHandler.java

@@ -0,0 +1,33 @@
+package org.sasanlabs.service.vulnerability.websocket;
+
+import org.springframework.stereotype.Component;
+import org.springframework.web.socket.TextMessage;
+import org.springframework.web.socket.WebSocketSession;
+import org.springframework.web.socket.handler.TextWebSocketHandler;
+
+/**
+ * Level 3: No Origin header validation (Cross-Site WebSocket Hijacking / CWE-346).
+ *
+ * <p>The server accepts WebSocket connections from any origin without checking the Origin header. A
+ * malicious page on a different domain could silently open a connection and communicate on behalf
+ * of a logged-in victim.
+ */
+@Component
+public class NoOriginValidationWebSocketHandler extends TextWebSocketHandler {
+
+    @Override
+    public void afterConnectionEstablished(WebSocketSession session) throws Exception {
+        String origin = session.getHandshakeHeaders().getFirst("Origin");
+        session.sendMessage(
+                new TextMessage(
+                        "Connected! Origin header was: "
+                                + origin
+                                + " — but it was never validated. Any origin is accepted."));
+    }
+
+    @Override
+    protected void handleTextMessage(WebSocketSession session, TextMessage message)
+            throws Exception {
+        session.sendMessage(new TextMessage("Echo: " + message.getPayload()));
+    }
+}

src/main/java/org/sasanlabs/service/vulnerability/websocket/SecureWebSocketHandler.java

@@ -0,0 +1,39 @@
+package org.sasanlabs.service.vulnerability.websocket;
+
+import org.apache.commons.text.StringEscapeUtils;
+import org.springframework.stereotype.Component;
+import org.springframework.web.socket.CloseStatus;
+import org.springframework.web.socket.TextMessage;
+import org.springframework.web.socket.WebSocketSession;
+import org.springframework.web.socket.handler.TextWebSocketHandler;
+
+/**
+ * Level 4: Secure WebSocket implementation.
+ *
+ * <p>Demonstrates the correct way to secure a WebSocket endpoint: token-based authentication via
+ * query parameter, origin restricted to the application's own host, and HTML sanitization of all
+ * incoming messages before echoing.
+ */
+@Component
+public class SecureWebSocketHandler extends TextWebSocketHandler {
+
+    static final String VALID_TOKEN = "secureToken123";
+
+    @Override
+    public void afterConnectionEstablished(WebSocketSession session) throws Exception {
+        String query = session.getUri() != null ? session.getUri().getQuery() : null;
+        if (query == null || !query.contains("token=" + VALID_TOKEN)) {
+            session.close(CloseStatus.POLICY_VIOLATION);
+            return;
+        }
+        session.sendMessage(
+                new TextMessage("Authenticated! Messages will be HTML-sanitized before echoing."));
+    }
+
+    @Override
+    protected void handleTextMessage(WebSocketSession session, TextMessage message)
+            throws Exception {
+        String sanitized = StringEscapeUtils.escapeHtml4(message.getPayload());
+        session.sendMessage(new TextMessage("Secure echo: " + sanitized));
+    }
+}

src/main/java/org/sasanlabs/service/vulnerability/websocket/UnauthenticatedWebSocketHandler.java

@@ -0,0 +1,42 @@
+package org.sasanlabs.service.vulnerability.websocket;
+
+import org.apache.logging.log4j.LogManager;
+import org.apache.logging.log4j.Logger;
+import org.springframework.stereotype.Component;
+import org.springframework.web.socket.CloseStatus;
+import org.springframework.web.socket.TextMessage;
+import org.springframework.web.socket.WebSocketSession;
+import org.springframework.web.socket.handler.TextWebSocketHandler;
+
+/**
+ * Level 1: WebSocket handler with no authentication. Any client can connect and exchange messages
+ * without providing any credentials or token.
+ */
+@Component
+public class UnauthenticatedWebSocketHandler extends TextWebSocketHandler {
+
+    private static final Logger LOGGER =
+            LogManager.getLogger(UnauthenticatedWebSocketHandler.class);
+
+    @Override
+    public void afterConnectionEstablished(WebSocketSession session) throws Exception {
+        LOGGER.info("WebSocket connection established without authentication: {}", session.getId());
+        session.sendMessage(
+                new TextMessage(
+                        "Connected! This endpoint accepts connections without any"
+                                + " authentication."));
+    }
+
+    @Override
+    protected void handleTextMessage(WebSocketSession session, TextMessage message)
+            throws Exception {
+        String payload = message.getPayload();
+        LOGGER.info("Message received on unauthenticated endpoint: {}", payload);
+        session.sendMessage(new TextMessage("Echo: " + payload));
+    }
+
+    @Override
+    public void afterConnectionClosed(WebSocketSession session, CloseStatus status) {
+        LOGGER.info("WebSocket connection closed: {}", session.getId());
+    }
+}

src/main/java/org/sasanlabs/service/vulnerability/websocket/WebSocketVulnerability.java

@@ -0,0 +1,70 @@
+package org.sasanlabs.service.vulnerability.websocket;
+
+import org.sasanlabs.internal.utility.LevelConstants;
+import org.sasanlabs.internal.utility.Variant;
+import org.sasanlabs.internal.utility.annotations.AttackVector;
+import org.sasanlabs.internal.utility.annotations.VulnerableAppRequestMapping;
+import org.sasanlabs.internal.utility.annotations.VulnerableAppRestController;
+import org.sasanlabs.service.vulnerability.bean.GenericVulnerabilityResponseBean;
+import org.sasanlabs.vulnerability.types.VulnerabilityType;
+import org.springframework.http.HttpStatus;
+import org.springframework.http.ResponseEntity;
+
+/**
+ * WebSocket Security Vulnerability module demonstrating insecure WebSocket implementations and
+ * their secure alternatives.
+ *
+ * <p>Issue: https://github.com/SasanLabs/VulnerableApp/issues/527
+ */
+@VulnerableAppRestController(
+        descriptionLabel = "WEBSOCKET_VULNERABILITY",
+        value = "WebSocketVulnerability")
+public class WebSocketVulnerability {
+
+    private static final String WS_LEVEL1_PATH = "/ws/websocket-vulnerability/level1";
+    private static final String WS_LEVEL2_PATH = "/ws/websocket-vulnerability/level2";
+    private static final String WS_LEVEL3_PATH = "/ws/websocket-vulnerability/level3";
+    private static final String WS_LEVEL4_PATH = "/ws/websocket-vulnerability/level4";
+
+    @AttackVector(
+            vulnerabilityExposed = VulnerabilityType.WEBSOCKET_MISSING_AUTHENTICATION,
+            description = "WEBSOCKET_NO_AUTHENTICATION",
+            payload = "WEBSOCKET_PAYLOAD_LEVEL_1")
+    @VulnerableAppRequestMapping(value = LevelConstants.LEVEL_1, htmlTemplate = "LEVEL_1/WebSocket")
+    public ResponseEntity<GenericVulnerabilityResponseBean<String>> level1() {
+        return new ResponseEntity<>(
+                new GenericVulnerabilityResponseBean<>(WS_LEVEL1_PATH, true), HttpStatus.OK);
+    }
+
+    @AttackVector(
+            vulnerabilityExposed = VulnerabilityType.WEBSOCKET_MESSAGE_INJECTION,
+            description = "WEBSOCKET_NO_INPUT_VALIDATION",
+            payload = "WEBSOCKET_PAYLOAD_LEVEL_2")
+    @VulnerableAppRequestMapping(value = LevelConstants.LEVEL_2, htmlTemplate = "LEVEL_1/WebSocket")
+    public ResponseEntity<GenericVulnerabilityResponseBean<String>> level2() {
+        return new ResponseEntity<>(
+                new GenericVulnerabilityResponseBean<>(WS_LEVEL2_PATH, true), HttpStatus.OK);
+    }
+
+    @AttackVector(
+            vulnerabilityExposed = VulnerabilityType.CROSS_ORIGIN_WEBSOCKET_HIJACKING,
+            description = "WEBSOCKET_NO_ORIGIN_VALIDATION",
+            payload = "WEBSOCKET_PAYLOAD_LEVEL_3")
+    @VulnerableAppRequestMapping(value = LevelConstants.LEVEL_3, htmlTemplate = "LEVEL_1/WebSocket")
+    public ResponseEntity<GenericVulnerabilityResponseBean<String>> level3() {
+        return new ResponseEntity<>(
+                new GenericVulnerabilityResponseBean<>(WS_LEVEL3_PATH, true), HttpStatus.OK);
+    }
+
+    @AttackVector(
+            vulnerabilityExposed = VulnerabilityType.WEBSOCKET_MISSING_AUTHENTICATION,
+            description = "WEBSOCKET_SECURE_IMPLEMENTATION")
+    @VulnerableAppRequestMapping(
+            value = LevelConstants.LEVEL_4,
+            htmlTemplate = "LEVEL_1/WebSocket",
+            variant = Variant.SECURE)
+    public ResponseEntity<GenericVulnerabilityResponseBean<String>> level4() {
+        return new ResponseEntity<>(
+                new GenericVulnerabilityResponseBean<>(WS_LEVEL4_PATH, true), HttpStatus.OK);
+    }
+}

src/main/java/org/sasanlabs/service/vulnerability/websocket/WebSocketVulnerabilityConfig.java

@@ -0,0 +1,38 @@
+package org.sasanlabs.service.vulnerability.websocket;
+
+import org.springframework.beans.factory.annotation.Autowired;
+import org.springframework.context.annotation.Configuration;
+import org.springframework.web.socket.config.annotation.EnableWebSocket;
+import org.springframework.web.socket.config.annotation.WebSocketConfigurer;
+import org.springframework.web.socket.config.annotation.WebSocketHandlerRegistry;
+
+/** Registers WebSocket handlers for each vulnerability level. */
+@Configuration
+@EnableWebSocket
+public class WebSocketVulnerabilityConfig implements WebSocketConfigurer {
+
+    @Autowired private UnauthenticatedWebSocketHandler unauthenticatedHandler;
+
+    @Autowired private InjectionWebSocketHandler injectionHandler;
+
+    @Autowired private NoOriginValidationWebSocketHandler noOriginHandler;
+
+    @Autowired private SecureWebSocketHandler secureHandler;
+
+    @Override
+    public void registerWebSocketHandlers(WebSocketHandlerRegistry registry) {
+        registry.addHandler(unauthenticatedHandler, "/ws/websocket-vulnerability/level1")
+                .setAllowedOriginPatterns("*");
+
+        registry.addHandler(injectionHandler, "/ws/websocket-vulnerability/level2")
+                .setAllowedOriginPatterns("*");
+
+        // Vulnerable: accepts connections from any origin without validation (CWE-346)
+        registry.addHandler(noOriginHandler, "/ws/websocket-vulnerability/level3")
+                .setAllowedOriginPatterns("*");
+
+        // Secure: restricts origin to application host only
+        registry.addHandler(secureHandler, "/ws/websocket-vulnerability/level4")
+                .setAllowedOrigins("http://localhost:9090");
+    }
+}

src/main/java/org/sasanlabs/vulnerability/types/VulnerabilityType.java

@@ -58,7 +58,12 @@ public enum VulnerabilityType {
     USERNAME_ENUMERATION(204, null),
 
     // LDAP Injection Vulnerability
-    LDAP_INJECTION(90, 29);
+    LDAP_INJECTION(90, 29),
+
+    // WebSocket Vulnerabilities
+    WEBSOCKET_MISSING_AUTHENTICATION(306, null),
+    WEBSOCKET_MESSAGE_INJECTION(20, null),
+    CROSS_ORIGIN_WEBSOCKET_HIJACKING(346, null);
 
     private Integer cweID;
     private Integer wascID;

src/main/resources/attackvectors/WebSocketVulnerabilityPayload.properties

@@ -0,0 +1,15 @@
+WEBSOCKET_PAYLOAD_LEVEL_1=Connect to the WebSocket endpoint without any authentication token. \
+The endpoint accepts all incoming connections - click <b>Connect</b> and observe that no credentials are required. \
+This demonstrates CWE-306 (Missing Authentication for Critical Function).
+WEBSOCKET_PAYLOAD_LEVEL_2=Connect to the WebSocket endpoint and send an injection payload such as \
+<code>&lt;script&gt;alert('XSS')&lt;/script&gt;</code> or <code>'; DROP TABLE users; --</code>. \
+The server reflects the message back without any sanitization, demonstrating how unsanitized \
+WebSocket messages can be exploited for injection attacks (CWE-20).
+WEBSOCKET_PAYLOAD_LEVEL_3=Connect from any origin and observe that the server accepts the connection \
+without validating the Origin header. The greeting message reveals the Origin that was sent. \
+In a real attack, a malicious page on evil.com could silently open a WebSocket to this server \
+and send messages on behalf of a logged-in victim (Cross-Site WebSocket Hijacking, CWE-346).
+WEBSOCKET_PAYLOAD_LEVEL_4=This level requires a valid token (<code>secureToken123</code>) as a \
+query parameter - enter it in the Token field and click Connect. Without the correct token the \
+connection is rejected. Messages are HTML-escaped before being echoed, preventing injection attacks. \
+Origin is restricted to localhost, blocking cross-origin hijacking attempts.

src/main/resources/i18n/messages.properties

@@ -370,6 +370,20 @@ LDAP_LEVEL_4_INPUT_SANITIZATION=This level demonstrates input encoding using a s
 LDAP_LEVEL_5_BLIND_INJECTION=Blind LDAP Injection leading to Authentication Bypass. The application does not display user data but reveals success or failure responses. Attackers can manipulate LDAP filters to bypass authentication or infer valid credentials.
 LDAP_LEVEL_6_SECURE=This level demonstrates a fully secure LDAP authentication implementation. \ User input is properly encoded, password verification is done using salted hashing, \ and no sensitive infformation is exposed in responses. LDAP Injection is not possible.
 
+### WebSocket Security Vulnerability
+WEBSOCKET_VULNERABILITY=<b>WebSocket Security Vulnerabilities</b> occur when applications fail to properly secure or validate WebSocket connections. Unlike HTTP, WebSocket provides a persistent, bidirectional communication channel, which introduces unique security challenges that traditional scanners may overlook.<br/><br/>\
+WebSocket vulnerabilities include missing authentication (allowing any client to connect), absence of input validation (enabling message injection), and missing Origin header validation (allowing cross-origin WebSocket hijacking).<br/><br/>\
+This module demonstrates the progression from an insecure WebSocket implementation to a hardened configuration.<br/><br/>\
+Important Links:<br/>\
+<ol> <li> <a href="https://owasp.org/www-community/attacks/Cross_Site_WebSocket_Hijacking" target="_blank">OWASP: Cross-Site WebSocket Hijacking</a> \
+    <li> <a href="https://portswigger.net/web-security/websockets" target="_blank">PortSwigger: WebSocket Vulnerabilities</a> \
+    <li> <a href="https://cwe.mitre.org/data/definitions/306.html" target="_blank">CWE-306: Missing Authentication for Critical Function</a> \
+</ol>
+WEBSOCKET_NO_AUTHENTICATION=The WebSocket endpoint accepts connections without any authentication, enabling any client to connect and interact with the server.
+WEBSOCKET_NO_INPUT_VALIDATION=User-controlled WebSocket messages are processed without input validation, allowing attackers to inject malicious payloads or manipulate application behavior.
+WEBSOCKET_NO_ORIGIN_VALIDATION=The server does not validate the Origin header during the WebSocket handshake, allowing malicious websites on any domain to establish WebSocket connections (Cross-Site WebSocket Hijacking, CWE-346).
+WEBSOCKET_SECURE_IMPLEMENTATION=Secure WebSocket implementation with token-based authentication via query parameter, origin restricted to the application domain, and HTML sanitization of all incoming messages before echoing.
+
 ### Authentication Vulnerability
 AUTHENTICATION_VULNERABILITY=<b>Authentication Vulnerabilities</b> occur when an application's login mechanism is implemented insecurely, \
 allowing attackers to compromise user accounts. Common weaknesses include SQL injection in the login query, logging sensitive data, storing passwords in plaintext or with weak hashing algorithms (MD5, SHA1), \
@@ -398,3 +412,4 @@ AUTH_LEVEL_8_WEAK_PASSWORD=Credentials are passed in the POST body to avoid URI
 AUTH_LEVEL_9_SECURE=Secure baseline: Credentials passed in the POST body, BCrypt hashing with unique salts, and a generic \"Invalid credentials\" error message regardless of the failure cause, preventing cracking and enumeration.
 AUTH_PAYLOAD_LEVEL_8=Try logging in with 'admin_weak' and a common weak password (hint: it starts with 'password'). Even with strong BCrypt hashing, a weak password can be easily guessed.
 AUTH_PAYLOAD_LEVEL_9=This is the secure implementation. Try to enumerate usernames or crack the hash \u2014 you will find it significantly more difficult!
+