llmgw/internal/web/oidc_test.go

package web

import (
	"context"
	"crypto/rand"
	"crypto/rsa"
	"encoding/json"
	"io"
	"log/slog"
	"net/http"
	"net/http/httptest"
	"net/url"
	"strings"
	"testing"
	"time"

	jose "github.com/go-jose/go-jose/v4"
	"github.com/go-jose/go-jose/v4/jwt"

	"gitea.mynodes.xyz/ilker/llmgw/internal/config"
	"gitea.mynodes.xyz/ilker/llmgw/internal/store"
)

// fakeIdP is an OpenID provider that issues an ID token with whatever
// claims the test sets.
type fakeIdP struct {
	srv    *httptest.Server
	key    *rsa.PrivateKey
	claims map[string]any
	nonce  string
}

func newFakeIdP(t *testing.T) *fakeIdP {
	t.Helper()
	k, err := rsa.GenerateKey(rand.Reader, 2048)
	if err != nil {
		t.Fatal(err)
	}
	f := &fakeIdP{key: k}
	mux := http.NewServeMux()
	f.srv = httptest.NewServer(mux)
	t.Cleanup(f.srv.Close)
	mux.HandleFunc("/.well-known/openid-configuration", func(w http.ResponseWriter, _ *http.Request) {
		_ = json.NewEncoder(w).Encode(map[string]any{"issuer": f.srv.URL,
			"authorization_endpoint": f.srv.URL + "/authorize", "token_endpoint": f.srv.URL + "/token",
			"jwks_uri": f.srv.URL + "/jwks", "id_token_signing_alg_values_supported": []string{"RS256"}})
	})
	mux.HandleFunc("/jwks", func(w http.ResponseWriter, _ *http.Request) {
		_ = json.NewEncoder(w).Encode(jose.JSONWebKeySet{Keys: []jose.JSONWebKey{
			{Key: &k.PublicKey, KeyID: "k1", Algorithm: "RS256", Use: "sig"}}})
	})
	mux.HandleFunc("/token", func(w http.ResponseWriter, r *http.Request) {
		_ = r.ParseForm()
		if r.Form.Get("code_verifier") == "" {
			http.Error(w, "pkce", 400)
			return
		}
		signer, _ := jose.NewSigner(jose.SigningKey{Algorithm: jose.RS256, Key: k},
			(&jose.SignerOptions{}).WithHeader("kid", "k1"))
		c := map[string]any{"iss": f.srv.URL, "aud": "llmgw", "sub": "s1", "nonce": f.nonce,
			"iat": time.Now().Unix(), "exp": time.Now().Add(time.Minute).Unix()}
		for k, v := range f.claims {
			c[k] = v
		}
		tok, _ := jwt.Signed(signer).Claims(c).Serialize()
		w.Header().Set("Content-Type", "application/json")
		_ = json.NewEncoder(w).Encode(map[string]any{"access_token": "a", "token_type": "Bearer", "id_token": tok})
	})
	return f
}

type oidcStoreMem struct {
	users  map[string]store.User
	roles  map[string]string
	joined []string
}

func (m *oidcStoreMem) UserByEmail(_ context.Context, e string) (store.User, error) {
	if u, ok := m.users[e]; ok {
		return u, nil
	}
	return store.User{}, store.ErrNotFound
}
func (m *oidcStoreMem) CreateUser(_ context.Context, u store.User) error {
	m.users[u.Email] = u
	return nil
}
func (m *oidcStoreMem) SetStaffRole(_ context.Context, id, role string) error {
	m.roles[id] = role
	return nil
}
func (m *oidcStoreMem) AddMembership(_ context.Context, user string, scope store.Scope, id, _, _ string, _ time.Time) error {
	m.joined = append(m.joined, user+"@"+string(scope)+":"+id)
	return nil
}
func (m *oidcStoreMem) AllOrgs(context.Context) ([]store.Org, error) { return nil, nil }
func (m *oidcStoreMem) AllTeams(context.Context) ([]store.TeamPath, error) {
	return []store.TeamPath{{ID: "t1", Name: "ml", OrgName: "Acme"}}, nil
}
func (m *oidcStoreMem) AllProjects(context.Context) ([]store.ProjectPath, error) { return nil, nil }
func (m *oidcStoreMem) ClusterSecret(_ context.Context, name string) (string, error) {
	return "the store's " + name, nil
}

type sessionAccounts struct{ fakeAccounts }

func (sessionAccounts) OpenSession(context.Context, string) (string, string, error) {
	return "sess", "csrf", nil
}

// signIn goes through /login/oidc and back, as a browser would.
func signIn(t *testing.T, mux http.Handler, idp *fakeIdP, tamper bool) *httptest.ResponseRecorder {
	t.Helper()
	return signInAcross(t, mux, mux, idp, tamper)
}

