// Echo the data received on the WebSocket. func NewWebSocketHandler(fn func([]byte) []byte) http.Handler { var server websocket.Server server.Handler = func(ws *websocket.Conn) { log.Printf("--- Received connection") buffer := make([]byte, 8000) var err error for { var nBytes int nBytes, err = ws.Read(buffer) if err != nil { break } frame := buffer[:nBytes] if response := fn(frame); response != nil { ws.Write(response) } } log.Printf("--- End connection (%v)", err) } return server }
func (h *Hrotti) AddListener(name string, config *ListenerConfig) error { listener := &internalListener{name: name, url: *config.URL} listener.stop = make(chan struct{}) h.listeners[name] = listener ln, err := net.Listen("tcp", listener.url.Host) if err != nil { ERROR.Println(err.Error()) return err } if listener.url.Scheme == "ws" && len(listener.url.Path) == 0 { listener.url.Path = "/" } h.listenersWaitGroup.Add(1) INFO.Println("Starting MQTT listener on", listener.url.String()) go func() { <-listener.stop INFO.Println("Listener", name, "is stopping...") ln.Close() }() //if this is a WebSocket listener if listener.url.Scheme == "ws" { var server websocket.Server //override the Websocket handshake to accept any protocol name server.Handshake = func(c *websocket.Config, req *http.Request) error { c.Origin, _ = url.Parse(req.RemoteAddr) c.Protocol = []string{"mqtt"} return nil } //set up the ws connection handler, ie what we do when we get a new websocket connection server.Handler = func(ws *websocket.Conn) { ws.PayloadType = websocket.BinaryFrame INFO.Println("New incoming websocket connection", ws.RemoteAddr()) listener.connections = append(listener.connections, ws) h.InitClient(ws) } //set the path that the http server will recognise as related to this websocket //server, needs to be configurable really. http.Handle(listener.url.Path, server) //ListenAndServe loops forever receiving connections and initiating the handler //for each one. go func(ln net.Listener) { defer h.listenersWaitGroup.Done() err := http.Serve(ln, nil) if err != nil { ERROR.Println(err.Error()) return } }(ln) } else { //loop forever accepting connections and launch InitClient as a goroutine with the connection go func() { defer h.listenersWaitGroup.Done() for { conn, err := ln.Accept() if err != nil { ERROR.Println(err.Error()) return } INFO.Println("New incoming connection", conn.RemoteAddr()) listener.connections = append(listener.connections, conn) go h.InitClient(conn) } }() } return nil }