snap/service.go

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))
}