81 lines
1.5 KiB
Go
81 lines
1.5 KiB
Go
package util
|
|
|
|
import (
|
|
"bufio"
|
|
"fmt"
|
|
"net"
|
|
"os"
|
|
"strings"
|
|
"sync"
|
|
)
|
|
|
|
// Hub is a helper to handle one to many chat
|
|
type Hub struct {
|
|
conns map[string]net.Conn
|
|
lock sync.RWMutex
|
|
}
|
|
|
|
// NewHub builds a new hub
|
|
func NewHub() *Hub {
|
|
return &Hub{conns: make(map[string]net.Conn)}
|
|
}
|
|
|
|
// Register adds a new conn to the Hub
|
|
func (h *Hub) Register(conn net.Conn) {
|
|
fmt.Printf("Connected to %s\n", conn.RemoteAddr())
|
|
h.lock.Lock()
|
|
defer h.lock.Unlock()
|
|
|
|
h.conns[conn.RemoteAddr().String()] = conn
|
|
|
|
go h.readLoop(conn)
|
|
}
|
|
|
|
func (h *Hub) readLoop(conn net.Conn) {
|
|
b := make([]byte, bufSize)
|
|
for {
|
|
n, err := conn.Read(b)
|
|
if err != nil {
|
|
h.unregister(conn)
|
|
return
|
|
}
|
|
fmt.Printf("Got message: %s\n", string(b[:n]))
|
|
}
|
|
}
|
|
|
|
func (h *Hub) unregister(conn net.Conn) {
|
|
h.lock.Lock()
|
|
defer h.lock.Unlock()
|
|
delete(h.conns, conn.RemoteAddr().String())
|
|
err := conn.Close()
|
|
if err != nil {
|
|
fmt.Println("Failed to disconnect", conn.RemoteAddr(), err)
|
|
} else {
|
|
fmt.Println("Disconnected ", conn.RemoteAddr())
|
|
}
|
|
}
|
|
|
|
func (h *Hub) broadcast(msg []byte) {
|
|
h.lock.RLock()
|
|
defer h.lock.RUnlock()
|
|
for _, conn := range h.conns {
|
|
_, err := conn.Write(msg)
|
|
if err != nil {
|
|
fmt.Printf("Failed to write message to %s: %v\n", conn.RemoteAddr(), err)
|
|
}
|
|
}
|
|
}
|
|
|
|
// Chat starts the stdin readloop to dispatch messages to the hub
|
|
func (h *Hub) Chat() {
|
|
reader := bufio.NewReader(os.Stdin)
|
|
for {
|
|
msg, err := reader.ReadString('\n')
|
|
Check(err)
|
|
if strings.TrimSpace(msg) == "exit" {
|
|
return
|
|
}
|
|
h.broadcast([]byte(msg))
|
|
}
|
|
}
|