ai-disclosure-standard/pkg/licenseclient/client.go

package licenseclient

import (
	"bytes"
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"net/http"
	"os"
	"path/filepath"
	"strings"
	"sync"
	"time"

	"github.com/b1tsblog/ai-disclosure-standard/pkg/licensekit"
)

type Status struct {
	Edition      string           `json:"edition"`
	Licensed     bool             `json:"licensed"`
	LicenseID    string           `json:"licenseId,omitempty"`
	Customer     string           `json:"customer,omitempty"`
	Product      string           `json:"product,omitempty"`
	Features     []string         `json:"features"`
	Limits       map[string]int64 `json:"limits,omitempty"`
	ExpiresAt    string           `json:"expiresAt,omitempty"`
	Mode         string           `json:"mode,omitempty"`
	Source       string           `json:"source,omitempty"`
	LastChecked  string           `json:"lastChecked,omitempty"`
	LeaseExpires string           `json:"leaseExpiresAt,omitempty"`
	Reason       string           `json:"reason,omitempty"`
}

type Config struct {
	Product        string
	Token          string
	TrustStore     licensekit.TrustStore
	BaseURL        string
	InstanceID     string
	Mode           licensekit.VerificationMode
	ServerURL      string
	CacheFile      string
	RefreshEvery   time.Duration
	RequestTimeout time.Duration
	ClientVersion  string
	HTTPClient     *http.Client
}

type Client struct {
	cfg      Config
	mu       sync.RWMutex
	status   Status
	claims   licensekit.Claims
	features map[string]bool
	stopOnce sync.Once
	stop     chan struct{}
}

type introspectRequest struct {
	Token         string `json:"token"`
	Product       string `json:"product"`
	BaseURL       string `json:"baseUrl"`
	Host          string `json:"host"`
	InstanceID    string `json:"instanceId,omitempty"`
	ClientVersion string `json:"clientVersion,omitempty"`
}

type introspectResponse struct {
	Valid      bool   `json:"valid"`
	LeaseToken string `json:"leaseToken,omitempty"`
	Reason     string `json:"reason,omitempty"`
}

type cacheDocument struct {
	LicenseID string `json:"licenseId"`
	Lease     string `json:"lease"`
	SavedAt   int64  `json:"savedAt"`
}

func New(ctx context.Context, cfg Config) *Client {
	if cfg.RefreshEvery <= 0 {
		cfg.RefreshEvery = 15 * time.Minute
	}
	if cfg.RequestTimeout <= 0 {
		cfg.RequestTimeout = 5 * time.Second
	}
	if cfg.HTTPClient == nil {
		cfg.HTTPClient = &http.Client{Timeout: cfg.RequestTimeout}
	}
	c := &Client{cfg: cfg, stop: make(chan struct{}), status: communityStatus(), features: map[string]bool{}}
	c.refresh(ctx)
	return c
}

func (c *Client) Start(ctx context.Context) {
	c.mu.RLock()
	mode := c.status.Mode
	c.mu.RUnlock()
	if mode == string(licensekit.ModeOffline) || strings.TrimSpace(c.cfg.Token) == "" {
		return
	}
	go func() {
		ticker := time.NewTicker(c.cfg.RefreshEvery)
		defer ticker.Stop()
		for {
			select {
			case <-ctx.Done():
				return
			case <-c.stop:
				return
			case <-ticker.C:
				refreshCtx, cancel := context.WithTimeout(context.Background(), c.cfg.RequestTimeout)
				c.refresh(refreshCtx)
				cancel()
			}
		}
	}()
}

func (c *Client) Close() { c.stopOnce.Do(func() { close(c.stop) }) }

func (c *Client) Refresh(ctx context.Context) Status {
	c.refresh(ctx)
	return c.Status()
}

func (c *Client) Has(feature string) bool {
	c.mu.RLock()
	defer c.mu.RUnlock()
	return c.isCurrentlyLicensedLocked(time.Now().UTC()) && c.features[feature]
}

func (c *Client) Limit(name string) (int64, bool) {
	c.mu.RLock()
	defer c.mu.RUnlock()
	if !c.isCurrentlyLicensedLocked(time.Now().UTC()) {
		return 0, false
	}
	value, ok := c.status.Limits[name]
	return value, ok
}

func (c *Client) Status() Status {
	c.mu.RLock()
	defer c.mu.RUnlock()
	out := c.status
	if out.Licensed && !c.isCurrentlyLicensedLocked(time.Now().UTC()) {
		out.Licensed = false
		out.Edition = "community"
		if out.Reason == "" {
			out.Reason = "license or online lease is no longer valid"
		}
	}
	out.Features = append([]string{}, c.status.Features...)
	if c.status.Limits != nil {
		out.Limits = make(map[string]int64, len(c.status.Limits))
		for key, value := range c.status.Limits {
			out.Limits[key] = value
		}
	}
	return out
}

