session_server.go

package qtunnul

import (
	"crypto/mlkem"
	"crypto/rand"
	"fmt"
	"sync"

	"golang.org/x/crypto/chacha20poly1305"
)

type ServerSession struct {
	attestationKey       []byte
	sharedSecretReceived []byte
	shaedSecretSendTo    []byte
	cipherText           []byte
	nonce                []byte
	nonceCount           uint64
	nonceCountLock       sync.Mutex
	bitmap               *Bitmap
}

func NewSession(attestationKeySeed []byte) (*ServerSession, error) {
	nonce := make([]byte, chacha20poly1305.NonceSizeX-8)
	if _, err := rand.Read(nonce); err != nil {
		return nil, fmt.Errorf("failed to generate nonce: %w", err)
	}

	return &ServerSession{
		attestationKey: attestationKeySeed,
		bitmap:         &Bitmap{},
		nonce:          nonce,
	}, nil
}

func boxAttestionResponse(cipherText, signature []byte) []byte {
	response := make([]byte, 4+len(cipherText)+len(signature))
	response[0] = byte(len(cipherText) >> 24)
	response[1] = byte(len(cipherText) >> 16)
	response[2] = byte(len(cipherText) >> 8)
	response[3] = byte(len(cipherText))
	copy(response[4:], cipherText)
	copy(response[4+len(cipherText):], signature)
	return response
}

func unboxAttestionResponse(response []byte) (cipherText, signature []byte, err error) {
	if len(response) < 4 {
		return nil, nil, fmt.Errorf("response too short")
	}
	cipherTextLen := (uint32(response[0]) << 24) | (uint32(response[1]) << 16) | (uint32(response[2]) << 8) | uint32(response[3])
	if len(response) < int(4+cipherTextLen) {
		return nil, nil, fmt.Errorf("response too short for ciphertext")
	}
	cipherText = response[4 : 4+cipherTextLen]
	signature = response[4+cipherTextLen:]
	return cipherText, signature, nil
}

func (ss *ServerSession) ProcessAttestationAndKeyExchange(challenge []byte) ([]byte, error) {
	if len(challenge) < mlkem.EncapsulationKeySize1024 {
		return nil, fmt.Errorf("challenge too short")
	}
	encapsulationKey, err := mlkem.NewEncapsulationKey1024(challenge[:mlkem.EncapsulationKeySize1024])
	if err != nil {
		return nil, fmt.Errorf("failed to create encapsulation key: %w", err)
	}

	sharedSecret, cipherText := encapsulationKey.Encapsulate()

	keys, err := DeriveKey(sharedSecret, []byte("qtunnul-session-keys"), 64)
	if err != nil {
		return nil, fmt.Errorf("failed to derive session keys: %w", err)
	}

	ss.sharedSecretReceived = keys[:32]
	ss.shaedSecretSendTo = keys[32:]
	ss.cipherText = cipherText

	signature, err := SignServerAttestation(ss.attestationKey, cipherText)
	if err != nil {
		return nil, fmt.Errorf("failed to sign attestation: %w", err)
	}

	return boxAttestionResponse(cipherText, signature), nil
}

func countToBytes(count uint64) []byte {
	b := make([]byte, 8)
	for i := uint(0); i < 8; i++ {
		b[7-i] = byte(count & 0xFF)
		count >>= 8
	}
	return b
}

func bytesToCount(b []byte) uint64 {
	var count uint64
	for i := 0; i < 8; i++ {
		count <<= 8
		count |= uint64(b[i])
	}
	return count
}

func (ss *ServerSession) NextNonce() ([]byte, error) {
	ss.nonceCountLock.Lock()
	count := ss.nonceCount
	ss.nonceCount++
	if ss.nonceCount == 0 {
		newNonce := make([]byte, chacha20poly1305.NonceSizeX-8)
		if _, err := rand.Read(newNonce); err != nil {
			ss.nonceCountLock.Unlock()
			return nil, fmt.Errorf("failed to generate new nonce: %w", err)
		}
		ss.nonce = newNonce
	}
	ss.nonceCountLock.Unlock()
	countBytes := countToBytes(count)

	nonce := make([]byte, chacha20poly1305.NonceSizeX)
	copy(nonce[:chacha20poly1305.NonceSizeX-8], ss.nonce)
	copy(nonce[chacha20poly1305.NonceSizeX-8:], countBytes)

	return nonce, nil
}

func (ss *ServerSession) Encrypt(plaintext []byte) ([]byte, error) {
	aead, err := chacha20poly1305.NewX(ss.shaedSecretSendTo)
	if err != nil {
		return nil, fmt.Errorf("failed to create AEAD: %w", err)
	}

	nonce, err := ss.NextNonce()
	if err != nil {
		return nil, fmt.Errorf("failed to get next nonce: %w", err)
	}
	ciphertext := aead.Seal(nil, nonce, plaintext, nil)

	result := make([]byte, len(nonce)+len(ciphertext))
	copy(result, nonce)
	copy(result[len(nonce):], ciphertext)

	return result, nil
}

func (ss *ServerSession) Decrypt(data []byte) ([]byte, error) {
	if len(data) < chacha20poly1305.NonceSizeX {
		return nil, fmt.Errorf("data too short")
	}

	nonce, ciphertext := data[:chacha20poly1305.NonceSizeX], data[chacha20poly1305.NonceSizeX:]

	count := bytesToCount(nonce[chacha20poly1305.NonceSizeX-8:])
	if !ss.bitmap.Add(count) {
		return nil, fmt.Errorf("replayed or old packet: %d", count)
	}

	aead, err := chacha20poly1305.NewX(ss.sharedSecretReceived)
	if err != nil {
		return nil, fmt.Errorf("failed to create AEAD: %w", err)
	}

	plaintext, err := aead.Open(nil, nonce, ciphertext, nil)
	if err != nil {
		return nil, fmt.Errorf("decryption failed: %w", err)
	}

	return plaintext, nil
}