polycule-connect/polyculeconnect/controller/auth/authcallback.go

package auth

import (
	"fmt"
	"net/http"

	"git.faercol.me/faercol/polyculeconnect/polyculeconnect/helpers"
	"git.faercol.me/faercol/polyculeconnect/polyculeconnect/internal/model"
	"git.faercol.me/faercol/polyculeconnect/polyculeconnect/internal/storage"
	"github.com/google/uuid"
	"github.com/zitadel/oidc/v3/pkg/client/rp"
	"github.com/zitadel/oidc/v3/pkg/oidc"
	"go.uber.org/zap"
)

const AuthCallbackRoute = "/callback"

type AuthCallbackController struct {
	l  *zap.SugaredLogger
	st *storage.Storage
}

func NewAuthCallbackController(l *zap.SugaredLogger, st *storage.Storage) *AuthCallbackController {
	return &AuthCallbackController{
		l:  l,
		st: st,
	}
}

func (c *AuthCallbackController) HandleUserInfoCallback(w http.ResponseWriter, r *http.Request, tokens *oidc.Tokens[*oidc.IDTokenClaims], state string, rp rp.RelyingParty, info *oidc.UserInfo) {
	requestID, err := uuid.Parse(state)
	if err != nil {
		c.l.Errorf("Invalid state, should be a request UUID, but got %s: %s", state, err)
		helpers.HandleResponse(w, r, http.StatusInternalServerError, []byte("failed to perform authentication"), c.l)
		return
	}

	c.l.Infof("Successful login from %s", info.Email)
	user := model.User{
		Subject:       info.Subject,
		Name:          info.Name,
		FamilyName:    info.FamilyName,
		GivenName:     info.GivenName,
		Picture:       info.Picture,
		UpdatedAt:     info.UpdatedAt.AsTime(),
		Email:         info.Email,
		EmailVerified: bool(info.EmailVerified),
	}

	err = c.st.LocalStorage.AuthRequestStorage().ValidateAuthRequest(r.Context(), requestID, user.Subject)
	if err != nil {
		c.l.Errorf("Failed to validate auth request from storage: %s", err)
		helpers.HandleResponse(w, r, http.StatusInternalServerError, []byte("failed to perform authentication"), c.l)
		return
	}
	if err := c.st.LocalStorage.UserStorage().AddUser(r.Context(), &user); err != nil {
		c.l.Errorf("Failed to add related user to storageL %w", err)
		helpers.HandleResponse(w, r, http.StatusInternalServerError, []byte("failed to perform authentication"), c.l)
		return
	}

	http.Redirect(w, r, "/authorize/callback?id="+state, http.StatusFound)
}

type CallbackDispatchController struct {
	l                *zap.SugaredLogger
	st               *storage.Storage
	callbackHandlers map[uuid.UUID]http.Handler
}

func NewCallbackDispatchController(l *zap.SugaredLogger, st *storage.Storage, handlers map[uuid.UUID]http.Handler) *CallbackDispatchController {
	return &CallbackDispatchController{
		l:                l,
		st:               st,
		callbackHandlers: handlers,
	}
}

func (c *CallbackDispatchController) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	errMsg := r.URL.Query().Get("error")
	if errMsg != "" {
		errorDesc := r.URL.Query().Get("error_description")
		c.l.Errorf("Failed to perform authentication: %s (%s)", errMsg, errorDesc)
		helpers.HandleResponse(w, r, http.StatusInternalServerError, []byte("failed to perform authentication"), c.l)
		return
	}

	state := r.URL.Query().Get("state")
	requestID, err := uuid.Parse(state)
	if err != nil {
		c.l.Errorf("Invalid state, should be a request UUID, but got %s: %s", state, err)
		helpers.HandleResponse(w, r, http.StatusInternalServerError, []byte("failed to perform authentication"), c.l)
		return
	}

	req, err := c.st.LocalStorage.AuthRequestStorage().GetAuthRequestByID(r.Context(), requestID)
	if err != nil {
		c.l.Errorf("Failed to get auth request from DB: %s", err)
		helpers.HandleResponse(w, r, http.StatusBadRequest, []byte("unknown request id"), c.l)
		return
	}

	if !req.Consent {
		c.l.Debug("Redirecting to consent endpoint")
		http.Redirect(w, r, fmt.Sprintf("/approval?state=%s&code=%s", state, r.URL.Query().Get("code")), http.StatusSeeOther)
		return
	}

	callbackHandler, ok := c.callbackHandlers[req.BackendID]
	if !ok {
		c.l.Errorf("Backend %s does not exist for request %s", req.ID, req.BackendID)
		helpers.HandleResponse(w, r, http.StatusNotFound, []byte("unknown backend"), c.l)
		return
	}
	callbackHandler.ServeHTTP(w, r)
}