nadir-agent/internal/rbac/middleware_test.go

package rbac

import (
	"context"
	"net/http"
	"path/filepath"
	"testing"

	"nadir/internal/auditlog"
	"nadir/internal/auth"

	"github.com/danielgtaylor/huma/v2"
	"github.com/danielgtaylor/huma/v2/adapters/humago"
	"github.com/danielgtaylor/huma/v2/humatest"
)

func TestRbacMiddleware(t *testing.T) {
	tempDir := t.TempDir()
	auditStore, err := auditlog.New(filepath.Join(tempDir, "audit.db"))
	if err != nil {
		t.Fatal(err)
	}
	defer auditStore.Close()

	sessions, err := auth.NewSessionStore(filepath.Join(tempDir, "sessions.db"))
	if err != nil {
		t.Fatal(err)
	}

	tokenStore, err := auth.NewTokenStore(filepath.Join(tempDir, "tokens.db"))
	if err != nil {
		t.Fatal(err)
	}
	tokenAuth := auth.NewTokenAuth(tokenStore)

	r := New()
	r.DefineRole(Role{
		Name: "test-role",
		ModuleGrants: map[string][]Permission{
			"test-mod": {Read, Write},
		},
	})
	r.AssignRole("alice", "test-role")
	// A machine credential is just another RBAC subject: the token name is
	// assigned a role exactly like a username.
	r.AssignRole("dash", "test-role")

	mux := http.NewServeMux()
	api := humatest.Wrap(t, humago.New(mux, huma.DefaultConfig("Test", "1.0.0")))

	api.UseMiddleware(RbacMiddleware(api, sessions, tokenAuth, r, auditStore))

	huma.Register(api, huma.Operation{
		OperationID: "public-get",
		Method:      "GET",
		Path:        "/public",
	}, func(ctx context.Context, _ *struct{}) (*struct{ Body string }, error) {
		return &struct{ Body string }{Body: "public"}, nil
	})

	huma.Register(api, huma.Operation{
		OperationID: "gated-get",
		Method:      "GET",
		Path:        "/gated-read",
		Metadata:    map[string]any{"module": "test-mod", "permission": "read"},
	}, func(ctx context.Context, _ *struct{}) (*struct{ Body string }, error) {
		return &struct{ Body string }{Body: "gated-read"}, nil
	})

	huma.Register(api, huma.Operation{
		OperationID: "gated-post",
		Method:      "POST",
		Path:        "/gated-write",
		Metadata:    map[string]any{"module": "test-mod", "permission": "write"},
	}, func(ctx context.Context, _ *struct{}) (*struct{ Body string }, error) {
		return &struct{ Body string }{Body: "gated-write"}, nil
	})

	// 1. Test public route
	resp := api.Get("/public")
	if resp.Code != http.StatusOK {
		t.Errorf("public GET: got status %d, want %d", resp.Code, http.StatusOK)
	}

	// 2. Test gated route without cookie -> 401 Unauthorized
	resp = api.Get("/gated-read")
	if resp.Code != http.StatusUnauthorized {
		t.Errorf("gated GET no cookie: got status %d, want %d", resp.Code, http.StatusUnauthorized)
	}

	// 3. Test gated route with invalid cookie -> 401 Unauthorized
	resp = api.Get("/gated-read", "Cookie: nadir_session_id=invalid")
	if resp.Code != http.StatusUnauthorized {
		t.Errorf("gated GET invalid cookie: got status %d, want %d", resp.Code, http.StatusUnauthorized)
	}

	// Create valid session
	token, err := sessions.Create("alice")
	if err != nil {
		t.Fatal(err)
	}

	// 4. Test gated route with valid cookie -> 200 OK
	resp = api.Get("/gated-read", "Cookie: nadir_session_id="+token)
	if resp.Code != http.StatusOK {
		t.Errorf("gated GET valid cookie: got status %d, want %d", resp.Code, http.StatusOK)
	}

	// 5. Test CSRF violation: POST with mismatched Origin header -> 403 Forbidden
	resp = api.Post("/gated-write", "Cookie: nadir_session_id="+token, "Origin: http://evil.com", "Host: example.com", struct{}{})
	if resp.Code != http.StatusForbidden {
		t.Errorf("CSRF mismatched Origin: got status %d, want %d", resp.Code, http.StatusForbidden)
	}

	// 6. Test CSRF success: POST with matching Origin header -> 200 OK
	resp = api.Post("/gated-write", "Cookie: nadir_session_id="+token, "Origin: http://example.com", "Host: example.com", struct{}{})
	if resp.Code != http.StatusOK {
		t.Errorf("CSRF matching Origin: got status %d, want %d", resp.Code, http.StatusOK)
	}

	// 7. Test gated route with unauthorized user -> 403 Forbidden
	tokenBob, err := sessions.Create("bob")
	if err != nil {
		t.Fatal(err)
	}
	resp = api.Get("/gated-read", "Cookie: nadir_session_id="+tokenBob)
	if resp.Code != http.StatusForbidden {
		t.Errorf("bob unauthorized GET: got status %d, want %d", resp.Code, http.StatusForbidden)
	}

	// 8. Bearer token for an assigned name -> 200 OK
	rawToken, err := tokenStore.Create("dash")
	if err != nil {
		t.Fatal(err)
	}
	resp = api.Get("/gated-read", "Authorization: Bearer "+rawToken)
	if resp.Code != http.StatusOK {
		t.Errorf("valid bearer GET: got status %d, want %d", resp.Code, http.StatusOK)
	}

	// 9. Bogus bearer token -> 401 Unauthorized
	resp = api.Get("/gated-read", "Authorization: Bearer nad_deadbeef")
	if resp.Code != http.StatusUnauthorized {
		t.Errorf("bogus bearer GET: got status %d, want %d", resp.Code, http.StatusUnauthorized)
	}

	// 10. Bearer token with no role assignment -> 403 Forbidden
	rawUnassigned, err := tokenStore.Create("orphan")
	if err != nil {
		t.Fatal(err)
	}
	resp = api.Get("/gated-read", "Authorization: Bearer "+rawUnassigned)
	if resp.Code != http.StatusForbidden {
		t.Errorf("unassigned bearer GET: got status %d, want %d", resp.Code, http.StatusForbidden)
	}
}