From 4fc0c438c829f9bd043f22cd12c92ded1bbc0543 Mon Sep 17 00:00:00 2001 From: Roger Ferdinan Date: Fri, 21 Nov 2025 16:24:33 +0700 Subject: [PATCH] feat: adding apiKey for authentication --- go.mod | 10 +++++- go.sum | 7 ++++ internal/error_response.go | 56 +++++++++++++++++++++++++++++++ v1/examples/client/main.go | 12 +++---- v1/examples/gobwas-server/main.go | 28 ++++++++++++++++ v1/examples/server/main.go | 1 + v1/server/server.go | 52 +++++++++++++++++++++------- 7 files changed, 147 insertions(+), 19 deletions(-) create mode 100644 internal/error_response.go create mode 100644 v1/examples/gobwas-server/main.go diff --git a/go.mod b/go.mod index 65d7e01..0dddb49 100644 --- a/go.mod +++ b/go.mod @@ -9,4 +9,12 @@ require ( github.com/google/uuid v1.6.0 ) -require git.neurocipta.com/rogerferdinan/safe-map v0.0.0-20251011004629-ab0b119a7c48 +require ( + git.neurocipta.com/rogerferdinan/safe-map v0.0.0-20251011004629-ab0b119a7c48 + github.com/gobwas/ws v1.4.0 +) + +require ( + github.com/gobwas/httphead v0.1.0 // indirect + github.com/gobwas/pool v0.2.1 // indirect +) diff --git a/go.sum b/go.sum index 7595de4..0b35889 100644 --- a/go.sum +++ b/go.sum @@ -2,7 +2,14 @@ git.neurocipta.com/rogerferdinan/custom-rwmutex v1.0.0 h1:KnNc40SrYsg0cksIIcQy/c git.neurocipta.com/rogerferdinan/custom-rwmutex v1.0.0/go.mod h1:9DvvHc2UZhBwEs63NgO4IhiuHnBNtTuBkTJgiMnnCss= git.neurocipta.com/rogerferdinan/safe-map v0.0.0-20251011004629-ab0b119a7c48 h1:4wXSbEuwFd2gycaaGP35bjUkKEEO6WcVfJ6cetEyT5s= git.neurocipta.com/rogerferdinan/safe-map v0.0.0-20251011004629-ab0b119a7c48/go.mod h1:QtIxG0BYCCq8a5qyklpSHA8qWUvKr+mfl42qF9QxTc0= +github.com/gobwas/httphead v0.1.0 h1:exrUm0f4YX0L7EBwZHuCF4GDp8aJfVeBrlLQrs6NqWU= +github.com/gobwas/httphead v0.1.0/go.mod h1:O/RXo79gxV8G+RqlR/otEwx4Q36zl9rqC5u12GKvMCM= +github.com/gobwas/pool v0.2.1 h1:xfeeEhW7pwmX8nuLVlqbzVc7udMDrwetjEv+TZIz1og= +github.com/gobwas/pool v0.2.1/go.mod h1:q8bcK0KcYlCgd9e7WYLm9LpyS+YeLd8JVDW6WezmKEw= +github.com/gobwas/ws v1.4.0 h1:CTaoG1tojrh4ucGPcoJFiAQUAsEWekEWvLy7GsVNqGs= +github.com/gobwas/ws v1.4.0/go.mod h1:G3gNqMNtPppf5XUz7O4shetPpcZ1VJ7zt18dlUeakrc= github.com/google/uuid v1.6.0 h1:NIvaJDMOsjHA8n1jAhLSgzrAzy1Hgr+hNrb57e+94F0= github.com/google/uuid v1.6.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo= github.com/gorilla/websocket v1.5.3 h1:saDtZ6Pbx/0u+bgYQ3q96pZgCzfhKXGPqt7kZ72aNNg= github.com/gorilla/websocket v1.5.3/go.mod h1:YR8l580nyteQvAITg2hZ9XVh4b55+EU/adAjf1fMHhE= +golang.org/x/sys v0.6.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= diff --git a/internal/error_response.go b/internal/error_response.go new file mode 100644 index 0000000..322ed2f --- /dev/null +++ b/internal/error_response.go @@ -0,0 +1,56 @@ +package internal + +import ( + "encoding/json" + "fmt" + "net/http" +) + +type StatusMessageBuilder struct { + statusCode *int + message *string +} + +func NewStatusMessage() *StatusMessageBuilder { + return &StatusMessageBuilder{} +} + +func (sm *StatusMessageBuilder) StatusCode(statusCode int) *StatusMessageBuilder { + sm.statusCode = &statusCode + return sm +} + +func (sm *StatusMessageBuilder) Message(message string) *StatusMessageBuilder { + sm.message = &message + return sm +} + +func (sm *StatusMessageBuilder) Build() *StatusMessage { + return &StatusMessage{ + StatusCode: sm.statusCode, + Message: sm.message, + } +} + +type StatusMessage struct { + StatusCode *int + Message *string `json:"msg"` +} + +func ErrorResponse(w http.ResponseWriter, statusMessage *StatusMessage) error { + if statusMessage.StatusCode == nil { + return fmt.Errorf("status code is missing") + } + if statusMessage.Message == nil { + return fmt.Errorf("message is missing") + } + w.Header().Add("Content-Type", "application/json") + w.WriteHeader(*statusMessage.StatusCode) + + b, err := json.Marshal(statusMessage) + if err != nil { + return fmt.Errorf("failed to marshal JSON") + } + w.Write(b) + return nil +} diff --git a/v1/examples/client/main.go b/v1/examples/client/main.go index e469801..8ae14c0 100644 --- a/v1/examples/client/main.go +++ b/v1/examples/client/main.go @@ -29,17 +29,17 @@ func main() { BasePort(8080). Path("/ws/test/data_1"). UseTLS(false). - ChannelSize(30). + ChannelSize(1). Build(ctx) if err != nil { log.Fatal(err) } - // go func() { - // for range wsClient.ReconnectChannel() { - // fmt.Println("Reconnection Success") - // } - // }() + go func() { + for range wsClient.ReconnectChannel() { + fmt.Println("Reconnection Success") + } + }() dataChannel := wsClient.DataChannel() diff --git a/v1/examples/gobwas-server/main.go b/v1/examples/gobwas-server/main.go new file mode 100644 index 0000000..cda344c --- /dev/null +++ b/v1/examples/gobwas-server/main.go @@ -0,0 +1,28 @@ +package main + +import ( + "syscall" +) + +func setMaxRLimit() { + var rLimit syscall.Rlimit + if err := syscall.Getrlimit(syscall.RLIMIT_NOFILE, &rLimit); err != nil { + panic(err) + } + rLimit.Cur = rLimit.Max + if err := syscall.Setrlimit(syscall.RLIMIT_NOFILE, &rLimit); err != nil { + panic(err) + } +} + +func main() { + setMaxRLimit() + // pooler, err := netpool.New() + + // http.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) { + // if _, _, _, err := ws.UpgradeHTTP(r, w); err != nil { + // log.Fatal(err) + // } + // }) + // epoller, err := MkEp +} diff --git a/v1/examples/server/main.go b/v1/examples/server/main.go index 90510ad..3841ea7 100644 --- a/v1/examples/server/main.go +++ b/v1/examples/server/main.go @@ -17,6 +17,7 @@ func main() { s, err := server.NewSafeWebsocketServerBuilder(). BaseHost("localhost"). BasePort(8080). + ApiKey(""). HandleFuncWebsocket("/ws/test/", "data_1", func(c chan []byte) { ticker := time.NewTicker(10 * time.Millisecond) for range ticker.C { diff --git a/v1/server/server.go b/v1/server/server.go index fa2b77f..4a17ccd 100644 --- a/v1/server/server.go +++ b/v1/server/server.go @@ -5,14 +5,27 @@ import ( "log" "net/http" "strings" + "syscall" "git.neurocipta.com/rogerferdinan/safe-web-socket/internal" "github.com/gorilla/websocket" ) +func setMaxRLimit() { + var rLimit syscall.Rlimit + if err := syscall.Getrlimit(syscall.RLIMIT_NOFILE, &rLimit); err != nil { + panic(err) + } + rLimit.Cur = rLimit.Max + if err := syscall.Setrlimit(syscall.RLIMIT_NOFILE, &rLimit); err != nil { + panic(err) + } +} + type SafeWebsocketServerBuilder struct { baseHost *string `nil_checker:"required"` basePort *uint16 `nil_checker:"required"` + apiKey *string `nil_checker:"required"` upgrader *websocket.Upgrader `nil_checker:"required"` mux *http.ServeMux `nil_checker:"required"` } @@ -40,16 +53,17 @@ func (b *SafeWebsocketServerBuilder) BasePort(basePort uint16) *SafeWebsocketSer return b } -func (b *SafeWebsocketServerBuilder) HandleFunc( - pattern string, fn func(http.ResponseWriter, *http.Request), -) *SafeWebsocketServerBuilder { +func (b *SafeWebsocketServerBuilder) ApiKey(apiKey string) *SafeWebsocketServerBuilder { + b.apiKey = &apiKey + 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 { +func (b *SafeWebsocketServerBuilder) HandleFuncWebsocket(pattern string, subscribedPath string, writeFunc func(chan []byte)) *SafeWebsocketServerBuilder { h := internal.NewHub() h.Run() @@ -79,22 +93,36 @@ func (b *SafeWebsocketServerBuilder) Build() (*SafeWebsocketServer, error) { return nil, err } + setMaxRLimit() safeServer := SafeWebsocketServer{ - url: fmt.Sprintf("%s:%d", *b.baseHost, *b.basePort), - mux: b.mux, + mux: b.mux, + url: fmt.Sprintf("%s:%d", *b.baseHost, *b.basePort), + apiKey: *b.apiKey, } return &safeServer, nil } type SafeWebsocketServer struct { - hub *internal.Hub - mux *http.ServeMux - url string + mux *http.ServeMux + url string + apiKey string +} + +func (s *SafeWebsocketServer) AuthMiddleware(next http.Handler) http.Handler { + return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.Header.Get("X-MBX-APIKEY") != s.apiKey { + internal.ErrorResponse(w, internal.NewStatusMessage(). + StatusCode(http.StatusForbidden). + Message("X-MBX-APIKEY is missing"). + Build()) + } + next.ServeHTTP(w, r) + }) } func (s *SafeWebsocketServer) ListenAndServe() error { log.Printf("HTTP serve on %s\n", s.url) - if err := http.ListenAndServe(s.url, s.mux); err != nil { + if err := http.ListenAndServe(s.url, s.AuthMiddleware(s.mux)); err != nil { return fmt.Errorf("failed to serve websocket: %w", err) } return nil