package snap import ( "io" "net/http" "github.com/gorilla/mux" ) type ServiceStatus struct { Message string } type ServiceResponse struct { Code int Message []byte } func (s *ServiceResponse) WriteResponse(w http.ResponseWriter) { w.WriteHeader(s.Code) w.Write(s.Message) } func (s *Server) serviceWrapper(handle func(c *ServiceRequest) *ServiceResponse) http.HandlerFunc { return func(w http.ResponseWriter, r *http.Request) { c := &ServiceRequest{ r: r, } if s.auth != nil { if rec, code := s.auth.DoAuth(w, r); code == http.StatusOK { c.auth = rec } } resp := handle(c) // discard the rest of the body content io.Copy(io.Discard, r.Body) defer r.Body.Close() resp.WriteResponse(w) } } func (s *Server) serviceWrapperAuthenticated(handle func(c *ServiceRequest) *ServiceResponse) http.HandlerFunc { return func(w http.ResponseWriter, r *http.Request) { defer func() { io.Copy(io.Discard, r.Body) r.Body.Close() }() if s.auth == nil { resp := ServiceResponse{ Code: http.StatusForbidden, Message: []byte("Not Authenticated"), } resp.WriteResponse(w) return } rec, code := s.auth.DoAuth(w, r) switch code { case http.StatusOK: c := &ServiceRequest{ r: r, auth: rec, } resp := handle(c) resp.WriteResponse(w) return case http.StatusUnauthorized: resp := ServiceResponse{ Code: http.StatusForbidden, Message: []byte("Not Authenticated"), } resp.WriteResponse(w) return default: resp := ServiceResponse{ Code: code, Message: []byte(http.StatusText(code)), } resp.WriteResponse(w) return } } } func (s *Server) HandleServiceFunc(path string, f func(c *ServiceRequest) *ServiceResponse) *mux.Route { return s.router.HandleFunc(path, s.serviceWrapper(f)) } func (s *Server) HandleServiceFuncWithAuth(path string, f func(c *ServiceRequest) *ServiceResponse) *mux.Route { return s.router.HandleFunc(path, s.serviceWrapperAuthenticated(f)) }