Harden session JWT revocation (#64)

* Harden session JWT revocation

* Confirm sessionjwt skew report is false positive

* test: fix session jwt claim parsing clock

---------

Co-authored-by: Cursor Agent <cursoragent@cursor.com>

Author: haasonsaas

db/migrations/0002_session_token_id.sql

@@ -0,0 +1,2 @@
+ALTER TABLE sessions
+    ADD COLUMN IF NOT EXISTS token_id TEXT NOT NULL DEFAULT '';

internal/app/service.go

@@ -131,6 +131,7 @@ func (s *Service) CreateSession(ctx context.Context, req *core.CreateSessionRequ
 		State:            core.SessionStateActive,
 		CreatedAt:        now,
 	}
+	session.TokenID = "tok_" + session.ID
 	if err := s.repo.SaveSession(ctx, session); err != nil {
 		return nil, fmt.Errorf("create session %q: save session: %w", session.ID, err)
 	}
@@ -446,6 +447,11 @@ func (s *Service) RevokeSession(ctx context.Context, req *core.RevokeSessionRequ
 		return fmt.Errorf("revoke session %q: save session: %w", session.ID, err)
 	}
 	s.metrics.recordSessionTransition(previousState, session.State, session.TenantID)
+	if s.runtime != nil && session.TokenID != "" {
+		if err := s.runtime.RevokeSessionToken(ctx, session.TokenID, session.ExpiresAt); err != nil {
+			return fmt.Errorf("revoke session %q: revoke session token %q: %w", session.ID, session.TokenID, err)
+		}
+	}
 
 	grants, err := s.repo.ListGrantsBySession(ctx, req.SessionID)
 	if err != nil {
@@ -836,6 +842,15 @@ func (s *Service) loadActiveSession(ctx context.Context, raw string) (*core.Sess
 	if err != nil {
 		return nil, nil, fmt.Errorf("load active session: verify session token: %w", err)
 	}
+	if s.runtime != nil && claims.TokenID != "" {
+		revoked, err := s.runtime.IsSessionTokenRevoked(ctx, claims.TokenID)
+		if err != nil {
+			return nil, nil, fmt.Errorf("load active session: check session token revocation: %w", err)
+		}
+		if revoked {
+			return nil, nil, fmt.Errorf("%w: session token revoked", core.ErrForbidden)
+		}
+	}
 	session, err := s.repo.GetSession(ctx, claims.SessionID)
 	if err != nil {
 		return nil, nil, fmt.Errorf("load active session %q: load session: %w", claims.SessionID, err)

internal/app/service_test.go

@@ -573,6 +573,93 @@ func TestService_RevokeSessionRevokesOutstandingGrants(t *testing.T) {
 	}
 }
 
+func TestService_RequestGrantRejectsRevokedSessionToken(t *testing.T) {
+	t.Parallel()
+
+	ctx := context.Background()
+	now := testNow()
+	repo := memstore.NewRepository()
+	runtime := memstore.NewRuntimeStore()
+	tools := toolregistry.New()
+	engine := policy.NewEngine()
+	signer := mustNewSigner(t)
+
+	mustPutTool(t, ctx, tools, core.Tool{
+		TenantID:             "t_acme",
+		Tool:                 "github",
+		ManifestHash:         "sha256:test",
+		RuntimeClass:         core.RuntimeClassHosted,
+		AllowedDeliveryModes: []core.DeliveryMode{core.DeliveryModeProxy},
+		AllowedCapabilities:  []string{"repo.read"},
+		TrustTags:            []string{"trusted", "github"},
+	})
+	mustPutPolicy(t, engine, core.Policy{
+		TenantID:             "t_acme",
+		Capability:           "repo.read",
+		ResourceKind:         core.ResourceKindGitHubRepo,
+		AllowedDeliveryModes: []core.DeliveryMode{core.DeliveryModeProxy},
+		DefaultTTL:           10 * time.Minute,
+		MaxTTL:               10 * time.Minute,
+		ApprovalMode:         core.ApprovalModeNone,
+		RequiredToolTags:     []string{"trusted", "github"},
+		Condition:            `request.tool == "github"`,
+	})
+
+	svc, err := app.NewService(app.Config{
+		Clock:         fixedClock(now),
+		IDs:           fixedIDs("sess_revoked"),
+		Repository:    repo,
+		Runtime:       runtime,
+		Verifier:      fakeVerifier{identity: workloadIdentity()},
+		SessionTokens: signer,
+		Policy:        engine,
+		Tools:         tools,
+		Connectors:    fakeConnectorResolver{connector: &fakeConnector{kind: "github"}},
+		Deliveries: map[core.DeliveryMode]core.DeliveryAdapter{
+			core.DeliveryModeProxy: &fakeDeliveryAdapter{
+				mode: core.DeliveryModeProxy,
+				delivery: &core.Delivery{
+					Kind:   core.DeliveryKindProxyHandle,
+					Handle: "ph_revoked",
+				},
+			},
+		},
+	})
+	if err != nil {
+		t.Fatalf("NewService() error = %v", err)
+	}
+
+	sessionResp, err := svc.CreateSession(ctx, &core.CreateSessionRequest{
+		TenantID:    "t_acme",
+		AgentID:     "agent_pr_reviewer",
+		RunID:       "run_revoked",
+		ToolContext: []string{"github"},
+		Attestation: &core.Attestation{Kind: core.AttestationKindK8SServiceAccountJWT, Token: "jwt"},
+	})
+	if err != nil {
+		t.Fatalf("CreateSession() error = %v", err)
+	}
+
+	claims, err := signer.Verify(sessionResp.SessionToken)
+	if err != nil {
+		t.Fatalf("Verify() error = %v", err)
+	}
+	if err := runtime.RevokeSessionToken(ctx, claims.TokenID, claims.ExpiresAt); err != nil {
+		t.Fatalf("RevokeSessionToken() error = %v", err)
+	}
+
+	_, err = svc.RequestGrant(ctx, &core.RequestGrantRequest{
+		SessionToken: sessionResp.SessionToken,
+		Tool:         "github",
+		Capability:   "repo.read",
+		ResourceRef:  "github:repo:acme/widgets",
+		DeliveryMode: core.DeliveryModeProxy,
+	})
+	if err == nil || !strings.Contains(err.Error(), "session token revoked") {
+		t.Fatalf("RequestGrant() error = %v, want revoked token failure", err)
+	}
+}
+
 func mustNewSigner(t *testing.T) core.SessionTokenManager {
 	t.Helper()
 

internal/bootstrap/service.go

@@ -403,18 +403,22 @@ func (v multiVerifier) Verify(ctx context.Context, in *core.Attestation) (*core.
 }
 
 func newSessionTokenManager() (core.SessionTokenManager, error) {
+	options := []sessionjwt.Option{}
+	if issuer := strings.TrimSpace(os.Getenv("ASB_SESSION_TOKEN_ISSUER")); issuer != "" {
+		options = append(options, sessionjwt.WithIssuer(issuer))
+	}
 	if path := os.Getenv("ASB_SESSION_SIGNING_PRIVATE_KEY_FILE"); path != "" {
 		privateKey, err := loadEd25519PrivateKey(path)
 		if err != nil {
 			return nil, err
 		}
-		return sessionjwt.NewManager(privateKey)
+		return sessionjwt.NewManager(privateKey, options...)
 	}
 	_, privateKey, err := ed25519.GenerateKey(rand.Reader)
 	if err != nil {
 		return nil, err
 	}
-	return sessionjwt.NewManager(privateKey)
+	return sessionjwt.NewManager(privateKey, options...)
 }
 
 func newDelegationValidator() (core.DelegationValidator, error) {

internal/core/types.go

@@ -133,6 +133,7 @@ type Session struct {
 	TenantID         string
 	AgentID          string
 	RunID            string
+	TokenID          string
 	WorkloadIdentity WorkloadIdentity
 	Delegation       *Delegation
 	ToolContext      []string
@@ -395,6 +396,7 @@ type SessionClaims struct {
 	TenantID     string
 	AgentID      string
 	RunID        string
+	TokenID      string
 	ToolContext  []string
 	WorkloadHash string
 	ExpiresAt    time.Time
@@ -504,6 +506,8 @@ type RuntimeStore interface {
 	CompleteProxyRequest(ctx context.Context, handle string, responseBytes int64) error
 	SaveRelaySession(ctx context.Context, relay *BrowserRelaySession) error
 	GetRelaySession(ctx context.Context, sessionID string) (*BrowserRelaySession, error)
+	RevokeSessionToken(ctx context.Context, tokenID string, expiresAt time.Time) error
+	IsSessionTokenRevoked(ctx context.Context, tokenID string) (bool, error)
 }
 
 type GitHubProxyExecutor interface {

internal/crypto/sessionjwt/manager.go

@@ -12,8 +12,18 @@ import (
 type Manager struct {
 	privateKey ed25519.PrivateKey
 	publicKey  ed25519.PublicKey
+	issuer     string
+	clockSkew  time.Duration
+	now        func() time.Time
 }
 
+type Option func(*Manager)
+
+const (
+	defaultIssuer    = "asb"
+	defaultClockSkew = 30 * time.Second
+)
+
 type claims struct {
 	SessionID    string   `json:"sid"`
 	TenantID     string   `json:"tenant_id"`
@@ -24,21 +34,56 @@ type claims struct {
 	jwt.RegisteredClaims
 }
 
-func NewManager(privateKey ed25519.PrivateKey) (*Manager, error) {
+func WithIssuer(issuer string) Option {
+	return func(manager *Manager) {
+		if issuer != "" {
+			manager.issuer = issuer
+		}
+	}
+}
+
+func WithClockSkew(clockSkew time.Duration) Option {
+	return func(manager *Manager) {
+		if clockSkew >= 0 {
+			manager.clockSkew = clockSkew
+		}
+	}
+}
+
+func WithNowFunc(now func() time.Time) Option {
+	return func(manager *Manager) {
+		if now != nil {
+			manager.now = now
+		}
+	}
+}
+
+func NewManager(privateKey ed25519.PrivateKey, options ...Option) (*Manager, error) {
 	if len(privateKey) == 0 {
 		return nil, fmt.Errorf("%w: private key is required", core.ErrInvalidRequest)
 	}
 	publicKey, ok := privateKey.Public().(ed25519.PublicKey)
 	if !ok {
 		return nil, fmt.Errorf("%w: private key public component is %T, want ed25519.PublicKey", core.ErrInvalidRequest, privateKey.Public())
 	}
-	return &Manager{
+	manager := &Manager{
 		privateKey: privateKey,
 		publicKey:  publicKey,
-	}, nil
+		issuer:     defaultIssuer,
+		clockSkew:  defaultClockSkew,
+		now:        time.Now,
+	}
+	for _, option := range options {
+		option(manager)
+	}
+	return manager, nil
 }
 
 func (m *Manager) Sign(session *core.Session) (string, error) {
+	tokenID := session.TokenID
+	if tokenID == "" {
+		tokenID = session.ID
+	}
 	token := jwt.NewWithClaims(jwt.SigningMethodEdDSA, claims{
 		SessionID:    session.ID,
 		TenantID:     session.TenantID,
@@ -49,7 +94,10 @@ func (m *Manager) Sign(session *core.Session) (string, error) {
 		RegisteredClaims: jwt.RegisteredClaims{
 			ExpiresAt: jwt.NewNumericDate(session.ExpiresAt),
 			IssuedAt:  jwt.NewNumericDate(session.CreatedAt),
+			NotBefore: jwt.NewNumericDate(session.CreatedAt.Add(-m.clockSkew)),
+			Issuer:    m.issuer,
 			Subject:   session.WorkloadIdentity.Subject,
+			ID:        tokenID,
 		},
 	})
 	return token.SignedString(m.privateKey)
@@ -61,7 +109,7 @@ func (m *Manager) Verify(raw string) (*core.SessionClaims, error) {
 			return nil, fmt.Errorf("%w: unexpected signing method %q", core.ErrUnauthorized, token.Method.Alg())
 		}
 		return m.publicKey, nil
-	})
+	}, jwt.WithIssuer(m.issuer), jwt.WithIssuedAt(), jwt.WithLeeway(m.clockSkew), jwt.WithTimeFunc(m.now))
 	if err != nil {
 		return nil, fmt.Errorf("%w: %v", core.ErrUnauthorized, err)
 	}
@@ -70,21 +118,27 @@ func (m *Manager) Verify(raw string) (*core.SessionClaims, error) {
 	if !ok || !parsed.Valid {
 		return nil, fmt.Errorf("%w: invalid session token", core.ErrUnauthorized)
 	}
+	if tokenClaims.ExpiresAt == nil {
+		return nil, fmt.Errorf("%w: missing exp", core.ErrUnauthorized)
+	}
+	if tokenClaims.NotBefore == nil {
+		return nil, fmt.Errorf("%w: missing nbf", core.ErrUnauthorized)
+	}
+	if tokenClaims.ID == "" {
+		return nil, fmt.Errorf("%w: missing jti", core.ErrUnauthorized)
+	}
+	if tokenClaims.SessionID == "" || tokenClaims.TenantID == "" {
+		return nil, fmt.Errorf("%w: missing required session claims", core.ErrUnauthorized)
+	}
 
 	return &core.SessionClaims{
 		SessionID:    tokenClaims.SessionID,
 		TenantID:     tokenClaims.TenantID,
 		AgentID:      tokenClaims.AgentID,
 		RunID:        tokenClaims.RunID,
+		TokenID:      tokenClaims.ID,
 		ToolContext:  append([]string(nil), tokenClaims.ToolContext...),
 		WorkloadHash: tokenClaims.WorkloadHash,
 		ExpiresAt:    tokenClaims.ExpiresAt.Time,
 	}, nil
 }
-
-func (c *claims) Valid() error {
-	if c.ExpiresAt == nil {
-		return fmt.Errorf("%w: missing exp", core.ErrUnauthorized)
-	}
-	return jwt.NewValidator(jwt.WithTimeFunc(time.Now)).Validate(c.RegisteredClaims)
-}