Warp/server/internal/clientapi/uia.go

package clientapi

import (
	"crypto/rand"
	"encoding/base64"
	"fmt"
	"net/http"
	"slices"
	"sync"
	"time"

	"github.com/dariuszjaniak/warp/server/internal/httputil"
)

// User-Interactive Authentication stage types. The spec's full list is long;
// these are the ones Warp implements. CAPTCHA, email and terms stages belong
// to deployments with public sign-up and arrive with the rest of that story.
const (
	stageDummy             = "m.login.dummy"
	stagePassword          = "m.login.password"
	stageRegistrationToken = "m.login.registration_token"

	// loginAppService is not an interactive stage but a login and
	// registration type: an application service's own token is the
	// credential.
	loginAppService = "m.login.application_service"
)

// uiaFlow is one route through the stages. A client may satisfy any flow.
type uiaFlow struct {
	Stages []string `json:"stages"`
}

// uiaChallenge is the 401 body of an incomplete User-Interactive
// Authentication: the flows on offer, the session to quote them against, and
// what has been done so far.
//
// When a stage has just been failed it also carries an error envelope, which
// is why errcode and error are part of this type rather than sent as a
// separate Matrix error: the spec puts both in one 401 body so the client
// can show the failure and retry the stage in one step.
type uiaChallenge struct {
	Flows     []uiaFlow      `json:"flows"`
	Params    map[string]any `json:"params"`
	Session   string         `json:"session"`
	Completed []string       `json:"completed,omitempty"`

	ErrCode string `json:"errcode,omitempty"`
	Error   string `json:"error,omitempty"`
}

// uiaSessionTTL is how long an unfinished authentication may sit idle.
// Clients complete a flow in seconds; the window only has to be long enough
// for a person to read a prompt and type a password.
const uiaSessionTTL = 15 * time.Minute

// uiaSession tracks one client's progress through a flow set.
type uiaSession struct {
	id        string
	completed []string
	expiresAt time.Time

	// registrationToken is the token a registration's
	// m.login.registration_token stage was passed with. It is spent when
	// the account is created, so it has to be remembered until then.
	registrationToken string
}

// uiaStore holds the in-flight authentication sessions.
//
// Sessions live in memory. They are worthless after a few minutes, they are
// only ever read by the process that wrote them (Warp is one process by
// design), and a restart costs a client exactly one extra round trip —
// whereas a table would need its own schema, its own sweeper and a write on
// every stage.
type uiaStore struct {
	ttl time.Duration
	now func() time.Time

	mu       sync.Mutex
	sessions map[string]*uiaSession
	sweptAt  time.Time
}

func newUIAStore(ttl time.Duration, now func() time.Time) *uiaStore {
	if now == nil {
		now = time.Now
	}
	return &uiaStore{ttl: ttl, now: now, sessions: map[string]*uiaSession{}}
}

// create starts a session.
func (s *uiaStore) create() (*uiaSession, error) {
	raw := make([]byte, 16)
	if _, err := rand.Read(raw); err != nil {
		return nil, fmt.Errorf("generate auth session: %w", err)
	}
	session := &uiaSession{
		id:        base64.RawURLEncoding.EncodeToString(raw),
		expiresAt: s.now().Add(s.ttl),
	}

	s.mu.Lock()
	defer s.mu.Unlock()
	s.sweep()
	s.sessions[session.id] = session
	return session, nil
}

// get returns a live session, or nil if it is unknown or expired.
func (s *uiaStore) get(id string) *uiaSession {
	if id == "" {
		return nil
	}
	s.mu.Lock()
	defer s.mu.Unlock()

	session, ok := s.sessions[id]
	if !ok {
		return nil
	}
	if s.now().After(session.expiresAt) {
		delete(s.sessions, id)
		return nil
	}
	// Each copy handed out is independent, so a caller cannot mutate the
	// stored progress by accident.
	return &uiaSession{
		id:                session.id,
		completed:         slices.Clone(session.completed),
		expiresAt:         session.expiresAt,
		registrationToken: session.registrationToken,
	}
}

// setRegistrationToken remembers the token a session's registration token
// stage was passed with.
func (s *uiaStore) setRegistrationToken(id, token string) {
	s.mu.Lock()
	defer s.mu.Unlock()
	if session, ok := s.sessions[id]; ok {
		session.registrationToken = token
	}
}

// complete records a finished stage and returns the session's new progress.
func (s *uiaStore) complete(id, stage string) []string {
	s.mu.Lock()
	defer s.mu.Unlock()

	session, ok := s.sessions[id]
	if !ok {
		return nil
	}
	if !slices.Contains(session.completed, stage) {
		session.completed = append(session.completed, stage)
	}
	// Finishing a stage is progress, so the window starts again: a flow
	// with several stages must not time out halfway through.
	session.expiresAt = s.now().Add(s.ttl)
	return slices.Clone(session.completed)
}

