package kdbx

import (
	"bytes"
	"fmt"
)

// Limits on what a file may demand of whoever opens it. The server opens
// files that its users upload — a merge decrypts both copies — so a header
// asking for a billion AES rounds, two gibibytes of Argon2 memory or a
// payload that inflates without end would otherwise be a way to stop it.
// The bounds sit well above what KeePassXC, KeePass and KeePassDX choose for
// a one-second unlock on a desktop, and far below what hurts a server.
var (
	// MaxAESRounds bounds AES-KDF.
	MaxAESRounds uint64 = 30_000_000
	// MaxArgon2MemoryKiB bounds Argon2's memory, per derivation.
	MaxArgon2MemoryKiB uint64 = 256 << 10
	// MaxArgon2Iterations bounds Argon2's passes.
	MaxArgon2Iterations uint64 = 64
	// MaxArgon2Work bounds memory × passes, in MiB-passes, so a file cannot
	// take the maximum of both at once.
	MaxArgon2Work uint64 = 4096
	// MaxArgon2Lanes bounds Argon2's parallelism.
	MaxArgon2Lanes uint64 = 16
	// MaxXMLBytes bounds the decompressed payload.
	MaxXMLBytes = 128 << 20
	// MaxXMLDepth and MaxXMLNodes bound the payload's shape. KeePass nests
	// groups a few levels deep and stores thousands of entries, not millions.
	MaxXMLDepth = 256
	MaxXMLNodes = 4_000_000
)

// checkKDF refuses key derivation parameters beyond the limits.
func checkKDF(kdf variantDict) error {
	uuid := kdf.bytes("$UUID")
	switch {
	case bytes.Equal(uuid, kdfAES):
		if r, ok := kdf.uint64("R"); ok && r > MaxAESRounds {
			return fmt.Errorf("kdbx: AES-KDF asks for %d rounds, refusing above %d", r, MaxAESRounds)
		}
	case bytes.Equal(uuid, kdfArgon2d), bytes.Equal(uuid, kdfArgon2id):
		mem, _ := kdf.uint64("M")
		iter, _ := kdf.uint64("I")
		par, _ := kdf.uint64("P")
		memKiB := mem / 1024
		switch {
		case memKiB > MaxArgon2MemoryKiB:
			return fmt.Errorf("kdbx: Argon2 asks for %d MiB, refusing above %d", memKiB>>10, MaxArgon2MemoryKiB>>10)
		case iter > MaxArgon2Iterations:
			return fmt.Errorf("kdbx: Argon2 asks for %d passes, refusing above %d", iter, MaxArgon2Iterations)
		case par > MaxArgon2Lanes:
			return fmt.Errorf("kdbx: Argon2 asks for %d lanes, refusing above %d", par, MaxArgon2Lanes)
		case (memKiB>>10)*iter > MaxArgon2Work:
			return fmt.Errorf("kdbx: Argon2 asks for %d MiB × %d passes, refusing above %d MiB-passes",
				memKiB>>10, iter, MaxArgon2Work)
		}
	}
	return nil
}