NianAIGC/backend/internal/identity/session.go

// Package identity implements the legacy zhinian_session cookie wire contract.
package identity

import (
	"crypto/hmac"
	"crypto/sha256"
	"encoding/base64"
	"encoding/json"
	"errors"
	"fmt"
	"net/url"
	"strings"
	"time"
)

const (
	SessionCookieName    = "zhinian_session"
	CookieChunkSize      = 3000
	CookieMaxChunks      = 20
	CookieMaxValueLength = CookieChunkSize * CookieMaxChunks
)

type AuthMode string

const (
	AuthModeUser  AuthMode = "user"
	AuthModeAdmin AuthMode = "admin"
)

// Session is the validated version-one session contract.
type Session struct {
	Version        int      `json:"version"`
	AuthMode       AuthMode `json:"authMode"`
	IssuedAt       int64    `json:"issuedAt"`
	ExpiresAt      int64    `json:"expiresAt"`
	SessionVersion *int     `json:"sessionVersion,omitempty"`
	AccessToken    string   `json:"accessToken,omitempty"`
	TokenType      string   `json:"tokenType,omitempty"`
	User           User     `json:"user"`
}

type User struct {
	ID               string   `json:"id"`
	Subject          string   `json:"subject"`
	Username         string   `json:"username,omitempty"`
	Phone            string   `json:"phone,omitempty"`
	DisplayName      string   `json:"displayName"`
	ClientID         string   `json:"clientId"`
	TenantID         string   `json:"tenantId,omitempty"`
	OrganizationID   string   `json:"organizationId,omitempty"`
	OrganizationName string   `json:"organizationName,omitempty"`
	Role             string   `json:"role,omitempty"`
	Status           string   `json:"status,omitempty"`
	Authorities      []string `json:"authorities"`
	Scope            []string `json:"scope"`
}

type CookieChunk struct {
	Name  string
	Value string
}

// CookieAttributes is a transport-neutral description of the legacy
// Set-Cookie attributes. Adapters can translate it to their HTTP framework.
type CookieAttributes struct {
	HTTPOnly      bool
	SameSite      string
	Secure        bool
	Path          string
	Expires       *time.Time
	MaxAgeSeconds *int
}

// CookieWrite describes one cookie mutation. A session write always returns
// all 20 names so stale chunks are cleared atomically with the new chunks.
type CookieWrite struct {
	Name       string
	Value      string
	Attributes CookieAttributes
}

var (
	ErrMalformedSession = errors.New("malformed session cookie")
	ErrInvalidSignature = errors.New("invalid session cookie signature")
	ErrInvalidSession   = errors.New("invalid session")
	ErrSessionTooLarge  = errors.New("session cookie exceeds maximum supported size")
)

// Sign preserves raw JSON bytes, base64url-encodes them without padding, and
// signs that encoded payload with HMAC-SHA256.
func Sign(rawJSON []byte, secret string) (string, error) {
	if !json.Valid(rawJSON) {
		return "", fmt.Errorf("%w: invalid JSON", ErrInvalidSession)
	}
	payload := base64.RawURLEncoding.EncodeToString(rawJSON)
	signature := signPayload(payload, secret)
	return payload + "." + signature, nil
}

// Parse authenticates and validates a version-one session cookie.
func Parse(value, secret string, now time.Time) (Session, error) {
	payload, signature, ok := strings.Cut(value, ".")
	if !ok || payload == "" || signature == "" || strings.Contains(signature, ".") {
		return Session{}, ErrMalformedSession
	}
	expected := signPayload(payload, secret)
	if !hmac.Equal([]byte(signature), []byte(expected)) {
		return Session{}, ErrInvalidSignature
	}
	rawJSON, err := base64.RawURLEncoding.DecodeString(payload)
	if err != nil {
		return Session{}, fmt.Errorf("%w: invalid payload encoding", ErrMalformedSession)
	}
	var session Session
	if err := json.Unmarshal(rawJSON, &session); err != nil {
		return Session{}, fmt.Errorf("%w: %v", ErrInvalidSession, err)
	}
	if session.Version != 1 || session.User.ID == "" || session.User.ClientID == "" || session.ExpiresAt <= now.Unix() {
		return Session{}, ErrInvalidSession
	}
	if session.AuthMode != AuthModeAdmin {
		session.AuthMode = AuthModeUser
	}
	if session.User.Authorities == nil {
		session.User.Authorities = []string{}
	}
	if session.User.Scope == nil {
		session.User.Scope = []string{}
	}
	return session, nil
}

