github.com/pion/dtls/v2@v2.2.12/examples/util/hub.go (about)

     1  // SPDX-FileCopyrightText: 2023 The Pion community <https://pion.ly>
     2  // SPDX-License-Identifier: MIT
     3  
     4  package util
     5  
     6  import (
     7  	"bufio"
     8  	"fmt"
     9  	"net"
    10  	"os"
    11  	"strings"
    12  	"sync"
    13  )
    14  
    15  // Hub is a helper to handle one to many chat
    16  type Hub struct {
    17  	conns map[string]net.Conn
    18  	lock  sync.RWMutex
    19  }
    20  
    21  // NewHub builds a new hub
    22  func NewHub() *Hub {
    23  	return &Hub{conns: make(map[string]net.Conn)}
    24  }
    25  
    26  // Register adds a new conn to the Hub
    27  func (h *Hub) Register(conn net.Conn) {
    28  	fmt.Printf("Connected to %s\n", conn.RemoteAddr())
    29  	h.lock.Lock()
    30  	defer h.lock.Unlock()
    31  
    32  	h.conns[conn.RemoteAddr().String()] = conn
    33  
    34  	go h.readLoop(conn)
    35  }
    36  
    37  func (h *Hub) readLoop(conn net.Conn) {
    38  	b := make([]byte, bufSize)
    39  	for {
    40  		n, err := conn.Read(b)
    41  		if err != nil {
    42  			h.unregister(conn)
    43  			return
    44  		}
    45  		fmt.Printf("Got message: %s\n", string(b[:n]))
    46  	}
    47  }
    48  
    49  func (h *Hub) unregister(conn net.Conn) {
    50  	h.lock.Lock()
    51  	defer h.lock.Unlock()
    52  	delete(h.conns, conn.RemoteAddr().String())
    53  	err := conn.Close()
    54  	if err != nil {
    55  		fmt.Println("Failed to disconnect", conn.RemoteAddr(), err)
    56  	} else {
    57  		fmt.Println("Disconnected ", conn.RemoteAddr())
    58  	}
    59  }
    60  
    61  func (h *Hub) broadcast(msg []byte) {
    62  	h.lock.RLock()
    63  	defer h.lock.RUnlock()
    64  	for _, conn := range h.conns {
    65  		_, err := conn.Write(msg)
    66  		if err != nil {
    67  			fmt.Printf("Failed to write message to %s: %v\n", conn.RemoteAddr(), err)
    68  		}
    69  	}
    70  }
    71  
    72  // Chat starts the stdin readloop to dispatch messages to the hub
    73  func (h *Hub) Chat() {
    74  	reader := bufio.NewReader(os.Stdin)
    75  	for {
    76  		msg, err := reader.ReadString('\n')
    77  		Check(err)
    78  		if strings.TrimSpace(msg) == "exit" {
    79  			return
    80  		}
    81  		h.broadcast([]byte(msg))
    82  	}
    83  }