func (c *Client) isCurrentlyLicensedLocked(now time.Time) bool {
	if !c.status.Licensed {
		return false
	}
	if c.claims.ExpiresAt > 0 && now.Unix() >= c.claims.ExpiresAt {
		return false
	}
	mode := licensekit.VerificationMode(c.status.Mode)
	if mode == licensekit.ModeOffline || c.status.LeaseExpires == "" {
		return true
	}
	leaseExpiry, err := time.Parse(time.RFC3339, c.status.LeaseExpires)
	if err != nil {
		return false
	}
	if mode == licensekit.ModeHybrid {
		leaseExpiry = leaseExpiry.Add(time.Duration(c.claims.Verification.OfflineGraceSeconds) * time.Second)
	}
	return now.Before(leaseExpiry)
}

func (c *Client) refresh(ctx context.Context) {
	now := time.Now().UTC()
	if strings.TrimSpace(c.cfg.Token) == "" {
		c.apply(communityStatus(), licensekit.Claims{})
		return
	}
	verified, err := licensekit.VerifyLicense(c.cfg.TrustStore, c.cfg.Token, now)
	if err != nil {
		c.apply(failedStatus(err.Error(), now), licensekit.Claims{})
		return
	}
	claims := verified.Claims
	if err := licensekit.ValidateLicenseContext(claims, c.cfg.Product, c.cfg.BaseURL, c.cfg.InstanceID); err != nil {
		c.apply(failedStatus(err.Error(), now), claims)
		return
	}
	mode := licensekit.StricterMode(claims.Verification.Mode, c.cfg.Mode)
	base := statusFromClaims(claims, mode, now)
	if mode == licensekit.ModeOffline {
		base.Licensed = true
		base.Source = "offline"
		c.apply(base, claims)
		return
	}
	if strings.TrimSpace(c.cfg.ServerURL) == "" {
		base.Edition = "community"
		base.Reason = "online verification is required but LICENSE_SERVER_URL is empty"
		c.apply(base, claims)
		return
	}
	lease, source, err := c.obtainLease(ctx, claims, mode, now)
	if err != nil {
		base.Edition = "community"
		base.Reason = err.Error()
		c.apply(base, claims)
		return
	}
	base.Licensed = true
	base.Source = source
	base.LeaseExpires = time.Unix(lease.ExpiresAt, 0).UTC().Format(time.RFC3339)
	base.Features = licensekit.UniqueSorted(intersection(claims.Features, lease.Features))
	c.apply(base, claims)
}

func (c *Client) obtainLease(ctx context.Context, claims licensekit.Claims, mode licensekit.VerificationMode, now time.Time) (licensekit.LeaseClaims, string, error) {
	leaseToken, err := c.requestLease(ctx)
	if err == nil {
		lease, verifyErr := licensekit.VerifyLease(c.cfg.TrustStore, leaseToken, now, 0)
		if verifyErr != nil {
			return licensekit.LeaseClaims{}, "", fmt.Errorf("online lease verification failed: %w", verifyErr)
		}
		if verifyErr = licensekit.ValidateLeaseContext(lease.Claims, claims, c.cfg.Product, c.cfg.BaseURL, c.cfg.InstanceID); verifyErr != nil {
			return licensekit.LeaseClaims{}, "", verifyErr
		}
		_ = c.writeCache(claims.LicenseID, leaseToken)
		return lease.Claims, "online", nil
	}
	if mode == licensekit.ModeOnline {
		return licensekit.LeaseClaims{}, "", fmt.Errorf("online verification failed: %w", err)
	}
	cached, cacheErr := c.readCache(claims, now)
	if cacheErr != nil {
		return licensekit.LeaseClaims{}, "", fmt.Errorf("online verification failed (%v) and no usable cached lease exists (%v)", err, cacheErr)
	}
	return cached, "cached-lease", nil
}

func (c *Client) requestLease(ctx context.Context) (string, error) {
	host, err := licensekit.HostFromBaseURL(c.cfg.BaseURL)
	if err != nil {
		return "", err
	}
	body, err := json.Marshal(introspectRequest{Token: c.cfg.Token, Product: c.cfg.Product, BaseURL: c.cfg.BaseURL, Host: host, InstanceID: c.cfg.InstanceID, ClientVersion: c.cfg.ClientVersion})
	if err != nil {
		return "", err
	}
	endpoint := strings.TrimRight(c.cfg.ServerURL, "/") + "/v1/introspect"
	req, err := http.NewRequestWithContext(ctx, http.MethodPost, endpoint, bytes.NewReader(body))
	if err != nil {
		return "", err
	}
	req.Header.Set("Content-Type", "application/json")
	req.Header.Set("Accept", "application/json")
	resp, err := c.cfg.HTTPClient.Do(req)
	if err != nil {
		return "", err
	}
	defer resp.Body.Close()
	payload, err := io.ReadAll(io.LimitReader(resp.Body, 1<<20))
	if err != nil {
		return "", err
	}
	var result introspectResponse
	if err := json.Unmarshal(payload, &result); err != nil {
		return "", fmt.Errorf("decode verification response: %w", err)
	}
	if resp.StatusCode != http.StatusOK || !result.Valid || result.LeaseToken == "" {
		if result.Reason == "" {
			result.Reason = resp.Status
		}
		return "", errors.New(result.Reason)
	}
	return result.LeaseToken, nil
}

