package api import ( "context" "fmt" "log" "net/http" "time" "git.freecumextremist.com/grumbulon/pomme/internal" "git.freecumextremist.com/grumbulon/pomme/internal/db" "github.com/go-chi/chi/v5" "github.com/go-chi/httplog" "github.com/go-chi/httprate" "github.com/go-chi/jwtauth/v5" "github.com/go-chi/render" "gorm.io/gorm" ) type key int const ( keyPrincipalContextID key = iota ) // setDBMiddleware is the http Handler func for the GORM middleware with context. func setDBMiddleware(next http.Handler) http.Handler { return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { var pommeDB *gorm.DB c, err := internal.ReadConfig() if err != nil { log.Printf("No config file defined: %v", err) } switch r.Header.Get("User-Agent") { case "pomme-api-test-slave": pommeDB = db.InitDb(c.TestDB) default: pommeDB = db.InitDb(c.DB) } timeoutContext, cancelContext := context.WithTimeout(context.Background(), time.Second) ctx := context.WithValue(r.Context(), keyPrincipalContextID, pommeDB.WithContext(timeoutContext)) defer cancelContext() next.ServeHTTP(w, r.WithContext(ctx)) }) } type GenericResponse[T map[string]any] struct { Response map[string]any `json:"response,omitempty"` } func APIError[T map[string]any](w http.ResponseWriter, r *http.Request, v map[string]any) { logger := httplog.NewLogger("Pomme", httplog.Options{ JSON: true, }) w.Header().Set("X-Content-Type-Options", "nosniff") w.Header().Set("Content-Type", "application/json; charset=utf-8") switch v["realm"] { case nil: w.Header().Add("API Error", v["message"].(string)) default: w.Header().Add("WWW-Authenticate", fmt.Sprintf(`realm="%s"`, v["realm"].(string))) w.Header().Add("API Error", v["message"].(string)) } w.WriteHeader(v["status"].(int)) render.JSON(w, r, v) switch v["error"] { case nil: logger.Info().Msg(fmt.Sprint(v["message"])) default: logger.Error().Msg(fmt.Sprint(v["error"])) } } // API subroute handler. func API() http.Handler { api := chi.NewRouter() // Protected routes api.Group(func(api chi.Router) { api.Use(httprate.Limit( 10, // requests 10*time.Second, // per duration httprate.WithKeyFuncs(httprate.KeyByIP, httprate.KeyByEndpoint), httprate.WithLimitHandler(func(w http.ResponseWriter, r *http.Request) { w.Header().Set("Content-Type", "application/json") w.WriteHeader(http.StatusTooManyRequests) resp := internal.Response{ Message: "API rate limit exceded", } render.JSON(w, r, resp) }), )) api.Use(jwtauth.Verifier(tokenAuth)) api.Use(jwtauth.Authenticator) api.With(setDBMiddleware).Post("/upload", ReceiveFile) api.With(setDBMiddleware).Post("/parse", ParseZoneFiles) }) // Open routes api.Group(func(api chi.Router) { api.Use(httprate.Limit( 5, // requests 5*time.Second, // per duration httprate.WithKeyFuncs(httprate.KeyByIP, httprate.KeyByEndpoint), httprate.WithLimitHandler(func(w http.ResponseWriter, r *http.Request) { w.Header().Set("Content-Type", "application/json") w.WriteHeader(http.StatusTooManyRequests) resp := internal.Response{ Message: "API rate limit exceded", } render.JSON(w, r, resp) }), )) api.Use(setDBMiddleware) api.With(setDBMiddleware).Post("/create", NewUser) api.With(setDBMiddleware).Post("/login", Login) api.Post("/logout", Logout) }) return api }