data-mesher/pkg/crypto/public_key.go

package crypto

import (
	"crypto/ed25519"
	"crypto/x509"
	"encoding/base64"
	"encoding/json"
	"encoding/pem"
	"errors"
	"fmt"
	"io"
	"os"
	"slices"
	"strings"
	"sync"

	p2pcrypto "github.com/libp2p/go-libp2p/core/crypto"
	"github.com/libp2p/go-libp2p/core/peer"
)

// PublicKey represents an ED25519 public key with added utility methods for encoding, comparison, and verification.
type PublicKey struct {
	key ed25519.PublicKey

	str     string
	strOnce *sync.Once
}

// Bytes returns the byte representation of the public key.
func (p *PublicKey) Bytes() []byte {
	if p == nil {
		return nil
	}

	return p.key
}

// Verify checks if the provided signature is valid for the given message using the public key.
// Returns true if valid.
func (p *PublicKey) Verify(msg []byte, sig []byte) bool {
	return ed25519.Verify(p.key, msg, sig)
}

// String returns the base64-encoded string representation of the public key.
// The result is computed once and cached for future calls.
func (p *PublicKey) String() string {
	if p == nil {
		return ""
	}

	p.strOnce.Do(func() {
		p.str = base64.StdEncoding.EncodeToString(p.key)
	})

	return p.str
}

// URLEncoded returns a filesystem-safe representation of the public key using unpadded base64url encoding (RFC 4648).
// The result contains only alphanumeric characters, hyphens, and underscores.
func (p *PublicKey) URLEncoded() string {
	if p == nil {
		return ""
	}

	return base64.RawURLEncoding.EncodeToString(p.key)
}

// PEM returns the PEM-encoded representation of the public key.
func (p *PublicKey) PEM() ([]byte, error) {
	der, err := x509.MarshalPKIXPublicKey(p.key)
	if err != nil {
		return nil, fmt.Errorf("failed to marshal public key: %w", err)
	}

	return pem.EncodeToMemory(&pem.Block{
		Type:  "PUBLIC KEY",
		Bytes: der,
	}), nil
}

// WriteToFile writes the PEM-encoded representation of the public key to the specified file path.
func (p *PublicKey) WriteToFile(path string) error {
	bytes, err := p.PEM()
	if err != nil {
		return fmt.Errorf("failed to encode public key to PEM: %w", err)
	}

	if err := os.WriteFile(path, bytes, 0o600); err != nil {
		return fmt.Errorf("failed to write public key to file: %w", err)
	}

	return nil
}

// ToLibP2P converts the public key to a libp2p crypto public key.
func (p *PublicKey) ToLibP2P() (p2pcrypto.PubKey, error) {
	result, err := p2pcrypto.UnmarshalEd25519PublicKey(p.key)
	if err != nil {
		return nil, fmt.Errorf("failed to convert key: %w", err)
	}

	return result, nil
}

// PeerID derives the libp2p peer ID from this public key.
func (p *PublicKey) PeerID() (peer.ID, error) {
	key, err := p.ToLibP2P()
	if err != nil {
		return "", err
	}

	id, err := peer.IDFromPublicKey(key)
	if err != nil {
		return "", fmt.Errorf("failed to derive peer ID: %w", err)
	}

	return id, nil
}

// Equal compares this PublicKey with another PublicKey and returns true if they represent the same key, false
// otherwise.
func (p *PublicKey) Equal(other *PublicKey) bool {
	if p == nil || other == nil {
		return p == other
	}

	return slices.Equal(p.key, other.key)
}

// MarshalJSON serializes the PublicKey to a JSON-encoded base64 representation of the key.
// It returns the serialized data or an error.
func (p *PublicKey) MarshalJSON() ([]byte, error) {
	buf, err := json.Marshal(p.String())
	if err != nil {
		return nil, fmt.Errorf("failed to marshal public key: %w", err)
	}

	return buf, nil
}

// UnmarshalJSON implements the json.Unmarshaler interface for the PublicKey type.
// It decodes a base64-encoded JSON string into the PublicKey's internal representation.
// Returns an error if decoding fails or if the decoded key length is incorrect.
func (p *PublicKey) UnmarshalJSON(b []byte) error {
	var (
		err     error
		encoded string
	)

	if err = json.Unmarshal(b, &encoded); err != nil {
		return fmt.Errorf("failed to unmarshal public key: %w", err)
	}

	// decode the base64 string
	if p.key, err = base64.StdEncoding.DecodeString(encoded); err != nil {
		return fmt.Errorf("failed to decode public key: %w", err)
	}

	// check the decoded bytes are the correct length
	if len(p.key) != ed25519.PublicKeySize {
		return fmt.Errorf("decoded public key has %d bytes, expected %d", len(p.key), ed25519.PublicKeySize)
	}

	// pre-set the string value since we have it here
	p.strOnce = &sync.Once{}
	p.strOnce.Do(func() {
		p.str = encoded
	})

	return nil
}

// MarshalBinary implements encoding.BinaryMarshaler for PublicKey.
// Returns the raw 32-byte Ed25519 public key.
func (p *PublicKey) MarshalBinary() ([]byte, error) {
	result := make([]byte, ed25519.PublicKeySize)
	copy(result, p.key)

	return result, nil
}