// remove forgets a session, which happens once its flow has been acted on.
// Leaving it behind would let a replayed request create a second account or
// delete a second device.
func (s *uiaStore) remove(id string) {
	s.mu.Lock()
	defer s.mu.Unlock()
	delete(s.sessions, id)
}

// sweep drops expired sessions. It runs under the caller's lock, at most
// once per TTL.
func (s *uiaStore) sweep() {
	now := s.now()
	if now.Sub(s.sweptAt) < s.ttl {
		return
	}
	s.sweptAt = now
	for id, session := range s.sessions {
		if now.After(session.expiresAt) {
			delete(s.sessions, id)
		}
	}
}

// uiaRequest is the `auth` dictionary a client sends. The stage-specific
// fields are flattened into it rather than nested, which is how the spec
// defines each stage's payload.
type uiaRequest struct {
	Type    string `json:"type"`
	Session string `json:"session"`

	// Token is the m.login.registration_token payload.
	Token string `json:"token"`

	// Identifier and Password are the m.login.password payload.
	Identifier loginIdentifier `json:"identifier"`
	Password   string          `json:"password"`

	// User is the deprecated top-level form of Identifier, still sent by
	// older clients.
	User string `json:"user"`
}

// flowSatisfied reports whether completed covers every stage of any flow.
func flowSatisfied(flows []uiaFlow, completed []string) bool {
	for _, flow := range flows {
		done := true
		for _, stage := range flow.Stages {
			if !slices.Contains(completed, stage) {
				done = false
				break
			}
		}
		if done {
			return true
		}
	}
	return false
}

// challenge writes a 401 asking the client to complete a stage.
//
// cause is optional: when a stage was attempted and failed, the client gets
// the failure and the flows in one body, which is what lets it re-prompt
// without starting over.
func (c *ClientAPI) challenge(w http.ResponseWriter, flows []uiaFlow, session *uiaSession, cause *httputil.MatrixError) {
	body := uiaChallenge{
		Flows: flows,
		// Always an object, never null: clients and Complement both expect
		// `params` to be present even when no stage needs parameters.
		Params:    map[string]any{},
		Session:   session.id,
		Completed: session.completed,
	}
	if cause != nil {
		body.ErrCode, body.Error = cause.Code, cause.Message
	}
	httputil.WriteJSON(w, http.StatusUnauthorized, body)
}

// stageHandler checks one stage of a flow.
//
// The two return values are different kinds of "no": a *MatrixError is a
// stage the client may retry, and is sent back inside the 401 alongside the
// flows; a plain error ends the request. Wrong password is the first kind,
// authenticating as somebody else is the second.
type stageHandler func(stage string, auth *uiaRequest) (*httputil.MatrixError, error)

// beginOrAdvanceUIA drives one request through a flow set.
//
// It returns done=true when the client has satisfied a flow and the caller
// may go ahead; otherwise it has already written the 401 challenge. The
// returned session ID is the one to forget once the operation succeeds.
func (c *ClientAPI) beginOrAdvanceUIA(
	w http.ResponseWriter,
	flows []uiaFlow,
	auth *uiaRequest,
	handleStage stageHandler,
) (sessionID string, done bool, err error) {
	// No auth at all, or an auth without the session the server handed out,
	// starts a new attempt. The spec requires the session to be quoted back,
	// so a request that drops it is not treated as a continuation.
	if auth == nil || auth.Type == "" || auth.Session == "" {
		session, createErr := c.uia.create()
		if createErr != nil {
			return "", false, createErr
		}
		c.challenge(w, flows, session, nil)
		return "", false, nil
	}

	session := c.uia.get(auth.Session)
	if session == nil {
		// An unknown or expired session: start a fresh one rather than
		// failing, so a client that was slow can simply carry on.
		session, createErr := c.uia.create()
		if createErr != nil {
			return "", false, createErr
		}
		c.challenge(w, flows, session,
			httputil.NewError(http.StatusUnauthorized, httputil.ErrCodeUnrecognized,
				"The authentication session has expired; start again"))
		return "", false, nil
	}

	if !stageOffered(flows, auth.Type) {
		return "", false, httputil.NewError(http.StatusBadRequest, httputil.ErrCodeUnrecognized,
			"Unknown authentication type %q", auth.Type)
	}

	retryable, fatal := handleStage(auth.Type, auth)
	if fatal != nil {
		return "", false, fatal
	}
	if retryable != nil {
		c.challenge(w, flows, session, retryable)
		return "", false, nil
	}

	session.completed = c.uia.complete(session.id, auth.Type)
	if !flowSatisfied(flows, session.completed) {
		c.challenge(w, flows, session, nil)
		return "", false, nil
	}
	return session.id, true, nil
}

// stageOffered reports whether stage appears in any of the flows, so that a
// client cannot complete a stage this operation never asked for.
func stageOffered(flows []uiaFlow, stage string) bool {
	for _, flow := range flows {
		if slices.Contains(flow.Stages, stage) {
			return true
		}
	}
	return false
}