data-mesher/pkg/libp2p/security/tls/transport.go

// Package tls implements a custom TLS security transport for libp2p that extends the standard libp2p TLS handshake with
// data-mesher CA certificate verification.
// It restricts cluster membership to nodes holding a certificate signed by a trusted CA.
//
// Adapted from https://github.com/libp2p/go-libp2p/blob/8a6fd5b923463676080a96c8429ce461c9b264bb/p2p/security/tls
package tls

import (
	"context"
	"crypto/tls"
	"errors"
	"fmt"
	"net"
	"os"
	"runtime/debug"

	dmcrypto "git.clan.lol/clan/data-mesher/pkg/crypto"
	ci "github.com/libp2p/go-libp2p/core/crypto"
	"github.com/libp2p/go-libp2p/core/network"
	"github.com/libp2p/go-libp2p/core/peer"
	"github.com/libp2p/go-libp2p/core/peerstore"
	"github.com/libp2p/go-libp2p/core/protocol"
	"github.com/libp2p/go-libp2p/core/sec"
	tptu "github.com/libp2p/go-libp2p/p2p/net/upgrader"
)

// ID is the protocol ID used when negotiating with multistream.
const ID = "/clan/data-mesher/tls/1.0.0"

// Transport constructs secure communication sessions for a peer, requiring CA-signed data-mesher certificates in
// addition to the standard libp2p TLS key binding.
type Transport struct {
	identity *Identity

	localPeer  peer.ID
	peerstore  peerstore.Peerstore
	privKey    ci.PrivKey
	muxers     []protocol.ID
	protocolID protocol.ID
}

var _ sec.SecureTransport = &Transport{}

// New returns a constructor closure compatible with libp2p.Security().
// The closure captures the CA keys and identity certificate, producing a Transport that verifies data-mesher
// certificates during TLS handshakes.
func New(
	caKeys []*dmcrypto.PublicKey,
	identityCert *dmcrypto.Certificate,
	peerstore peerstore.Peerstore,
) func(id protocol.ID, key ci.PrivKey, muxers []tptu.StreamMuxer) (*Transport, error) {
	return func(id protocol.ID, key ci.PrivKey, muxers []tptu.StreamMuxer) (*Transport, error) {
		localPeer, err := peer.IDFromPrivateKey(key)
		if err != nil {
			return nil, fmt.Errorf("failed to derive peer ID from private key: %w", err)
		}

		muxerIDs := make([]protocol.ID, 0, len(muxers))
		for _, m := range muxers {
			muxerIDs = append(muxerIDs, m.ID)
		}

		identity, err := NewIdentity(key, identityCert, caKeys)
		if err != nil {
			return nil, err
		}

		return &Transport{
			identity:   identity,
			protocolID: id,
			localPeer:  localPeer,
			peerstore:  peerstore,
			privKey:    key,
			muxers:     muxerIDs,
		}, nil
	}
}

// SecureInbound runs the TLS handshake as a server.
// If p is empty, connections from any peer are accepted.
func (t *Transport) SecureInbound(ctx context.Context, insecure net.Conn, p peer.ID) (sec.SecureConn, error) {
	config, keyCh := t.identity.ConfigForPeer(p, t.peerstore)

	muxers := make([]string, 0, len(t.muxers))
	for _, muxer := range t.muxers {
		muxers = append(muxers, string(muxer))
	}

	// TLS' ALPN selection prefers the server's list. We want to prefer the client's preference, so we reorder in
	// GetConfigForClient.
	getConfigForClient := config.GetConfigForClient

	config.GetConfigForClient = func(info *tls.ClientHelloInfo) (*tls.Config, error) {
	alpnLoop:
		for _, proto := range info.SupportedProtos {
			for _, m := range muxers {
				if m == proto {
					config.NextProtos = []string{proto}
					break alpnLoop
				}
			}
		}

		if getConfigForClient != nil {
			return getConfigForClient(info)
		}

		return config, nil
	}

	config.NextProtos = append(muxers, config.NextProtos...)

	cs, err := t.handshake(ctx, tls.Server(insecure, config), keyCh)
	if err != nil {
		_ = insecure.Close()
	}

	return cs, err
}

// SecureOutbound runs the TLS handshake as a client.
func (t *Transport) SecureOutbound(ctx context.Context, insecure net.Conn, p peer.ID) (sec.SecureConn, error) {
	config, keyCh := t.identity.ConfigForPeer(p, t.peerstore)

	muxers := make([]string, 0, len(t.muxers))
	for _, muxer := range t.muxers {
		muxers = append(muxers, string(muxer))
	}

	config.NextProtos = append(muxers, config.NextProtos...)

	cs, err := t.handshake(ctx, tls.Client(insecure, config), keyCh)
	if err != nil {
		_ = insecure.Close()
	}

	return cs, err
}

// ID returns the protocol ID for this transport.
func (t *Transport) ID() protocol.ID {
	return t.protocolID
}

func (t *Transport) handshake(
	ctx context.Context,
	tlsConn *tls.Conn,
	keyCh <-chan ci.PubKey,
) (_ sec.SecureConn, err error) {
	defer func() {
		if recoverErr := recover(); recoverErr != nil {
			_, _ = fmt.Fprintf(os.Stderr, "panic in TLS handshake: %s\n%s\n", recoverErr, debug.Stack())
			err = fmt.Errorf("panic in TLS handshake: %s", recoverErr)
		}
	}()

	if err = tlsConn.HandshakeContext(ctx); err != nil {
		return nil, fmt.Errorf("TLS handshake failed: %w", err)
	}

	var remotePubKey ci.PubKey
	select {
	case remotePubKey = <-keyCh:
	default:
	}

	if remotePubKey == nil {
		return nil, errors.New("expected remote pub key to be set after handshake")
	}

	return t.setupConn(tlsConn, remotePubKey)
}

func (t *Transport) setupConn(
	tlsConn *tls.Conn,
	remotePubKey ci.PubKey,
) (sec.SecureConn, error) {
	remotePeerID, err := peer.IDFromPublicKey(remotePubKey)
	if err != nil {
		return nil, fmt.Errorf("failed to derive peer ID from public key: %w", err)
	}

	nextProto := tlsConn.ConnectionState().NegotiatedProtocol
	if nextProto == "libp2p" {
		nextProto = ""
	}

	return &conn{
		Conn:         tlsConn,
		localPeer:    t.localPeer,
		remotePeer:   remotePeerID,
		remotePubKey: remotePubKey,
		connectionState: network.ConnectionState{
			StreamMultiplexer:         protocol.ID(nextProto),
			UsedEarlyMuxerNegotiation: nextProto != "",
		},
	}, nil
}