3 files changed, 326 insertions, 0 deletions

diff --git a/internal/handler/grpc/server.go b/internal/handler/grpc/server.go
index 74857f6..b65b527 100644
--- a/internal/handler/grpc/server.go
+++ b/internal/handler/grpc/server.go
@@ -28,6 +28,10 @@ func NewServer(store EventStore) *Server {
         }
 }
 
+func (s *Server) SetSubscriptionManager(mgr *subscription.Manager) {
+        s.subs = mgr
+}
+
 func (s *Server) PublishEvent(ctx context.Context, req *pb.PublishEventRequest) (*pb.PublishEventResponse, error) {
         if req.Event == nil {
                 return &pb.PublishEventResponse{
diff --git a/internal/handler/websocket/convert.go b/internal/handler/websocket/convert.go
new file mode 100644
index 0000000..0458ee4
--- /dev/null
+++ b/internal/handler/websocket/convert.go
@@ -0,0 +1,76 @@
+package websocket
+
+import (
+        pb "northwest.io/nostr-grpc/api/nostr/v1"
+        "northwest.io/nostr-grpc/internal/nostr"
+)
+
+func NostrToPB(n *nostr.Event) *pb.Event {
+        tags := make([]*pb.Tag, len(n.Tags))
+        for i, tag := range n.Tags {
+                tags[i] = &pb.Tag{Values: tag}
+        }
+
+        return &pb.Event{
+                Id:        n.ID,
+                Pubkey:    n.PubKey,
+                CreatedAt: n.CreatedAt,
+                Kind:      int32(n.Kind),
+                Tags:      tags,
+                Content:   n.Content,
+                Sig:       n.Sig,
+        }
+}
+
+func PBToNostr(e *pb.Event) *nostr.Event {
+        tags := make(nostr.Tags, len(e.Tags))
+        for i, tag := range e.Tags {
+                tags[i] = tag.Values
+        }
+
+        return &nostr.Event{
+                ID:        e.Id,
+                PubKey:    e.Pubkey,
+                CreatedAt: e.CreatedAt,
+                Kind:      int(e.Kind),
+                Tags:      tags,
+                Content:   e.Content,
+                Sig:       e.Sig,
+        }
+}
+
+func NostrFilterToPB(f *nostr.Filter) *pb.Filter {
+        pbFilter := &pb.Filter{
+                Ids:     f.IDs,
+                Authors: f.Authors,
+                Kinds:   make([]int32, len(f.Kinds)),
+        }
+
+        for i, kind := range f.Kinds {
+                pbFilter.Kinds[i] = int32(kind)
+        }
+
+        if f.Since != nil {
+                since := int64(*f.Since)
+                pbFilter.Since = &since
+        }
+
+        if f.Until != nil {
+                until := int64(*f.Until)
+                pbFilter.Until = &until
+        }
+
+        if f.Limit > 0 {
+                limit := int32(f.Limit)
+                pbFilter.Limit = &limit
+        }
+
+        if len(f.Tags) > 0 {
+                pbFilter.TagFilters = make(map[string]*pb.TagFilter)
+                for tagName, values := range f.Tags {
+                        pbFilter.TagFilters[tagName] = &pb.TagFilter{Values: values}
+                }
+        }
+
+        return pbFilter
+}
diff --git a/internal/handler/websocket/handler.go b/internal/handler/websocket/handler.go
new file mode 100644
index 0000000..cef83dd
--- /dev/null
+++ b/internal/handler/websocket/handler.go
@@ -0,0 +1,246 @@
+package websocket
+
+import (
+        "context"
+        "encoding/json"
+        "fmt"
+        "log"
+        "net/http"
+
+        pb "northwest.io/nostr-grpc/api/nostr/v1"
+        "northwest.io/nostr-grpc/internal/nostr"
+        "northwest.io/nostr-grpc/internal/storage"
+        "northwest.io/nostr-grpc/internal/subscription"
+        "northwest.io/nostr-grpc/internal/websocket"
+)
+
+type EventStore interface {
+        StoreEvent(context.Context, *storage.EventData) error
+        QueryEvents(context.Context, []*pb.Filter, *storage.QueryOptions) ([]*pb.Event, error)
+}
+
+type Handler struct {
+        store EventStore
+        subs  *subscription.Manager
+}
+
+func NewHandler(store EventStore, subs *subscription.Manager) *Handler {
+        return &Handler{
+                store: store,
+                subs:  subs,
+        }
+}
+
+func (h *Handler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
+        conn, err := websocket.Accept(w, r)
+        if err != nil {
+                log.Printf("WebSocket accept failed: %v", err)
+                return
+        }
+        defer conn.Close(websocket.StatusNormalClosure, "")
+
+        ctx := r.Context()
+        clientSubs := make(map[string]*subscription.Subscription)
+        defer func() {
+                for subID := range clientSubs {
+                        h.subs.Remove(subID)
+                }
+        }()
+
+        for {
+                _, data, err := conn.Read(ctx)
+                if err != nil {
+                        return
+                }
+
+                if err := h.handleMessage(ctx, conn, data, clientSubs); err != nil {
+                        log.Printf("Message handling error: %v", err)
+                        h.sendNotice(ctx, conn, err.Error())
+                }
+        }
+}
+
+func (h *Handler) handleMessage(ctx context.Context, conn *websocket.Conn, data []byte, clientSubs map[string]*subscription.Subscription) error {
+        var raw []json.RawMessage
+        if err := json.Unmarshal(data, &raw); err != nil {
+                return fmt.Errorf("invalid JSON")
+        }
+
+        if len(raw) == 0 {
+                return fmt.Errorf("empty message")
+        }
+
+        var msgType string
+        if err := json.Unmarshal(raw[0], &msgType); err != nil {
+                return fmt.Errorf("invalid message type")
+        }
+
+        switch msgType {
+        case "EVENT":
+                return h.handleEvent(ctx, conn, raw)
+        case "REQ":
+                return h.handleReq(ctx, conn, raw, clientSubs)
+        case "CLOSE":
+                return h.handleClose(raw, clientSubs)
+        default:
+                return fmt.Errorf("unknown message type: %s", msgType)
+        }
+}
+
+func (h *Handler) handleEvent(ctx context.Context, conn *websocket.Conn, raw []json.RawMessage) error {
+        if len(raw) != 2 {
+                return fmt.Errorf("EVENT expects 2 elements")
+        }
+
+        var event nostr.Event
+        if err := json.Unmarshal(raw[1], &event); err != nil {
+                return fmt.Errorf("invalid event: %w", err)
+        }
+
+        if !event.CheckID() {
+                h.sendOK(ctx, conn, event.ID, false, "invalid: event ID mismatch")
+                return nil
+        }
+
+        if !event.Verify() {
+                h.sendOK(ctx, conn, event.ID, false, "invalid: signature verification failed")
+                return nil
+        }
+
+        pbEvent := NostrToPB(&event)
+        canonicalJSON := event.Serialize()
+
+        eventData := &storage.EventData{
+                Event:         pbEvent,
+                CanonicalJSON: canonicalJSON,
+        }
+
+        err := h.store.StoreEvent(ctx, eventData)
+        if err == storage.ErrEventExists {
+                h.sendOK(ctx, conn, event.ID, true, "duplicate: already have this event")
+                return nil
+        }
+        if err != nil {
+                h.sendOK(ctx, conn, event.ID, false, fmt.Sprintf("error: %v", err))
+                return nil
+        }
+
+        h.subs.MatchAndFan(pbEvent)
+
+        h.sendOK(ctx, conn, event.ID, true, "")
+        return nil
+}
+
+func (h *Handler) handleReq(ctx context.Context, conn *websocket.Conn, raw []json.RawMessage, clientSubs map[string]*subscription.Subscription) error {
+        if len(raw) < 3 {
+                return fmt.Errorf("REQ expects at least 3 elements")
+        }
+
+        var subID string
+        if err := json.Unmarshal(raw[1], &subID); err != nil {
+                return fmt.Errorf("invalid subscription ID")
+        }
+
+        var filters []*pb.Filter
+        for i := 2; i < len(raw); i++ {
+                var nostrFilter nostr.Filter
+                if err := json.Unmarshal(raw[i], &nostrFilter); err != nil {
+                        return fmt.Errorf("invalid filter: %w", err)
+                }
+
+                pbFilter := NostrFilterToPB(&nostrFilter)
+                filters = append(filters, pbFilter)
+        }
+
+        if existing, ok := clientSubs[subID]; ok {
+                h.subs.Remove(existing.ID)
+                delete(clientSubs, subID)
+        }
+
+        storedEvents, err := h.store.QueryEvents(ctx, filters, &storage.QueryOptions{Limit: 0})
+        if err != nil {
+                return fmt.Errorf("query failed: %w", err)
+        }
+
+        for _, pbEvent := range storedEvents {
+                event := PBToNostr(pbEvent)
+                h.sendEvent(ctx, conn, subID, event)
+        }
+
+        h.sendEOSE(ctx, conn, subID)
+
+        sub := &subscription.Subscription{
+                ID:      subID,
+                Filters: filters,
+                Events:  make(chan *pb.Event, 100),
+        }
+        sub.InitDone()
+
+        h.subs.Add(sub)
+        clientSubs[subID] = sub
+
+        go h.streamEvents(ctx, conn, sub)
+
+        return nil
+}
+
+func (h *Handler) handleClose(raw []json.RawMessage, clientSubs map[string]*subscription.Subscription) error {
+        if len(raw) != 2 {
+                return fmt.Errorf("CLOSE expects 2 elements")
+        }
+
+        var subID string
+        if err := json.Unmarshal(raw[1], &subID); err != nil {
+                return fmt.Errorf("invalid subscription ID")
+        }
+
+        if sub, ok := clientSubs[subID]; ok {
+                h.subs.Remove(sub.ID)
+                delete(clientSubs, subID)
+        }
+
+        return nil
+}
+
+func (h *Handler) streamEvents(ctx context.Context, conn *websocket.Conn, sub *subscription.Subscription) {
+        for {
+                select {
+                case pbEvent, ok := <-sub.Events:
+                        if !ok {
+                                return
+                        }
+                        event := PBToNostr(pbEvent)
+                        h.sendEvent(ctx, conn, sub.ID, event)
+
+                case <-ctx.Done():
+                        return
+
+                case <-sub.Done():
+                        return
+                }
+        }
+}
+
+func (h *Handler) sendEvent(ctx context.Context, conn *websocket.Conn, subID string, event *nostr.Event) error {
+        msg := []interface{}{"EVENT", subID, event}
+        data, _ := json.Marshal(msg)
+        return conn.Write(ctx, websocket.MessageText, data)
+}
+
+func (h *Handler) sendOK(ctx context.Context, conn *websocket.Conn, eventID string, accepted bool, message string) error {
+        msg := []interface{}{"OK", eventID, accepted, message}
+        data, _ := json.Marshal(msg)
+        return conn.Write(ctx, websocket.MessageText, data)
+}
+
+func (h *Handler) sendEOSE(ctx context.Context, conn *websocket.Conn, subID string) error {
+        msg := []interface{}{"EOSE", subID}
+        data, _ := json.Marshal(msg)
+        return conn.Write(ctx, websocket.MessageText, data)
+}
+
+func (h *Handler) sendNotice(ctx context.Context, conn *websocket.Conn, notice string) error {
+        msg := []interface{}{"NOTICE", notice}
+        data, _ := json.Marshal(msg)
+        return conn.Write(ctx, websocket.MessageText, data)
+}