meowlib/server/invitation.go

package server

import (
	"crypto/rand"
	"crypto/subtle"
	"errors"
	"fmt"
	"math/big"
	"strings"
	"time"

	"github.com/go-redis/redis"
)

func (r *RedisRouter) StoreInvitation(invitation []byte, timeout int, password string, serverTimeout int, urlLen int) (string, time.Time, error) {
	id, err := r.createShortId(urlLen)
	if err != nil {
		return "", time.Time{}, fmt.Errorf("failed to create invitation ID: %w", err)
	}
	if timeout > serverTimeout {
		timeout = serverTimeout
	}
	r.Client.Set("mwiv:"+id, invitation, 0) //, time.Duration(timeout*1000000))
	if len(password) > 0 {
		r.Client.Set("mwpw:"+id, password, 0) //, time.Duration(timeout*1000000))
	}
	return id, time.Now().Add(time.Duration(timeout * 1000000)).UTC(), nil
}

func (r *RedisRouter) GetInvitation(id string, password string) ([]byte, error) {
	// Check failed attempts counter
	failedAttemptsKey := "mwfa:" + id
	failedAttempts := 0 // Default when key doesn't exist
	val, err := r.Client.Get(failedAttemptsKey).Int()
	if err == nil {
		failedAttempts = val
	} else if err != redis.Nil {
		return nil, fmt.Errorf("failed to check attempts: %w", err)
	}
	// If err == redis.Nil, key doesn't exist, so failedAttempts stays 0

	// If already hit the limit, delete invitation and fail
	if failedAttempts >= 3 {
		r.deleteInvitation(id)
		return nil, errors.New("invitation locked due to too many failed attempts")
	}

	// Check if password is required
	passRequired := false
	expectedpass, err := r.Client.Get("mwpw:" + id).Result()
	if err != nil {
		passRequired = false
	} else {
		passRequired = true
	}

	// Validate password with constant-time comparison
	if passRequired {
		if subtle.ConstantTimeCompare([]byte(password), []byte(expectedpass)) != 1 {
			// Increment failed attempts
			newCount := failedAttempts + 1
			r.Client.Set(failedAttemptsKey, newCount, 0)

			// If this was the 3rd attempt, delete invitation
			if newCount >= 3 {
				r.deleteInvitation(id)
				return nil, errors.New("auth failed - invitation destroyed after 3 attempts")
			}
			return nil, errors.New("auth failed")
		}
	}

	// Success - get invitation data
	mwiv, err := r.Client.Get("mwiv:" + id).Result()
	if err != nil {
		return nil, err
	}

	// Clear failed attempts counter on successful access
	r.Client.Del(failedAttemptsKey)

	return []byte(mwiv), nil
}

// deleteInvitation removes all invitation-related keys from Redis
func (r *RedisRouter) deleteInvitation(id string) {
	r.Client.Del("mwiv:" + id)  // invitation data
	r.Client.Del("mwpw:" + id)  // password
	r.Client.Del("mwfa:" + id)  // failed attempts
	r.Client.Del("mwan:" + id)  // answer to invitation
}

func (r *RedisRouter) StoreAnswerToInvitation(id string, timeout int, invitation []byte, serverTimeout int) time.Time {
	if timeout > serverTimeout {
		timeout = serverTimeout
	}
	r.Client.Set("mwan:"+id, invitation, time.Duration(timeout*1000000))
	return time.Now().Add(time.Duration(timeout * 1000000)).UTC()
}

func (r *RedisRouter) GetAnswerToInvitation(id string) ([]byte, error) {
	mwan, err := r.Client.Get("mwan:" + id).Result()
	if err != nil {
		return nil, err
	}
	return []byte(mwan), nil
}

func (r *RedisRouter) createShortId(length int) (string, error) {
	alphabet := "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789"
	alphabetLen := big.NewInt(int64(len(alphabet)))

	// for not in redis
	for {
		var id strings.Builder
		id.Grow(length)

		for i := 0; i < length; i++ {
			n, err := rand.Int(rand.Reader, alphabetLen)
			if err != nil {
				return "", fmt.Errorf("random generation failed: %w", err)
			}
			id.WriteByte(alphabet[n.Int64()])
		}

		idStr := id.String()
		if r.Client.Get("mwiv:"+idStr).Err() == redis.Nil {
			return idStr, nil
		}
	}
}