partila

Author: zeripath

go.mod

@@ -51,6 +51,7 @@ require (
 	github.com/google/uuid v1.3.0
 	github.com/gorilla/feeds v1.1.1
 	github.com/gorilla/sessions v1.2.1
+	github.com/gorilla/websocket v1.4.2
 	github.com/hashicorp/go-version v1.4.0
 	github.com/hashicorp/golang-lru v0.5.4
 	github.com/huandu/xstrings v1.3.2
@@ -194,7 +195,6 @@ require (
 	github.com/gorilla/handlers v1.5.1 // indirect
 	github.com/gorilla/mux v1.8.0 // indirect
 	github.com/gorilla/securecookie v1.1.1 // indirect
-	github.com/gorilla/websocket v1.4.2 // indirect
 	github.com/grpc-ecosystem/go-grpc-middleware v1.3.0 // indirect
 	github.com/grpc-ecosystem/go-grpc-prometheus v1.2.0 // indirect
 	github.com/grpc-ecosystem/grpc-gateway v1.16.0 // indirect

routers/web/events/websocket.go

@@ -0,0 +1,126 @@
+// Copyright 2022 The Gitea Authors. All rights reserved.
+// Use of this source code is governed by a MIT-style
+// license that can be found in the LICENSE file.
+
+package events
+
+import (
+	"code.gitea.io/gitea/modules/context"
+	"code.gitea.io/gitea/modules/eventsource"
+	"code.gitea.io/gitea/modules/graceful"
+	"code.gitea.io/gitea/modules/json"
+	"code.gitea.io/gitea/modules/log"
+	"code.gitea.io/gitea/routers/web/auth"
+	"github.com/gorilla/websocket"
+)
+
+type readMessage struct {
+	messageType int
+	message     []byte
+	err         error
+}
+
+// Events listens for events
+func Websocket(ctx *context.Context) {
+	upgrader := websocket.Upgrader{
+		ReadBufferSize:  1024,
+		WriteBufferSize: 1024,
+	}
+
+	conn, err := upgrader.Upgrade(ctx.Resp, ctx.Req, nil)
+	if err != nil {
+		log.Error("Unable to upgrade due to error: %v", err)
+		return
+	}
+	defer conn.Close()
+
+	notify := ctx.Done()
+	shutdownCtx := graceful.GetManager().ShutdownContext()
+
+	eventChan := make(<-chan *eventsource.Event)
+	uid := int64(0)
+	unregister := func() {}
+	if ctx.IsSigned {
+		uid = ctx.Doer.ID
+		eventChan = eventsource.GetManager().Register(uid)
+		unregister = func() {
+			go func() {
+				eventsource.GetManager().Unregister(uid, eventChan)
+				// ensure the messageChan is closed
+				for {
+					_, ok := <-eventChan
+					if !ok {
+						break
+					}
+				}
+			}()
+		}
+	}
+	defer unregister()
+
+	readChan := make(chan readMessage, 20)
+	go func() {
+		for {
+			messageType, message, err := conn.ReadMessage()
+			readChan <- readMessage{
+				messageType: messageType,
+				message:     message,
+				err:         err,
+			}
+			if err != nil {
+				close(readChan)
+				return
+			}
+		}
+	}()
+
+	for {
+		select {
+		case <-notify:
+			return
+		case <-shutdownCtx.Done():
+			return
+		case _, ok := <-readChan:
+			if !ok {
+				break
+			}
+		case event, ok := <-eventChan:
+			if !ok {
+				break
+			}
+			if event.Name == "logout" {
+				if ctx.Session.ID() == event.Data {
+					_, _ = (&eventsource.Event{
+						Name: "logout",
+						Data: "here",
+					}).WriteTo(ctx.Resp)
+					ctx.Resp.Flush()
+					go unregister()
+					auth.HandleSignOut(ctx)
+					break
+				}
+				// Replace the event - we don't want to expose the session ID to the user
+				event = &eventsource.Event{
+					Name: "logout",
+					Data: "elsewhere",
+				}
+			}
+
+			w, err := conn.NextWriter(websocket.TextMessage)
+			if err != nil {
+				log.Warn("Unable to get writer for websocket %v", err)
+				return
+			}
+
+			if err := json.NewEncoder(w).Encode(event); err != nil {
+				log.Error("Unable to create encoder for %v %v", event, err)
+				return
+			}
+			if err := w.Close(); err != nil {
+				log.Warn("Unable to close writer for websocket %v", err)
+				return
+			}
+
+		}
+	}
+}