Безопасность IoT WebSocket сервера (Go)

// ❌ ОПАСНО — любой сайт может подключиться к твоему WS
var upgrader = websocket.Upgrader{
    CheckOrigin: func(r *http.Request) bool {
        return true
    },
}

// ✅ ПРАВИЛЬНО — проверяй origin как в CORS
var allowedOrigins = map[string]bool{
    "http://localhost:5173":    true,
    "https://grf-bs.ru":        true,
    "https://erp.e-mqtt.ru":    true,
}

var upgrader = websocket.Upgrader{
    CheckOrigin: func(r *http.Request) bool {
        origin := r.Header.Get("Origin")
        return allowedOrigins[origin]
    },
}

// ❌ Любой может подключиться к /api/v1/ws и читать данные IoT

// ✅ Проверяй JWT/API-ключ перед апгрейдом
func wsHandler(w http.ResponseWriter, r *http.Request) {
    token := r.URL.Query().Get("token") // или Header Authorization
    if !isValidToken(token) {
        http.Error(w, "Unauthorized", http.StatusUnauthorized)
        return
    }
    conn, err := upgrader.Upgrade(w, r, nil)
    // ...
}

// ❌ map без мьютекса — паника при конкурентном доступе
var wsClients = make(map[*websocket.Conn]bool)

// ✅ Защити мьютексом
var (
    wsClients   = make(map[*websocket.Conn]bool)
    wsClientsMu sync.RWMutex
)

// При добавлении:
wsClientsMu.Lock()
wsClients[conn] = true
wsClientsMu.Unlock()

// При удалении:
wsClientsMu.Lock()
delete(wsClients, conn)
wsClientsMu.Unlock()

// ❌ /api/v1/objects/{imei}/commands/{cmd} — команды на устройства без авторизации!

// ✅ Middleware для проверки токена
func authMiddleware(next http.Handler) http.Handler {
    return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
        token := r.Header.Get("Authorization")
        if !isValidToken(strings.TrimPrefix(token, "Bearer ")) {
            http.Error(w, "Unauthorized", http.StatusUnauthorized)
            return
        }
        next.ServeHTTP(w, r)
    })
}

// Применяй к роутеру:
r.HandleFunc("/api/v1/objects/{imei}/commands/{cmd}", CommandHandler).Methods("POST")
handler := c.Handler(authMiddleware(r))

// ❌ Клиент может прислать гигабайтное сообщение

// ✅ Ограничь размер
conn.SetReadLimit(512 * 1024) // 512 KB максимум

// ✅ Устанавливай дедлайн на чтение и обрабатывай pong
const pongWait = 60 * time.Second
const pingPeriod = 50 * time.Second

conn.SetReadDeadline(time.Now().Add(pongWait))
conn.SetPongHandler(func(string) error {
    conn.SetReadDeadline(time.Now().Add(pongWait))
    return nil
})

// ❌ ws://  — данные IoT передаются в открытом виде
// ✅ wss:// — используй TLS (nginx reverse proxy + Let's Encrypt)

// ❌ /api/v1/objects/../../../etc — path traversal / некорректные данные
// ✅ Валидируй IMEI — только цифры, 15 символов
func isValidIMEI(imei string) bool {
    matched, _ := regexp.MatchString(`^\d{15}$`, imei)
    return matched
}

package main

import (
    "log"
    "net/http"
    "regexp"
    "strings"
    "sync"
    "time"

    "github.com/gorilla/websocket"
)

// --- Константы ---
const (
    pongWait       = 60 * time.Second
    pingPeriod     = (pongWait * 9) / 10 // чуть меньше pongWait
    maxMessageSize = 512 * 1024           // 512 KB
)

// --- Разрешённые origins ---
var allowedOrigins = map[string]bool{
    "http://localhost:5173": true,
    "https://grf-bs.ru":     true,
    "https://erp.e-mqtt.ru": true,
}

// --- Потокобезопасное хранилище клиентов ---
type ClientHub struct {
    mu      sync.RWMutex
    clients map[*websocket.Conn]bool
}

var hub = &ClientHub{
    clients: make(map[*websocket.Conn]bool),
}

func (h *ClientHub) add(conn *websocket.Conn) {
    h.mu.Lock()
    defer h.mu.Unlock()
    h.clients[conn] = true
}

func (h *ClientHub) remove(conn *websocket.Conn) {
    h.mu.Lock()
    defer h.mu.Unlock()
    delete(h.clients, conn)
}

func (h *ClientHub) broadcast(message []byte) {
    h.mu.RLock()
    defer h.mu.RUnlock()
    for conn := range h.clients {
        // каждый conn.WriteMessage тоже лучше защищать отдельным мьютексом на запись,
        // если broadcast вызывается из разных горутин
        if err := conn.WriteMessage(websocket.TextMessage, message); err != nil {
            log.Println("Ошибка broadcast:", err)
        }
    }
}

