simple server relay with clients

cmd/server/main.go

@@ -0,0 +1,83 @@
+package main
+
+import (
+	"encoding/json"
+	"log"
+	"net/http"
+	"sync"
+
+	"gossip/pkg/protocol"
+
+	"github.com/gorilla/websocket"
+)
+
+type Server struct {
+	clients  map[string]*websocket.Conn
+	mu       sync.Mutex
+	upgrader websocket.Upgrader
+}
+
+func main() {
+	srv := &Server{
+		clients: make(map[string]*websocket.Conn),
+		upgrader: websocket.Upgrader{
+			ReadBufferSize:  1024,
+			WriteBufferSize: 1024,
+			CheckOrigin:     func(r *http.Request) bool { return true },
+		},
+	}
+
+	http.HandleFunc("/ws", srv.handleWS)
+	log.Println("Relay Server listening on :8080")
+	log.Fatal(http.ListenAndServe(":8080", nil))
+}
+
+func (s *Server) handleWS(w http.ResponseWriter, r *http.Request) {
+	conn, err := s.upgrader.Upgrade(w, r, nil)
+	if err != nil {
+		log.Println("Upgrade error:", err)
+		return
+	}
+	defer conn.Close()
+
+	var myPubKey string
+
+	for {
+		_, data, err := conn.ReadMessage()
+		if err != nil {
+			break
+		}
+
+		var msg protocol.Message
+		if err := json.Unmarshal(data, &msg); err != nil {
+			continue
+		}
+
+		switch msg.Type {
+		case "login":
+			s.mu.Lock()
+			s.clients[msg.Sender] = conn
+			s.mu.Unlock()
+			myPubKey = msg.Sender
+			log.Printf("Client connected: %s...", myPubKey[:8])
+		case "msg":
+			s.mu.Lock()
+			targetConn, ok := s.clients[msg.Target]
+			s.mu.Unlock()
+
+			if ok {
+				err = targetConn.WriteMessage(websocket.TextMessage, data)
+				if err != nil {
+					log.Printf("Failed to relay to %s", msg.Target[:8])
+				}
+			}
+		}
+	}
+
+	if myPubKey != "" {
+		s.mu.Lock()
+		delete(s.clients, myPubKey)
+		s.mu.Unlock()
+		log.Printf("Client disconnected: %s...", myPubKey[:8])
+	}
+}