Compare commits

...

32 Commits

Author SHA1 Message Date
e011931436 fix: messageType typo 2025-10-01 15:09:06 +07:00
70d9a37a4e fix: pong data race 2025-10-01 14:41:23 +07:00
2ab2e07b9b fix: pong not sent 2025-10-01 14:00:22 +07:00
9a9c65c24c fix: ping not sent 2025-10-01 13:45:48 +07:00
8a11ce0103 fix: extend read deadline wait 2025-09-30 13:13:29 +07:00
cf63683c9c fix: fixing panic on pong write message 2025-09-30 10:43:39 +07:00
92bd56aac0 fix: fixing data race 2025-09-30 10:33:30 +07:00
b4e7238b0b fix: fixing data race 2025-09-30 09:59:46 +07:00
b092e36987 fix: adding reconnect when deadline met 2025-09-30 08:16:58 +07:00
4f956b8fe9 fix: read issues 2025-09-30 07:24:54 +07:00
967b8a98b3 fix: write & reader posiiton 2025-09-30 07:12:38 +07:00
d09d389011 fix: initialize write channel 2025-09-30 06:59:01 +07:00
c550701dfa fix: adding lock for cancel function slices 2025-09-29 22:53:29 +07:00
2225391fc3 fix: changing to one concurrent writer & reader 2025-09-29 22:47:33 +07:00
9f8ee49b5c fix: adding SafeWebSocketClient parameter 2025-09-29 18:48:38 +07:00
735f55858e fix: adding nil checker 2025-09-29 17:58:19 +07:00
5e5df9090f fix: adding back read mutex for data race 2025-09-29 17:55:24 +07:00
1537e58444 fix: fixing deadlock caused by read mutex 2025-09-29 16:38:08 +07:00
66b27082b9 feat: adding reconnect channel for authentication purposes 2025-09-29 15:45:02 +07:00
9d74e72ede feat: add writeJSON function 2025-09-29 14:40:07 +07:00
dfbe2f2808 fix: fixing client deadlock condition & reconnection issues 2025-09-29 09:06:31 +07:00
4792d638c1 feat: adding channelSize to safe websocket client 2025-09-27 05:27:57 +07:00
b96e574726 feat: adding client close function 2025-09-27 04:52:47 +07:00
752804cc58 fix: restructuring safe websocket client 2025-09-27 04:48:01 +07:00
e686e69a24 fix: fixing client nil pointer 2025-09-26 16:10:19 +07:00
034e5b68e0 fix: client reconnection data race fix 2025-09-26 16:07:20 +07:00
6fb7cea1fa fix client & server experimental fix 2025-09-26 15:02:54 +07:00
7457604e0f fix: client data race hotfix 2025-09-26 14:18:53 +07:00
f32d97eac4 chore: cleaning unused mutex 2025-09-26 09:16:35 +07:00
4111dd0f25 fix: removing debugger for writePump 2025-09-25 17:33:00 +07:00
f91e24cc02 feat: adding rawQuery to safeWebsocketClient 2025-09-25 14:51:44 +07:00
ef98404caf feat: safe websocket client implementation 2025-09-25 14:14:24 +07:00
9 changed files with 555 additions and 113 deletions

View File

@@ -1,2 +1,22 @@
# safe-web-socket
> A secure, production-ready WebSocket wrapper for Go with built-in validation, rate limiting, authentication, and connection management.
## Overview
`SafeWebSocket` is a Go library designed to simplify the creation of **secure, scalable, and resilient WebSocket servers**. Built on top of `gorilla/websocket`, it adds essential safety layers — including authentication, input validation, connection limits, rate limiting, and graceful shutdown — so you can focus on your application logic, not security pitfalls.
Whether youre building real-time dashboards, chat apps, or live data feeds, `SafeWebSocket` ensures your WebSocket endpoints are protected against common attacks (e.g., DoS, injection, unauthorized access).
## Features
-**Graceful Shutdown & Cleanup**
-**Automatic Reconnection Handling (Client-side helpers)**
-**WebSocket Ping/Pong Health Checks**
## Installation
```bash
go get git.neurocipta.com/rogerferdinan/safe-web-socket
```
> Requires Go 1.24+

