properly handling serving websocket and frontend at the same time

config.example.yml

@@ -1,7 +1,5 @@
 mongodb:
   uri: "mongodb://localhost:27017/"
   database: "grain"
-relay:
+server:
   port: ":8080"
-web:
-  port: ":8181"

main.go

@@ -27,23 +27,33 @@ func main() {
 	}
 	defer db.DisconnectDB()
 
-	// Run the WebSocket server in a goroutine
-	go func() {
-		fmt.Printf("WebSocket server is running on ws://localhost%s\n", config.Relay.Port)
-		err := http.ListenAndServe(config.Relay.Port, websocket.Handler(relay.Listener))
-		if err != nil {
-			fmt.Println("Error starting WebSocket server:", err)
-		}
-	}()
-
-	// Run the HTTP server for serving static files and home page
+	// Create a new ServeMux
 	mux := http.NewServeMux()
-	mux.HandleFunc("/", web.RootHandler)
+
+	// Handle the root path
+	mux.HandleFunc("/", ListenAndServe)
+
+	// Serve static files
 	mux.Handle("/static/", http.StripPrefix("/static/", http.FileServer(http.Dir("web/static"))))
 
-	fmt.Printf("Http server is running on http://localhost%s\n", config.Web.Port)
-	err = http.ListenAndServe(config.Web.Port, mux)
+	// Serve the favicon
+	mux.HandleFunc("/favicon.ico", func(w http.ResponseWriter, r *http.Request) {
+		http.ServeFile(w, r, "web/static/img/favicon.ico")
+	})
+
+	// Start the Relay
+	fmt.Printf("Server is running on http://localhost%s\n", config.Server.Port)
+	err = http.ListenAndServe(config.Server.Port, mux)
 	if err != nil {
-		fmt.Println("Error starting web server:", err)
+		fmt.Println("Error starting server:", err)
+	}
+}
+
+// Listener serves both WebSocket and HTML
+func ListenAndServe(w http.ResponseWriter, r *http.Request) {
+	if r.Header.Get("Upgrade") == "websocket" {
+		websocket.Handler(relay.WebSocketHandler).ServeHTTP(w, r)
+	} else {
+		web.RootHandler(w, r)
 	}
 }

relay/relay.go

@@ -8,7 +8,8 @@ import (
 	"golang.org/x/net/websocket"
 )
 
-func Listener(ws *websocket.Conn) {
+// WebSocketHandler handles incoming WebSocket connections
+func WebSocketHandler(ws *websocket.Conn) {
 	defer ws.Close()
 
 	var msg string

relay/utils/loadConfig.go

@@ -11,12 +11,9 @@ type Config struct {
 		URI      string `yaml:"uri"`
 		Database string `yaml:"database"`
 	} `yaml:"mongodb"`
-	Relay struct {
+	Server struct {
 		Port string `yaml:"port"`
-	} `yaml:"relay"`
-	Web struct {
-		Port string `yaml:"port"`
-	} `yaml:"web"`
+	} `yaml:"server"`
 }
 
 func LoadConfig(filename string) (*Config, error) {