From 963f4020b726b6b5212025e39b649540de43ba0e Mon Sep 17 00:00:00 2001
From: Daniel Ponte <amigan@gmail.com>
Date: Tue, 25 Oct 2022 21:18:50 -0400
Subject: [PATCH] Improve routing

---
 pkg/auth/flow.go            | 70 +++++++++++++++++++++++++------------
 pkg/server/config/config.go |  7 ++--
 pkg/server/server.go        | 20 ++++++-----
 3 files changed, 62 insertions(+), 35 deletions(-)

diff --git a/pkg/auth/flow.go b/pkg/auth/flow.go
index 45db5cc..2f9c72b 100644
--- a/pkg/auth/flow.go
+++ b/pkg/auth/flow.go
@@ -69,6 +69,10 @@ func (fs FlowStore) register(f *Flow) {
 	fs[f.ID] = f
 }
 
+func (fs FlowStore) Remove(f *Flow) {
+	delete(fs, f.ID)
+}
+
 const cullAge = time.Minute * 30
 
 func (fs FlowStore) cull() {
@@ -142,6 +146,8 @@ func (f *Flow) progress(a *Authenticator, c echo.Context) error {
 		err = a.Check(f, rm)
 		switch err {
 		case nil:
+			// TODO: setup the session. delete the flow.
+			a.Flows.Remove(f)
 			return c.String(http.StatusOK, "login success!")
 		case ErrInvalidHandler:
 			return c.String(http.StatusNotFound, err.Error())
@@ -170,35 +176,53 @@ func (a *Authenticator) LoginFlowDeleteHandler(c echo.Context) error {
 	return c.String(http.StatusOK, "deleted")
 }
 
-func (a *Authenticator) LoginFlowHandler(c echo.Context) error {
+func setJSON(c echo.Context) {
 	if c.Request().Method == http.MethodPost && strings.HasPrefix(c.Request().Header.Get(echo.HeaderContentType), "text/plain") {
 		// hack around the content-type, Context.JSON refuses to work otherwise
 		c.Request().Header.Set(echo.HeaderContentType, echo.MIMEApplicationJSONCharsetUTF8)
 	}
+}
+
+func (a *Authenticator) BeginLoginFlowHandler(c echo.Context) error {
+	setJSON(c)
+
+	var flowReq FlowRequest
+	err := c.Bind(&flowReq)
+	if err != nil {
+		return c.String(http.StatusBadRequest, err.Error())
+	}
+
+	resp := a.NewFlow(&flowReq)
+
+	if resp == nil {
+		return c.String(http.StatusBadRequest, "no such handler")
+	}
+
+	return c.JSON(http.StatusOK, resp)
+}
+
+func (a *Authenticator) LoginFlowHandler(c echo.Context) error {
+	setJSON(c)
 
 	flowID := c.Param("flow_id")
 
-	switch flowID {
-	case "": // new
-		var flowReq FlowRequest
-		err := c.Bind(&flowReq)
-		if err != nil {
-			return c.String(http.StatusBadRequest, err.Error())
-		}
-
-		resp := a.NewFlow(&flowReq)
-
-		if resp == nil {
-			return c.String(http.StatusBadRequest, "no such handler")
-		}
-
-		return c.JSON(http.StatusOK, resp)
-	default:
-		flow := a.Flows.Get(FlowID(flowID))
-		if flow == nil {
-			return c.String(http.StatusNotFound, "no such flow")
-		}
-
-		return flow.progress(a, c)
+	flow := a.Flows.Get(FlowID(flowID))
+	if flow == nil {
+		return c.String(http.StatusNotFound, "no such flow")
 	}
+
+	return flow.progress(a, c)
+}
+
+func (a *Authenticator) InstallRoutes(e *echo.Echo) {
+	authG := e.Group("/auth")
+	authG.GET("/authorize", a.AuthorizeHandler)
+	authG.GET("/providers", a.ProvidersHandler)
+
+	authG.POST("/login_flow", a.BeginLoginFlowHandler)
+
+	loginFlow := authG.Group("/login_flow") // TODO: add IP address affinity middleware
+	loginFlow.POST("/:flow_id", a.LoginFlowHandler)
+	loginFlow.DELETE("/:flow_id", a.LoginFlowDeleteHandler)
+
 }
diff --git a/pkg/server/config/config.go b/pkg/server/config/config.go
index 170a0ca..bc86567 100644
--- a/pkg/server/config/config.go
+++ b/pkg/server/config/config.go
@@ -34,7 +34,8 @@ func (ips *IPSource) UnmarshalYAML(val *yaml.Node) error {
 }
 
 type Config struct {
-	Bind           string   `yaml:"bind"`
-	IPSource       IPSource `yaml:"ip_source"`
-	TrustedProxies []string `yaml:"trusted_proxies"`
+	Bind             string   `yaml:"bind"`
+	IPSource         IPSource `yaml:"ip_source"`
+	LogRequestErrors bool     `yaml:"log_req_errors"`
+	TrustedProxies   []string `yaml:"trusted_proxies"`
 }
diff --git a/pkg/server/server.go b/pkg/server/server.go
index 7b037df..bca6892 100644
--- a/pkg/server/server.go
+++ b/pkg/server/server.go
@@ -2,7 +2,6 @@ package server
 
 import (
 	"context"
-	"io/fs"
 	"net/http"
 	"os"
 	"os/signal"
@@ -24,19 +23,14 @@ type Server struct {
 	*blas.Blas
 	*echo.Echo
 	auth.Authenticator
-	rootFS fs.FS
-	wg     sync.WaitGroup
+	wg sync.WaitGroup
 }
 
 func (s *Server) installRoutes() {
 	s.GET("/", echo.WrapHandler(frontend.FSHandler))
 	s.GET("/api/websocket", s.wsHandler)
-	s.GET("/auth/authorize", s.AuthorizeHandler)
-	s.GET("/auth/providers", s.ProvidersHandler)
 
-	s.POST("/auth/login_flow", s.LoginFlowHandler)
-	s.POST("/auth/login_flow/:flow_id", s.LoginFlowHandler)
-	s.DELETE("/auth/login_flow/:flow_id", s.LoginFlowDeleteHandler)
+	s.Authenticator.InstallRoutes(s.Echo)
 }
 
 func New(cfg *config.Config) (s *Server, err error) {
@@ -56,7 +50,15 @@ func New(cfg *config.Config) (s *Server, err error) {
 
 	s.Echo.Debug = true
 	s.Echo.HideBanner = true
-	s.Echo.Logger = lecho.From(log.Logger)
+	logger := lecho.From(log.Logger)
+	s.Echo.Logger = logger
+
+	if cfg.Server.LogRequestErrors {
+		s.Echo.Use(lecho.Middleware(lecho.Config{
+			Logger: logger,
+		}))
+	}
+
 	s.Echo.HidePort = true
 
 	ipext := map[conf.IPSource]echo.IPExtractor{