Complete registration flow

1 files changed, 71 insertions, 4 deletions

diff --git a/app/controllers/register.go b/app/controllers/register.go
index 8698bda..7c1a0f3 100644
--- a/app/controllers/register.go
+++ b/app/controllers/register.go
@@ -1,7 +1,13 @@
 package controllers
 
 import (
+        "bytes"
+        "context"
+        "encoding/json"
+        "fmt"
+        "io"
         "net/http"
+        "time"
 
         "code.crute.us/mcrute/golib/echo/session"
         "code.crute.us/mcrute/ssh-proxy/app"
@@ -19,13 +25,45 @@ type RegisterController[T app.AppSession] struct {
         Webauthn     *webauthn.WebAuthn
 }
 
+func (a *RegisterController[T]) validateRequest(ctx context.Context, u *models.User, code string) (*models.AuthSession, error) {
+        if code == "" {
+                return nil, fmt.Errorf("Code not passed in request")
+        }
+
+        authSession, err := a.AuthSessions.GetByUserCode(ctx, code)
+        if err != nil {
+                return nil, fmt.Errorf("No auth session exists")
+        }
+
+        if time.Now().After(authSession.Expires) {
+                return nil, fmt.Errorf("Session is expired")
+        }
+
+        if !authSession.IsRegistration {
+                return nil, fmt.Errorf("Session is not an invitation to register")
+        }
+
+        if authSession.UserId != u.Username {
+                return nil, fmt.Errorf("Session not valid for this user")
+        }
+
+        return authSession, nil
+}
+
 func (a *RegisterController[T]) HandleStart(c echo.Context) error {
-        user, err := a.Users.Get(c.Request().Context(), c.Param("username"))
+        ctx := c.Request().Context()
+
+        user, err := a.Users.Get(ctx, c.Param("username"))
         if err != nil {
                 a.Logger.Errorf("Error getting user: %s", err)
                 return c.NoContent(http.StatusNotFound)
         }
 
+        if _, err := a.validateRequest(ctx, user, c.QueryParam("code")); err != nil {
+                a.Logger.Errorf("Error creating registration request: %s", err)
+                return c.NoContent(http.StatusNotFound)
+        }
+
         request, sessionData, err := a.Webauthn.BeginRegistration(user)
         if err != nil {
                 a.Logger.Errorf("Error creating webauthn request: %s", err)
@@ -41,13 +79,42 @@ func (a *RegisterController[T]) HandleStart(c echo.Context) error {
 }
 
 func (a *RegisterController[T]) HandleFinish(c echo.Context) error {
-        user, err := a.Users.Get(c.Request().Context(), c.Param("username"))
+        ctx := c.Request().Context()
+
+        body, err := io.ReadAll(c.Request().Body)
+        if err != nil {
+                a.Logger.Errorf("Error reading request body:", err)
+                return c.NoContent(http.StatusInternalServerError)
+        }
+
+        user, err := a.Users.Get(ctx, c.Param("username"))
         if err != nil {
                 a.Logger.Errorf("Error getting user: %s", err)
                 return c.NoContent(http.StatusNotFound)
         }
 
-        response, err := protocol.ParseCredentialCreationResponseBody(c.Request().Body)
+        var code struct {
+                Code string `json:"code"`
+        }
+        if err := json.Unmarshal(body, &code); err != nil {
+                a.Logger.Errorf("Error decoding json body")
+                return c.NoContent(http.StatusBadRequest)
+        }
+
+        authSession, err := a.validateRequest(ctx, user, code.Code)
+        if err != nil {
+                a.Logger.Errorf("Error finishing register request: %s", err)
+                return c.NoContent(http.StatusNotFound)
+        }
+
+        // Delete before anything else to avoid allowing double use of an auth
+        // session in case of other errors
+        if err := a.AuthSessions.Delete(ctx, authSession); err != nil {
+                a.Logger.Errorf("Error deleting auth session: %s", err)
+                return c.NoContent(http.StatusInternalServerError)
+        }
+
+        response, err := protocol.ParseCredentialCreationResponseBody(bytes.NewBuffer(body))
         if err != nil {
                 a.Logger.Errorf("Error parsing credential response: %s", err)
                 return c.NoContent(http.StatusBadRequest)
@@ -69,7 +136,7 @@ func (a *RegisterController[T]) HandleFinish(c echo.Context) error {
 
         user.Fido2Credentials = append(user.Fido2Credentials, *credential)
 
-        if err := a.Users.Upsert(c.Request().Context(), user); err != nil {
+        if err := a.Users.Upsert(ctx, user); err != nil {
                 a.Logger.Errorf("Error saving user: %s", err)
                 return c.NoContent(http.StatusInternalServerError)
         }