108 lines
2.9 KiB
Go
108 lines
2.9 KiB
Go
package main
|
|
|
|
import (
|
|
"context"
|
|
_ "embed"
|
|
"errors"
|
|
"flag"
|
|
"fmt"
|
|
"net/http/httptest"
|
|
"os"
|
|
"strings"
|
|
"time"
|
|
|
|
omni "git.yuntuo.hk/k2safe/go-omni"
|
|
"git.yuntuo.hk/k2safe/go-omni/httpx"
|
|
"git.yuntuo.hk/k2safe/go-omni/ws"
|
|
"git.yuntuo.hk/k2safe/go-omni/ws/hub"
|
|
wsomni "git.yuntuo.hk/k2safe/go-omni/ws/omni"
|
|
)
|
|
|
|
//go:embed index.html
|
|
var page []byte
|
|
|
|
func main() {
|
|
serve := flag.Bool("serve", false, "serve the browser demo on 127.0.0.1:8088")
|
|
flag.Parse()
|
|
if err := run(*serve); err != nil {
|
|
fmt.Fprintln(os.Stderr, err)
|
|
os.Exit(1)
|
|
}
|
|
}
|
|
|
|
func run(serve bool) (runErr error) {
|
|
application := omni.New(omni.WithLogger(false))
|
|
defer func() { runErr = errors.Join(runErr, application.Close()) }()
|
|
connections, err := hub.New(hub.Config{})
|
|
if err != nil {
|
|
return err
|
|
}
|
|
application.OnShutdown(connections.Close)
|
|
server, err := ws.NewServer(ws.ServerConfig{})
|
|
if err != nil {
|
|
return err
|
|
}
|
|
if err = wsomni.Mount(application, "/ws", server, ws.Handlers{
|
|
OnOpen: func(_ context.Context, s *ws.Session) error {
|
|
if err := connections.Add(s); err != nil {
|
|
return err
|
|
}
|
|
return connections.Join(s.ID(), "demo")
|
|
},
|
|
OnMessage: func(ctx context.Context, s *ws.Session, m ws.Message) error {
|
|
if m.Type == ws.Text && string(m.Data) == "broadcast" {
|
|
result, err := connections.SendGroup("demo", ws.Message{Type: ws.Text, Data: []byte("hello group")})
|
|
if err != nil {
|
|
return err
|
|
}
|
|
if len(result.Errors) > 0 {
|
|
return s.SendJSON(ctx, map[string]int{"queued": result.Queued, "failed": len(result.Errors)})
|
|
}
|
|
return nil
|
|
}
|
|
return s.Send(ctx, m)
|
|
},
|
|
}); err != nil {
|
|
return err
|
|
}
|
|
application.GET("/", func(c *httpx.Context) error {
|
|
c.Raw().Header().Set("Content-Type", "text/html; charset=utf-8")
|
|
_, err := c.Raw().Write(page)
|
|
return err
|
|
})
|
|
if serve {
|
|
fmt.Println("Open http://127.0.0.1:8088 in two browser tabs; Ctrl+C stops cleanly.")
|
|
return application.Run("127.0.0.1:8088")
|
|
}
|
|
httpServer := httptest.NewServer(application)
|
|
defer httpServer.Close()
|
|
client, err := ws.NewClient(ws.ClientConfig{ProxyMode: "direct"})
|
|
if err != nil {
|
|
return err
|
|
}
|
|
defer func() {
|
|
ctx, cancel := context.WithTimeout(context.Background(), time.Second)
|
|
defer cancel()
|
|
runErr = errors.Join(runErr, client.Shutdown(ctx))
|
|
}()
|
|
ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second)
|
|
defer cancel()
|
|
replies := make(chan string, 2)
|
|
session, err := client.Dial(ctx, "ws"+strings.TrimPrefix(httpServer.URL, "http")+"/ws", ws.Handlers{OnMessage: func(_ context.Context, _ *ws.Session, m ws.Message) error { replies <- string(m.Data); return nil }})
|
|
if err != nil {
|
|
return err
|
|
}
|
|
for _, message := range []string{"hello websocket", "broadcast"} {
|
|
if err = session.Send(ctx, ws.Message{Type: ws.Text, Data: []byte(message)}); err != nil {
|
|
return err
|
|
}
|
|
select {
|
|
case reply := <-replies:
|
|
fmt.Println(reply)
|
|
case <-ctx.Done():
|
|
return ctx.Err()
|
|
}
|
|
}
|
|
return nil
|
|
}
|