2
go.mod
View File

@@ -3,3 +3,5 @@ module git.neurocipta.com/rogerferdinan/safe-web-socket
go 1.24.5
require github.com/gorilla/websocket v1.5.3
require git.neurocipta.com/rogerferdinan/custom-rwmutex v1.0.0 // indirect

2
go.sum
View File

@@ -1,2 +1,4 @@
git.neurocipta.com/rogerferdinan/custom-rwmutex v1.0.0 h1:KnNc40SrYsg0cksIIcQy/ca6bunkGADQOs1u7O/E+iY=
git.neurocipta.com/rogerferdinan/custom-rwmutex v1.0.0/go.mod h1:9DvvHc2UZhBwEs63NgO4IhiuHnBNtTuBkTJgiMnnCss=
github.com/gorilla/websocket v1.5.3 h1:saDtZ6Pbx/0u+bgYQ3q96pZgCzfhKXGPqt7kZ72aNNg=
github.com/gorilla/websocket v1.5.3/go.mod h1:YR8l580nyteQvAITg2hZ9XVh4b55+EU/adAjf1fMHhE=

View File

@@ -1,7 +1,6 @@
package internal
import (
"fmt"
"log"
"time"
@@ -11,7 +10,7 @@ import (
const (
writeWait = 10 * time.Second
pongWait = 60 * time.Second
pingPeriod = 55 * time.Second
pingPeriod = 25 * time.Second
)
type Client struct {
@@ -19,7 +18,6 @@ type Client struct {
Send chan []byte
SubscribedPath string
done chan struct{}
mu *CustomRwMutex
}
func NewClient(conn *websocket.Conn, subscribedPath string) *Client {
@@ -28,7 +26,6 @@ func NewClient(conn *websocket.Conn, subscribedPath string) *Client {
Send: make(chan []byte, 64),
SubscribedPath: subscribedPath,
done: make(chan struct{}),
mu: NewCustomRwMutex(),
}
}
@@ -37,8 +34,6 @@ type Hub struct {
Broadcast chan []byte
Register chan *Client
Unregister chan *Client
writeMu *CustomRwMutex
readMu *CustomRwMutex
}
func NewHub() *Hub {
@@ -47,7 +42,6 @@ func NewHub() *Hub {
Register: make(chan *Client),
Unregister: make(chan *Client),
Clients: make(map[*Client]bool),
writeMu: NewCustomRwMutex(),
}
}
@@ -63,6 +57,7 @@ func (h *Hub) Run() {
delete(h.Clients, c)
close(c.Send)
}
log.Println("Client Unregistered")
case message := <-h.Broadcast:
for client := range h.Clients {
select {
@@ -99,7 +94,7 @@ func WritePump(c *Client, h *Hub) {
}
case <-pingTicker.C:
c.Conn.SetWriteDeadline(time.Now().Add(writeWait))
if err := c.Conn.WriteMessage(websocket.PingMessage, nil); err != nil {
if err := c.Conn.WriteMessage(websocket.PingMessage, []byte{}); err != nil {
return
}
}
@@ -122,14 +117,14 @@ func ReadPump(c *Client, h *Hub) {
for {
messageType, message, err := c.Conn.ReadMessage()
if err != nil {
if websocket.IsUnexpectedCloseError(err, websocket.CloseGoingAway, websocket.CloseAbnormalClosure) {
if !websocket.IsUnexpectedCloseError(err, websocket.CloseGoingAway, websocket.CloseAbnormalClosure) {
log.Printf("WebSocket error: %v", err)
}
break
}
if messageType == websocket.TextMessage {
fmt.Printf("Received: %s\n", message)
log.Printf("Received: %s\n", message)
}
}
}

View File

@@ -1,33 +0,0 @@
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
}

View File

@@ -1,82 +1,471 @@
package client
// import (
// "context"
// "fmt"
// "time"
import (
"context"
"fmt"
"log"
"net/url"
"strings"
"sync"
"time"
// "git.neurocipta.com/rogerferdinan/safe-web-socket/internal"
// "github.com/gorilla/websocket"
// )
custom_rwmutex "git.neurocipta.com/rogerferdinan/custom-rwmutex"
"git.neurocipta.com/rogerferdinan/safe-web-socket/internal"
"github.com/gorilla/websocket"
)
// const (
// pingPeriod = 30 * time.Second
// )
const (
pingPeriod = 10 * time.Second
readDeadline = 30 * time.Second
)
// type SafeWebsocketClientBuilder struct {
// baseHost *string `nil_checker:"required"`
// basePort *uint16 `nil_checker:"required"`
// }
type SafeMap[K comparable, V any] struct {
m sync.Map
}
// func NewSafeWebsocketClientBuilder() *SafeWebsocketClientBuilder {
// return &SafeWebsocketClientBuilder{}
// }
func NewSafeMap[K comparable, V any]() *SafeMap[K, V] {
return &SafeMap[K, V]{
m: sync.Map{},
}
}
// func (b *SafeWebsocketClientBuilder) BaseHost(host string) *SafeWebsocketClientBuilder {
// b.baseHost = &host
// return b
// }
func (sm *SafeMap[K, V]) Store(key K, value V) {
sm.m.Store(key, value)
}
// func (b *SafeWebsocketClientBuilder) BasePort(port uint16) *SafeWebsocketClientBuilder {
// b.basePort = &port
// return b
// }
func (sm *SafeMap[K, V]) Load(key K) (value V, ok bool) {
val, loaded := sm.m.Load(key)
if !loaded {
return *new(V), false
}
return val.(V), true
}
// func (b *SafeWebsocketClientBuilder) Build() (*SafeWebsocketClient, error) {
// if err := internal.NilChecker(b); err != nil {
// return nil, err
// }
func (sm *SafeMap[K, V]) Delete(key K) {
sm.m.Delete(key)
}
// ctx, cancel := context.WithCancel(context.Background())
func (sm *SafeMap[K, V]) Range(f func(K, V) bool) {
sm.m.Range(func(key, value any) bool {
k, ok1 := key.(K)
v, ok2 := value.(V)
if !ok1 || !ok2 {
return true
}
return f(k, v)
})
}
// wsClient := SafeWebsocketClient{
// baseHost: b.baseHost,
// basePort: b.basePort,
// ctx: ctx,
// cancel: cancel,
// reconnectCh: make(chan struct{}, 1),
// isConnected: false,
// }
func (sm *SafeMap[K, V]) Len() int {
count := 0
sm.Range(func(_ K, _ V) bool {
count++
return true
})
return count
}
// if err := wsClient.connect(); err != nil {
// cancel()
// return nil, fmt.Errorf("failed to establish initial connection: %v", err)
// }
type MessageType uint
// wsClient.startPingTicker()
// wsClient.startReceiveHandler()
const (
MessageTypeText MessageType = websocket.TextMessage
MessageTypePing MessageType = websocket.PingMessage
MessageTypePong MessageType = websocket.PongMessage
MessageTypeClose MessageType = websocket.CloseMessage
)
// return &wsClient, nil
// }
type Message struct {
MessageType MessageType
Data []byte
}
// type SafeWebsocketClient struct {
// baseHost *string
// basePort *uint16
// mu *internal.CustomRwMutex
// ctx context.Context
// cancel context.CancelFunc
// reconnectCh chan struct{}
// isConnected bool
// }
type SafeWebsocketClientBuilder struct {
baseHost *string `nil_checker:"required"`
basePort *uint16 `nil_checker:"required"`
path *string
rawQuery *string
useTLS *bool
channelSize *int64
authenticateFn func(*SafeWebsocketClient) error
}
// func (wsClient *SafeWebsocketClient) connect() error {
// url := fmt.Sprintf("%s:%d", *wsClient.baseHost, *wsClient.basePort)
// conn, _, err := websocket.DefaultDialer.Dial(url, nil)
// if err != nil {
// return fmt.Errorf("failed to connect to %s: %w", *wsClient.baseHost, err)
// }
func NewSafeWebsocketClientBuilder() *SafeWebsocketClientBuilder {
return &SafeWebsocketClientBuilder{}
}
// conn.SetPingHandler(func(pingData string) error {
// conn.WriteMessage(websocket.PongMessage, []byte(pingData))
// })
// }
func (b *SafeWebsocketClientBuilder) BaseHost(host string) *SafeWebsocketClientBuilder {
b.baseHost = &host
return b
}
func (b *SafeWebsocketClientBuilder) BasePort(port uint16) *SafeWebsocketClientBuilder {
b.basePort = &port
return b
}
func (b *SafeWebsocketClientBuilder) UseTLS(useTLS bool) *SafeWebsocketClientBuilder {
b.useTLS = &useTLS
return b
}
func (b *SafeWebsocketClientBuilder) AuthenticateFn(authenticateFn func(*SafeWebsocketClient) error) *SafeWebsocketClientBuilder {
b.authenticateFn = authenticateFn
return b
}
func (b *SafeWebsocketClientBuilder) Path(path string) *SafeWebsocketClientBuilder {
b.path = &path
return b
}
func (b *SafeWebsocketClientBuilder) RawQuery(rawQuery string) *SafeWebsocketClientBuilder {
b.rawQuery = &rawQuery
return b
}
func (b *SafeWebsocketClientBuilder) ChannelSize(channelSize int64) *SafeWebsocketClientBuilder {
b.channelSize = &channelSize
return b
}
func (b *SafeWebsocketClientBuilder) Build(ctx context.Context) (*SafeWebsocketClient, error) {
if err := internal.NilChecker(b); err != nil {
return nil, err
}
var useTLS bool
if b.useTLS != nil {
useTLS = *b.useTLS
}
if b.channelSize == nil {
channelSize := int64(1)
b.channelSize = &channelSize
}
wsClient := SafeWebsocketClient{
baseHost: *b.baseHost,
basePort: *b.basePort,
useTLS: useTLS,
path: b.path,
rawQuery: b.rawQuery,
dataChannel: make(chan []byte, *b.channelSize),
mu: custom_rwmutex.NewCustomRwMutex(),
ctx: ctx,
reconnectCh: make(chan struct{}, 1),
isConnected: false,
doneMap: NewSafeMap[string, chan struct{}](),
writeChan: make(chan Message),
}
if b.authenticateFn != nil {
wsClient.authenticateFn = b.authenticateFn
}
go wsClient.reconnectHandler()
if err := wsClient.connect(); err != nil {
return nil, fmt.Errorf("failed to establish initial connection: %v", err)
}
return &wsClient, nil
}
type SafeWebsocketClient struct {
baseHost string
basePort uint16
useTLS bool
path *string
rawQuery *string
dataChannel chan []byte
mu *custom_rwmutex.CustomRwMutex
conn *websocket.Conn
cancelFuncs []context.CancelFunc
ctx context.Context
Cancel context.CancelFunc
reconnectCh chan struct{}
reconnectChans []chan struct{}
isConnected bool
doneMap *SafeMap[string, chan struct{}]
authenticateFn func(*SafeWebsocketClient) error
writeChan chan Message
pongChan chan error
}
func (wsClient *SafeWebsocketClient) connect() error {
var scheme string
if wsClient.useTLS {
scheme = "wss"
} else {
scheme = "ws"
}
newURL := url.URL{
Scheme: scheme,
Host: fmt.Sprintf("%s:%d", wsClient.baseHost, wsClient.basePort),
}
if wsClient.path != nil && strings.TrimSpace(*wsClient.path) != "" {
newURL.Path = *wsClient.path
}
if wsClient.rawQuery != nil && strings.TrimSpace(*wsClient.rawQuery) != "" {
newURL.RawQuery = *wsClient.rawQuery
}
conn, _, err := websocket.DefaultDialer.Dial(newURL.String(), nil)
if err != nil {
return fmt.Errorf("failed to connect to %s: %w", wsClient.baseHost, err)
}
pingCtx, pingCancel := context.WithCancel(context.Background())
wsClient.mu.WriteHandler(func() error {
wsClient.cancelFuncs = append(wsClient.cancelFuncs, pingCancel)
return nil
})
go wsClient.startPingTicker(pingCtx)
if wsClient.conn != nil {
wsClient.conn.Close()
}
wsClient.conn = conn
wsClient.isConnected = true
if wsClient.authenticateFn != nil {
if err := wsClient.authenticateFn(wsClient); err != nil {
return err
}
}
go func() {
ctx, cancel := context.WithCancel(context.Background())
wsClient.mu.WriteHandler(func() error {
wsClient.cancelFuncs = append(wsClient.cancelFuncs, cancel)
return nil
})
var c *websocket.Conn
wsClient.mu.ReadHandler(func() error {
c = conn
return nil
})
for {
select {
case <-ctx.Done():
log.Println("Writer stopped due to client shutdown")
return
case data := <-wsClient.writeChan:
if c == nil {
wsClient.triggerReconnect()
return
}
if err := c.WriteMessage(int(data.MessageType), data.Data); err != nil {
log.Printf("error on write message: %v\n", err)
wsClient.triggerReconnect()
if data.MessageType == MessageTypePong {
wsClient.pongChan <- err
}
return
}
}
}
}()
go func() {
ctx, cancel := context.WithCancel(context.Background())
wsClient.mu.WriteHandler(func() error {
wsClient.cancelFuncs = append(wsClient.cancelFuncs, cancel)
return nil
})
var c *websocket.Conn
wsClient.mu.ReadHandler(func() error {
c = conn
return nil
})
if c == nil {
wsClient.triggerReconnect()
return
}
for {
select {
case <-ctx.Done():
log.Println("Reader stopped due to client shutdown")
return
default:
if err := c.SetReadDeadline(time.Now().Add(readDeadline)); err != nil {
log.Printf("error on read deadline: %v\n", err)
return
}
messageType, data, err := c.ReadMessage()
if err != nil {
log.Printf("error on read message: %v\n", err)
wsClient.triggerReconnect()
return
}
if messageType != websocket.TextMessage {
continue
}
select {
case wsClient.dataChannel <- data:
case <-ctx.Done():
return
default:
log.Println("Data channel full, dropping message")
}
}
}
}()
conn.SetPingHandler(func(pingData string) error {
wsClient.writeChan <- Message{
MessageType: MessageTypePong,
Data: []byte(pingData),
}
select {
case err := <-wsClient.pongChan:
return err
default:
}
return nil
})
return nil
}
func (wsClient *SafeWebsocketClient) startPingTicker(ctx context.Context) {
ticker := time.NewTicker(pingPeriod)
defer ticker.Stop()
for {
select {
case <-ctx.Done():
log.Println("ping ticker stopped")
return
case <-ticker.C:
wsClient.writeChan <- Message{
MessageType: websocket.PingMessage,
Data: []byte{},
}
}
}
}
func (wsClient *SafeWebsocketClient) triggerReconnect() {
select {
case wsClient.reconnectCh <- struct{}{}:
default:
}
}
func (wsClient *SafeWebsocketClient) reconnectHandler() {
backoff := 1 * time.Second
maxBackoff := 30 * time.Second
for {
select {
case <-wsClient.reconnectCh:
log.Println("Reconnect triggered")
wsClient.mu.ReadHandler(func() error {
if wsClient.cancelFuncs != nil {
for _, cancel := range wsClient.cancelFuncs {
cancel()
}
}
return nil
})
wsClient.isConnected = false
time.Sleep(100 * time.Millisecond)
isInnerLoop := true
for isInnerLoop {
log.Printf("Attempting reconnect in %v...", backoff)
select {
case <-time.After(backoff):
if err := wsClient.connect(); err != nil {
log.Printf("Reconnect failed: %v", err)
if backoff < maxBackoff {
backoff *= 2
}
continue
}
log.Println("Reconnected successfully")
backoff = 1 * time.Second
isInnerLoop = false
continue
case <-wsClient.ctx.Done():
log.Println("reconnect handler stopped due to client shutdown")
wsClient.Close()
return
}
}
if wsClient.reconnectChans != nil {
for _, reconnectCh := range wsClient.reconnectChans {
reconnectCh <- struct{}{}
}
}
case <-wsClient.ctx.Done():
log.Println("reconnect handler stopped due to client shutdown")
wsClient.Close()
return
}
}
}
func (wsClient *SafeWebsocketClient) ReconnectChannel() <-chan struct{} {
reconnectCh := make(chan struct{})
wsClient.mu.WriteHandler(func() error {
wsClient.reconnectChans = append(wsClient.reconnectChans, reconnectCh)
return nil
})
return reconnectCh
}
func (wsClient *SafeWebsocketClient) DataChannel() <-chan []byte {
return wsClient.dataChannel
}
func (wsClient *SafeWebsocketClient) Write(data []byte) error {
wsClient.writeChan <- Message{
MessageType: MessageTypeText,
Data: data,
}
return nil
}
func (wsClient *SafeWebsocketClient) Close() error {
wsClient.mu.ReadHandler(func() error {
if wsClient.cancelFuncs != nil {
for _, cancel := range wsClient.cancelFuncs {
cancel()
}
}
return nil
})
if wsClient.reconnectChans != nil {
for _, reconnectChan := range wsClient.reconnectChans {
close(reconnectChan)
}
}
if wsClient.conn != nil {
wsClient.conn.Close()
}
wsClient.isConnected = false
// close(wsClient.dataChannel)
return nil
}

View File

@@ -0,0 +1,50 @@
package main
import (
"context"
"fmt"
"log"
"os"
"os/signal"
"syscall"
"git.neurocipta.com/rogerferdinan/safe-web-socket/v1/client"
)
func main() {
sigChan := make(chan os.Signal, 1)
signal.Notify(sigChan, syscall.SIGINT, syscall.SIGTERM)
ctx, cancel := context.WithCancel(context.Background())
defer cancel()
go func() {
<-sigChan
fmt.Println("\nReceived interrupt signal. Shutting down gracefully...")
cancel()
}()
wsClient, err := client.NewSafeWebsocketClientBuilder().
BaseHost("localhost").
BasePort(8080).
Path("/ws/test/data_1").
UseTLS(false).
ChannelSize(30).
Build(ctx)
if err != nil {
log.Fatal(err)
}
go func() {
for range wsClient.ReconnectChannel() {
fmt.Println("Reconnection Success")
}
}()
dataChannel := wsClient.DataChannel()
for data := range dataChannel {
// _ = data
fmt.Println(string(data))
}
}

View File

@@ -1,12 +1,18 @@
package main
import (
"encoding/json"
"log"
"time"
"git.neurocipta.com/rogerferdinan/safe-web-socket/v1/server"
)
type ExampleData struct {
Time time.Time `json:"time"`
Data string `json:"data"`
}
func main() {
s, err := server.NewSafeWebsocketServerBuilder().
BaseHost("localhost").
@@ -14,13 +20,27 @@ func main() {
HandleFuncWebsocket("/ws/test/", "data_1", func(c chan []byte) {
ticker := time.NewTicker(10 * time.Millisecond)
for range ticker.C {
c <- []byte(time.Now().Format("2006-01-02 15:04:05") + "_data_1")
jsonBytes, err := json.Marshal(ExampleData{
Time: time.Now(),
Data: "data_1",
})
if err != nil {
continue
}
c <- jsonBytes
}
}).
HandleFuncWebsocket("/ws/test/", "data_2", func(c chan []byte) {
ticker := time.NewTicker(10 * time.Millisecond)
for range ticker.C {
c <- []byte(time.Now().Format("2006-01-02 15:04:05") + "_data_2")
jsonBytes, err := json.Marshal(ExampleData{
Time: time.Now(),
Data: "data_2",
})
if err != nil {
continue
}
c <- jsonBytes
}
}).
Build()

View File

@@ -18,9 +18,6 @@ type SafeWebsocketServerBuilder struct {
}
func NewSafeWebsocketServerBuilder() *SafeWebsocketServerBuilder {
h := internal.NewHub()
h.Run()
return &SafeWebsocketServerBuilder{
upgrader: &websocket.Upgrader{
ReadBufferSize: 1024,