WebSocket

package main

import (
    "encoding/json"
    "fmt"
    "log"
    "net/http"
    "sync"
    "time"

    "github.com/gorilla/websocket"
)

// 1. WebSocket upgrader
var upgrader = websocket.Upgrader{
    ReadBufferSize:  1024,
    WriteBufferSize: 1024,
    CheckOrigin: func(r *http.Request) bool {
        return true // Allow all origins for demo
    },
}

// 2. Message types
type Message struct {
    Type    string      `json:"type"`
    Content interface{} `json:"content"`
    From    string      `json:"from"`
    To      string      `json:"to,omitempty"`
    Time    time.Time   `json:"time"`
}

// 3. Client structure
type Client struct {
    hub      *Hub
    conn     *websocket.Conn
    send     chan Message
    username string
    mu       sync.Mutex
}

func newClient(hub *Hub, conn *websocket.Conn) *Client {
    return &Client{
        hub:  hub,
        conn: conn,
        send: make(chan Message, 256),
    }
}

// 4. Client message reader
func (c *Client) readPump() {
    defer func() {
        c.hub.unregister <- c
        c.conn.Close()
    }()

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

    for {
        _, data, err := c.conn.ReadMessage()
        if err != nil {
            if websocket.IsUnexpectedCloseError(err,
                websocket.CloseGoingAway,
                websocket.CloseAbnormalClosure) {
                log.Printf("error: %v", err)
            }
            break
        }

        var msg Message
        if err := json.Unmarshal(data, &msg); err != nil {
            log.Printf("error unmarshaling message: %v", err)
            continue
        }

        msg.Time = time.Now()
        msg.From = c.username

        c.hub.broadcast <- msg
    }
}

// 5. Client message writer
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
            }

            data, err := json.Marshal(message)
            if err != nil {
                return
            }
            w.Write(data)

            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
            }
        }
    }
}

// 6. Hub structure
type Hub struct {
    clients    map[*Client]bool
    broadcast  chan Message
    register   chan *Client
    unregister chan *Client
    mu         sync.RWMutex
}

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

// 7. Hub run loop
func (h *Hub) run() {
    for {
        select {
        case client := <-h.register:
            h.mu.Lock()
            h.clients[client] = true
            h.mu.Unlock()

            // Notify others
            h.broadcast <- Message{
                Type:    "system",
                Content: fmt.Sprintf("%s joined", client.username),
                Time:    time.Now(),
            }

        case client := <-h.unregister:
            h.mu.Lock()
            if _, ok := h.clients[client]; ok {
                delete(h.clients, client)
                close(client.send)
            }
            h.mu.Unlock()

            // Notify others
            h.broadcast <- Message{
                Type:    "system",
                Content: fmt.Sprintf("%s left", client.username),
                Time:    time.Now(),
            }

        case message := <-h.broadcast:
            h.mu.RLock()
            for client := range h.clients {
                // Skip private messages not intended for this client
                if message.To != "" &&
                    message.To != client.username &&
                    message.From != client.username {
                    continue
                }

                select {
                case client.send <- message:
                default:
                    close(client.send)
                    delete(h.clients, client)
                }
            }
            h.mu.RUnlock()
        }
    }
}

// 8. Constants
const (
    writeWait      = 10 * time.Second
    pongWait       = 60 * time.Second
    pingPeriod     = (pongWait * 9) / 10
    maxMessageSize = 512
)

// 9. WebSocket handler
func serveWs(hub *Hub, w http.ResponseWriter, r *http.Request) {
    username := r.URL.Query().Get("username")
    if username == "" {
        http.Error(w, "Username is required",
            http.StatusBadRequest)
        return
    }

    conn, err := upgrader.Upgrade(w, r, nil)
    if err != nil {
        log.Println(err)
        return
    }

    client := newClient(hub, conn)
    client.username = username
    client.hub.register <- client

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

// 10. HTML template
const homeTemplate = `
<!DOCTYPE html>
<html>
<head>
    <title>WebSocket Chat</title>
    <style>
        body { margin: 0; padding: 20px; font-family: Arial; }
        #messages { height: 400px; overflow-y: scroll; border: 1px solid #ccc; padding: 10px; }
        #form { margin-top: 20px; }
        input[type="text"] { width: 70%; padding: 5px; }
        button { padding: 5px 10px; }
        .system { color: #666; font-style: italic; }
        .private { color: #0066cc; }
    </style>
</head>
<body>
    <div id="messages"></div>
    <form id="form">
        <input type="text" id="msg" placeholder="Message">
        <input type="text" id="to" placeholder="To (optional)">
        <button type="submit">Send</button>
    </form>
    <script>
        const username = prompt('Enter your username:');
        if (!username) {
            alert('Username is required!');
            window.location.reload();
        }

        const ws = new WebSocket('ws://localhost:8080/ws?username=' +
            encodeURIComponent(username));
        const messages = document.getElementById('messages');
        const form = document.getElementById('form');
        const msgInput = document.getElementById('msg');
        const toInput = document.getElementById('to');

        ws.onmessage = function(e) {
            const msg = JSON.parse(e.data);
            const div = document.createElement('div');
            
            if (msg.type === 'system') {
                div.className = 'system';
                div.textContent = msg.content;
            } else {
                div.className = msg.to ? 'private' : '';
                const time = new Date(msg.time).toLocaleTimeString();
                if (msg.to) {
                    div.textContent = `[${time}] (Private) ${msg.from}: ${msg.content}`;
                } else {
                    div.textContent = `[${time}] ${msg.from}: ${msg.content}`;
                }
            }
            
            messages.appendChild(div);
            messages.scrollTop = messages.scrollHeight;
        };

        form.onsubmit = function(e) {
            e.preventDefault();
            if (msgInput.value) {
                const msg = {
                    type: 'message',
                    content: msgInput.value,
                    to: toInput.value
                };
                ws.send(JSON.stringify(msg));
                msgInput.value = '';
            }
        };
    </script>
</body>
</html>
`

func serveHome(w http.ResponseWriter, r *http.Request) {
    if r.URL.Path != "/" {
        http.Error(w, "Not found", http.StatusNotFound)
        return
    }
    if r.Method != http.MethodGet {
        http.Error(w, "Method not allowed",
            http.StatusMethodNotAllowed)
        return
    }
    w.Write([]byte(homeTemplate))
}

func main() {
    hub := newHub()
    go hub.run()

    http.HandleFunc("/", serveHome)
    http.HandleFunc("/ws", func(w http.ResponseWriter,
        r *http.Request) {
        serveWs(hub, w, r)
    })

    log.Println("Server starting on :8080")
    err := http.ListenAndServe(":8080", nil)
    if err != nil {
        log.Fatal("ListenAndServe: ", err)
    }
}

go get github.com/gorilla/websocket

go run main.go