package router import ( "net/http" // just for status codes "sync" "github.com/kataras/iris/context" ) func statusCodeSuccessful(statusCode int) bool { return !context.StatusCodeNotSuccessful(statusCode) } // ErrorCodeHandler is the entry // of the list of all http error code handlers. type ErrorCodeHandler struct { StatusCode int Handlers context.Handlers mu sync.Mutex } // Fire executes the specific an error http error status. // it's being wrapped to make sure that the handler // will render correctly. func (ch *ErrorCodeHandler) Fire(ctx context.Context) { // if we can reset the body if w, ok := ctx.IsRecording(); ok { if statusCodeSuccessful(w.StatusCode()) { // if not an error status code w.WriteHeader(ch.StatusCode) // then set it manually here, otherwise it should be set via ctx.StatusCode(...) } // reset if previous content and it's recorder, keep the status code. w.ClearHeaders() w.ResetBody() } else if w, ok := ctx.ResponseWriter().(*context.GzipResponseWriter); ok { // reset and disable the gzip in order to be an expected form of http error result w.ResetBody() w.Disable() } else { // if we can't reset the body and the body has been filled // which means that the status code already sent, // then do not fire this custom error code. if ctx.ResponseWriter().Written() > 0 { // != -1, rel: context/context.go#EndRequest return } } // ctx.StopExecution() // not uncomment this, is here to remember why to. // note for me: I don't stopping the execution of the other handlers // because may the user want to add a fallback error code // i.e // users := app.Party("/users") // users.Done(func(ctx context.Context){ if ctx.StatusCode() == 400 { /* custom error code for /users */ }}) // use .HandlerIndex // that sets the current handler index to zero // in order to: // ignore previous runs that may changed the handler index, // via ctx.Next or ctx.StopExecution, if any. // // use .Do // that overrides the existing handlers and sets and runs these error handlers. // in order to: // ignore the route's after-handlers, if any. ctx.HandlerIndex(0) ctx.Do(ch.Handlers) } func (ch *ErrorCodeHandler) updateHandlers(handlers context.Handlers) { ch.mu.Lock() ch.Handlers = handlers ch.mu.Unlock() } // ErrorCodeHandlers contains the http error code handlers. // User of this struct can register, get // a status code handler based on a status code or // fire based on a receiver context. type ErrorCodeHandlers struct { handlers []*ErrorCodeHandler } func defaultErrorCodeHandlers() *ErrorCodeHandlers { chs := new(ErrorCodeHandlers) // register some common error handlers. // Note that they can be registered on-fly but // we don't want to reduce the performance even // on the first failed request. for _, statusCode := range []int{ http.StatusNotFound, http.StatusMethodNotAllowed, http.StatusInternalServerError, } { chs.Register(statusCode, statusText(statusCode)) } return chs } func statusText(statusCode int) context.Handler { return func(ctx context.Context) { ctx.WriteString(http.StatusText(statusCode)) } } // Get returns an http error handler based on the "statusCode". // If not found it returns nil. func (s *ErrorCodeHandlers) Get(statusCode int) *ErrorCodeHandler { for i, n := 0, len(s.handlers); i < n; i++ { if h := s.handlers[i]; h.StatusCode == statusCode { return h } } return nil } // Register registers an error http status code // based on the "statusCode" < 200 || >= 400 (`context.StatusCodeNotSuccessful`). // The handler is being wrapepd by a generic // handler which will try to reset // the body if recorder was enabled // and/or disable the gzip if gzip response recorder // was active. func (s *ErrorCodeHandlers) Register(statusCode int, handlers ...context.Handler) *ErrorCodeHandler { if statusCodeSuccessful(statusCode) { return nil } h := s.Get(statusCode) if h == nil { // create new and add it ch := &ErrorCodeHandler{ StatusCode: statusCode, Handlers: handlers, } s.handlers = append(s.handlers, ch) return ch } // otherwise update the handlers h.updateHandlers(handlers) return h } // Fire executes an error http status code handler // based on the context's status code. // // If a handler is not already registered, // then it creates & registers a new trivial handler on the-fly. func (s *ErrorCodeHandlers) Fire(ctx context.Context) { statusCode := ctx.GetStatusCode() if statusCodeSuccessful(statusCode) { return } ch := s.Get(statusCode) if ch == nil { ch = s.Register(statusCode, statusText(statusCode)) } ch.Fire(ctx) }