package main

import (
    "log"
    "net/http"
    "github.com/gorilla/websocket"
)

var upgrader = websocket.Upgrader{
    ReadBufferSize:  1024,
    WriteBufferSize: 1024,
    CheckOrigin: func(r *http.Request) bool {
        // In production, replace this with proper origin validation
        return true
    },
}

func handleWebSocket(w http.ResponseWriter, r *http.Request) {
    conn, err := upgrader.Upgrade(w, r, nil)
    if err != nil {
        log.Printf("Failed to upgrade connection: %v", err)
        return
    }
    defer conn.Close()

    // Implement connection handling logic here
    for {
        messageType, message, err := conn.ReadMessage()
        if err != nil {
            log.Printf("Error reading message: %v", err)
            break
        }

        // Echo the message back
        if err := conn.WriteMessage(messageType, message); err != nil {
            log.Printf("Error writing message: %v", err)
            break
        }
    }
}

func main() {
    http.HandleFunc("/ws", handleWebSocket)

    // Load TLS certificates
    certFile := "path/to/cert.pem"
    keyFile := "path/to/key.pem"

    log.Printf("Starting secure WebSocket server on :8443")
    err := http.ListenAndServeTLS(":8443", certFile, keyFile, nil)
    if err != nil {
        log.Fatal("ListenAndServeTLS: ", err)
    }
}

var upgrader = websocket.Upgrader{
    CheckOrigin: func(r *http.Request) bool {
        origin := r.Header.Get("Origin")
        return origin == "https://yourtrustedorigin.com"
    },
}

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

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

type Client struct {
    conn      *websocket.Conn
    lastPing  time.Time
    msgCount  int
    rateLimit time.Duration
}

func (c *Client) checkRateLimit() bool {
    now := time.Now()
    if now.Sub(c.lastPing) < c.rateLimit {
        c.msgCount++
        if c.msgCount > 100 { // Max 100 messages per rate limit window
            return false
        }
    } else {
        c.lastPing = now
        c.msgCount = 0
    }
    return true
}