add OnDisconnect() handlers.

go.mod

@@ -44,6 +44,7 @@ require (
 	github.com/mattn/go-sqlite3 v1.14.17 // indirect
 	github.com/pkg/errors v0.9.1 // indirect
 	github.com/savsgio/gotils v0.0.0-20230208104028-c358bd845dee // indirect
+	github.com/sebest/xff v0.0.0-20210106013422-671bd2870b3a // indirect
 	github.com/tidwall/gjson v1.14.4 // indirect
 	github.com/tidwall/match v1.1.1 // indirect
 	github.com/tidwall/pretty v1.2.0 // indirect

go.sum

@@ -102,6 +102,8 @@ github.com/rs/cors v1.7.0 h1:+88SsELBHx5r+hZ8TCkggzSstaWNbDvThkVK8H6f9ik=
 github.com/rs/cors v1.7.0/go.mod h1:gFx+x8UowdsKA9AchylcLynDq+nNFfI8FkUZdN/jGCU=
 github.com/savsgio/gotils v0.0.0-20230208104028-c358bd845dee h1:8Iv5m6xEo1NR1AvpV+7XmhI4r39LGNzwUL4YpMuL5vk=
 github.com/savsgio/gotils v0.0.0-20230208104028-c358bd845dee/go.mod h1:qwtSXrKuJh/zsFQ12yEE89xfCrGKK63Rr7ctU/uCo4g=
+github.com/sebest/xff v0.0.0-20210106013422-671bd2870b3a h1:iLcLb5Fwwz7g/DLK89F+uQBDeAhHhwdzB5fSlVdhGcM=
+github.com/sebest/xff v0.0.0-20210106013422-671bd2870b3a/go.mod h1:wozgYq9WEBQBaIJe4YZ0qTSFAMxmcwBhQH0fO0R34Z0=
 github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME=
 github.com/stretchr/testify v1.4.0/go.mod h1:j7eGeouHqKxXV5pUuKE4zz7dFj8WfuZ+81PSLYec5m4=
 github.com/stretchr/testify v1.8.4 h1:CcVxjf3Q8PM0mHUKJCdn+eZZtm5yQwehR5yeSVQQcUk=

handlers.go

@@ -61,6 +61,10 @@ func (rl *Relay) HandleWebsocket(w http.ResponseWriter, r *http.Request) {
 	)
 
 	kill := func() {
+		for _, ondisconnect := range rl.OnDisconnect {
+			ondisconnect(ctx)
+		}
+
 		ticker.Stop()
 		cancel()
 		if _, ok := rl.clients.Load(conn); ok {
@@ -96,7 +100,7 @@ func (rl *Relay) HandleWebsocket(w http.ResponseWriter, r *http.Request) {
 				) {
 					rl.Log.Printf("unexpected close error from %s: %v\n", r.Header.Get("X-Forwarded-For"), err)
 				}
-				break
+				return
 			}
 
 			if typ == websocket.PingMessage {
@@ -195,7 +199,6 @@ func (rl *Relay) HandleWebsocket(w http.ResponseWriter, r *http.Request) {
 					if pubkey, ok := nip42.ValidateAuthEvent(&env.Event, ws.Challenge, wsBaseUrl); ok {
 						ws.AuthedPublicKey = pubkey
 						close(ws.Authed)
-						ctx = context.WithValue(ctx, AUTH_CONTEXT_KEY, pubkey)
 						ws.WriteJSON(nostr.OKEnvelope{EventID: env.Event.ID, OK: true})
 					} else {
 						ws.WriteJSON(nostr.OKEnvelope{EventID: env.Event.ID, OK: false, Reason: "error: failed to authenticate"})

relay.go

@@ -56,6 +56,7 @@ type Relay struct {
 	CountEvents               []func(ctx context.Context, filter nostr.Filter) (int64, error)
 	OnAuth                    []func(ctx context.Context, pubkey string)
 	OnConnect                 []func(ctx context.Context)
+	OnDisconnect              []func(ctx context.Context)
 	OnEventSaved              []func(ctx context.Context, event *nostr.Event)
 
 	// editing info will affect

utils.go

@@ -2,6 +2,9 @@ package khatru
 
 import (
 	"context"
+
+	"github.com/nbd-wtf/go-nostr"
+	"github.com/sebest/xff"
 )
 
 func GetConnection(ctx context.Context) *WebSocket {
@@ -9,9 +12,21 @@ func GetConnection(ctx context.Context) *WebSocket {
 }
 
 func GetAuthed(ctx context.Context) string {
-	authedPubkey := ctx.Value(AUTH_CONTEXT_KEY)
-	if authedPubkey == nil {
-		return ""
-	}
-	return authedPubkey.(string)
+	return GetConnection(ctx).AuthedPublicKey
+}
+
+func GetIP(ctx context.Context) string {
+	return xff.GetRemoteAddr(GetConnection(ctx).Request)
+}
+
+func GetOpenSubscriptions(ctx context.Context) []nostr.Filter {
+	if listeners, ok := listeners.Load(GetConnection(ctx)); ok {
+		res := make([]nostr.Filter, 0, listeners.Size()*2)
+		listeners.Range(func(_ string, listener *Listener) bool {
+			res = append(res, listener.filters...)
+			return true
+		})
+		return res
+	}
+	return nil
 }