kedge/internal/shellstream/handshake.go

package shellstream

import (
	"context"
	"fmt"
	"net"

	"go.uber.org/zap"

	"github.com/guildhouse-co/kedge/internal/config"
	"github.com/guildhouse-co/kedge/internal/quartermaster"
)

// Listener accepts incoming Shellstream connections and performs the
// 3-way attestation handshake (ATTEST-INIT → ATTEST-VERIFY → ATTEST-CONFIRM).
type Listener struct {
	cfg config.ShellstreamConfig
	qm  *quartermaster.Client
	log *zap.SugaredLogger
}

// NewListener creates a new Shellstream handshake listener.
func NewListener(cfg config.ShellstreamConfig, qm *quartermaster.Client, log *zap.SugaredLogger) *Listener {
	return &Listener{cfg: cfg, qm: qm, log: log}
}

// Run starts listening for incoming Shellstream connections.
func (l *Listener) Run(ctx context.Context) error {
	listener, err := net.Listen("tcp", l.cfg.ListenAddr)
	if err != nil {
		return fmt.Errorf("failed to listen on %s: %w", l.cfg.ListenAddr, err)
	}
	defer listener.Close()

	l.log.Infof("shellstream listener started on %s", l.cfg.ListenAddr)

	go func() {
		<-ctx.Done()
		listener.Close()
	}()

	for {
		conn, err := listener.Accept()
		if err != nil {
			select {
			case <-ctx.Done():
				return nil
			default:
				l.log.Warnw("accept error", "error", err)
				continue
			}
		}

		go l.handleConnection(ctx, conn)
	}
}

func (l *Listener) handleConnection(ctx context.Context, conn net.Conn) {
	defer conn.Close()

	// Phase 1: Receive ATTEST-INIT from remote.
	initMsg, err := readAttestInit(conn)
	if err != nil {
		l.log.Warnw("failed to read ATTEST-INIT", "remote", conn.RemoteAddr(), "error", err)
		return
	}

	// Validate the SAT token.
	if err := ValidateSAT(initMsg.SATToken, l.cfg.TrustBundlePath); err != nil {
		l.log.Warnw("SAT validation failed", "remote", conn.RemoteAddr(), "error", err)
		return
	}

	// Evaluate capability request against local accord policy.
	grant, err := EvaluateCapability(initMsg.Capabilities, l.cfg.AccordPolicyPath)
	if err != nil {
		l.log.Warnw("capability evaluation failed", "remote", conn.RemoteAddr(), "error", err)
		return
	}

	// Phase 2: Send ATTEST-VERIFY with granted capabilities.
	if err := writeAttestVerify(conn, grant); err != nil {
		l.log.Warnw("failed to write ATTEST-VERIFY", "error", err)
		return
	}

	// Phase 3: Receive ATTEST-CONFIRM.
	if err := readAttestConfirm(conn); err != nil {
		l.log.Warnw("failed to read ATTEST-CONFIRM", "error", err)
		return
	}

	// Record session transit artifact.
	artifact := quartermaster.SessionTransitArtifact{
		SourceCluster:     initMsg.SourceCluster,
		DestCluster:       l.cfg.ClusterID,
		TargetDevice:      initMsg.TargetDevice,
		GrantedMode:       grant.Mode,
		GrantedOperations: grant.Operations,
	}
	if err := l.qm.SubmitSessionTransit(ctx, &artifact); err != nil {
		l.log.Warnw("failed to submit session transit", "error", err)
	}

	l.log.Infow("session established",
		"remote", conn.RemoteAddr(),
		"mode", grant.Mode,
		"target", initMsg.TargetDevice,
	)
}

// AttestInitMsg represents the ATTEST-INIT message from the remote peer.
type AttestInitMsg struct {
	SATToken      []byte
	SourceCluster string
	TargetDevice  string
	TargetSubnet  string
	Capabilities  CapabilityRequest
}

// CapabilityRequest describes what mode and operations the remote peer is requesting.
type CapabilityRequest struct {
	Mode       string   // "overlay", "underlay", or "both"
	Targets    []string // Target device addresses
	Operations []string // "read", "mutate"
}

// CapabilityGrant describes what was actually granted after accord evaluation.
type CapabilityGrant struct {
	Mode       string
	Operations []string
}

func readAttestInit(conn net.Conn) (*AttestInitMsg, error) {
	// TODO: Implement Shellstream wire protocol parsing.
	_ = conn
	return nil, fmt.Errorf("not yet implemented")
}

func writeAttestVerify(conn net.Conn, grant *CapabilityGrant) error {
	// TODO: Implement Shellstream wire protocol serialization.
	_ = conn
	_ = grant
	return fmt.Errorf("not yet implemented")
}

func readAttestConfirm(conn net.Conn) error {
	// TODO: Implement Shellstream wire protocol parsing.
	_ = conn
	return fmt.Errorf("not yet implemented")
}