evcc-io/server/socket.go
2020-03-06 08:24:05 +01:00

141 lines
3.1 KiB
Go

package server
import (
"fmt"
"net/http"
"time"
"github.com/andig/evcc/core"
"github.com/gorilla/websocket"
)
const (
// Time allowed to write a message to the peer
writeTimeout = 10 * time.Second
)
var upgrader = websocket.Upgrader{
ReadBufferSize: 1024,
WriteBufferSize: 1024,
CheckOrigin: func(r *http.Request) bool { return true },
}
// SocketClient is a middleman between the websocket connection and the hub.
type SocketClient struct {
hub *SocketHub
// The websocket connection.
conn *websocket.Conn
// Buffered channel of outbound messages.
send chan []byte
}
// writePump pumps messages from the hub to the websocket connection.
func (c *SocketClient) writePump() {
defer func() {
c.conn.Close()
}()
for {
msg := <-c.send
if err := c.conn.SetWriteDeadline(time.Now().Add(writeTimeout)); err != nil {
return
}
if err := c.conn.WriteMessage(websocket.TextMessage, msg); err != nil {
return
}
}
}
// ServeWebsocket handles websocket requests from the peer.
func ServeWebsocket(hub *SocketHub, w http.ResponseWriter, r *http.Request) {
conn, err := upgrader.Upgrade(w, r, nil)
if err != nil {
log.ERROR.Println(err)
return
}
client := &SocketClient{hub: hub, conn: conn, send: make(chan []byte, 256)}
client.hub.register <- client
// run writing to client in goroutine
go client.writePump()
}
// SocketHub maintains the set of active clients and broadcasts messages to the
// clients.
type SocketHub struct {
// Registered clients.
clients map[*SocketClient]bool
// Register requests from the clients.
register chan *SocketClient
// Unregister requests from clients.
unregister chan *SocketClient
}
// NewSocketHub creates a web socket hub that distributes meter status and
// query results for the ui or other clients
func NewSocketHub() *SocketHub {
return &SocketHub{
register: make(chan *SocketClient),
unregister: make(chan *SocketClient),
clients: make(map[*SocketClient]bool),
}
}
func (h *SocketHub) encode(v core.Param) ([]byte, error) {
switch val := v.Val.(type) {
case string:
s := fmt.Sprintf("{\"%s\": \"%s\"}", v.Key, val)
return []byte(s), nil
case float64:
s := fmt.Sprintf("{\"%s\": %.3f}", v.Key, val)
return []byte(s), nil
case time.Duration:
s := fmt.Sprintf("{\"%s\": %d}", v.Key, int64(val.Seconds()))
return []byte(s), nil
default:
s := fmt.Sprintf("{\"%s\": %v}", v.Key, val)
return []byte(s), nil
}
}
func (h *SocketHub) broadcast(i core.Param) {
if len(h.clients) > 0 {
message, err := h.encode(i)
if err != nil {
log.FATAL.Fatal(err)
}
for client := range h.clients {
select {
case client.send <- message:
default:
close(client.send)
delete(h.clients, client)
}
}
}
}
// Run starts data and status distribution
func (h *SocketHub) Run(in <-chan core.Param) {
for {
select {
case client := <-h.register:
h.clients[client] = true
case client := <-h.unregister:
if _, ok := h.clients[client]; ok {
delete(h.clients, client)
close(client.send)
}
case obj, ok := <-in:
if !ok {
return // break if channel closed
}
h.broadcast(obj)
}
}
}