data-mesher/pkg/crypto/private_key.go

package crypto

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

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

// PrivateKey represents a cryptographic private key using the ED25519 algorithm.
// It contains the corresponding PublicKey for signature verification.
type PrivateKey struct {
	key ed25519.PrivateKey

	Public *PublicKey
}

func (p *PrivateKey) ToLibP2P() (p2pcrypto.PrivKey, error) {
	result, err := p2pcrypto.UnmarshalEd25519PrivateKey(p.key)
	if err != nil {
		return nil, fmt.Errorf("failed to convert key: %w", err)
	}

	return result, nil
}

func (p *PrivateKey) PeerID() (peer.ID, error) {
	key, err := p.ToLibP2P()
	if err != nil {
		return "", err
	}

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

	return id, nil
}

// Sign generates an ED25519 signature for the given message using the private key.
// Returns the generated signature.
func (p *PrivateKey) Sign(msg []byte) []byte {
	return ed25519.Sign(p.key, msg)
}

// String returns the base64 representation of the associated public key to prevent exposing sensitive private key data.
func (p *PrivateKey) String() string {
	// we return the string value of the public key to avoid leaking secrets
	return p.Public.String()
}

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

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

// UnmarshalJSON implements the json.Unmarshaler interface for deserializing a PrivateKey from a JSON-encoded value.
// UnmarshalJSON accepts a base64-encoded string, decodes it, and validates the private key length.
func (p *PrivateKey) UnmarshalJSON(b []byte) error {
	var (
		err     error
		encoded string
	)

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

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

	if len(p.key) != ed25519.PrivateKeySize {
		return fmt.Errorf("decoded private key has %d bytes, expected %d", len(p.key), ed25519.PrivateKeySize)
	}

	return nil
}

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

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

	return nil
}

// newPrivateKey creates a new PrivateKey instance from the provided byte slice if it meets the required length.
// Returns an error if the byte slice length is invalid or if a PublicKey cannot be created.
func newPrivateKey(buf []byte) (*PrivateKey, error) {
	if len(buf) != ed25519.PrivateKeySize {
		return nil, fmt.Errorf("buf has %d bytes, expected %d", len(buf), ed25519.PrivateKeySize)
	}

	pub, err := newPublicKey(buf[32:])
	if err != nil {
		return nil, fmt.Errorf("failed to create public key from private key: %w", err)
	}

	return &PrivateKey{
		key:    buf,
		Public: pub,
	}, nil
}

// GenerateKey generates a new ED25519 private key using the provided random source.
// It returns a pointer to the created PrivateKey or an error if key generation fails.
func GenerateKey(r io.Reader) (*PrivateKey, error) {
	_, privateKey, err := ed25519.GenerateKey(r)
	if err != nil {
		return nil, fmt.Errorf("failed to generate keypair: %w", err)
	}

	return newPrivateKey(privateKey)
}

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

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

	return ParsePrivateKey(buf)
}

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

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

// ParsePrivateKey parses a PEM-encoded ED25519 private key and returns a PrivateKey instance or an error.
//
//nolint:nestif
func ParsePrivateKey(buf []byte) (*PrivateKey, error) {
	var parsed ed25519.PrivateKey

	if len(buf) == 88 {
		// we assume we're dealing with a simpled base64 encoded key, as 64 bytes in base64 is 88
		parsed = make([]byte, ed25519.PrivateKeySize)

		_, err := base64.StdEncoding.Decode(parsed, buf)
		if err != nil {
			return nil, fmt.Errorf("failed to decode private key: %w", err)
		}
	} else {
		// 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 PRIVATE KEY-----") {
			buf = []byte("-----BEGIN PRIVATE KEY-----\n" + string(buf) + "\n-----END PRIVATE KEY-----")
		}

		block, _ := pem.Decode(buf)

		if block == nil {
			return nil, errors.New("failed to decode PEM block")
		}

		if block.Type != "PRIVATE KEY" {
			return nil, fmt.Errorf("unexpected block type: %s, expected PRIVATE KEY", block.Type)
		}

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

		// Type assert to ED25519 private key
		ed25519Key, ok := key.(ed25519.PrivateKey)
		if !ok {
			return nil, errors.New("key is not an ED25519 private key")
		}

		parsed = ed25519Key
	}

	return newPrivateKey(parsed)
}