// signInAcross begins a sign-in on one gateway and comes back from the
// provider to another, as a balancer in front of two may send it.
func signInAcross(t *testing.T, mux, back http.Handler, idp *fakeIdP, tamper bool) *httptest.ResponseRecorder {
	t.Helper()
	start := httptest.NewRecorder()
	mux.ServeHTTP(start, httptest.NewRequest("GET", "/login/oidc", nil))
	if start.Code != http.StatusFound {
		t.Fatalf("start: %d %s", start.Code, start.Body.String())
	}
	to, _ := url.Parse(start.Header().Get("Location"))
	if !strings.HasPrefix(to.String(), idp.srv.URL+"/authorize") || to.Query().Get("code_challenge") == "" {
		t.Fatalf("redirected to %s", to)
	}
	idp.nonce = to.Query().Get("nonce")
	state := to.Query().Get("state")
	if tamper {
		state = "someone-elses"
	}
	cb := httptest.NewRequest("GET", "/login/oidc/callback?code=c1&state="+state, nil)
	for _, c := range start.Result().Cookies() {
		cb.AddCookie(c)
	}
	rec := httptest.NewRecorder()
	back.ServeHTTP(rec, cb)
	return rec
}

func TestSigningInThroughAProvider(t *testing.T) {
	idp := newFakeIdP(t)
	st := &oidcStoreMem{users: map[string]store.User{}, roles: map[string]string{}}
	p := NewAccountPages(sessionAccounts{}, nil, nil, "v0", slog.New(slog.NewTextHandler(io.Discard, nil)))
	mux := http.NewServeMux()
	if err := p.RegisterOIDC(mux, config.OIDC{Issuer: idp.srv.URL, ClientID: "llmgw", CreateAccounts: true,
		Roles:       map[string][]string{"operator": {"ops"}, "admin": {"root"}},
		Memberships: map[string]string{"ml": "team:Acme/ml"}}, st); err != nil {
		t.Fatal(err)
	}

	idp.claims = map[string]any{"email": "Ann@Example.com", "email_verified": true, "name": "Ann", "groups": []string{"ops", "ml"}}
	rec := signIn(t, mux, idp, false)
	if rec.Code != http.StatusSeeOther || rec.Header().Get("Location") != "/profile" {
		t.Fatalf("callback: %d %s", rec.Code, rec.Body.String())
	}
	u, ok := st.users["ann@example.com"]
	if !ok || u.State != store.UserActive || st.roles[u.ID] != "operator" || len(st.joined) != 1 || !strings.HasSuffix(st.joined[0], "team:t1") {
		t.Errorf("account %+v roles %v joined %v", u, st.roles, st.joined)
	}
	if !strings.Contains(rec.Header().Get("Set-Cookie")+strings.Join(rec.Header().Values("Set-Cookie"), ";"), "sess") {
		t.Error("no session cookie")
	}

	// Refused: a state that is not this browser's, an unverified address,
	// and nobody the gateway knows when it does not make accounts.
	if rec := signIn(t, mux, idp, true); rec.Code != http.StatusForbidden {
		t.Errorf("a tampered state: %d", rec.Code)
	}
	idp.claims = map[string]any{"email": "bob@example.com", "email_verified": false}
	if rec := signIn(t, mux, idp, false); rec.Code != http.StatusForbidden || !strings.Contains(rec.Body.String(), "not verified") {
		t.Errorf("unverified: %d", rec.Code)
	}
	p2 := NewAccountPages(sessionAccounts{}, nil, nil, "v0", slog.New(slog.NewTextHandler(io.Discard, nil)))
	mux2 := http.NewServeMux()
	_ = p2.RegisterOIDC(mux2, config.OIDC{Issuer: idp.srv.URL, ClientID: "llmgw"}, st)
	idp.claims = map[string]any{"email": "carol@example.com"}
	if rec := signIn(t, mux2, idp, false); rec.Code != http.StatusForbidden || !strings.Contains(rec.Body.String(), "no account for carol@example.com") {
		t.Errorf("unknown without create_accounts: %d %s", rec.Code, rec.Body.String())
	}
}

func TestASignInBegunOnOneControlPlaneEndsOnTheOther(t *testing.T) {
	// The cookie carrying a sign-in across the provider was signed with a
	// key each process made for itself, so with two control planes behind a
	// balancer a callback reaching the other one was refused as tampered.
	idp := newFakeIdP(t)
	st := &oidcStoreMem{users: map[string]store.User{}, roles: map[string]string{}}
	cfg := config.OIDC{Issuer: idp.srv.URL, ClientID: "llmgw", CreateAccounts: true}
	var muxes []*http.ServeMux
	for range 2 {
		p := NewAccountPages(sessionAccounts{}, nil, nil, "v0", slog.New(slog.NewTextHandler(io.Discard, nil)))
		mux := http.NewServeMux()
		if err := p.RegisterOIDC(mux, cfg, st); err != nil {
			t.Fatal(err)
		}
		muxes = append(muxes, mux)
	}
	idp.claims = map[string]any{"email": "dan@example.com", "email_verified": true}
	if rec := signInAcross(t, muxes[0], muxes[1], idp, false); rec.Code != http.StatusSeeOther {
		t.Fatalf("a callback on the other control plane: %d %s", rec.Code, rec.Body.String())
	}
	if rec := signInAcross(t, muxes[0], muxes[1], idp, true); rec.Code != http.StatusForbidden {
		t.Errorf("a tampered state on the other control plane: %d", rec.Code)
	}
}