mayhem-party/src/device/input/raw/ws.go

149 lines
2.8 KiB
Go

package raw
import (
"context"
"embed"
_ "embed"
"fmt"
"log"
"net/http"
"net/http/httputil"
"net/url"
"os"
"path"
"strings"
"github.com/gorilla/websocket"
)
var (
FlagWSProxy = os.Getenv("RAW_WS_PROXY_URL")
FlagWSDebug = os.Getenv("RAW_WS_DEBUG") != ""
)
type WS struct {
ctx context.Context
can context.CancelFunc
ch chan []byte
}
func NewWS(ctx context.Context, port int) WS {
ctx, can := context.WithCancel(ctx)
ws := WS{ctx: ctx, can: can, ch: make(chan []byte, 256)}
go ws.listen(port)
return ws
}
func (ws WS) Read() []byte {
select {
case v := <-ws.ch:
return v
case <-ws.ctx.Done():
return nil
}
}
func (ws WS) Close() {
ws.can()
}
func (ws WS) listen(port int) {
server := &http.Server{
Addr: fmt.Sprintf(":%d", port),
Handler: ws,
}
go func() {
if err := server.ListenAndServe(); err != nil && ws.ctx.Err() == nil {
panic(err)
}
}()
log.Println("WS on", port)
<-ws.ctx.Done()
server.Close()
}
var upgrader = websocket.Upgrader{
ReadBufferSize: 1024,
WriteBufferSize: 1024,
}
func (ws WS) ServeHTTP(w http.ResponseWriter, r *http.Request) {
r = r.WithContext(ws.ctx)
if err := ws.serveHTTP(w, r); err != nil {
log.Println(err)
http.Error(w, err.Error(), http.StatusInternalServerError)
}
}
func (ws WS) serveHTTP(w http.ResponseWriter, r *http.Request) error {
switch r.URL.Path {
case "/api/ws":
return ws.serveWS(w, r)
}
if strings.HasPrefix(r.URL.Path, "/proxy") {
return ws.serveProxy(w, r)
}
return ws.serveStaticFile(w, r)
}
func (ws WS) serveProxy(w http.ResponseWriter, r *http.Request) error {
u, err := url.Parse(FlagWSProxy)
if err != nil {
return err
}
r.URL.Path = strings.TrimPrefix(r.URL.Path, "/proxy")
if r.URL.Path == "" {
r.URL.Path = "/"
}
proxy := httputil.NewSingleHostReverseProxy(u)
proxy.ServeHTTP(w, r)
return nil
}
//go:embed public/*
var staticFiles embed.FS
func (ws WS) serveStaticFile(w http.ResponseWriter, r *http.Request) error {
if FlagWSDebug {
b, _ := os.ReadFile("src/device/input/raw/public/root.html")
w.Write(b)
return nil
}
if r.URL.Path == "/" {
r.URL.Path = "root.html"
}
r.URL.Path = path.Join("public", r.URL.Path)
http.FileServer(http.FS(staticFiles)).ServeHTTP(w, r)
return nil
}
func (ws WS) serveWS(w http.ResponseWriter, r *http.Request) error {
if err := ws._serveWS(w, r); err != nil {
log.Println("_serveWS:", err)
}
return nil
}
func (ws WS) _serveWS(w http.ResponseWriter, r *http.Request) error {
conn, err := upgrader.Upgrade(w, r, nil)
if err != nil {
return err
}
defer conn.Close()
for ws.ctx.Err() == nil {
msgType, p, err := conn.ReadMessage()
if err != nil {
if websocket.IsCloseError(err) || websocket.IsUnexpectedCloseError(err) {
return nil
}
return err
}
if msgType == websocket.TextMessage {
log.Println(string(p))
ws.ch <- p
}
}
return nil
}