func (c *Client) readCache(claims licensekit.Claims, now time.Time) (licensekit.LeaseClaims, error) {
	if strings.TrimSpace(c.cfg.CacheFile) == "" {
		return licensekit.LeaseClaims{}, errors.New("cache file is not configured")
	}
	data, err := os.ReadFile(c.cfg.CacheFile)
	if err != nil {
		return licensekit.LeaseClaims{}, err
	}
	var doc cacheDocument
	if err := json.Unmarshal(data, &doc); err != nil {
		return licensekit.LeaseClaims{}, err
	}
	if doc.LicenseID != claims.LicenseID {
		return licensekit.LeaseClaims{}, errors.New("cached lease belongs to another license")
	}
	grace := time.Duration(claims.Verification.OfflineGraceSeconds) * time.Second
	verified, err := licensekit.VerifyLease(c.cfg.TrustStore, doc.Lease, now, grace)
	if err != nil {
		return licensekit.LeaseClaims{}, err
	}
	if err := licensekit.ValidateLeaseContext(verified.Claims, claims, c.cfg.Product, c.cfg.BaseURL, c.cfg.InstanceID); err != nil {
		return licensekit.LeaseClaims{}, err
	}
	return verified.Claims, nil
}

func (c *Client) writeCache(licenseID, lease string) error {
	if strings.TrimSpace(c.cfg.CacheFile) == "" {
		return nil
	}
	if err := os.MkdirAll(filepath.Dir(c.cfg.CacheFile), 0o700); err != nil && filepath.Dir(c.cfg.CacheFile) != "." {
		return err
	}
	data, err := json.Marshal(cacheDocument{LicenseID: licenseID, Lease: lease, SavedAt: time.Now().UTC().Unix()})
	if err != nil {
		return err
	}
	temp := c.cfg.CacheFile + ".tmp"
	if err := os.WriteFile(temp, data, 0o600); err != nil {
		return err
	}
	return os.Rename(temp, c.cfg.CacheFile)
}

func (c *Client) apply(status Status, claims licensekit.Claims) {
	status.Features = licensekit.UniqueSorted(status.Features)
	c.mu.Lock()
	defer c.mu.Unlock()
	c.status = status
	c.claims = claims
	c.features = make(map[string]bool, len(status.Features))
	for _, feature := range status.Features {
		c.features[feature] = true
	}
}

func communityStatus() Status {
	return Status{Edition: "community", Features: []string{}, Limits: map[string]int64{}}
}

func failedStatus(reason string, now time.Time) Status {
	return Status{Edition: "community", Features: []string{}, Limits: map[string]int64{}, Reason: reason, LastChecked: now.Format(time.RFC3339)}
}

func statusFromClaims(claims licensekit.Claims, mode licensekit.VerificationMode, now time.Time) Status {
	limits := map[string]int64{}
	for key, value := range claims.Limits {
		limits[key] = value
	}
	return Status{
		Edition: claims.Edition, LicenseID: claims.LicenseID, Customer: claims.Customer, Product: claims.Product,
		Features: append([]string(nil), claims.Features...), Limits: limits,
		ExpiresAt: time.Unix(claims.ExpiresAt, 0).UTC().Format(time.RFC3339), Mode: string(mode), LastChecked: now.Format(time.RFC3339),
	}
}

func intersection(a, b []string) []string {
	allowed := make(map[string]bool, len(b))
	for _, value := range b {
		allowed[value] = true
	}
	out := make([]string, 0, len(a))
	for _, value := range a {
		if allowed[value] {
			out = append(out, value)
		}
	}
	return out
}

// NewDevelopment returns an in-memory licensed client for local development.
// Production applications should not expose this path without an explicit development switch.
func NewDevelopment(product, edition string, features []string) *Client {
	status := Status{Edition: edition, Licensed: true, Customer: "development", Product: product, Features: licensekit.UniqueSorted(features), Limits: map[string]int64{}, Mode: string(licensekit.ModeOffline), Source: "development", Reason: "insecure development override"}
	c := &Client{status: status, features: map[string]bool{}, stop: make(chan struct{})}
	for _, feature := range status.Features {
		c.features[feature] = true
	}
	return c
}