wireguard-vpn/app/stats.go

package main

import (
    "encoding/json"
    "log"
    "net/http"
    "time"
    "math/rand"
    "github.com/gorilla/websocket"
)

// Client represents a single WebSocket connection.
type Client struct {
    hub  *Hub
    conn *websocket.Conn
    send chan []byte
}

// Hub maintains the set of active clients and broadcasts messages to them.
type Hub struct {
    clients    map[*Client]bool
    broadcast  chan []byte
    register   chan *Client
    unregister chan *Client
}

func newHub() *Hub {
    return &Hub{
        clients:    make(map[*Client]bool),
        broadcast:  make(chan []byte),
        register:   make(chan *Client),
        unregister: make(chan *Client),
    }
}

func (h *Hub) run() {
    for {
        select {
        case c := <-h.register:
            h.clients[c] = true
        case c := <-h.unregister:
            if _, ok := h.clients[c]; ok {
                delete(h.clients, c)
                close(c.send)
            }
        case message := <-h.broadcast:
            for c := range h.clients {
                select {
                case c.send <- message:
                default:
                    close(c.send)
                    delete(h.clients, c)
                }
            }
        }
    }
}

const (
    writeWait  = 10 * time.Second
    pongWait   = 60 * time.Second
    pingPeriod = (pongWait * 9) / 10
)

var upgrader = websocket.Upgrader{
    CheckOrigin: func(r *http.Request) bool { return true },
}

func (c *Client) writePump() {
    ticker := time.NewTicker(pingPeriod)
    defer func() {
        ticker.Stop()
        c.conn.Close()
    }()
    for {
        select {
        case message, ok := <-c.send:
            c.conn.SetWriteDeadline(time.Now().Add(writeWait))
            if !ok {
                c.conn.WriteMessage(websocket.CloseMessage, []byte{})
                return
            }
            w, err := c.conn.NextWriter(websocket.TextMessage)
            if err != nil {
                return
            }
            w.Write(message)
            if err := w.Close(); err != nil {
                return
            }
        case <-ticker.C:
            c.conn.SetWriteDeadline(time.Now().Add(writeWait))
            if err := c.conn.WriteMessage(websocket.PingMessage, nil); err != nil {
                return
            }
        }
    }
}

func (c *Client) readPump() {
    defer func() {
        c.hub.unregister <- c
        c.conn.Close()
    }()
    c.conn.SetReadLimit(5120)
    c.conn.SetReadDeadline(time.Now().Add(pongWait))
    c.conn.SetPongHandler(func(string) error { c.conn.SetReadDeadline(time.Now().Add(pongWait)); return nil })
    for {
        _, _, err := c.conn.ReadMessage()
        if err != nil {
            break
        }
    }
}

func serveWs(hub *Hub, w http.ResponseWriter, r *http.Request) {
    ws, err := upgrader.Upgrade(w, r, nil)
    if err != nil {
        log.Println("WebSocket upgrade error:", err)
        return
    }
    client := &Client{hub: hub, conn: ws, send: make(chan []byte, 256)}
    client.hub.register <- client

    go client.writePump()
    go client.readPump()
}

// Stats payload shape
type Stats struct {
    TotalRx      int `json:"total_rx"`
    TotalTx      int `json:"total_tx"`
    PeersOnline  int `json:"peers_online"`
    PeersOffline int `json:"peers_offline"`
}

// startStatsBroadcast periodically emits mock stats to all connected clients.
func startStatsBroadcast(hub *Hub) {
    rand.Seed(time.Now().UnixNano())
    var totalRx, totalTx int
    ticker := time.NewTicker(5 * time.Second)
    go func() {
        for {
            <-ticker.C
            totalRx += rand.Intn(120)
            totalTx += rand.Intn(150)
            online := rand.Intn(5) + 1
            offline := 5 - online
            s := Stats{TotalRx: totalRx, TotalTx: totalTx, PeersOnline: online, PeersOffline: offline}
            payload, err := json.Marshal(s)
            if err != nil {
                log.Println("Stats marshal error:", err)
                continue
            }
            hub.broadcast <- payload
        }
    }()
}