jiggablend/internal/auth/jobtoken.go

package auth

import (
	"crypto/hmac"
	"crypto/rand"
	"crypto/sha256"
	"encoding/base64"
	"encoding/json"
	"fmt"
	"sync"
	"time"
)

// DefaultJobTokenDuration is the minimum validity period for job tokens.
const DefaultJobTokenDuration = 1 * time.Hour

// JobTokenSkew is extra lifetime added beyond configured task timeouts.
const JobTokenSkew = 15 * time.Minute

// JobTokenDuration is the current validity period for newly issued job tokens.
// Prefer JobTokenTTL() / SetJobTokenTTL; this var remains for backward-compatible reads.
var (
	jobTokenTTL   = DefaultJobTokenDuration
	jobTokenTTLMu sync.RWMutex
)

// JobTokenTTL returns the current job token lifetime used when minting tokens.
func JobTokenTTL() time.Duration {
	jobTokenTTLMu.RLock()
	defer jobTokenTTLMu.RUnlock()
	return jobTokenTTL
}

// SetJobTokenTTL sets the lifetime for newly issued job tokens.
// Values below DefaultJobTokenDuration are raised to the default.
func SetJobTokenTTL(d time.Duration) {
	if d < DefaultJobTokenDuration {
		d = DefaultJobTokenDuration
	}
	jobTokenTTLMu.Lock()
	jobTokenTTL = d
	jobTokenTTLMu.Unlock()
}

// ConfigureJobTokenTTLFromTimeouts sets token lifetime from the longest of the
// given task timeouts (seconds) plus JobTokenSkew, floored at DefaultJobTokenDuration.
func ConfigureJobTokenTTLFromTimeouts(timeoutSeconds ...int) time.Duration {
	maxSec := 0
	for _, s := range timeoutSeconds {
		if s > maxSec {
			maxSec = s
		}
	}
	d := time.Duration(maxSec)*time.Second + JobTokenSkew
	SetJobTokenTTL(d)
	return JobTokenTTL()
}

// JobTokenClaims represents the claims in a job token
type JobTokenClaims struct {
	JobID    int64 `json:"job_id"`
	RunnerID int64 `json:"runner_id"`
	TaskID   int64 `json:"task_id"`
	Exp      int64 `json:"exp"` // Unix timestamp
}

// jobTokenSecret is the secret used to sign job tokens
// Generated once at startup and kept in memory
var jobTokenSecret []byte

func init() {
	// Generate a random secret for signing job tokens
	// This means tokens are invalidated on server restart, which is acceptable
	// for short-lived job tokens
	jobTokenSecret = make([]byte, 32)
	if _, err := rand.Read(jobTokenSecret); err != nil {
		panic(fmt.Sprintf("failed to generate job token secret: %v", err))
	}
}

// GenerateJobToken creates a new job token for a specific job/runner/task combination
func GenerateJobToken(jobID, runnerID, taskID int64) (string, error) {
	claims := JobTokenClaims{
		JobID:    jobID,
		RunnerID: runnerID,
		TaskID:   taskID,
		Exp:      time.Now().Add(JobTokenTTL()).Unix(),
	}

	// Encode claims to JSON
	claimsJSON, err := json.Marshal(claims)
	if err != nil {
		return "", fmt.Errorf("failed to marshal claims: %w", err)
	}

	// Create HMAC signature
	h := hmac.New(sha256.New, jobTokenSecret)
	h.Write(claimsJSON)
	signature := h.Sum(nil)

	// Combine claims and signature: base64(claims).base64(signature)
	token := base64.RawURLEncoding.EncodeToString(claimsJSON) + "." +
		base64.RawURLEncoding.EncodeToString(signature)

	return token, nil
}

// ValidateJobToken validates a job token and returns the claims if valid
func ValidateJobToken(token string) (*JobTokenClaims, error) {
	// Split token into claims and signature
	var claimsB64, sigB64 string
	dotIdx := -1
	for i := len(token) - 1; i >= 0; i-- {
		if token[i] == '.' {
			dotIdx = i
			break
		}
	}
	if dotIdx == -1 {
		return nil, fmt.Errorf("invalid token format")
	}
	claimsB64 = token[:dotIdx]
	sigB64 = token[dotIdx+1:]

	// Decode claims
	claimsJSON, err := base64.RawURLEncoding.DecodeString(claimsB64)
	if err != nil {
		return nil, fmt.Errorf("invalid token encoding: %w", err)
	}

	// Decode signature
	signature, err := base64.RawURLEncoding.DecodeString(sigB64)
	if err != nil {
		return nil, fmt.Errorf("invalid signature encoding: %w", err)
	}

	// Verify signature
	h := hmac.New(sha256.New, jobTokenSecret)
	h.Write(claimsJSON)
	expectedSig := h.Sum(nil)
	if !hmac.Equal(signature, expectedSig) {
		return nil, fmt.Errorf("invalid signature")
	}

	// Parse claims
	var claims JobTokenClaims
	if err := json.Unmarshal(claimsJSON, &claims); err != nil {
		return nil, fmt.Errorf("invalid claims: %w", err)
	}

	// Check expiration
	if time.Now().Unix() > claims.Exp {
		return nil, fmt.Errorf("token expired")
	}

	return &claims, nil
}