# Milestone 10: Network - WebSocket

**Status**: Complete
**Goal**: Secure WebSocket connections with same validation as HTTP.

---

## Overview

This milestone extends the network system with WebSocket support:
- Reuse HttpValidator for URL/domain validation (wss:// only)
- Connection limits per app
- Message size limits
- Idle timeout handling
- Clean disconnection on app stop

### Key Deliverables

1. **WebSocketManager class** - Connection pool management
2. **WebSocket class** - Individual connection wrapper
3. **Lua websocket API** - `network.websocket()` function
4. **Mock implementation** - For desktop testing

---

## File Structure

```
src/main/cpp/sandbox/
├── http_validator.h        # Existing - add WSS support
├── http_validator.cpp      # Existing - add WSS support
├── websocket_manager.h     # NEW - WebSocket pool
└── websocket_manager.cpp   # NEW - Connection management
```

---

## Implementation Details

### 1. WebSocketManager Class

```cpp
// websocket_manager.h
#pragma once

#include <string>
#include <map>
#include <memory>
#include <mutex>
#include <functional>
#include <vector>
#include "http_validator.h"

struct lua_State;

namespace mosis {

struct WebSocketLimits {
    int max_connections_per_app = 5;
    size_t max_message_size = 1 * 1024 * 1024;     // 1 MB
    int idle_timeout_ms = 5 * 60 * 1000;           // 5 minutes
    int connect_timeout_ms = 30000;                 // 30 seconds
};

enum class WebSocketState {
    Connecting,
    Open,
    Closing,
    Closed
};

class WebSocket {
public:
    using MessageCallback = std::function<void(const std::string& data, bool binary)>;
    using CloseCallback = std::function<void(int code, const std::string& reason)>;
    using ErrorCallback = std::function<void(const std::string& error)>;
    using OpenCallback = std::function<void()>;

    WebSocket(int id, const std::string& url);
    ~WebSocket();

    int GetId() const { return m_id; }
    const std::string& GetUrl() const { return m_url; }
    WebSocketState GetState() const { return m_state; }

    // Send message (returns false if not connected or message too large)
    bool Send(const std::string& data, bool binary = false);

    // Close connection
    void Close(int code = 1000, const std::string& reason = "");

    // Event callbacks
    void SetOnOpen(OpenCallback cb) { m_on_open = std::move(cb); }
    void SetOnMessage(MessageCallback cb) { m_on_message = std::move(cb); }
    void SetOnClose(CloseCallback cb) { m_on_close = std::move(cb); }
    void SetOnError(ErrorCallback cb) { m_on_error = std::move(cb); }

    // For mock mode - simulate events
    void SimulateOpen();
    void SimulateMessage(const std::string& data, bool binary);
    void SimulateClose(int code, const std::string& reason);
    void SimulateError(const std::string& error);

private:
    int m_id;
    std::string m_url;
    WebSocketState m_state;
    size_t m_max_message_size;

    OpenCallback m_on_open;
    MessageCallback m_on_message;
    CloseCallback m_on_close;
    ErrorCallback m_on_error;
};

class WebSocketManager {
public:
    WebSocketManager(const std::string& app_id, const WebSocketLimits& limits = WebSocketLimits{});
    ~WebSocketManager();

    // Configure domain restrictions (reuses HttpValidator)
    void SetAllowedDomains(const std::vector<std::string>& domains);
    void ClearDomainRestrictions();

    // Create new WebSocket connection
    // Returns connection ID on success, -1 on failure (sets error)
    std::shared_ptr<WebSocket> Connect(const std::string& url, std::string& error);

    // Get connection by ID
    std::shared_ptr<WebSocket> GetConnection(int id);

    // Close specific connection
    void CloseConnection(int id, int code = 1000, const std::string& reason = "");

    // Close all connections (called on app stop)
    void CloseAll();

    // Stats
    int GetActiveConnectionCount() const;

    // Access validator for testing
    HttpValidator& GetValidator() { return m_validator; }

    // For testing: set mock mode
    void SetMockMode(bool enabled) { m_mock_mode = enabled; }
    bool IsMockMode() const { return m_mock_mode; }

private:
    std::string m_app_id;
    WebSocketLimits m_limits;
    HttpValidator m_validator;
    std::map<int, std::shared_ptr<WebSocket>> m_connections;
    int m_next_id = 1;
    mutable std::mutex m_mutex;
    bool m_mock_mode = true;

    bool ValidateUrl(const std::string& url, std::string& error);
};

// Register websocket APIs
void RegisterWebSocketAPI(lua_State* L, WebSocketManager* manager);

} // namespace mosis
```

### 2. URL Validation Changes

The HttpValidator needs to accept both `https://` and `wss://` schemes:

```cpp
// In Validate():
if (parsed->scheme != "https" && parsed->scheme != "wss") {
    error = "HTTPS or WSS required, got: " + parsed->scheme;
    return std::nullopt;
}
```

### 3. Lua API

```lua
-- Create WebSocket connection
local ws, err = network.websocket("wss://api.example.com/socket")
if not ws then
    print("Failed to connect:", err)
    return
end

-- Event handlers
ws:on("open", function()
    print("Connected!")
    ws:send("Hello server")
end)

ws:on("message", function(data, binary)
    print("Received:", data)
end)

ws:on("close", function(code, reason)
    print("Closed:", code, reason)
end)

ws:on("error", function(err)
    print("Error:", err)
end)

-- Send message
ws:send("Hello")
ws:send(binaryData, true)  -- binary mode

-- Close connection
ws:close()
ws:close(1000, "Normal closure")

-- Get state
local state = ws:state()  -- "connecting", "open", "closing", "closed"
```

### 4. Connection Lifecycle

```
Connect(url) ──► Validating ──► Connecting ──► Open
                    │              │             │
                    │ error        │ error       │ message
                    ▼              ▼             ▼
                 Closed         Closed        Handler
                                               │
                                               │ close/error
                                               ▼
                                            Closed
```

### 5. Limits Enforcement

| Limit | Default | Description |
|-------|---------|-------------|
| max_connections_per_app | 5 | Maximum concurrent WebSocket connections |
| max_message_size | 1 MB | Maximum message size (send or receive) |
| idle_timeout_ms | 5 min | Close idle connections |
| connect_timeout_ms | 30 sec | Connection establishment timeout |

---

## Test Cases

### Test 1: WebSocket URL Validation

```cpp
bool Test_WebSocketUrlValidation(std::string& error_msg) {
    mosis::WebSocketManager manager("test.app");
    manager.ClearDomainRestrictions();

    std::string err;

    // WSS should be allowed
    auto ws = manager.Connect("wss://example.com/socket", err);
    // In mock mode, connection will fail but validation should pass
    EXPECT_TRUE(err.find("mock") != std::string::npos ||
                err.find("disabled") != std::string::npos ||
                ws != nullptr);

    // WS (plain) should be blocked
    ws = manager.Connect("ws://example.com/socket", err);
    EXPECT_TRUE(ws == nullptr);
    EXPECT_TRUE(err.find("WSS") != std::string::npos ||
                err.find("HTTPS") != std::string::npos ||
                err.find("required") != std::string::npos);

    return true;
}
```

### Test 2: Connection Limits

```cpp
bool Test_WebSocketConnectionLimits(std::string& error_msg) {
    mosis::WebSocketLimits limits;
    limits.max_connections_per_app = 2;

    mosis::WebSocketManager manager("test.app", limits);
    manager.ClearDomainRestrictions();

    std::string err;

    // Create max connections
    auto ws1 = manager.Connect("wss://example.com/socket1", err);
    auto ws2 = manager.Connect("wss://example.com/socket2", err);

    // Third should fail (or be limited)
    auto ws3 = manager.Connect("wss://example.com/socket3", err);
    EXPECT_TRUE(ws3 == nullptr);
    EXPECT_TRUE(err.find("limit") != std::string::npos ||
                err.find("connections") != std::string::npos);

    return true;
}
```

### Test 3: WebSocket Blocks Private IPs

```cpp
bool Test_WebSocketBlocksPrivateIP(std::string& error_msg) {
    mosis::WebSocketManager manager("test.app");
    manager.ClearDomainRestrictions();

    std::string err;

    std::vector<std::string> private_urls = {
        "wss://127.0.0.1/socket",
        "wss://localhost/socket",
        "wss://10.0.0.1/socket",
        "wss://192.168.1.1/socket",
        "wss://169.254.169.254/socket"
    };

    for (const auto& url : private_urls) {
        auto ws = manager.Connect(url, err);
        EXPECT_TRUE(ws == nullptr);
    }

    return true;
}
```

### Test 4: Domain Whitelist

```cpp
bool Test_WebSocketDomainWhitelist(std::string& error_msg) {
    mosis::WebSocketManager manager("test.app");
    manager.SetAllowedDomains({"api.example.com"});

    std::string err;

    // Allowed domain - should pass validation (may fail in mock mode for other reasons)
    auto ws1 = manager.Connect("wss://api.example.com/socket", err);
    bool allowed_passed = (ws1 != nullptr) ||
                          (err.find("mock") != std::string::npos) ||
                          (err.find("disabled") != std::string::npos);
    EXPECT_TRUE(allowed_passed);

    // Disallowed domain - should fail validation
    auto ws2 = manager.Connect("wss://evil.com/socket", err);
    EXPECT_TRUE(ws2 == nullptr);
    EXPECT_TRUE(err.find("allowed") != std::string::npos ||
                err.find("whitelist") != std::string::npos ||
                err.find("Domain") != std::string::npos);

    return true;
}
```

### Test 5: Message Size Limits

```cpp
bool Test_WebSocketMessageLimits(std::string& error_msg) {
    mosis::WebSocketLimits limits;
    limits.max_message_size = 1024;  // 1 KB for testing

    mosis::WebSocketManager manager("test.app", limits);
    manager.ClearDomainRestrictions();

    // Create a mock WebSocket directly to test send limits
    mosis::WebSocket ws(1, "wss://example.com/socket");

    // Small message should work (if connected)
    // Large message should fail
    std::string large_message(2048, 'X');  // 2 KB
    bool send_result = ws.Send(large_message);
    EXPECT_FALSE(send_result);  // Should fail - not connected and/or too large

    return true;
}
```

### Test 6: Close All Connections

```cpp
bool Test_WebSocketCloseAll(std::string& error_msg) {
    mosis::WebSocketManager manager("test.app");
    manager.ClearDomainRestrictions();

    std::string err;

    // Create some connections (may fail in mock but that's ok)
    manager.Connect("wss://example.com/socket1", err);
    manager.Connect("wss://example.com/socket2", err);

    // Close all
    manager.CloseAll();

    // Should have no active connections
    EXPECT_TRUE(manager.GetActiveConnectionCount() == 0);

    return true;
}
```

### Test 7: Lua Integration

```cpp
bool Test_WebSocketLuaIntegration(std::string& error_msg) {
    SandboxContext ctx = TestContext();
    LuaSandbox sandbox(ctx);

    mosis::WebSocketManager manager("test.app");
    manager.ClearDomainRestrictions();
    mosis::RegisterWebSocketAPI(sandbox.GetState(), &manager);

    std::string script = R"lua(
        -- Test that network.websocket exists
        if not network then
            error("network global not found")
        end
        if not network.websocket then
            error("network.websocket not found")
        end

        -- Test validation rejection (private IP)
        local ws, err = network.websocket("wss://127.0.0.1/socket")
        if ws then
            error("expected private IP to be blocked")
        end
    )lua";

    bool ok = sandbox.LoadString(script, "websocket_test");
    if (!ok) {
        error_msg = "Lua test failed: " + sandbox.GetLastError();
        return false;
    }
    return true;
}
```

---

## Acceptance Criteria

All tests must pass:

- [x] `Test_WebSocketUrlValidation` - WSS required, WS blocked
- [x] `Test_WebSocketConnectionLimits` - Per-app connection limits enforced
- [x] `Test_WebSocketBlocksPrivateIP` - SSRF prevention works
- [x] `Test_WebSocketDomainWhitelist` - Domain restrictions work
- [x] `Test_WebSocketMessageLimits` - Message size limits enforced
- [x] `Test_WebSocketCloseAll` - Cleanup works
- [x] `Test_WebSocketLuaIntegration` - Lua API works

---

## Dependencies

- Milestone 1 (LuaSandbox)
- Milestone 9 (HttpValidator)

---

## Notes

### Desktop vs Android Implementation

For desktop testing, WebSocketManager operates in mock mode:
- URL validation runs normally
- Connection limits are tracked
- Actual WebSocket connections are not made
- Events can be simulated for testing

On Android, the real implementation would:
1. Use Java WebSocket client through JNI
2. Handle background thread for socket I/O
3. Marshal events back to Lua thread

### Security Considerations

1. **Same-origin**: WSS URLs validated same as HTTPS
2. **Connection limits**: Prevent resource exhaustion
3. **Message limits**: Prevent memory exhaustion
4. **Idle timeout**: Automatic cleanup of abandoned connections
5. **Clean shutdown**: All connections closed on app stop

---

## Next Steps

After Milestone 10 passes:
1. Milestone 11: Virtual Hardware - Camera