photoprism/internal/api/websocket.go

194 lines
3.7 KiB
Go
Raw Normal View History

package api
import (
"encoding/json"
"net/http"
"sync"
"time"
"github.com/gin-gonic/gin"
"github.com/gorilla/websocket"
"github.com/photoprism/photoprism/internal/config"
"github.com/photoprism/photoprism/internal/entity"
"github.com/photoprism/photoprism/internal/event"
"github.com/photoprism/photoprism/internal/service"
"github.com/photoprism/photoprism/pkg/rnd"
)
var wsConnection = websocket.Upgrader{
CheckOrigin: func(r *http.Request) bool {
return true
},
}
var wsTimeout = 90 * time.Second
type clientInfo struct {
SessionToken string `json:"session"`
JsHash string `json:"js"`
CssHash string `json:"css"`
ManifestHash string `json:"manifest"`
Version string `json:"version"`
}
var wsAuth = struct {
user map[string]entity.User
mutex sync.RWMutex
}{user: make(map[string]entity.User)}
func wsReader(ws *websocket.Conn, writeMutex *sync.Mutex, connId string, conf *config.Config) {
defer ws.Close()
ws.SetReadLimit(512)
ws.SetReadDeadline(time.Now().Add(wsTimeout))
ws.SetPongHandler(func(string) error { ws.SetReadDeadline(time.Now().Add(wsTimeout)); return nil })
for {
_, m, err := ws.ReadMessage()
if err != nil {
break
}
var info clientInfo
if err := json.Unmarshal(m, &info); err != nil {
// Do nothing.
} else {
if sess := Session(info.SessionToken); sess.Valid() {
wsAuth.mutex.Lock()
wsAuth.user[connId] = sess.User
wsAuth.mutex.Unlock()
var clientConfig config.ClientConfig
if sess.User.Guest() {
clientConfig = conf.GuestConfig()
} else if sess.User.Registered() {
clientConfig = conf.UserConfig()
} else {
clientConfig = conf.PublicConfig()
}
writeMutex.Lock()
ws.SetWriteDeadline(time.Now().Add(30 * time.Second))
if err := ws.WriteJSON(gin.H{"event": "config.updated", "data": event.Data{"config": clientConfig}}); err != nil {
// Do nothing.
}
writeMutex.Unlock()
}
}
}
}
func wsWriter(ws *websocket.Conn, writeMutex *sync.Mutex, connId string) {
pingTicker := time.NewTicker(15 * time.Second)
s := event.Subscribe(
"log.*",
"notify.*",
"index.*",
"upload.*",
"import.*",
"config.*",
"count.*",
"photos.*",
"cameras.*",
"lenses.*",
"countries.*",
"albums.*",
"labels.*",
"sync.*",
)
defer func() {
pingTicker.Stop()
event.Unsubscribe(s)
ws.Close()
wsAuth.mutex.Lock()
wsAuth.user[connId] = entity.UnknownUser
wsAuth.mutex.Unlock()
}()
for {
select {
case <-pingTicker.C:
writeMutex.Lock()
ws.SetWriteDeadline(time.Now().Add(30 * time.Second))
if err := ws.WriteMessage(websocket.PingMessage, []byte{}); err != nil {
writeMutex.Unlock()
return
}
writeMutex.Unlock()
case msg := <-s.Receiver:
wsAuth.mutex.RLock()
user := entity.UnknownUser
if hit, ok := wsAuth.user[connId]; ok {
user = hit
}
wsAuth.mutex.RUnlock()
if user.Registered() {
writeMutex.Lock()
ws.SetWriteDeadline(time.Now().Add(30 * time.Second))
if err := ws.WriteJSON(gin.H{"event": msg.Name, "data": msg.Fields}); err != nil {
writeMutex.Unlock()
return
}
writeMutex.Unlock()
}
}
}
}
// GET /api/v1/ws
func Websocket(router *gin.RouterGroup) {
if router == nil {
return
}
conf := service.Config()
if conf == nil {
return
}
router.GET("/ws", func(c *gin.Context) {
w := c.Writer
r := c.Request
ws, err := wsConnection.Upgrade(w, r, nil)
if err != nil {
return
}
var writeMutex sync.Mutex
defer ws.Close()
connId := rnd.UUID()
// Init connection.
wsAuth.mutex.Lock()
if conf.Public() {
wsAuth.user[connId] = entity.Admin
} else {
wsAuth.user[connId] = entity.UnknownUser
}
wsAuth.mutex.Unlock()
go wsWriter(ws, &writeMutex, connId)
wsReader(ws, &writeMutex, connId, conf)
})
}