forked from Gavinpub/ws
-
Notifications
You must be signed in to change notification settings - Fork 0
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
- Loading branch information
Showing
9 changed files
with
392 additions
and
3 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Binary file not shown.
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,87 @@ | ||
| package main | ||
|
|
||
| import ( | ||
| "encoding/json" | ||
| "fmt" | ||
| "net/url" | ||
| "os" | ||
| "os/signal" | ||
| "strings" | ||
|
|
||
| "github.com/gorilla/websocket" | ||
| "github.com/slytomcat/ws/server" | ||
| ) | ||
|
|
||
| type msg struct { | ||
| Type string `json:"type,omitempty"` | ||
| Payload string `json:"payload,omitempty"` | ||
| ListenerCount *int `json:"listenerCount,omitempty"` | ||
| } | ||
|
|
||
| func sendMsg(m msg, c *websocket.Conn) { | ||
| response, _ := json.Marshal(m) | ||
| c.WriteMessage(websocket.TextMessage, response) | ||
| } | ||
|
|
||
| func main() { | ||
| if len(os.Args) < 2 { | ||
| fmt.Printf("Usage: %s ws://host[:port][/path]", os.Args[0]) | ||
| os.Exit(1) | ||
| } | ||
| u, err := url.Parse(os.Args[1]) | ||
| if err != nil { | ||
| fmt.Printf("url parsing error: %v", err) | ||
| os.Exit(1) | ||
| } | ||
| srv := server.NewServer(u.Host) | ||
| if !strings.HasPrefix(u.Path, "/") { | ||
| u.Path = "/" + u.Path | ||
| } | ||
| srv.WSHandleFunc(u.Path, func(conn *websocket.Conn) { | ||
| in := msg{} | ||
| id := conn.RemoteAddr().String() | ||
| for { | ||
| _, message, err := conn.ReadMessage() | ||
| if err != nil { | ||
| if websocket.IsUnexpectedCloseError(err, websocket.CloseGoingAway, websocket.CloseNormalClosure, websocket.CloseServiceRestart) { | ||
| fmt.Printf("echoHandler for %s: websocket reading error: %v\n", id, err) | ||
| } else { | ||
| fmt.Printf("echoHandler for %s: %s\n", id, err) | ||
| } | ||
| return | ||
| } | ||
| fmt.Printf("echoHandler for %s: handle message: %s\n", id, message) | ||
| if err := json.Unmarshal(message, &in); err != nil { | ||
| sendMsg(msg{"error", fmt.Sprintf("message parsing error: %v", err), nil}, conn) | ||
| continue | ||
| } | ||
| switch in.Type { | ||
| case "echo": | ||
| sendMsg(msg{in.Type, in.Payload, nil}, conn) | ||
| case "broadcast": | ||
| count := 0 | ||
| out, _ := json.Marshal(msg{in.Type, in.Payload, nil}) | ||
| srv.ForEachConnection(func(c *websocket.Conn) bool { | ||
| c.WriteMessage(websocket.TextMessage, out) | ||
| count++ | ||
| return true | ||
| }) | ||
| sendMsg(msg{"broadcastResult", in.Payload, &count}, conn) | ||
| default: | ||
| sendMsg(msg{"error", "unknown type", nil}, conn) | ||
| } | ||
| } | ||
| }) | ||
| go func() { | ||
| sig := make(chan os.Signal, 2) | ||
| signal.Notify(sig, os.Interrupt, os.Kill) | ||
| fmt.Printf("\n%s signal received, exiting...\n", <-sig) | ||
| srv.Close() | ||
| }() | ||
| fmt.Printf("starting echo server on %s...\n", u.Host) | ||
| err = srv.ListenAndServe() | ||
| if err.Error() != "http: Server closed" { | ||
| fmt.Println(err) | ||
| os.Exit(1) | ||
| } | ||
| } |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -1,19 +1,23 @@ | ||
| module github.com/slytomcat/ws | ||
|
|
||
| go 1.21.6 | ||
| go 1.22.1 | ||
|
|
||
| require ( | ||
| github.com/chzyer/readline v1.5.1 | ||
| github.com/fatih/color v1.16.0 | ||
| github.com/gorilla/websocket v1.5.1 | ||
| github.com/spf13/cobra v1.8.0 | ||
| github.com/stretchr/testify v1.9.0 | ||
| ) | ||
|
|
||
| require ( | ||
| github.com/davecgh/go-spew v1.1.1 // indirect | ||
| github.com/inconshreveable/mousetrap v1.1.0 // indirect | ||
| github.com/mattn/go-colorable v0.1.13 // indirect | ||
| github.com/mattn/go-isatty v0.0.20 // indirect | ||
| github.com/pmezard/go-difflib v1.0.0 // indirect | ||
| github.com/spf13/pflag v1.0.5 // indirect | ||
| golang.org/x/net v0.20.0 // indirect | ||
| golang.org/x/sys v0.16.0 // indirect | ||
| gopkg.in/yaml.v3 v3.0.1 // indirect | ||
| ) |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,25 @@ | ||
| package server | ||
|
|
||
| import ( | ||
| "fmt" | ||
|
|
||
| "github.com/gorilla/websocket" | ||
| ) | ||
|
|
||
| // EchoHandler is a handler that sends back all received messages | ||
| func EchoHandler(conn *websocket.Conn) { | ||
| id := conn.RemoteAddr().String() | ||
| for { | ||
| mt, message, err := conn.ReadMessage() | ||
| if err != nil { | ||
| if websocket.IsUnexpectedCloseError(err, websocket.CloseGoingAway, websocket.CloseNormalClosure, websocket.CloseServiceRestart) { | ||
| fmt.Printf("echoHandler for %s: websocket reading error: %v\n", id, err) | ||
| } else { | ||
| fmt.Printf("echoHandler for %s: %s\n", id, err) | ||
| } | ||
| return | ||
| } | ||
| fmt.Printf("echoHandler for %s: handle message: %s\n", id, message) | ||
| conn.WriteMessage(mt, message) | ||
| } | ||
| } |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,103 @@ | ||
| package server | ||
|
|
||
| import ( | ||
| "fmt" | ||
| "net/http" | ||
| "strings" | ||
| "sync" | ||
| "time" | ||
|
|
||
| "github.com/gorilla/websocket" | ||
| ) | ||
|
|
||
| var upgrader = websocket.Upgrader{ | ||
| HandshakeTimeout: time.Second, | ||
| Subprotocols: []string{}, | ||
| CheckOrigin: func(r *http.Request) bool { | ||
| return true | ||
| }, | ||
| } | ||
|
|
||
| // Server is a websocket/http server. It is wrapper for standard http.Server with additional functionality for websocket request handling. | ||
| type Server struct { | ||
| http.Server | ||
| connections sync.Map | ||
| handle func(*websocket.Conn) | ||
| mux *http.ServeMux | ||
| Upgrader websocket.Upgrader | ||
| } | ||
|
|
||
| // NewServer creates new websocket/http server. It is configured to start on provided addr and creates the standard serve mux and sets it as server handler. | ||
| // Use WSHandleFunc and HandleFunc for setting handlers on desired paths and start Server via ListenAndServe/ListenAndServeTLS. | ||
| // ForEachConnection allows to iterate over currently active WS connections. | ||
| // For example you can send some broadcast message via s.ForEachConnection(func(c *websocket.Conn){c.WriteMessage(websocket.TextMessage, message)}) | ||
| func NewServer(addr string) *Server { | ||
| mux := http.NewServeMux() | ||
| return &Server{ | ||
| mux: mux, | ||
| Upgrader: upgrader, | ||
| Server: http.Server{ | ||
| Addr: addr, | ||
| Handler: mux, | ||
| }, | ||
| } | ||
| } | ||
|
|
||
| // WSHandleFunc setups new WS handler for path | ||
| func (s *Server) WSHandleFunc(path string, handler func(*websocket.Conn)) { | ||
| s.mux.HandleFunc(path, s.serve(handler)) | ||
| } | ||
|
|
||
| // HandleFunc setups new regular http handler for path | ||
| func (s *Server) HandleFunc(path string, handler func(w http.ResponseWriter, r *http.Request)) { | ||
| s.mux.HandleFunc(path, handler) | ||
| } | ||
|
|
||
| // Close correctly closes all active ws connections and close the server | ||
| func (s *Server) Close() error { | ||
| s.ForEachConnection(func(c *websocket.Conn) bool { | ||
| if err := TryCloseNormally(c, "server going down"); err != nil { | ||
| fmt.Printf("server: closing connection from %s error: %v\n", c.RemoteAddr(), err) | ||
| } | ||
| c.Close() | ||
| return true | ||
| }) | ||
| return s.Server.Close() | ||
| } | ||
|
|
||
| // TryCloseNormally tries to close websocket connection normally i.e. according to RFC | ||
| // NOTE It doesn't close underlying connection as socket reader have to read and handle close response. | ||
| func TryCloseNormally(conn *websocket.Conn, message string) error { | ||
| closeMessage := websocket.FormatCloseMessage(websocket.CloseNormalClosure, message) | ||
| if err := conn.WriteControl(websocket.CloseMessage, closeMessage, time.Now().Add(time.Second)); err != nil { | ||
| if strings.Contains(err.Error(), "close sent") { | ||
| return nil | ||
| } | ||
| return err | ||
| } | ||
| return nil | ||
| } | ||
|
|
||
| func (s *Server) serve(handler func(*websocket.Conn)) func(w http.ResponseWriter, r *http.Request) { | ||
| return func(w http.ResponseWriter, r *http.Request) { | ||
| connection, err := upgrader.Upgrade(w, r, nil) | ||
| if err != nil { | ||
| fmt.Printf("server: connection upgrade error: %v\n", err) | ||
| return | ||
| } | ||
| addr := connection.RemoteAddr().String() | ||
| fmt.Printf("server: new WS connection from %v\n", addr) | ||
| s.connections.Store(connection, nil) | ||
| handler(connection) | ||
| s.connections.Delete(connection) | ||
| fmt.Printf("server: WS connection from %v closed\n", addr) | ||
| connection.Close() | ||
| } | ||
| } | ||
|
|
||
| // ForEachConnection allow to iterate over all active connections | ||
| func (s *Server) ForEachConnection(f func(*websocket.Conn) bool) { | ||
| s.connections.Range(func(key, value any) bool { | ||
| return f(key.(*websocket.Conn)) | ||
| }) | ||
| } |
Oops, something went wrong.