diff --git a/tests/e2e/distributed/cluster/admin.go b/tests/e2e/distributed/cluster/admin.go new file mode 100644 index 000000000..b1371f532 --- /dev/null +++ b/tests/e2e/distributed/cluster/admin.go @@ -0,0 +1,168 @@ +package cluster + +import ( + "bytes" + "encoding/json" + "fmt" + "io" + "net/http" + "net/http/cookiejar" + "net/url" + "time" + + "github.com/mudler/LocalAI/pkg/httpclient" +) + +const ( + // adminPassword must satisfy core/http/auth's policy (>= 12 chars and a + // zxcvbn score of 3 against hints that include "admin" and "localai"). + // This one scores 3 today; acknowledgeWeakPassword is sent alongside it so + // a future tightening of the policy cannot silently break every failover + // spec at setup time. + adminPassword = "e2e-admin-password" + // sessionCookieName mirrors the unexported constant in core/http/auth. + // The register handler returns 201 both for "user created, here is your + // session" and for "this email already exists" (a deliberate account + // enumeration defence), so the status code alone cannot tell the two + // apart: the presence of this cookie is the only reliable signal. + sessionCookieName = "session" + // authRequestTimeout bounds one register/login round trip. + authRequestTimeout = 30 * time.Second + // bodyExcerptLimit caps how much of an error response is quoted back. + bodyExcerptLimit = 512 +) + +// ForTestingEmpty returns a Cluster with no processes. It exists so the package's +// own argument-validation specs do not need to start anything. +func ForTestingEmpty() *Cluster { + return &Cluster{} +} + +// AdminSession registers the admin user on frontend i and returns a client +// carrying the resulting session cookie. The email matches LOCALAI_ADMIN_EMAIL, +// which core/http/auth exempts from the approval gate and assigns the admin +// role, so registration alone yields an active admin session. +// +// Call this ONCE per cluster and share the client. Two reasons: the auth +// endpoints are rate limited to 5 requests per minute per client IP, and every +// e2e request arrives from 127.0.0.1; and the returned client is already good +// for every frontend, because sessions live in the shared Postgres auth DB and +// Go's cookie jar keys cookies by host without the port. +func (c *Cluster) AdminSession(i int) (*http.Client, error) { + base, err := c.frontendBaseURL(i) + if err != nil { + return nil, err + } + + jar, err := cookiejar.New(nil) + if err != nil { + return nil, fmt.Errorf("creating cookie jar: %w", err) + } + // httpclient hardens the transport and refuses redirects; the jar is the one + // thing it does not configure, and a session cookie is the whole point here. + client := httpclient.NewWithTimeout(authRequestTimeout) + client.Jar = jar + + credentials := map[string]any{ + "email": c.opts.AdminEmail, + "password": adminPassword, + } + registration := map[string]any{ + "email": c.opts.AdminEmail, + "password": adminPassword, + "name": "E2E Admin", + "acknowledge_weak_password": true, + } + + registerStatus, registerBody, err := postJSON(client, base+"/api/auth/register", registration) + if err != nil { + return nil, fmt.Errorf("registering admin on frontend %d: %w", i, err) + } + if hasSessionCookie(jar, base) { + return client, nil + } + + // No cookie means the user already existed (a repeat call against the same + // Postgres), or registration was rejected. Log in; on failure the + // registration response is the diagnosis, so carry it into the error. + loginStatus, loginBody, err := postJSON(client, base+"/api/auth/login", credentials) + if err != nil { + return nil, fmt.Errorf("logging in admin on frontend %d: %w", i, err) + } + if loginStatus != http.StatusOK { + return nil, fmt.Errorf( + "admin login on frontend %d returned %d (%s); registration had returned %d (%s)", + i, loginStatus, loginBody, registerStatus, registerBody) + } + if !hasSessionCookie(jar, base) { + return nil, fmt.Errorf("admin login on frontend %d returned 200 but set no %q cookie: %s", i, sessionCookieName, loginBody) + } + return client, nil +} + +// GetJSON performs an authenticated GET against a frontend and decodes the body. +func (c *Cluster) GetJSON(client *http.Client, frontend int, path string, out any) error { + base, err := c.frontendBaseURL(frontend) + if err != nil { + return err + } + resp, err := client.Get(base + path) + if err != nil { + return fmt.Errorf("GET %s on frontend %d: %w", path, frontend, err) + } + defer func() { _ = resp.Body.Close() }() + if resp.StatusCode != http.StatusOK { + return fmt.Errorf("GET %s on frontend %d returned %d: %s", path, frontend, resp.StatusCode, excerpt(resp.Body)) + } + if err := json.NewDecoder(resp.Body).Decode(out); err != nil { + return fmt.Errorf("decoding %s from frontend %d: %w", path, frontend, err) + } + return nil +} + +// frontendBaseURL validates the index before FrontendURL indexes the slice: a +// bare index panic in a helper every failover spec calls is far harder to read +// than a named error. +func (c *Cluster) frontendBaseURL(i int) (string, error) { + if i < 0 || i >= len(c.frontends) { + return "", fmt.Errorf("frontend %d out of range (cluster has %d)", i, len(c.frontends)) + } + return c.FrontendURL(i), nil +} + +// postJSON sends body as JSON and returns the status plus an excerpt of the +// response, closing the body in every path. +func postJSON(client *http.Client, endpoint string, body any) (int, string, error) { + encoded, err := json.Marshal(body) + if err != nil { + return 0, "", fmt.Errorf("marshalling request body: %w", err) + } + resp, err := client.Post(endpoint, "application/json", bytes.NewReader(encoded)) + if err != nil { + return 0, "", err + } + defer func() { _ = resp.Body.Close() }() + return resp.StatusCode, excerpt(resp.Body), nil +} + +// hasSessionCookie reports whether the jar holds a usable session for base. +func hasSessionCookie(jar *cookiejar.Jar, base string) bool { + u, err := url.Parse(base) + if err != nil { + return false + } + for _, cookie := range jar.Cookies(u) { + if cookie.Name == sessionCookieName && cookie.Value != "" { + return true + } + } + return false +} + +func excerpt(r io.Reader) string { + data, err := io.ReadAll(io.LimitReader(r, bodyExcerptLimit)) + if err != nil { + return fmt.Sprintf("", err) + } + return string(bytes.TrimSpace(data)) +} diff --git a/tests/e2e/distributed/cluster/cluster_test.go b/tests/e2e/distributed/cluster/cluster_test.go index 4d48f1ecb..807c275c1 100644 --- a/tests/e2e/distributed/cluster/cluster_test.go +++ b/tests/e2e/distributed/cluster/cluster_test.go @@ -3,7 +3,9 @@ package cluster_test import ( "os" "path/filepath" + "time" + "github.com/mudler/LocalAI/pkg/httpclient" "github.com/mudler/LocalAI/tests/e2e/distributed/cluster" . "github.com/onsi/ginkgo/v2" @@ -35,3 +37,23 @@ var _ = Describe("Cluster options", Label("Distributed"), func() { Expect(err.Error()).To(ContainSubstring("at least one frontend")) }) }) + +// The HTTP flow inside AdminSession and GetJSON cannot run here: it needs a +// built local-ai plus real Postgres and NATS, which arrive with the failover +// suites. These specs cover the argument validation that would otherwise panic +// on an out-of-range slice index inside a helper every later spec calls. +var _ = Describe("Admin session", Label("Distributed"), func() { + It("reports a clear error when the frontend index is out of range", func() { + c := cluster.ForTestingEmpty() + _, err := c.AdminSession(3) + Expect(err).To(HaveOccurred()) + Expect(err.Error()).To(ContainSubstring("frontend 3")) + }) + + It("reports a clear error when GetJSON names a frontend that does not exist", func() { + c := cluster.ForTestingEmpty() + err := c.GetJSON(httpclient.NewWithTimeout(time.Second), 1, "/api/nodes", &struct{}{}) + Expect(err).To(HaveOccurred()) + Expect(err.Error()).To(ContainSubstring("frontend 1")) + }) +})