controller options to set a OnSocketConnect,Disconnect callback function (#56)

Author: adnaan

controller.go

@@ -28,9 +28,11 @@ type Controller interface {
 }
 
 type opt struct {
-	channelFunc       func(r *http.Request, viewID string) *string
-	pathParamsFunc    func(r *http.Request) PathParams
-	websocketUpgrader websocket.Upgrader
+	onSocketConnect    func(userOrSessionID string) error
+	onSocketDisconnect func(userOrSessionID string)
+	channelFunc        func(r *http.Request, viewID string) *string
+	pathParamsFunc     func(r *http.Request) PathParams
+	websocketUpgrader  websocket.Upgrader
 
 	disableTemplateCache  bool
 	disableWebsocket      bool
@@ -146,6 +148,27 @@ func WithDropDuplicateInterval(interval time.Duration) ControllerOption {
 	}
 }
 
+// WithOnSocketConnect takes a function that is called when a new websocket connection is established.
+// The function should return an error if the connection should be rejected.
+// The user or fir's browser session id is passed to the function.
+// user must be set in request.Context with the key UserKey by a developer supplied authentication mechanism.
+// It can be used to track user connections and disconnections.
+// It can be be used to reject connections based on user or session id.
+// It can be used to refresh the page data when a user re-connects.
+func WithOnSocketConnect(f func(userOrSessionID string) error) ControllerOption {
+	return func(o *opt) {
+		o.onSocketConnect = f
+	}
+}
+
+// WithOnSocketDisconnect takes a function that is called when a websocket connection is disconnected.
+func WithOnSocketDisconnect(f func(userOrSessionID string)) ControllerOption {
+	return func(o *opt) {
+		o.onSocketDisconnect = f
+	}
+
+}
+
 // DisableTemplateCache is an option to disable template caching. This is useful for development.
 func DisableTemplateCache() ControllerOption {
 	return func(o *opt) {

examples/counter-ticker/main.go

@@ -116,7 +116,17 @@ func (i *index) updated(ctx fir.RouteContext) error {
 
 func main() {
 	pubsubAdapter := pubsub.NewInmem()
-	controller := fir.NewController("counter_app", fir.DevelopmentMode(true), fir.WithPubsubAdapter(pubsubAdapter))
+	controller := fir.NewController("counter_app",
+		fir.DevelopmentMode(true),
+		fir.WithPubsubAdapter(pubsubAdapter),
+		fir.WithOnSocketConnect(func(userOrSessionID string) error {
+			fmt.Printf("socket connected for user %s\n", userOrSessionID)
+			return nil
+		}),
+		fir.WithOnSocketDisconnect(func(userOrSessionID string) {
+			fmt.Printf("socket disconnected for user %s\n", userOrSessionID)
+		}),
+	)
 	http.Handle("/", controller.Route(NewCounterIndex(pubsubAdapter)))
 	http.ListenAndServe(":9867", nil)
 }

websocket.go

@@ -96,6 +96,20 @@ func onWebsocket(w http.ResponseWriter, r *http.Request, cntrl *controller) {
 
 	user := getUserFromRequestContext(r)
 
+	connectedUser := user
+	if user == "" {
+		connectedUser = sessionID
+	}
+	if cntrl.onSocketConnect != nil {
+		err := cntrl.onSocketConnect(connectedUser)
+		if err != nil {
+			return
+		}
+	}
+	if cntrl.onSocketDisconnect != nil {
+		defer cntrl.onSocketDisconnect(connectedUser)
+	}
+
 	send := make(chan []byte, 100)
 
 	ctx := context.Background()