add 404 handler (#57)

1 files changed, 7 insertions, 0 deletions

diff --git a/internal/router/router.go b/internal/router/router.go
index e575b112f..1b8d899fa 100644
--- a/internal/router/router.go
+++ b/internal/router/router.go
@@ -42,6 +42,8 @@ type Router interface {
 	AttachHandler(method string, path string, f gin.HandlerFunc)
 	// Attach a gin middleware to the router that will be used globally
 	AttachMiddleware(handler gin.HandlerFunc)
+	// Attach 404 NoRoute handler
+	AttachNoRouteHandler(handler gin.HandlerFunc)
 	// Start the router
 	Start()
 	// Stop the router
@@ -109,6 +111,11 @@ func (r *router) AttachMiddleware(middleware gin.HandlerFunc) {
 	r.engine.Use(middleware)
 }
 
+// AttachNoRouteHandler attaches a gin.HandlerFunc to NoRoute to handle 404's
+func (r *router) AttachNoRouteHandler(handler gin.HandlerFunc) {
+	r.engine.NoRoute(handler)
+}
+
 // New returns a new Router with the specified configuration, using the given logrus logger.
 func New(config *config.Config, logger *logrus.Logger) (Router, error) {
 	lvl, err := logrus.ParseLevel(config.LogLevel)