feat: safe websocket server implementation
This commit is contained in:
5
go.mod
Normal file
5
go.mod
Normal file
@@ -0,0 +1,5 @@
|
||||
module git.neurocipta.com/rogerferdinan/safe-web-socket
|
||||
|
||||
go 1.24.5
|
||||
|
||||
require github.com/gorilla/websocket v1.5.3
|
||||
2
go.sum
Normal file
2
go.sum
Normal file
@@ -0,0 +1,2 @@
|
||||
github.com/gorilla/websocket v1.5.3 h1:saDtZ6Pbx/0u+bgYQ3q96pZgCzfhKXGPqt7kZ72aNNg=
|
||||
github.com/gorilla/websocket v1.5.3/go.mod h1:YR8l580nyteQvAITg2hZ9XVh4b55+EU/adAjf1fMHhE=
|
||||
188
internal/hub.go
Normal file
188
internal/hub.go
Normal file
@@ -0,0 +1,188 @@
|
||||
package internal
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"log"
|
||||
"time"
|
||||
|
||||
"github.com/gorilla/websocket"
|
||||
)
|
||||
|
||||
type Client struct {
|
||||
Conn *websocket.Conn
|
||||
Send chan []byte
|
||||
SubscribedPath string
|
||||
done chan struct{}
|
||||
mu *CustomRwMutex
|
||||
}
|
||||
|
||||
func NewClient(conn *websocket.Conn, subscribedPath string) *Client {
|
||||
return &Client{
|
||||
Conn: conn,
|
||||
Send: make(chan []byte, 1024),
|
||||
SubscribedPath: subscribedPath,
|
||||
done: make(chan struct{}),
|
||||
mu: NewCustomRwMutex(),
|
||||
}
|
||||
}
|
||||
|
||||
type Hub struct {
|
||||
Clients map[*Client]bool
|
||||
Broadcast chan []byte
|
||||
Register chan *Client
|
||||
Unregister chan *Client
|
||||
ClientData map[string]chan []byte
|
||||
writeMu *CustomRwMutex
|
||||
readMu *CustomRwMutex
|
||||
}
|
||||
|
||||
func NewHub() *Hub {
|
||||
return &Hub{
|
||||
Broadcast: make(chan []byte),
|
||||
Register: make(chan *Client),
|
||||
Unregister: make(chan *Client),
|
||||
Clients: make(map[*Client]bool),
|
||||
ClientData: make(map[string]chan []byte),
|
||||
writeMu: NewCustomRwMutex(),
|
||||
}
|
||||
}
|
||||
|
||||
func (h *Hub) AddDataChannel(dataID string) chan []byte {
|
||||
ch := make(chan []byte, 256)
|
||||
h.writeMu.WriteHandler(func() error {
|
||||
if innerCh, ok := h.ClientData[dataID]; ok {
|
||||
ch = innerCh
|
||||
return nil
|
||||
}
|
||||
|
||||
h.ClientData[dataID] = ch
|
||||
log.Printf("Created data channel for: %s\n", dataID)
|
||||
return nil
|
||||
})
|
||||
|
||||
return ch
|
||||
}
|
||||
|
||||
func (h *Hub) GetDataChannel(dataID string) (chan []byte, bool) {
|
||||
var ch chan []byte
|
||||
var ok bool
|
||||
h.writeMu.ReadHandler(func() error {
|
||||
innerCh, innerOk := h.ClientData[dataID]
|
||||
ch = innerCh
|
||||
ok = innerOk
|
||||
return nil
|
||||
})
|
||||
|
||||
return ch, ok
|
||||
}
|
||||
|
||||
func (h *Hub) RemoveDataChannel(dataID string) {
|
||||
h.writeMu.WriteHandler(func() error {
|
||||
if ch, ok := h.ClientData[dataID]; ok {
|
||||
close(ch)
|
||||
delete(h.ClientData, dataID)
|
||||
log.Printf("Removed data channel for: %s\n", dataID)
|
||||
}
|
||||
return nil
|
||||
})
|
||||
}
|
||||
|
||||
func (h *Hub) Run() {
|
||||
go func() {
|
||||
for {
|
||||
select {
|
||||
case c := <-h.Register:
|
||||
h.Clients[c] = true
|
||||
log.Println("Client registered")
|
||||
case c := <-h.Unregister:
|
||||
if _, ok := h.Clients[c]; ok {
|
||||
delete(h.Clients, c)
|
||||
close(c.Send)
|
||||
c.Conn.Close()
|
||||
log.Println("Client unregistered")
|
||||
}
|
||||
case message := <-h.Broadcast:
|
||||
for c := range h.Clients {
|
||||
select {
|
||||
case c.Send <- message:
|
||||
default:
|
||||
close(c.Send)
|
||||
delete(h.Clients, c)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}()
|
||||
}
|
||||
|
||||
func WritePump(c *Client, h *Hub) {
|
||||
go func() {
|
||||
defer func() {
|
||||
h.Unregister <- c
|
||||
c.Conn.Close()
|
||||
}()
|
||||
|
||||
c.Conn.SetReadLimit(1024)
|
||||
c.Conn.SetPongHandler(func(string) error {
|
||||
if err := c.Conn.WriteMessage(websocket.PongMessage, []byte{}); err != nil {
|
||||
return fmt.Errorf("failed to send pong: %v", err)
|
||||
}
|
||||
return nil
|
||||
})
|
||||
|
||||
ticker := time.NewTicker(30 * time.Second)
|
||||
defer ticker.Stop()
|
||||
|
||||
for {
|
||||
select {
|
||||
case message, ok := <-c.Send:
|
||||
if err := c.mu.WriteHandler(func() error {
|
||||
if !ok {
|
||||
c.Conn.WriteMessage(websocket.CloseMessage, []byte{})
|
||||
return fmt.Errorf("message not ok")
|
||||
}
|
||||
|
||||
w, err := c.Conn.NextWriter(websocket.TextMessage)
|
||||
if err != nil {
|
||||
return fmt.Errorf("failed to get writer: %q", err)
|
||||
}
|
||||
|
||||
w.Write(message)
|
||||
|
||||
if err := w.Close(); err != nil {
|
||||
return err
|
||||
}
|
||||
return nil
|
||||
}); err == nil {
|
||||
continue
|
||||
}
|
||||
case <-ticker.C:
|
||||
if err := c.mu.WriteHandler(func() error {
|
||||
if err := c.Conn.WriteMessage(websocket.PingMessage, nil); err != nil {
|
||||
return err
|
||||
}
|
||||
return nil
|
||||
}); err != nil {
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
}()
|
||||
}
|
||||
|
||||
func ReadPump(c *Client) {
|
||||
go func() {
|
||||
for {
|
||||
_, message, err := c.Conn.ReadMessage()
|
||||
if err != nil {
|
||||
if websocket.IsUnexpectedCloseError(err, websocket.CloseGoingAway, websocket.CloseAbnormalClosure) {
|
||||
log.Printf("WebSocket error: %v", err)
|
||||
}
|
||||
break
|
||||
}
|
||||
|
||||
fmt.Println(string(message))
|
||||
}
|
||||
}()
|
||||
|
||||
}
|
||||
33
internal/mutex.go
Normal file
33
internal/mutex.go
Normal file
@@ -0,0 +1,33 @@
|
||||
package internal
|
||||
|
||||
import (
|
||||
"sync"
|
||||
)
|
||||
|
||||
type CustomRwMutex struct {
|
||||
mu *sync.RWMutex
|
||||
}
|
||||
|
||||
func NewCustomRwMutex() *CustomRwMutex {
|
||||
return &CustomRwMutex{
|
||||
mu: &sync.RWMutex{},
|
||||
}
|
||||
}
|
||||
|
||||
func (rwMu *CustomRwMutex) WriteHandler(fn func() error) error {
|
||||
rwMu.mu.Lock()
|
||||
defer rwMu.mu.Unlock()
|
||||
if err := fn(); err != nil {
|
||||
return err
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (rwMu *CustomRwMutex) ReadHandler(fn func() error) error {
|
||||
rwMu.mu.RLock()
|
||||
defer rwMu.mu.RUnlock()
|
||||
if err := fn(); err != nil {
|
||||
return err
|
||||
}
|
||||
return nil
|
||||
}
|
||||
40
internal/nil_checker.go
Normal file
40
internal/nil_checker.go
Normal file
@@ -0,0 +1,40 @@
|
||||
package internal
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"reflect"
|
||||
"strings"
|
||||
)
|
||||
|
||||
func NilChecker(data any) error {
|
||||
val := reflect.ValueOf(data)
|
||||
|
||||
if val.Kind() == reflect.Ptr {
|
||||
val = val.Elem()
|
||||
}
|
||||
|
||||
valType := val.Type()
|
||||
|
||||
if val.Kind() != reflect.Struct {
|
||||
return fmt.Errorf("data is not a struct")
|
||||
}
|
||||
|
||||
nilFields := []string{}
|
||||
for i := range val.NumField() {
|
||||
field := val.Field(i)
|
||||
fieldType := valType.Field(i)
|
||||
tagValue := fieldType.Tag.Get("nil_checker")
|
||||
|
||||
if tagValue == "required" {
|
||||
if field.Kind() == reflect.Ptr {
|
||||
if field.IsNil() {
|
||||
nilFields = append(nilFields, fieldType.Name)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
if len(nilFields) > 0 {
|
||||
return fmt.Errorf("%s is empty", strings.Join(nilFields, ","))
|
||||
}
|
||||
return nil
|
||||
}
|
||||
32
v1/example/server/main.go
Normal file
32
v1/example/server/main.go
Normal file
@@ -0,0 +1,32 @@
|
||||
package main
|
||||
|
||||
import (
|
||||
"log"
|
||||
"time"
|
||||
|
||||
"git.neurocipta.com/rogerferdinan/safe-web-socket/v1/server"
|
||||
)
|
||||
|
||||
func main() {
|
||||
s, err := server.NewSafeWebsocketServerBuilder().
|
||||
BaseHost("localhost").
|
||||
BasePort(8080).
|
||||
HandleFuncWebsocket("/ws/test/", "data_1", func(c chan []byte) {
|
||||
ticker := time.NewTicker(100 * time.Millisecond)
|
||||
for range ticker.C {
|
||||
c <- []byte(time.Now().Format("2006-01-02 15:04:05") + "_data_1")
|
||||
}
|
||||
}).
|
||||
HandleFuncWebsocket("/ws/test/", "data_2", func(c chan []byte) {
|
||||
ticker := time.NewTicker(100 * time.Millisecond)
|
||||
for range ticker.C {
|
||||
c <- []byte(time.Now().Format("2006-01-02 15:04:05") + "_data_2")
|
||||
}
|
||||
}).
|
||||
Build()
|
||||
|
||||
if err != nil {
|
||||
log.Fatal(err)
|
||||
}
|
||||
s.ListenAndServe()
|
||||
}
|
||||
107
v1/server/server.go
Normal file
107
v1/server/server.go
Normal file
@@ -0,0 +1,107 @@
|
||||
package server
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"log"
|
||||
"net/http"
|
||||
"strings"
|
||||
|
||||
"git.neurocipta.com/rogerferdinan/safe-web-socket/internal"
|
||||
"github.com/gorilla/websocket"
|
||||
)
|
||||
|
||||
type SafeWebsocketServerBuilder struct {
|
||||
baseHost *string `nil_checker:"required"`
|
||||
basePort *uint16 `nil_checker:"required"`
|
||||
upgrader *websocket.Upgrader `nil_checker:"required"`
|
||||
mux *http.ServeMux `nil_checker:"required"`
|
||||
}
|
||||
|
||||
func NewSafeWebsocketServerBuilder() *SafeWebsocketServerBuilder {
|
||||
h := internal.NewHub()
|
||||
h.Run()
|
||||
|
||||
return &SafeWebsocketServerBuilder{
|
||||
upgrader: &websocket.Upgrader{
|
||||
ReadBufferSize: 1024,
|
||||
WriteBufferSize: 1024,
|
||||
CheckOrigin: func(r *http.Request) bool {
|
||||
return true
|
||||
},
|
||||
},
|
||||
mux: http.NewServeMux(),
|
||||
}
|
||||
}
|
||||
|
||||
func (b *SafeWebsocketServerBuilder) BaseHost(baseHost string) *SafeWebsocketServerBuilder {
|
||||
b.baseHost = &baseHost
|
||||
return b
|
||||
}
|
||||
|
||||
func (b *SafeWebsocketServerBuilder) BasePort(basePort uint16) *SafeWebsocketServerBuilder {
|
||||
b.basePort = &basePort
|
||||
return b
|
||||
}
|
||||
|
||||
func (b *SafeWebsocketServerBuilder) HandleFunc(
|
||||
pattern string, fn func(http.ResponseWriter, *http.Request),
|
||||
) *SafeWebsocketServerBuilder {
|
||||
b.mux.HandleFunc(pattern, fn)
|
||||
return b
|
||||
}
|
||||
|
||||
func (b *SafeWebsocketServerBuilder) HandleFuncWebsocket(
|
||||
pattern string, subscribedPath string, writeFunc func(chan []byte),
|
||||
) *SafeWebsocketServerBuilder {
|
||||
h := internal.NewHub()
|
||||
h.Run()
|
||||
|
||||
b.mux.HandleFunc(pattern+subscribedPath, func(w http.ResponseWriter, r *http.Request) {
|
||||
conn, err := b.upgrader.Upgrade(w, r, nil)
|
||||
if err != nil {
|
||||
http.Error(w, "upgrade failed", http.StatusBadRequest)
|
||||
return
|
||||
}
|
||||
|
||||
subscribedPath := strings.TrimPrefix(r.URL.Path, pattern)
|
||||
fmt.Println(subscribedPath)
|
||||
if subscribedPath == "" {
|
||||
http.Error(w, "invalid path", http.StatusBadRequest)
|
||||
return
|
||||
}
|
||||
c := internal.NewClient(conn, subscribedPath)
|
||||
h.Register <- c
|
||||
internal.WritePump(c, h)
|
||||
internal.ReadPump(c)
|
||||
writeFunc(h.Broadcast)
|
||||
})
|
||||
return b
|
||||
}
|
||||
|
||||
func (b *SafeWebsocketServerBuilder) Build() (*SafeWebsocketServer, error) {
|
||||
if err := internal.NilChecker(b); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
safeServer := SafeWebsocketServer{
|
||||
url: fmt.Sprintf("%s:%d", *b.baseHost, *b.basePort),
|
||||
mux: b.mux,
|
||||
mu: internal.NewCustomRwMutex(),
|
||||
}
|
||||
return &safeServer, nil
|
||||
}
|
||||
|
||||
type SafeWebsocketServer struct {
|
||||
hub *internal.Hub
|
||||
mux *http.ServeMux
|
||||
url string
|
||||
mu *internal.CustomRwMutex
|
||||
}
|
||||
|
||||
func (s *SafeWebsocketServer) ListenAndServe() error {
|
||||
log.Printf("HTTP serve on %s\n", s.url)
|
||||
if err := http.ListenAndServe(s.url, s.mux); err != nil {
|
||||
return fmt.Errorf("failed to serve websocket: %w", err)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
Reference in New Issue
Block a user