// UnmarshalBinary implements encoding.BinaryUnmarshaler for PublicKey.
// Expects exactly 32 bytes of raw Ed25519 public key data.
func (p *PublicKey) UnmarshalBinary(b []byte) error {
	if len(b) != ed25519.PublicKeySize {
		return fmt.Errorf("public key has %d bytes, expected %d", len(b), ed25519.PublicKeySize)
	}

	p.key = make([]byte, ed25519.PublicKeySize)
	copy(p.key, b)
	p.strOnce = &sync.Once{}

	return nil
}

// newPublicKey creates a new PublicKey instance from the provided byte slice if it satisfies the required length.
// Returns an error if the byte slice length does not match the expected ed25519.PublicKeySize.
func newPublicKey(buf []byte) (*PublicKey, error) {
	if len(buf) != ed25519.PublicKeySize {
		return nil, fmt.Errorf("buf has %d bytes, expected %d", len(buf), ed25519.PublicKeySize)
	}

	return &PublicKey{
		key:     buf,
		strOnce: &sync.Once{},
	}, nil
}

// PublicKeyList represents a list of pointers to PublicKey objects, enabling batch operations on multiple keys.
type PublicKeyList []*PublicKey

// String returns a comma-separated string representation of the PublicKeyList.
// Returns an empty string if the list is empty.
func (p PublicKeyList) String() string {
	if len(p) == 0 {
		return ""
	}

	var b strings.Builder
	for _, key := range p {
		b.WriteString(",")
		b.WriteString(key.String())
	}

	return b.String()[1:]
}

// LoadPublicKey loads a public key from a file path or parses it from a base64/PEM string.
func LoadPublicKey(pathOrData string) (*PublicKey, error) {
	// first try to read as a file
	data, err := os.ReadFile(pathOrData) //nolint:gosec
	if err == nil {
		// file exists, parse its contents
		return ParsePublicKey(data)
	}

	// not a file, try to parse as base64/PEM string
	return ParsePublicKey([]byte(pathOrData))
}

// ReadPublicKey reads an ED25519 public key from a file at the specified path and returns a PublicKey or an error.
func ReadPublicKey(path string) (*PublicKey, error) {
	f, err := os.Open(path) //nolint:gosec
	if err != nil {
		return nil, fmt.Errorf("failed to open public key file: %w", err)
	}
	defer f.Close() //nolint:errcheck

	buf, err := io.ReadAll(f)
	if err != nil {
		return nil, fmt.Errorf("failed to read public key: %w", err)
	}

	return ParsePublicKey(buf)
}

// ParsePublicKey parses a byte slice into a PublicKey.
// It wraps data in a PEM block if not already formatted as one.
// It ensures the key is ED25519 and returns an error for other types, invalid PEM, or incorrect key format.
func ParsePublicKey(buf []byte) (*PublicKey, error) {
	var parsed ed25519.PublicKey

	switch {
	case len(buf) == ed25519.PublicKeySize:
		// raw 32-byte key
		return newPublicKey(buf)
	case len(buf) == 44:
		// we assume we're dealing with a simple base64 encoded key, as 32 bytes encoded in base64 is 44
		parsed = make([]byte, ed25519.PublicKeySize)

		_, err := base64.StdEncoding.Decode(parsed, buf)
		if err != nil {
			return nil, fmt.Errorf("failed to decode public key: %w", err)
		}
	default:
		// otherwise, we assume the key was generated directly with openssl in PEM format
		// wrap in a PEM block if it isn't already
		// this is a convenience allowing users to just copy the main line from the key file
		if !strings.HasPrefix(string(buf), "-----BEGIN PUBLIC KEY-----") {
			buf = []byte("-----BEGIN PUBLIC KEY-----\n" + string(buf) + "\n-----END PUBLIC KEY-----")
		}

		// Decode the PEM block
		block, _ := pem.Decode(buf)
		if block == nil {
			return nil, errors.New("failed to decode PEM block")
		}

		// Ensure the block type is "PUBLIC KEY"
		if block.Type != "PUBLIC KEY" {
			return nil, fmt.Errorf("unexpected block type: %s, expected PUBLIC KEY", block.Type)
		}

		// Parse the public key
		key, err := x509.ParsePKIXPublicKey(block.Bytes)
		if err != nil {
			return nil, fmt.Errorf("failed to parse public key: %w", err)
		}

		// Ensure the key is an ED25519 public key
		ed25519Key, ok := key.(ed25519.PublicKey)
		if !ok {
			return nil, errors.New("key is not an ED25519 public key")
		}

		parsed = ed25519Key
	}

	return newPublicKey(parsed)
}

// ParsePublicKeyURLEncoded decodes an unpadded base64url-encoded string (as produced by URLEncoded) into a PublicKey.
func ParsePublicKeyURLEncoded(encoded string) (*PublicKey, error) {
	buf, err := base64.RawURLEncoding.DecodeString(encoded)
	if err != nil {
		return nil, fmt.Errorf("failed to decode public key dir name: %w", err)
	}

	return newPublicKey(buf)
}

// UnmarshalParam implements the echo.BindUnmarshaler interface for binding from headers/query params.
func (p *PublicKey) UnmarshalParam(param string) error {
	buf, err := base64.StdEncoding.DecodeString(param)
	if err != nil {
		return fmt.Errorf("failed to decode public key: %w", err)
	}

	if len(buf) != ed25519.PublicKeySize {
		return fmt.Errorf("decoded public key has %d bytes, expected %d", len(buf), ed25519.PublicKeySize)
	}

	p.key = buf
	p.strOnce = &sync.Once{}
	p.strOnce.Do(func() {
		p.str = param
	})

	return nil
}