fhe-coprocessor/pkg/fhe/adapter.go

// Package fhe provides a coprocessor-facing adapter over the canonical
// github.com/luxfi/fhe homomorphic-encryption library.
//
// The canonical library exposes a radix-integer API: encryption produces a
// *fhe.RadixCiphertext, and homomorphic operations are methods on a
// *fhe.IntegerEvaluator constructed from public parameters plus a
// *fhe.BootstrapKey (the public evaluation key — the FHE analogue of a
// "server key"). This adapter wraps that API behind the byte-oriented contract
// the worker relies on: ciphertexts travel as length-prefixed bytes and
// operations are selected by OperationType.
//
// Parameter set: PN10QP27 with 2-bit radix blocks. This is the repository-wide
// default (every cmd/* binary and the precompile use it) and must match the
// parameters used to generate the BootstrapKey passed to NewAdapter; ciphertext
// bytes produced under different parameters will not evaluate correctly.
package fhe

import (
	"context"
	"errors"
	"fmt"

	fhe "github.com/luxfi/fhe"
)

// Common errors.
var (
	ErrInvalidCiphertext = errors.New("invalid ciphertext")
	ErrOperationFailed   = errors.New("FHE operation failed")
	ErrKeyNotFound       = errors.New("key not found")
	ErrUnsupportedOp     = errors.New("unsupported operation")
)

// radixBlockBits is the radix block width used across the coprocessor. It must
// match the parameters used to generate the bootstrap and secret keys.
const radixBlockBits = 2

// OperationType defines supported FHE operations.
type OperationType uint8

const (
	OpAdd OperationType = iota
	OpSub
	OpMul
	OpNot
	OpAnd
	OpOr
	OpXor
	OpEq
	OpNe
	OpGt
	OpGe
	OpLt
	OpLe
	OpShl
	OpShr
	OpRotl
	OpRotr
	OpMin
	OpMax
	OpNeg
)

// Adapter wraps canonical FHE integer operations with context support. It holds
// the integer parameters and a bootstrap-key-backed evaluator; it carries no
// secret material.
type Adapter struct {
	params *fhe.IntegerParams
	eval   *fhe.IntegerEvaluator
}

// NewAdapter creates a new adapter from the serialized public evaluation key
// (a marshaled *fhe.BootstrapKey).
func NewAdapter(serverKeyBytes []byte) (*Adapter, error) {
	params, err := integerParams()
	if err != nil {
		return nil, err
	}

	bsk := new(fhe.BootstrapKey)
	if err := bsk.UnmarshalBinary(serverKeyBytes); err != nil {
		return nil, fmt.Errorf("deserialize bootstrap key: %w", err)
	}

	return &Adapter{
		params: params,
		eval:   fhe.NewIntegerEvaluator(params, bsk),
	}, nil
}

// Ciphertext represents an encrypted integer value. Data is a marshaled
// *fhe.RadixCiphertext; BitLen records the integer width (used to choose the
// FHE type when (re)constructing the radix ciphertext).
type Ciphertext struct {
	Data   []byte
	BitLen uint8
}

// Execute performs an FHE operation on one or two ciphertexts.
func (a *Adapter) Execute(ctx context.Context, op OperationType, lhs, rhs *Ciphertext) (*Ciphertext, error) {
	select {
	case <-ctx.Done():
		return nil, ctx.Err()
	default:
	}

	if lhs == nil {
		return nil, ErrInvalidCiphertext
	}

	lhsCt, err := decodeRadix(lhs.Data)
	if err != nil {
		return nil, fmt.Errorf("decode lhs: %w", err)
	}

	var result *fhe.RadixCiphertext

	switch op {
	case OpNot, OpNeg:
		result, err = a.executeUnary(op, lhsCt)
	case OpShl, OpShr, OpRotl, OpRotr:
		// The canonical API models shifts/rotates as a ciphertext plus a
		// plaintext shift amount (Shl/Shr(a, shift int)); rotates are not
		// implemented at all. None can be honestly expressed against an
		// opaque ciphertext right-hand side, so reject rather than fabricate
		// a shift derived from a second ciphertext.
		return nil, fmt.Errorf("%w: %v requires a plaintext shift amount", ErrUnsupportedOp, op)
	default:
		if rhs == nil {
			return nil, ErrInvalidCiphertext
		}
		rhsCt, decErr := decodeRadix(rhs.Data)
		if decErr != nil {
			return nil, fmt.Errorf("decode rhs: %w", decErr)
		}
		result, err = a.executeBinary(op, lhsCt, rhsCt)
	}

	if err != nil {
		return nil, fmt.Errorf("%w: %v", ErrOperationFailed, err)
	}

	data, err := result.MarshalBinary()
	if err != nil {
		return nil, fmt.Errorf("marshal result: %w", err)
	}

	return &Ciphertext{Data: data, BitLen: lhs.BitLen}, nil
}

