pxmon/internal/cluster/transport.go

package cluster

import (
	"context"
	"crypto/sha256"
	"crypto/tls"
	"crypto/x509"
	"encoding/hex"
	"errors"
	"fmt"
	"net"
	"net/http"
	"strconv"
	"strings"
	"time"

	"golang.org/x/crypto/ssh"
)

// agentClient wraps an *http.Client pointed at the correct base URL for a
// cluster's agent (direct or SSH-tunneled).
type agentClient struct {
	http   *http.Client
	target string
	close  func()
}

func (a *agentClient) Close() {
	if a != nil && a.close != nil {
		a.close()
	}
}

// newAgentClient builds the right HTTP client for reaching a cluster's agent.
//
// For TransportDirect it returns a plain client talking to cluster.Host:port.
//
// For TransportIPFabric it reuses a cached ssh.Client from the service pool
// and returns a client whose Transport routes every TCP connection through
// ssh.Client.Dial to 127.0.0.1:port. The SSH connection stays pooled after
// Close() — only the HTTP transport's idle conns are released.
func (s *Service) newAgentClient(ctx context.Context, c Cluster, timeout time.Duration) (*agentClient, error) {
	if c.Agent.Port == 0 {
		return nil, errors.New("agent port is not set")
	}
	if timeout <= 0 {
		timeout = 8 * time.Second
	}

	scheme := agentScheme(c)
	tlsCfg, err := agentTLSConfig(c)
	if err != nil {
		return nil, err
	}

	switch normalizeTransport(c.Transport) {
	case TransportIPFabric:
		sshClient, err := s.acquireTunnelClient(ctx, c)
		if err != nil {
			return nil, err
		}
		tr := &http.Transport{
			TLSClientConfig: tlsCfg,
			DialContext: func(dctx context.Context, network, _ string) (net.Conn, error) {
				addr := net.JoinHostPort("127.0.0.1", strconv.Itoa(c.Agent.Port))
				return sshDialWithContext(dctx, sshClient, network, addr)
			},
			DisableKeepAlives:     true,
			IdleConnTimeout:       30 * time.Second,
			ResponseHeaderTimeout: timeout,
			ExpectContinueTimeout: 1 * time.Second,
		}
		httpClient := &http.Client{
			Timeout:   timeout,
			Transport: tr,
		}
		return &agentClient{
			http:   httpClient,
			target: scheme + "://127.0.0.1:" + strconv.Itoa(c.Agent.Port),
			close: func() {
				tr.CloseIdleConnections()
			},
		}, nil

	default:
		tr := &http.Transport{
			TLSClientConfig:       tlsCfg,
			ResponseHeaderTimeout: timeout,
			ExpectContinueTimeout: 1 * time.Second,
		}
		httpClient := &http.Client{Timeout: timeout, Transport: tr}
		return &agentClient{
			http:   httpClient,
			target: scheme + "://" + net.JoinHostPort(c.Host, strconv.Itoa(c.Agent.Port)),
			close: func() {
				tr.CloseIdleConnections()
			},
		}, nil
	}
}

func agentScheme(c Cluster) string {
	if c.Agent.TLSEnabled {
		return "https"
	}
	return "http"
}

func agentTLSConfig(c Cluster) (*tls.Config, error) {
	if !c.Agent.TLSEnabled {
		return nil, nil
	}
	fp := strings.ToLower(strings.TrimSpace(c.Agent.TLSFingerprint))
	if fp == "" {
		return nil, errors.New("agent TLS is enabled but certificate fingerprint is missing")
	}
	fp = strings.ReplaceAll(fp, ":", "")
	want, err := hex.DecodeString(fp)
	if err != nil {
		return nil, fmt.Errorf("invalid agent TLS fingerprint: %w", err)
	}
	return &tls.Config{
		MinVersion:         tls.VersionTLS12,
		InsecureSkipVerify: true, // verified via explicit fingerprint pinning below
		VerifyPeerCertificate: func(rawCerts [][]byte, _ [][]*x509.Certificate) error {
			if len(rawCerts) == 0 {
				return errors.New("agent TLS: peer certificate is missing")
			}
			sum := sha256.Sum256(rawCerts[0])
			if len(want) != len(sum) {
				return errors.New("agent TLS: fingerprint length mismatch")
			}
			if !hmacEqual(sum[:], want) {
				return errors.New("agent TLS: fingerprint mismatch")
			}
			return nil
		},
	}, nil
}

