material_manage_system/internal/service/auth.go

package service

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

	"github.com/wmms/wmms/internal/model"
	"github.com/wmms/wmms/internal/repository"
	"github.com/wmms/wmms/pkg/database"
	"golang.org/x/crypto/bcrypt"
)

var (
	ErrInvalidCredentials = errors.New("invalid username or password")
	ErrUserNotFound       = errors.New("user not found")
)

type AuthService struct {
	userRepo *repository.UserRepository
}

func NewAuthService() *AuthService {
	return &AuthService{
		userRepo: repository.NewUserRepository(),
	}
}

// SimpleToken represents a simple JWT-like token
type SimpleToken struct {
	SessionID string   `json:"session_id,omitempty"`
	UserID    int64    `json:"user_id"`
	Username  string   `json:"username"`
	Roles     []string `json:"roles"`
	Exp       int64    `json:"exp"`
}

// GenerateToken creates a simple token for the user
func (s *AuthService) GenerateToken(user *model.User) (string, error) {
	nonce := make([]byte, 32)
	if _, err := rand.Read(nonce); err != nil {
		return "", err
	}
	exp := time.Now().Add(24 * time.Hour).Unix()

	roleCodes := make([]string, 0, len(user.Roles))
	for _, role := range user.Roles {
		if role.Enabled {
			roleCodes = append(roleCodes, role.Code)
		}
	}
	token := SimpleToken{
		SessionID: base64.RawURLEncoding.EncodeToString(nonce),
		UserID:    user.ID,
		Username:  user.Username,
		Roles:     roleCodes,
		Exp:       exp,
	}

	// Create header
	header := base64.RawURLEncoding.EncodeToString([]byte(`{"alg":"HS256","typ":"JWT"}`))

	// Create payload
	payloadBytes, err := json.Marshal(token)
	if err != nil {
		return "", err
	}
	payload := base64.RawURLEncoding.EncodeToString(payloadBytes)

	// Create signature using HMAC-SHA256 (cryptographically secure)
	signatureInput := header + "." + payload
	mac := hmac.New(sha256.New, []byte("warehouse-mms-secret-key"))
	mac.Write([]byte(signatureInput))
	signature := base64.RawURLEncoding.EncodeToString(mac.Sum(nil))

	return header + "." + payload + "." + signature, nil
}

// ValidateToken validates a token and returns the user info
func (s *AuthService) ValidateToken(tokenString string) (*SimpleToken, error) {
	parts := splitToken(tokenString)
	if len(parts) != 3 {
		return nil, errors.New("invalid token format")
	}

	header, payload, signature := parts[0], parts[1], parts[2]

	// Verify signature using HMAC-SHA256 (cryptographically secure)
	signatureInput := header + "." + payload
	mac := hmac.New(sha256.New, []byte("warehouse-mms-secret-key"))
	mac.Write([]byte(signatureInput))
	expectedSig := base64.RawURLEncoding.EncodeToString(mac.Sum(nil))

	if !hmac.Equal([]byte(signature), []byte(expectedSig)) {
		return nil, errors.New("invalid token signature")
	}

	// Decode payload
	payloadBytes, err := base64.RawURLEncoding.DecodeString(payload)
	if err != nil {
		return nil, err
	}

	var token SimpleToken
	if err := json.Unmarshal(payloadBytes, &token); err != nil {
		return nil, err
	}

	// Check expiration
	if time.Now().Unix() >= token.Exp {
		return nil, errors.New("token expired")
	}

	if err := ValidateSession(TokenSessionID(tokenString), token.Exp); err != nil {
		return nil, err
	}
	return &token, nil
}

// The single-server session cache is restored from revoked_tokens on startup.
// Keep only a fingerprint, never a raw bearer token, and expire entries with JWTs.
var revokedSessions = struct {
	sync.Mutex
	entries map[string]int64
}{entries: make(map[string]int64)}

func TokenSessionID(token string) string {
	sum := sha256.Sum256([]byte(token))
	return base64.RawURLEncoding.EncodeToString(sum[:])
}

func ValidateSession(id string, expiresAt int64) error {
	if id == "" || time.Now().Unix() >= expiresAt {
		return errors.New("session expired")
	}
	revokedSessions.Lock()
	defer revokedSessions.Unlock()
	if _, revoked := revokedSessions.entries[id]; revoked {
		return errors.New("session revoked")
	}
	return nil
}