func (a *Adapter) executeUnary(op OperationType, ct *fhe.RadixCiphertext) (*fhe.RadixCiphertext, error) {
	switch op {
	case OpNot:
		return a.eval.Not(ct)
	case OpNeg:
		return a.eval.Neg(ct)
	default:
		return nil, fmt.Errorf("%w: unary %d", ErrUnsupportedOp, op)
	}
}

func (a *Adapter) executeBinary(op OperationType, lhs, rhs *fhe.RadixCiphertext) (*fhe.RadixCiphertext, error) {
	switch op {
	case OpAdd:
		return a.eval.Add(lhs, rhs)
	case OpSub:
		return a.eval.Sub(lhs, rhs)
	case OpMul:
		return a.eval.Mul(lhs, rhs)
	case OpAnd:
		return a.eval.And(lhs, rhs)
	case OpOr:
		return a.eval.Or(lhs, rhs)
	case OpXor:
		return a.eval.Xor(lhs, rhs)
	case OpEq:
		return a.eval.Eq(lhs, rhs)
	case OpNe:
		return a.eval.Ne(lhs, rhs)
	case OpGt:
		return a.eval.Gt(lhs, rhs)
	case OpGe:
		return a.eval.Ge(lhs, rhs)
	case OpLt:
		return a.eval.Lt(lhs, rhs)
	case OpLe:
		return a.eval.Le(lhs, rhs)
	case OpMin:
		return a.eval.Min(lhs, rhs)
	case OpMax:
		return a.eval.Max(lhs, rhs)
	default:
		return nil, fmt.Errorf("%w: binary %d", ErrUnsupportedOp, op)
	}
}

// Decrypt decrypts a ciphertext using the client key (for testing only). The
// clientKeyBytes argument is a marshaled *fhe.SecretKey.
func Decrypt(clientKeyBytes []byte, ct *Ciphertext) (uint64, error) {
	params, sk, err := decryptorParts(clientKeyBytes)
	if err != nil {
		return 0, err
	}

	rc, err := decodeRadix(ct.Data)
	if err != nil {
		return 0, fmt.Errorf("decode ciphertext: %w", err)
	}

	return fhe.NewIntegerDecryptor(params, sk).DecryptUint64(rc), nil
}

// Encrypt encrypts a plaintext value (for testing only). The clientKeyBytes
// argument is a marshaled *fhe.SecretKey.
func Encrypt(clientKeyBytes []byte, value uint64, bitLen uint8) (*Ciphertext, error) {
	params, sk, err := decryptorParts(clientKeyBytes)
	if err != nil {
		return nil, err
	}

	rc, err := fhe.NewIntegerEncryptor(params, sk).EncryptUint64(value, fheTypeFor(bitLen))
	if err != nil {
		return nil, fmt.Errorf("encrypt: %w", err)
	}

	data, err := rc.MarshalBinary()
	if err != nil {
		return nil, fmt.Errorf("marshal: %w", err)
	}

	return &Ciphertext{Data: data, BitLen: bitLen}, nil
}

// integerParams builds the fixed PN10QP27 / 2-bit-block integer parameters.
func integerParams() (*fhe.IntegerParams, error) {
	base, err := fhe.NewParametersFromLiteral(fhe.PN10QP27)
	if err != nil {
		return nil, fmt.Errorf("fhe parameters: %w", err)
	}
	params, err := fhe.NewIntegerParams(base, radixBlockBits)
	if err != nil {
		return nil, fmt.Errorf("integer parameters: %w", err)
	}
	return params, nil
}

// decryptorParts builds the integer parameters and secret key used by the
// test-only Encrypt/Decrypt helpers.
func decryptorParts(clientKeyBytes []byte) (*fhe.IntegerParams, *fhe.SecretKey, error) {
	params, err := integerParams()
	if err != nil {
		return nil, nil, err
	}
	sk := new(fhe.SecretKey)
	if err := sk.UnmarshalBinary(clientKeyBytes); err != nil {
		return nil, nil, fmt.Errorf("deserialize secret key: %w", err)
	}
	return params, sk, nil
}

// decodeRadix unmarshals a radix ciphertext from its byte form.
func decodeRadix(data []byte) (*fhe.RadixCiphertext, error) {
	if len(data) == 0 {
		return nil, ErrInvalidCiphertext
	}
	rc := new(fhe.RadixCiphertext)
	if err := rc.UnmarshalBinary(data); err != nil {
		return nil, err
	}
	return rc, nil
}

// fheTypeFor maps an integer bit width to the canonical FheUintType, rounding
// up to the next supported width.
func fheTypeFor(bitLen uint8) fhe.FheUintType {
	switch {
	case bitLen <= 1:
		return fhe.FheBool
	case bitLen <= 4:
		return fhe.FheUint4
	case bitLen <= 8:
		return fhe.FheUint8
	case bitLen <= 16:
		return fhe.FheUint16
	case bitLen <= 32:
		return fhe.FheUint32
	case bitLen <= 64:
		return fhe.FheUint64
	case bitLen <= 128:
		return fhe.FheUint128
	default:
		return fhe.FheUint256
	}
}