// --- Upgrader с проверкой Origin ---
var upgrader = websocket.Upgrader{
    ReadBufferSize:  1024,
    WriteBufferSize: 1024,
    CheckOrigin: func(r *http.Request) bool {
        origin := r.Header.Get("Origin")
        return allowedOrigins[origin]
    },
}

// --- Валидация IMEI ---
var imeiRegex = regexp.MustCompile(`^\d{15}$`)

func isValidIMEI(imei string) bool {
    return imeiRegex.MatchString(imei)
}

// --- Валидация токена ---
func isValidToken(token string) bool {
    // Замени на свою логику: JWT-проверку, lookup в Redis и т.д.
    // Пример с фиксированным API-ключом (только для примера!):
    return token != "" && token == getExpectedToken()
}

func getExpectedToken() string {
    // В реальном коде — читай из env переменной или проверяй JWT
    return "" // TODO: os.Getenv("WS_API_KEY") или JWT-валидация
}

// --- Middleware аутентификации для HTTP ---
func authMiddleware(next http.Handler) http.Handler {
    return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
        authHeader := r.Header.Get("Authorization")
        token := strings.TrimPrefix(authHeader, "Bearer ")

        if !isValidToken(token) {
            http.Error(w, `{"error":"unauthorized"}`, http.StatusUnauthorized)
            return
        }
        next.ServeHTTP(w, r)
    })
}

// --- WebSocket Handler ---
func wsHandler(w http.ResponseWriter, r *http.Request) {
    // 1. Аутентификация — токен передаём в query param (т.к. WS не поддерживает кастомные заголовки из браузера)
    token := r.URL.Query().Get("token")
    if !isValidToken(token) {
        http.Error(w, `{"error":"unauthorized"}`, http.StatusUnauthorized)
        log.Println("WS: неавторизованная попытка подключения от", r.RemoteAddr)
        return
    }

    // 2. Апгрейд соединения
    conn, err := upgrader.Upgrade(w, r, nil)
    if err != nil {
        log.Println("WS: ошибка апгрейда:", err)
        return
    }
    defer func() {
        hub.remove(conn)
        conn.Close()
        log.Println("WS: клиент отключился", r.RemoteAddr)
    }()

    // 3. Настройки безопасности соединения
    conn.SetReadLimit(maxMessageSize)
    conn.SetReadDeadline(time.Now().Add(pongWait))
    conn.SetPongHandler(func(string) error {
        // Сбрасываем дедлайн при получении pong
        conn.SetReadDeadline(time.Now().Add(pongWait))
        return nil
    })

    // 4. Регистрируем клиента
    hub.add(conn)
    log.Println("WS: новый клиент подключился", r.RemoteAddr)

    // 5. Горутина для отправки ping-ов
    go func() {
        ticker := time.NewTicker(pingPeriod)
        defer ticker.Stop()
        for range ticker.C {
            conn.SetWriteDeadline(time.Now().Add(10 * time.Second))
            if err := conn.WriteMessage(websocket.PingMessage, nil); err != nil {
                log.Println("WS: ошибка ping, отключаем клиента:", err)
                conn.Close()
                return
            }
        }
    }()

    // 6. Основной цикл чтения
    for {
        _, msg, err := conn.ReadMessage()
        if err != nil {
            if websocket.IsUnexpectedCloseError(err,
                websocket.CloseGoingAway,
                websocket.CloseAbnormalClosure,
            ) {
                log.Println("WS: неожиданное закрытие:", err)
            }
            break
        }

        // Обработка входящего сообщения (при необходимости)
        log.Printf("WS: получено сообщение от %s: %s", r.RemoteAddr, msg)
    }
}

r := mux.NewRouter()
r.HandleFunc("/", HomeHandler)

// WS — аутентификация внутри самого wsHandler (через query param ?token=...)
r.HandleFunc("/api/v1/ws", wsHandler)

// Все REST эндпоинты защищены middleware
api := r.PathPrefix("/api").Subrouter()
api.Use(authMiddleware)
api.HandleFunc("/v1/objects/{imei}/commands/{cmd}", func(w http.ResponseWriter, r *http.Request) {
    vars := mux.Vars(r)
    if !isValidIMEI(vars["imei"]) {
        http.Error(w, `{"error":"invalid imei"}`, http.StatusBadRequest)
        return
    }
    CommandHandler(w, r)
}).Methods("POST")
api.HandleFunc("/v1/objects/{imei}", HTTPIotLatestOneScooterData).Methods("GET")
api.HandleFunc("/v1/objects/{imei}/config", HTTPScooterConfig).Methods("GET")
api.HandleFunc("/v1/objects", HTTPIotDataScootersV2).Methods("GET")
api.HandleFunc("/v2/objects", HTTPIotDataScootersV2).Methods("GET")
api.HandleFunc("/v1/objects/{imei}/journal", HTTPIotOneScooterCmdsJournal).Methods("GET")
api.HandleFunc("/v1/objects/{imei}/packets", HTTPIotScooterPackets).Methods("GET")
api.HandleFunc("/v1/objects-multiple", HTTPIotLatestMultipleScootersData).Methods("GET")