func (s *AuthService) RevokeSession(value string) error {
	token, err := s.ValidateToken(value)
	if err != nil {
		return err
	}
	if _, err := database.DB.Exec(`INSERT INTO revoked_tokens (jti, user_id, expires_at) VALUES (?, ?, ?) ON DUPLICATE KEY UPDATE expires_at=VALUES(expires_at)`, TokenSessionID(value), token.UserID, time.Unix(token.Exp, 0)); err != nil {
		return err
	}
	revokedSessions.Lock()
	defer revokedSessions.Unlock()
	for id, expires := range revokedSessions.entries {
		if expires <= time.Now().Unix() {
			delete(revokedSessions.entries, id)
		}
	}
	revokedSessions.entries[TokenSessionID(value)] = token.Exp
	return nil
}

// LoadRevokedSessions runs before serving requests, so restarting the process
// cannot reactivate a logged-out token. Tickets remain local to this server.
func LoadRevokedSessions() error {
	// Bind the cutoff through the same driver location as the stored DATETIME.
	// MySQL NOW() may use a different session timezone and omit live revocations.
	rows, err := database.DB.Query(`SELECT jti, expires_at FROM revoked_tokens WHERE expires_at > ?`, time.Now())
	if err != nil {
		return err
	}
	defer rows.Close()
	entries := make(map[string]int64)
	for rows.Next() {
		var id string
		var expiry time.Time
		if err := rows.Scan(&id, &expiry); err != nil {
			return err
		}
		entries[id] = expiry.Unix()
	}
	if err := rows.Err(); err != nil {
		return err
	}
	revokedSessions.Lock()
	revokedSessions.entries = entries
	revokedSessions.Unlock()
	return nil
}

func splitToken(token string) []string {
	var parts []string
	start := 0
	for i := 0; i < len(token); i++ {
		if token[i] == '.' {
			parts = append(parts, token[start:i])
			start = i + 1
		}
	}
	parts = append(parts, token[start:])
	return parts
}

// Login authenticates a user and returns a token
func (s *AuthService) Login(username, password string) (*model.LoginResponse, error) {
	user, err := s.userRepo.GetByUsername(username)
	if err != nil {
		return nil, ErrInvalidCredentials
	}

	// Verify password
	if err := bcrypt.CompareHashAndPassword([]byte(user.Password), []byte(password)); err != nil {
		return nil, ErrInvalidCredentials
	}

	// Generate token
	token, err := s.GenerateToken(user)
	if err != nil {
		return nil, err
	}

	return &model.LoginResponse{
		Token: token,
		User:  *user,
	}, nil
}

// GetUserByID returns a user by ID
func (s *AuthService) GetUserByID(id int64) (*model.User, error) {
	return s.userRepo.GetByID(id)
}

// ListUsers returns all users with pagination
func (s *AuthService) ListUsers(search string, limit, offset int) ([]model.User, int64, error) {
	return s.userRepo.List(search, limit, offset)
}

// CreateUser creates a new user
func (s *AuthService) CreateUser(username, password, description string, roleIDs []int64, organization ...model.UserOrganizationInput) (*model.User, error) {
	// Hash password
	hashedPassword, err := bcrypt.GenerateFromPassword([]byte(password), bcrypt.DefaultCost)
	if err != nil {
		return nil, err
	}

	user := &model.User{
		Username:    username,
		Password:    string(hashedPassword),
		Description: description,
	}

	// Service callers always validate company access. Repository calls without
	// input are reserved for the legacy bootstrap users created before migration.
	input := model.UserOrganizationInput{}
	if len(organization) > 0 {
		input = organization[0]
	}
	if err := s.userRepo.Create(user, roleIDs, input); err != nil {
		return nil, err
	}
	return s.userRepo.GetByID(user.ID)
}

// UpdateUser updates a user (admin function)
func (s *AuthService) UpdateUser(id int64, description string, roleIDs []int64, newPassword string, organization ...model.UserOrganizationInput) error {
	user, err := s.userRepo.GetByID(id)
	if err != nil {
		return err
	}

	user.Description = description

	if newPassword != "" {
		hashedPassword, err := bcrypt.GenerateFromPassword([]byte(newPassword), bcrypt.DefaultCost)
		if err != nil {
			return err
		}
		user.Password = string(hashedPassword)
		return s.userRepo.UpdateWithPassword(user, roleIDs, organization...)
	}

	return s.userRepo.Update(user, roleIDs, organization...)
}

// DeleteUser deletes a user
func (s *AuthService) DeleteUser(id int64) error {
	return s.userRepo.Delete(id)
}

// HashPassword hashes a password using bcrypt
func HashPassword(password string) (string, error) {
	hash, err := bcrypt.GenerateFromPassword([]byte(password), bcrypt.DefaultCost)
	if err != nil {
		return "", err
	}
	return string(hash), nil
}