iop/apps/node/internal/transport/client.go

package transport

import (
	"context"
	"crypto/ecdh"
	"crypto/tls"
	"errors"
	"fmt"
	"net"
	"strconv"
	"time"

	toki "git.toki-labs.com/toki/proto-socket/go"
	"go.uber.org/zap"

	"iop/packages/go/auth"
	"iop/packages/go/config"
	"iop/packages/go/credentiallease"
	iop "iop/proto/gen/iop"
)

// connectFailureClass classifies a DialEdge failure so the connectivity
// supervisor can distinguish transient failures it should retry from fatal ones
// that must terminate the Node.
type connectFailureClass int

const (
	// connectFailureRetryable marks a transient failure (dial refused, network
	// unreachable, timeout, transient register transport failure) that the
	// supervisor should retry under the reconnect policy.
	connectFailureRetryable connectFailureClass = iota
	// connectFailureFatal marks a non-retryable failure (invalid address or
	// authoritative registration rejection) that retrying cannot recover.
	connectFailureFatal
)

// ConnectError wraps a DialEdge failure with a retry classification while
// preserving the underlying cause for errors.Is/errors.As and any transport
// status it carries. The supervisor consumes the classification through
// IsFatalConnectError rather than comparing error strings.
type ConnectError struct {
	class connectFailureClass
	cause error
}

// Error returns the wrapped cause's message unchanged.
func (e *ConnectError) Error() string {
	if e.cause == nil {
		return "transport: connect error"
	}
	return e.cause.Error()
}

// Unwrap exposes the original cause so errors.Is/errors.As reach it.
func (e *ConnectError) Unwrap() error { return e.cause }

// Retryable reports whether the failure is a transient one the supervisor
// should retry.
func (e *ConnectError) Retryable() bool { return e.class == connectFailureRetryable }

func wrapConnectError(class connectFailureClass, cause error) *ConnectError {
	return &ConnectError{class: class, cause: cause}
}

// IsFatalConnectError reports whether err (or any wrapped cause) is a fatal,
// non-retryable connect failure. Errors that are not classified ConnectErrors
// are treated as retryable transient failures.
func IsFatalConnectError(err error) bool {
	var ce *ConnectError
	if errors.As(err, &ce) {
		return !ce.Retryable()
	}
	return false
}

const (
	// Provider tunnel traffic can legitimately delay heartbeat processing while a
	// long prompt is in prefill or a streamed response is applying backpressure.
	// Keep the transport liveness window above those request-local stalls so an
	// otherwise healthy field node is not disconnected mid-generation.
	heartbeatIntervalSec = 30
	// heartbeatWaitSec is kept above heartbeatIntervalSec as defence in depth:
	// the library wait-timer callback already self-heals stale state (see
	// proto-socket go/base_client.go sendHeartBeat), but a larger wait window
	// gives the peer's next heartbeat an extra chance to overwrite any stray
	// timer on slow or jittery links before it fires.
	heartbeatWaitSec    = 45
	tcpWriteTimeout     = 10 * time.Second
	registerTimeout     = 10 * time.Second
	registerInitialWait = 100 * time.Millisecond
	registerRetryWait   = 250 * time.Millisecond
	registerAttempts    = 3
)

type writeDeadlineConn struct {
	net.Conn
	timeout time.Duration
}

func (c *writeDeadlineConn) Write(b []byte) (int, error) {
	if c.timeout > 0 {
		if err := c.Conn.SetWriteDeadline(time.Now().Add(c.timeout)); err != nil {
			return 0, err
		}
		defer c.Conn.SetWriteDeadline(time.Time{})
	}
	return c.Conn.Write(b)
}

// RegisterResult is returned by DialEdge after successful registration.
type RegisterResult struct {
	Session *Session
	NodeID  string
	Alias   string
	Config  *iop.NodeConfigPayload
}

// DialEdge connects to edge, performs the registration handshake, and returns
// a RegisterResult. Call result.Session.SetHandler after creating node.Node.
func DialEdge(ctx context.Context, addr, token string, logger *zap.Logger) (*RegisterResult, error) {
	return DialEdgeTLS(ctx, addr, token, nil, logger)
}

// DialEdgeConfig loads the configured mTLS identity before dialing. Local TLS
// configuration failures are classified fatal so reconnect supervision cannot
// loop forever on an invalid or missing certificate.
func DialEdgeConfig(ctx context.Context, nodeConfig *config.NodeConfig, logger *zap.Logger) (*RegisterResult, error) {
	if nodeConfig == nil {
		return nil, wrapConnectError(connectFailureFatal, fmt.Errorf("transport: node config is required"))
	}
	conf := nodeConfig.Transport
	var tlsConfig *tls.Config
	if conf.TLS.Enabled {
		if err := conf.TLS.ValidateClient("edge"); err != nil {
			return nil, wrapConnectError(connectFailureFatal, fmt.Errorf("transport: tls config: %w", err))
		}
		loaded, err := auth.LoadClientTLSWithIdentity(conf.TLS.Cert, conf.TLS.Key, conf.TLS.CA, conf.TLS.ServerName, conf.TLS.EffectivePeerRole("edge"), conf.TLS.PeerName)
		if err != nil {
			return nil, wrapConnectError(connectFailureFatal, fmt.Errorf("transport: load tls: %w", err))
		}
		tlsConfig = loaded
	}
	var recipientKeyID string
	var recipientPublic []byte
	if nodeConfig.CredentialPlane.Enabled {
		private, err := credentiallease.LoadPrivateKeyFile(nodeConfig.CredentialPlane.RecipientPrivateKey, 32)
		if err != nil {
			return nil, wrapConnectError(connectFailureFatal, fmt.Errorf("transport: load credential recipient key: %w", err))
		}
		defer zeroTransportKey(private)
		key, err := ecdh.X25519().NewPrivateKey(private)
		if err != nil {
			return nil, wrapConnectError(connectFailureFatal, fmt.Errorf("transport: parse credential recipient key: %w", err))
		}
		recipientKeyID = nodeConfig.CredentialPlane.RecipientKeyID
		recipientPublic = key.PublicKey().Bytes()
	}
	return dialEdgeTLSIdentity(ctx, conf.EdgeAddr, conf.Token, tlsConfig, recipientKeyID, recipientPublic, logger)
}

