From af2b62f782d9ce120bdeb7d81b525b35e9c60b99 Mon Sep 17 00:00:00 2001
From: Marvin Blum <marvin.blum@emvi.com>
Date: Tue, 4 Jun 2024 20:29:38 +0200
Subject: [PATCH] Fixed router setup.

---
 CHANGELOG.md      |  4 ++++
 README.md         |  2 +-
 cmd/shifu/main.go | 13 +++++++++++--
 pkg/server.go     | 41 +++++++++++++++++++++++++----------------
 pkg/version.go    |  2 +-
 5 files changed, 42 insertions(+), 20 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 390dc90..0a74a01 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -1,5 +1,9 @@
 # Changelog
 
+## 0.6.2
+
+* fixed router setup
+
 ## 0.6.1
 
 * fixed router setup
diff --git a/README.md b/README.md
index 4b462a3..bb006cc 100644
--- a/README.md
+++ b/README.md
@@ -259,7 +259,7 @@ func main() {
     })
 	
 	// Start the server. The router and cancel function are optional.
-	if err := server.Start(nil, nil); err != nil {
+	if err := server.Start(nil); err != nil {
 		slog.Error("Error starting Shifu", "error", err)
 	}
 }
diff --git a/cmd/shifu/main.go b/cmd/shifu/main.go
index 2a4aa40..bfec787 100644
--- a/cmd/shifu/main.go
+++ b/cmd/shifu/main.go
@@ -2,7 +2,9 @@ package main
 
 import (
 	shifu "github.com/emvi/shifu/pkg"
+	"github.com/go-chi/chi/v5"
 	"log/slog"
+	"net/http"
 	"os"
 	"strings"
 )
@@ -29,9 +31,16 @@ func main() {
 
 	switch cmd {
 	case "run":
-		server := shifu.NewServer(dir, shifu.ServerOptions{})
+		r := chi.NewRouter()
+		r.HandleFunc("/debug", func(w http.ResponseWriter, _ *http.Request) {
+			w.Write([]byte("HELLO WORLD!"))
+		})
 
-		if err := server.Start(nil, nil); err != nil {
+		server := shifu.NewServer(dir, shifu.ServerOptions{
+			Router: r,
+		})
+
+		if err := server.Start(nil); err != nil {
 			slog.Error("Error starting Shifu", "error", err)
 		}
 	case "init":
diff --git a/pkg/server.go b/pkg/server.go
index 39e2131..a8a182f 100644
--- a/pkg/server.go
+++ b/pkg/server.go
@@ -39,10 +39,7 @@ type Server struct {
 	// Content is the CMS content.
 	Content *cms.CMS
 
-	// Router is the router for the server.
-	// It has a CORS and gzip middleware, as well as a route for robots.txt, the static directory, and any unmatched path (/*) by default.
-	Router chi.Router
-
+	router  chi.Router
 	dir     string
 	funcMap template.FuncMap
 }
@@ -51,7 +48,7 @@ type Server struct {
 // The second argument is an optional template.FuncMap that will be merged with Shifu's funcmap.
 func NewServer(dir string, options ServerOptions) *Server {
 	return &Server{
-		Router:  options.Router,
+		router:  options.Router,
 		dir:     dir,
 		funcMap: options.FuncMap,
 	}
@@ -59,7 +56,7 @@ func NewServer(dir string, options ServerOptions) *Server {
 
 // Start starts the Shifu server.
 // The context.CancelFunc is optional and will be called on server shutdown or error if set.
-func (server *Server) Start(router chi.Router, cancel context.CancelFunc) error {
+func (server *Server) Start(cancel context.CancelFunc) error {
 	slog.Info("Starting Shifu", "version", version, "directory", server.dir)
 	server.funcMap = cms.Merge(server.funcMap)
 	ctx, cancelServer := context.WithCancel(context.Background())
@@ -112,28 +109,40 @@ func (server *Server) Start(router chi.Router, cancel context.CancelFunc) error
 	})
 	analytics.Init()
 	server.setupRouter(server.dir, server.Content, sm)
-	<-server.startServer(server.Router, stop)
+	<-server.startServer(server.router, stop)
 	return nil
 }
 
 func (server *Server) setupRouter(dir string, cms *cms.CMS, sm *sitemap.Sitemap) {
-	var router chi.Router
-
-	if server.Router != nil {
-		router = server.Router
-	} else {
-		router = chi.NewRouter()
-	}
-
+	router := chi.NewRouter()
 	router.Use(
 		middleware.Cors(),
 		middleware.Gzip(),
 	)
+
+	if server.router != nil {
+		slog.Info("Merging router with Shifu router...")
+
+		for _, route := range server.router.Routes() {
+			for method, handler := range route.Handlers {
+				if method == "*" {
+					router.Handle(route.Pattern, handler)
+				} else {
+					router.Method(method, route.Pattern, handler)
+				}
+			}
+		}
+	}
+
 	sm.Serve(router)
 	server.serveRobotsTxt(router)
 	server.serveStaticDir(router, dir)
 	router.Handle("/*", http.HandlerFunc(cms.Serve))
-	server.Router = router
+	server.router = router
+
+	for _, route := range router.Routes() {
+		slog.Info("Added route", "route", route.Pattern)
+	}
 }
 
 func (server *Server) serveRobotsTxt(router chi.Router) {
diff --git a/pkg/version.go b/pkg/version.go
index 90d5bd9..70887fc 100644
--- a/pkg/version.go
+++ b/pkg/version.go
@@ -1,7 +1,7 @@
 package pkg
 
 const (
-	version = "0.6.1"
+	version = "0.6.2"
 )
 
 // Version returns the Shifu version number.