1 files changed, 30 insertions, 6 deletions

diff --git a/cmd/relay/main.go b/cmd/relay/main.go
index 3db466d..53296b9 100644
--- a/cmd/relay/main.go
+++ b/cmd/relay/main.go
@@ -4,20 +4,26 @@ import (
         "flag"
         "log"
         "net"
+        "net/http"
         "os"
         "os/signal"
         "syscall"
 
+        "context"
+
         "google.golang.org/grpc"
 
         pb "northwest.io/nostr-grpc/api/nostr/v1"
         grpchandler "northwest.io/nostr-grpc/internal/handler/grpc"
+        wshandler "northwest.io/nostr-grpc/internal/handler/websocket"
         "northwest.io/nostr-grpc/internal/storage"
+        "northwest.io/nostr-grpc/internal/subscription"
 )
 
 func main() {
         var (
                 grpcAddr = flag.String("grpc-addr", ":50051", "gRPC server address")
+                wsAddr   = flag.String("ws-addr", ":8080", "WebSocket server address")
                 dbPath   = flag.String("db", "relay.db", "SQLite database path")
         )
         flag.Parse()
@@ -28,17 +34,28 @@ func main() {
         }
         defer store.Close()
 
-        handler := grpchandler.NewServer(store)
+        subManager := subscription.NewManager()
+
+        grpcHandler := grpchandler.NewServer(store)
+        grpcHandler.SetSubscriptionManager(subManager)
 
-        lis, err := net.Listen("tcp", *grpcAddr)
+        wsHandler := wshandler.NewHandler(store, subManager)
+
+        grpcLis, err := net.Listen("tcp", *grpcAddr)
         if err != nil {
-                log.Fatalf("failed to listen: %v", err)
+                log.Fatalf("failed to listen on gRPC port: %v", err)
         }
 
         grpcServer := grpc.NewServer()
-        pb.RegisterNostrRelayServer(grpcServer, handler)
+        pb.RegisterNostrRelayServer(grpcServer, grpcHandler)
+
+        httpServer := &http.Server{
+                Addr:    *wsAddr,
+                Handler: wsHandler,
+        }
 
         log.Printf("gRPC server listening on %s", *grpcAddr)
+        log.Printf("WebSocket server listening on %s", *wsAddr)
         log.Printf("Database: %s", *dbPath)
 
         sigChan := make(chan os.Signal, 1)
@@ -48,9 +65,16 @@ func main() {
                 <-sigChan
                 log.Println("Shutting down...")
                 grpcServer.GracefulStop()
+                httpServer.Shutdown(context.Background())
+        }()
+
+        go func() {
+                if err := grpcServer.Serve(grpcLis); err != nil {
+                        log.Fatalf("gRPC server failed: %v", err)
+                }
         }()
 
-        if err := grpcServer.Serve(lis); err != nil {
-                log.Fatalf("failed to serve: %v", err)
+        if err := httpServer.ListenAndServe(); err != nil && err != http.ErrServerClosed {
+                log.Fatalf("WebSocket server failed: %v", err)
         }
 }