// DialEdgeTLS connects with the supplied TLS configuration. A non-nil config
// is used for the initial connection and every caller-driven reconnect; no
// plaintext fallback is attempted after a TLS failure.
func DialEdgeTLS(ctx context.Context, addr, token string, tlsConfig *tls.Config, logger *zap.Logger) (*RegisterResult, error) {
	return dialEdgeTLSIdentity(ctx, addr, token, tlsConfig, "", nil, logger)
}

func dialEdgeTLSIdentity(ctx context.Context, addr, token string, tlsConfig *tls.Config, recipientKeyID string, recipientPublic []byte, logger *zap.Logger) (*RegisterResult, error) {
	host, portStr, err := net.SplitHostPort(addr)
	if err != nil {
		// Local address preparation error: retrying cannot recover it.
		return nil, wrapConnectError(connectFailureFatal, fmt.Errorf("transport: invalid addr %q: %w", addr, err))
	}
	port, err := strconv.Atoi(portStr)
	if err != nil {
		return nil, wrapConnectError(connectFailureFatal, fmt.Errorf("transport: invalid port %q: %w", portStr, err))
	}
	if port < 1 || port > 65535 {
		// A numeric port outside the TCP range is still a local configuration
		// error. Preserve a standard range cause so callers can inspect it with
		// errors.Is/errors.As instead of relying on the message.
		rangeErr := &strconv.NumError{Func: "Atoi", Num: portStr, Err: strconv.ErrRange}
		return nil, wrapConnectError(connectFailureFatal, fmt.Errorf("transport: invalid port %q: %w", portStr, rangeErr))
	}

	dialer := net.Dialer{KeepAlive: 15 * time.Second}
	var conn net.Conn
	if tlsConfig != nil {
		tlsDialer := tls.Dialer{NetDialer: &dialer, Config: tlsConfig}
		conn, err = tlsDialer.DialContext(ctx, "tcp", net.JoinHostPort(host, strconv.Itoa(port)))
	} else {
		conn, err = dialer.DialContext(ctx, "tcp", net.JoinHostPort(host, strconv.Itoa(port)))
	}
	if err != nil {
		// Edge unavailable / network transient failure: retryable.
		return nil, wrapConnectError(connectFailureRetryable, fmt.Errorf("transport: dial edge %s: %w", addr, err))
	}
	client := toki.NewTcpClient(&writeDeadlineConn{Conn: conn, timeout: tcpWriteTimeout}, heartbeatIntervalSec, heartbeatWaitSec, nodeParserMap())

	resp, err := registerWithEdgeIdentity(ctx, client, token, recipientKeyID, recipientPublic, logger)
	if err != nil {
		_ = client.Close()
		// Register request transport failure or context cancellation: retryable.
		return nil, wrapConnectError(connectFailureRetryable, fmt.Errorf("transport: register: %w", err))
	}
	if !resp.GetAccepted() {
		_ = client.Close()
		// Authoritative registration rejection (bad token/credential): fatal.
		return nil, wrapConnectError(connectFailureFatal, fmt.Errorf("transport: register rejected: %s", resp.GetReason()))
	}

	sess := newSession(client, logger, resp.GetNodeId(), resp.GetAlias())
	logger.Info("registered with edge",
		zap.String("node_id", resp.GetNodeId()),
		zap.String("alias", resp.GetAlias()),
	)
	return &RegisterResult{
		Session: sess,
		NodeID:  resp.GetNodeId(),
		Alias:   resp.GetAlias(),
		Config:  resp.GetConfig(),
	}, nil
}

func registerWithEdge(ctx context.Context, client *toki.TcpClient, token string, logger *zap.Logger) (*iop.RegisterResponse, error) {
	return registerWithEdgeIdentity(ctx, client, token, "", nil, logger)
}

func registerWithEdgeIdentity(ctx context.Context, client *toki.TcpClient, token, recipientKeyID string, recipientPublic []byte, logger *zap.Logger) (*iop.RegisterResponse, error) {
	timer := time.NewTimer(registerInitialWait)
	select {
	case <-ctx.Done():
		timer.Stop()
		return nil, ctx.Err()
	case <-timer.C:
	}

	var lastErr error
	for attempt := 1; attempt <= registerAttempts; attempt++ {
		resp, err := toki.SendRequestTyped[*iop.RegisterRequest, *iop.RegisterResponse](
			&client.Communicator,
			&iop.RegisterRequest{Token: token, CredentialRecipientKeyId: recipientKeyID, CredentialRecipientPublicKey: append([]byte(nil), recipientPublic...)},
			registerTimeout,
		)
		if err == nil {
			return resp, nil
		}
		lastErr = err
		if attempt == registerAttempts || !client.IsAlive() {
			break
		}
		if logger != nil {
			logger.Warn("register request failed, retrying",
				zap.Int("attempt", attempt),
				zap.Error(err),
			)
		}
		timer := time.NewTimer(registerRetryWait)
		select {
		case <-ctx.Done():
			timer.Stop()
			return nil, ctx.Err()
		case <-timer.C:
		}
	}
	return nil, lastErr
}

func zeroTransportKey(value []byte) {
	for i := range value {
		value[i] = 0
	}
}