func hmacEqual(a, b []byte) bool {
	if len(a) != len(b) {
		return false
	}
	var v byte
	for i := 0; i < len(a); i++ {
		v |= a[i] ^ b[i]
	}
	return v == 0
}

// acquireTunnelClient returns a pooled ssh.Client for the cluster, creating
// one if necessary. Credential changes invalidate the cached entry via the
// fingerprint field. Dead clients are evicted lazily: when a DialContext
// through a stale client fails, the caller invokes CloseTunnelClient and the
// next acquire re-dials.
func (s *Service) acquireTunnelClient(ctx context.Context, c Cluster) (*ssh.Client, error) {
	fp := credentialFingerprint(c)

	s.tunnelMu.Lock()
	entry, ok := s.tunnelPool[c.ID]
	if ok && entry.fp != fp {
		_ = entry.client.Close()
		delete(s.tunnelPool, c.ID)
		entry = nil
		ok = false
	}
	if ok {
		s.tunnelMu.Unlock()
		return entry.client, nil
	}
	s.tunnelMu.Unlock()

	client, err := s.dialSSH(ctx, c, "", "")
	if err != nil {
		return nil, err
	}

	s.tunnelMu.Lock()
	if existing, ok := s.tunnelPool[c.ID]; ok && existing.fp == fp {
		// Another goroutine won the race; drop ours.
		s.tunnelMu.Unlock()
		_ = client.Close()
		return existing.client, nil
	}
	s.tunnelPool[c.ID] = &tunneledSSH{
		client: client,
		fp:     fp,
	}
	s.tunnelMu.Unlock()
	return client, nil
}

// CloseTunnelClient drops a pooled SSH tunnel for a cluster. Safe to call if
// no entry exists.
func (s *Service) CloseTunnelClient(clusterID string) {
	s.tunnelMu.Lock()
	entry, ok := s.tunnelPool[clusterID]
	if ok {
		delete(s.tunnelPool, clusterID)
	}
	s.tunnelMu.Unlock()
	if ok && entry != nil && entry.client != nil {
		_ = entry.client.Close()
	}
}

// CloseAllTunnelClients tears down every pooled SSH tunnel.
func (s *Service) CloseAllTunnelClients() {
	s.tunnelMu.Lock()
	pool := s.tunnelPool
	s.tunnelPool = make(map[string]*tunneledSSH)
	s.tunnelMu.Unlock()
	for _, e := range pool {
		if e != nil && e.client != nil {
			_ = e.client.Close()
		}
	}
}

// credentialFingerprint returns a short hash over the fields that affect how
// we'd reconnect. If any of these change we must not reuse a cached client.
func credentialFingerprint(c Cluster) string {
	h := sha256.New()
	h.Write([]byte(c.Host))
	h.Write([]byte{'|'})
	h.Write([]byte(strconv.Itoa(c.Port)))
	h.Write([]byte{'|'})
	h.Write([]byte(c.User))
	h.Write([]byte{'|'})
	h.Write([]byte(c.AuthMethod))
	h.Write([]byte{'|'})
	h.Write([]byte(c.Password))
	h.Write([]byte{'|'})
	h.Write([]byte(c.KeyPath))
	h.Write([]byte{'|'})
	h.Write([]byte(c.KeyPassphrase))
	h.Write([]byte{'|'})
	h.Write([]byte(c.KeyPassphraseFile))
	return hex.EncodeToString(h.Sum(nil)[:8])
}

// sshDialWithContext wraps ssh.Client.Dial so it respects ctx cancellation.
// ssh.Client has no context-aware dial, so we fall back to a watcher goroutine
// that closes the connection if ctx fires before the dial returns.
func sshDialWithContext(ctx context.Context, client *ssh.Client, network, addr string) (net.Conn, error) {
	type result struct {
		conn net.Conn
		err  error
	}
	ch := make(chan result, 1)
	go func() {
		conn, err := client.Dial(network, addr)
		ch <- result{conn: conn, err: err}
	}()

	select {
	case <-ctx.Done():
		go func() {
			r := <-ch
			if r.conn != nil {
				_ = r.conn.Close()
			}
		}()
		return nil, ctx.Err()
	case r := <-ch:
		return r.conn, r.err
	}
}