diff --git a/tower.go b/tower.go index 1735c47..021f9ef 100644 --- a/tower.go +++ b/tower.go @@ -49,6 +49,22 @@ func (t *Tower) SetErrorHandler(errorHandler ErrorHandler) { t.errorHandler = errorHandler } +func (t *Tower) SetNotFoundHandler(notFoundHandler Handler) { + t.router.NotFound = http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + ctx := Context{ + request: r, + response: &responseWriter{ResponseWriter: w}, + params: httprouter.Params{}, + } + + if err := notFoundHandler(&ctx); err != nil { + if err := t.errorHandler(&ctx, err); err != nil { + log.Printf("[Error] %v", err) + } + } + }) +} + func (t *Tower) Use(globalMiddlewares ...Middleware) { t.globalMiddlewares = append(t.globalMiddlewares, globalMiddlewares...) }