fhe-coprocessor/pkg/fhe/adapter.go

// Package fhe provides a coprocessor-facing adapter over the canonical
// github.com/luxfi/fhe homomorphic-encryption library.
//
// luxfi/fhe is a lattigo-style library built from boolean gates: an integer is
// carried as a *fhe.BitCiphertext (a vector of encrypted bits that records its
// own width), homomorphic operations are methods on a *fhe.BitwiseEvaluator,
// and the evaluator is constructed from public Parameters plus a
// *fhe.BootstrapKey — the public evaluation key, which is the FHE analogue of a
// classic "server key". This adapter wraps that API behind the byte-oriented
// contract the worker relies on: ciphertexts travel as marshaled bytes and
// operations are selected by OperationType.
//
// BitCiphertext is the integer carrier here (rather than RadixCiphertext)
// because it is the only multi-bit ciphertext luxfi/fhe can MarshalBinary /
// UnmarshalBinary, and the adapter's []byte contract requires serialization.
// The evaluator never holds secret material: the secret-key parameter of
// NewBitwiseEvaluator is deprecated and ignored by the library, so it is passed
// nil here.
//
// Parameters: PN10QP27, the library's ~128-bit-security default. They must
// match the parameters used to generate the BootstrapKey passed to NewAdapter
// and the secret key used by the test-only Encrypt/Decrypt helpers; 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")
)

// 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 public parameters and a bootstrap-key-backed evaluator; it carries no
// secret material.
type Adapter struct {
	params fhe.Parameters
	eval   *fhe.BitwiseEvaluator
}

// NewAdapter creates a new adapter from the serialized public evaluation key
// (a marshaled *fhe.BootstrapKey).
func NewAdapter(serverKeyBytes []byte) (*Adapter, error) {
	params, err := defaultParameters()
	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)
	}

	// The third argument (secret key) is deprecated and ignored by the
	// library; the evaluator operates purely on public material.
	return &Adapter{
		params: params,
		eval:   fhe.NewBitwiseEvaluator(params, bsk, nil),
	}, nil
}

// Ciphertext represents an encrypted integer value. Data is a marshaled
// *fhe.BitCiphertext (which encodes the true bit width); BitLen records the
// nominal integer width carried alongside it.
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 := decodeBits(lhs.Data)
	if err != nil {
		return nil, fmt.Errorf("decode lhs: %w", err)
	}

	var result *fhe.BitCiphertext
	bitLen := lhs.BitLen

	switch op {
	case OpNot, OpNeg:
		result, err = a.executeUnary(op, lhsCt)
	case OpShl, OpShr, OpRotl, OpRotr:
		// luxfi/fhe models shifts as (ciphertext, plaintext-int) and exposes
		// no rotate primitive at all. A plaintext shift amount cannot be
		// honestly recovered from an encrypted right-hand side without the
		// secret key, so reject rather than fabricate a result.
		return nil, fmt.Errorf("%w: op %d needs a plaintext shift amount unavailable in luxfi/fhe", ErrUnsupportedOp, op)
	default:
		if rhs == nil {
			return nil, ErrInvalidCiphertext
		}
		rhsCt, decErr := decodeBits(rhs.Data)
		if decErr != nil {
			return nil, fmt.Errorf("decode rhs: %w", decErr)
		}
		result, err = a.executeBinary(op, lhsCt, rhsCt)
		if isComparison(op) {
			bitLen = 1 // comparisons yield a single encrypted boolean bit
		}
	}

	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: bitLen}, nil
}

func (a *Adapter) executeUnary(op OperationType, ct *fhe.BitCiphertext) (*fhe.BitCiphertext, error) {
	switch op {
	case OpNot:
		return a.eval.Not(ct), nil
	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.BitCiphertext) (*fhe.BitCiphertext, 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 OpMin:
		return a.eval.Min(lhs, rhs)
	case OpMax:
		return a.eval.Max(lhs, rhs)
	case OpEq:
		return boolResult(a.eval.Eq(lhs, rhs))
	case OpNe:
		// luxfi/fhe has no Ne; logically negate the equality bit.
		bit, eqErr := a.eval.Eq(lhs, rhs)
		if eqErr != nil {
			return nil, eqErr
		}
		return a.eval.Not(fhe.WrapBoolCiphertext(bit)), nil
	case OpGt:
		return boolResult(a.eval.Gt(lhs, rhs))
	case OpGe:
		return boolResult(a.eval.Ge(lhs, rhs))
	case OpLt:
		return boolResult(a.eval.Lt(lhs, rhs))
	case OpLe:
		return boolResult(a.eval.Le(lhs, rhs))
	default:
		return nil, fmt.Errorf("%w: binary %d", ErrUnsupportedOp, op)
	}
}

// boolResult wraps a single-bit comparison ciphertext as a 1-bit
// BitCiphertext so it can travel through the byte-oriented contract.
func boolResult(bit *fhe.Ciphertext, err error) (*fhe.BitCiphertext, error) {
	if err != nil {
		return nil, err
	}
	return fhe.WrapBoolCiphertext(bit), nil
}

// isComparison reports whether op yields a single boolean bit.
func isComparison(op OperationType) bool {
	switch op {
	case OpEq, OpNe, OpGt, OpGe, OpLt, OpLe:
		return true
	default:
		return false
	}
}

// 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 := secretParts(clientKeyBytes)
	if err != nil {
		return 0, err
	}

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

	return fhe.NewBitwiseDecryptor(params, sk).DecryptUint64(bc), 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 := secretParts(clientKeyBytes)
	if err != nil {
		return nil, err
	}

	bc := fhe.NewBitwiseEncryptor(params, sk).EncryptUint64(value, fheTypeFor(bitLen))

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

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

// defaultParameters builds the fixed PN10QP27 parameter set.
func defaultParameters() (fhe.Parameters, error) {
	params, err := fhe.NewParametersFromLiteral(fhe.PN10QP27)
	if err != nil {
		return fhe.Parameters{}, fmt.Errorf("fhe parameters: %w", err)
	}
	return params, nil
}

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

// decodeBits unmarshals a BitCiphertext from its byte form.
func decodeBits(data []byte) (*fhe.BitCiphertext, error) {
	if len(data) == 0 {
		return nil, ErrInvalidCiphertext
	}
	bc := new(fhe.BitCiphertext)
	if err := bc.UnmarshalBinary(data); err != nil {
		return nil, err
	}
	return bc, 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
	}
}