func signPayload(payload, secret string) string {
	mac := hmac.New(sha256.New, []byte(secret))
	_, _ = mac.Write([]byte(payload))
	return base64.RawURLEncoding.EncodeToString(mac.Sum(nil))
}

// Chunk splits a value into legacy cookie chunks named base, base.1, base.2,
// ... and rejects values the legacy reader cannot fully reassemble.
func Chunk(baseName, value string) ([]CookieChunk, error) {
	if len(value) > CookieMaxValueLength {
		return nil, ErrSessionTooLarge
	}
	if value == "" {
		return []CookieChunk{{Name: baseName}}, nil
	}
	chunks := make([]CookieChunk, 0, (len(value)+CookieChunkSize-1)/CookieChunkSize)
	for start, index := 0, 0; start < len(value); start, index = start+CookieChunkSize, index+1 {
		end := start + CookieChunkSize
		if end > len(value) {
			end = len(value)
		}
		chunks = append(chunks, CookieChunk{Name: chunkName(baseName, index), Value: value[start:end]})
	}
	return chunks, nil
}

// CookieNames returns every cookie name read or cleared by the legacy session
// contract, from zhinian_session through zhinian_session.19.
func CookieNames() []string {
	names := make([]string, CookieMaxChunks)
	for index := range names {
		names[index] = chunkName(SessionCookieName, index)
	}
	return names
}

// SetSessionCookies creates writes for the current chunks and deletion writes
// for every remaining legacy chunk name.
func SetSessionCookies(value string, expires time.Time, secure bool) ([]CookieWrite, error) {
	chunks, err := Chunk(SessionCookieName, value)
	if err != nil {
		return nil, err
	}
	writes := make([]CookieWrite, 0, CookieMaxChunks)
	for _, chunk := range chunks {
		expiresCopy := expires
		writes = append(writes, CookieWrite{
			Name:  chunk.Name,
			Value: chunk.Value,
			Attributes: CookieAttributes{
				HTTPOnly: true,
				SameSite: "lax",
				Secure:   secure,
				Path:     "/",
				Expires:  &expiresCopy,
			},
		})
	}
	for index := len(chunks); index < CookieMaxChunks; index++ {
		writes = append(writes, clearCookieWrite(chunkName(SessionCookieName, index), secure))
	}
	return writes, nil
}

// ClearSessionCookies clears the base cookie and every possible legacy chunk.
func ClearSessionCookies(secure bool) []CookieWrite {
	writes := make([]CookieWrite, 0, CookieMaxChunks)
	for _, name := range CookieNames() {
		writes = append(writes, clearCookieWrite(name, secure))
	}
	return writes
}

func clearCookieWrite(name string, secure bool) CookieWrite {
	maxAge := 0
	return CookieWrite{
		Name: name,
		Attributes: CookieAttributes{
			HTTPOnly:      true,
			SameSite:      "lax",
			Secure:        secure,
			Path:          "/",
			MaxAgeSeconds: &maxAge,
		},
	}
}

// ResolveSecureCookie mirrors the legacy environment precedence: a recognized
// explicit setting wins, then the configured public base URL wins over the
// request URL, and only HTTPS enables Secure.
func ResolveSecureCookie(explicit, publicBaseURL, requestURL string) bool {
	switch strings.ToLower(strings.TrimSpace(explicit)) {
	case "1", "true", "yes", "on":
		return true
	case "0", "false", "no", "off":
		return false
	}
	candidate := strings.TrimSpace(publicBaseURL)
	if candidate == "" {
		candidate = strings.TrimSpace(requestURL)
	}
	parsed, err := url.Parse(candidate)
	return err == nil && strings.EqualFold(parsed.Scheme, "https")
}

// Reassemble reads at most 20 contiguous legacy cookie chunks.
func Reassemble(baseName string, getValue func(name string) (string, bool)) (string, bool) {
	first, ok := getValue(baseName)
	if !ok || first == "" {
		return "", false
	}
	var value strings.Builder
	value.WriteString(first)
	for index := 1; index < CookieMaxChunks; index++ {
		chunk, found := getValue(chunkName(baseName, index))
		if !found || chunk == "" {
			break
		}
		value.WriteString(chunk)
	}
	return value.String(), true
}

func chunkName(baseName string, index int) string {
	if index == 0 {
		return baseName
	}
	return fmt.Sprintf("%s.%d", baseName, index)
}