fix: Ffixes to console

server/internal/services/consoleproxy.go

@@ -21,18 +21,30 @@ import (
 // so this is fatal rather than a degraded mode.
 var ErrAgentOffline = errors.New("agent is not connected")
 
-// A console session spans two processes once there is more than one replica.
+// A console session spans up to three processes once there is more than one
+// replica, and no two of them can be assumed to be the same one:
 //
-// The browser's WebSocket lands on an arbitrary pod. The agent's ProxyStream
-// lands on the pod holding that agent's command stream. The relay listener has
-// to be on the latter — that is the only process that can match an incoming
-// ProxyStream to a waiting listener — while guacd is dialled from the former.
+//	the browser's WebSocket   lands on an arbitrary pod
+//	the agent's CommandStream lands on the pod holding presence for it
+//	the agent's ProxyStream   lands on an arbitrary pod
 //
-// So the WebSocket's pod asks, over the bus, for a relay to be bound on the
-// agent's pod, and gets back an address to hand to guacd. That address is the
-// owner pod's own, which is why it must resolve to a single pod (POD_IP under
-// Kubernetes) rather than to the Service, which would send guacd to a pod
-// holding no listener roughly (n-1)/n of the time.
+// That third line is the one that is easy to get wrong. A ProxyStream is a
+// separate HTTP/2 request, and an L7 proxy (Traefik, which the chart's gRPC
+// ingress uses) balances requests rather than connections — so it does not
+// follow the command stream. Binding the relay listener on the command
+// stream's pod therefore fails roughly (n-1)/n of the time with "proxy session
+// not found": the stream arrives at a pod whose registry is empty.
+//
+// So the listener is bound by whichever pod receives the ProxyStream, at the
+// moment it receives it, and that pod announces its own address on
+// ProxyAddrChannel. The WebSocket's pod subscribes before dispatching and
+// hands the announced address to guacd. The address is the announcing pod's
+// own, which is why it must resolve to a single pod (POD_IP under Kubernetes)
+// rather than to the Service.
+//
+// Authorisation cannot live in that pod's memory either, so a pending record
+// in Redis (bus.SetPendingProxy) carries the instance and server a proxy_id
+// was minted for, and is consumed atomically on first claim.
 //
 // Teardown needs no message of its own. When the browser goes away guac closes
 // its connection to the relay, the relay sees the read end, and the session
@@ -46,6 +58,15 @@ var ErrAgentOffline = errors.New("agent is not connected")
 // waited for longer.
 const proxyEndGrace = 2 * time.Second
 
+// How long a minted proxy_id stays claimable, and how long the WebSocket's pod
+// waits for the relay's address to be announced. The TTL is the longer of the
+// two on purpose: a record that expired while its own opener was still waiting
+// would turn a slow agent into an unexplained refusal.
+const (
+	proxyPendingTTL = 30 * time.Second
+	proxyAddrWait   = 15 * time.Second
+)
+
 // ConsoleProxy is a relay as seen by the pod serving the WebSocket.
 type ConsoleProxy struct {
 	ProxyID string
@@ -114,51 +135,55 @@ func guacdHosts(addr string) []string {
 	return ips
 }
 
-// localRelay is a listener bound by this process on behalf of a remote request.
-type localRelay struct {
-	proxyID string
-	host    string
-	port    int
-	session *proxy.Session
+// proxyAddr is what a relay's binding pod announces: the address guacd should
+// dial to reach the listener it has just bound.
+type proxyAddr struct {
+	Host string `json:"host"`
+	Port int    `json:"port"`
 }
 
-// openLocalRelay binds a listener here and registers it, so the agent's
-// ProxyStream — which will arrive at this process — can be matched to it.
-// Called on the owner pod, from the dispatch handler.
-func openLocalRelay(instanceID, serverID, proxyID string) (*localRelay, error) {
+// ClaimProxyStream authorises an incoming ProxyStream, binds a relay listener
+// for it on this pod, and announces the address to whichever pod is serving the
+// browser's WebSocket. It is called from the gRPC handler, on whichever replica
+// the stream happened to reach.
+//
+// instanceID and serverID are the *authenticated* identity of the calling
+// agent; they must match the pending record or the claim is refused, so an
+// agent cannot relay a console session minted for another server.
+func ClaimProxyStream(instanceID, serverID, proxyID string) (*proxy.Session, error) {
+	ctx, cancel := context.WithTimeout(context.Background(), dispatchAckTimeout)
+	defer cancel()
+
+	wantInstance, wantServer := bus.ClaimPendingProxy(ctx, proxyID)
+	if wantInstance == "" {
+		return nil, proxy.ErrNotFound
+	}
+	if wantInstance != instanceID || wantServer != serverID {
+		return nil, proxy.ErrForbidden
+	}
+
 	sess, err := proxy.NewSession(proxyListenHost(), guacdHosts(guacdAddr()))
 	if err != nil {
 		return nil, err
 	}
 
 	sess.OnEnd(func(reason string) {
-		ctx, cancel := context.WithTimeout(context.Background(), dispatchAckTimeout)
-		defer cancel()
-		if _, err := bus.Publish(ctx, bus.ProxyEndChannel+proxyID, proxyEnd{Reason: reason}); err != nil {
+		endCtx, endCancel := context.WithTimeout(context.Background(), dispatchAckTimeout)
+		defer endCancel()
+		if _, err := bus.Publish(endCtx, bus.ProxyEndChannel+proxyID, proxyEnd{Reason: reason}); err != nil {
 			log.Printf("proxy: publish end for %s: %v", proxyID, err)
 		}
 	})
 
-	proxy.Default.Add(&proxy.Entry{
-		ProxyID:    proxyID,
-		InstanceID: instanceID,
-		ServerID:   serverID,
-		Session:    sess,
-	})
+	addr := proxyAddr{Host: proxyAdvertiseHost(), Port: sess.Port()}
+	if _, err := bus.Publish(ctx, bus.ProxyAddrChannel+proxyID, addr); err != nil {
+		// Nobody will ever dial this listener, so it is closed now rather than
+		// left to sit out its rendezvous timeout.
+		sess.Close("announce_failed")
+		return nil, fmt.Errorf("announce relay address: %w", err)
+	}
 
-	return &localRelay{
-		proxyID: proxyID,
-		host:    proxyAdvertiseHost(),
-		port:    sess.Port(),
-		session: sess,
-	}, nil
-}
-
-// abandon tears down a relay that was bound but whose command never reached the
-// agent, so the listener does not sit out its rendezvous timeout for nothing.
-func (r *localRelay) abandon() {
-	proxy.Default.Remove(r.proxyID)
-	r.session.Close("dispatch_failed")
+	return sess, nil
 }
 
 // OpenConsoleProxy asks the pod holding serverID's stream to bind a relay and
@@ -173,44 +198,85 @@ func OpenConsoleProxy(instanceID, serverID string, targetPort int) (*ConsoleProx
 		return nil, fmt.Errorf("generate proxy id: %w", err)
 	}
 
-	// Subscribed before the relay is asked for: a relay that fails immediately
-	// (the agent never claims it, the dial is refused) publishes its reason at
-	// once, and that reason is the whole content of the audit event.
+	// Both subscriptions are established before the command is dispatched: a
+	// fast agent binds and announces its relay within milliseconds, and a relay
+	// that fails immediately publishes its reason just as quickly. Either
+	// arriving before the subscriber is in place would be lost.
 	cp := &ConsoleProxy{ProxyID: proxyID, serverID: serverID, ended: make(chan struct{})}
 	endCtx, endCancel := context.WithCancel(context.Background())
-	ends, unsub, err := bus.Subscribe(endCtx, bus.ProxyEndChannel+proxyID)
+	ends, unsubEnd, err := bus.Subscribe(endCtx, bus.ProxyEndChannel+proxyID)
 	if err != nil {
 		endCancel()
 		return nil, fmt.Errorf("subscribe relay end: %w", err)
 	}
+	addrs, unsubAddr, err := bus.Subscribe(endCtx, bus.ProxyAddrChannel+proxyID)
+	if err != nil {
+		endCancel()
+		unsubEnd()
+		return nil, fmt.Errorf("subscribe relay address: %w", err)
+	}
 	cp.stop = func() {
 		endCancel()
-		unsub()
+		unsubAddr()
+		unsubEnd()
 	}
 	go cp.watchEnd(ends)
 
-	ack, err := Dispatcher.send(CommandEnvelope{
+	// The pending record authorises the ProxyStream the agent is about to open,
+	// and is written before the command so it cannot lose the race with it.
+	pendCtx, pendCancel := context.WithTimeout(context.Background(), dispatchAckTimeout)
+	if err := bus.SetPendingProxy(pendCtx, proxyID, instanceID, serverID, proxyPendingTTL); err != nil {
+		pendCancel()
+		cp.stop()
+		return nil, fmt.Errorf("register pending relay: %w", err)
+	}
+	pendCancel()
+
+	abandon := func() {
+		ctx, cancel := context.WithTimeout(context.Background(), dispatchAckTimeout)
+		bus.ClearPendingProxy(ctx, proxyID)
+		cancel()
+		cp.stop()
+	}
+
+	if _, err := Dispatcher.send(CommandEnvelope{
 		ServerID: serverID,
 		Command: &pb.ServerCommand{
 			CommandId: proxyID,
 			OpenProxy: &pb.OpenProxyCmd{ProxyId: proxyID, Port: uint32(targetPort)},
 		},
-		Proxy: &ProxyRelayRequest{InstanceID: instanceID, ProxyID: proxyID},
-	})
-	if err != nil {
-		cp.stop()
+	}); err != nil {
+		abandon()
 		if errors.Is(err, ErrAgentNotConnected) {
 			return nil, ErrAgentOffline
 		}
 		return nil, err
 	}
-	if ack.ProxyHost == "" || ack.ProxyPort == 0 {
-		cp.stop()
+
+	// The command has reached the agent; the relay's address arrives only once
+	// the agent has actually opened its ProxyStream somewhere in the fleet.
+	var addr proxyAddr
+	select {
+	case b, ok := <-addrs:
+		if !ok {
+			abandon()
+			return nil, fmt.Errorf("relay address subscription closed")
+		}
+		if err := json.Unmarshal(b, &addr); err != nil {
+			abandon()
+			return nil, fmt.Errorf("undecodable relay address: %w", err)
+		}
+	case <-time.After(proxyAddrWait):
+		abandon()
+		return nil, fmt.Errorf("agent did not open a relay for %s", serverID)
+	}
+	if addr.Host == "" || addr.Port == 0 {
+		abandon()
 		return nil, fmt.Errorf("relay opened without an address")
 	}
 
-	cp.Host = ack.ProxyHost
-	cp.Port = ack.ProxyPort
+	cp.Host = addr.Host
+	cp.Port = addr.Port
 	return cp, nil
 }
 

server/internal/services/dispatch.go

@@ -38,11 +38,10 @@ const (
 // CommandEnvelope is what actually crosses the bus. It is the command plus the
 // small amount of context the owning pod needs to act on it locally.
 type CommandEnvelope struct {
-	ServerID string             `json:"server_id"`
-	Command  *pb.ServerCommand  `json:"command"`
-	ReplyTo  string             `json:"reply_to"`
-	Log      *LogRequest        `json:"log,omitempty"`
-	Proxy    *ProxyRelayRequest `json:"proxy,omitempty"`
+	ServerID string            `json:"server_id"`
+	Command  *pb.ServerCommand `json:"command"`
+	ReplyTo  string            `json:"reply_to"`
+	Log      *LogRequest       `json:"log,omitempty"`
 }
 
 // LogRequest asks the owner pod to open a step log before it dispatches.
@@ -56,25 +55,17 @@ type LogRequest struct {
 	Mask     []string `json:"mask,omitempty"`
 }
 
-// ProxyRelayRequest asks the owner pod to bind a console relay listener and
-// register it before dispatching OpenProxyCmd.
+// CommandAck is the owner pod's answer. It reports only that the command
+// reached the agent's stream.
 //
-// The listener has to live on the owner pod: the agent's ProxyStream arrives
-// there, and only there can it be matched to a waiting listener. The pod
-// serving the browser's WebSocket learns the address from the ack and hands
-// that to guacd.
-type ProxyRelayRequest struct {
-	InstanceID string `json:"instance_id"`
-	ProxyID    string `json:"proxy_id"`
-}
-
-// CommandAck is the owner pod's answer.
+// A console relay listener used to be bound here and its address returned in
+// this ack. It no longer is: the agent's ProxyStream does not necessarily
+// arrive at the pod holding its command stream, so the listener is bound by
+// whichever pod receives that stream and announced on bus.ProxyAddrChannel.
 type CommandAck struct {
-	OK        bool   `json:"ok"`
-	Error     string `json:"error,omitempty"`
-	Node      string `json:"node,omitempty"`
-	ProxyHost string `json:"proxy_host,omitempty"`
-	ProxyPort int    `json:"proxy_port,omitempty"`
+	OK    bool   `json:"ok"`
+	Error string `json:"error,omitempty"`
+	Node  string `json:"node,omitempty"`
 }
 
 type commandDispatcher struct{}
@@ -169,27 +160,10 @@ func (d *commandDispatcher) handleEnvelope(ctx context.Context, raw []byte, out
 		}
 	}
 
-	var relay *localRelay
-	if env.Proxy != nil {
-		r, err := openLocalRelay(env.Proxy.InstanceID, env.ServerID, env.Proxy.ProxyID)
-		if err != nil {
-			ack = CommandAck{OK: false, Error: err.Error(), Node: bus.NodeID()}
-		} else {
-			relay = r
-			ack.ProxyHost = r.host
-			ack.ProxyPort = r.port
-		}
-	}
-
-	if ack.OK {
-		select {
-		case out <- env.Command:
-		default:
-			ack = CommandAck{OK: false, Error: "command queue full", Node: bus.NodeID()}
-			if relay != nil {
-				relay.abandon()
-			}
-		}
+	select {
+	case out <- env.Command:
+	default:
+		ack = CommandAck{OK: false, Error: "command queue full", Node: bus.NodeID()}
 	}
 
 	if err := bus.Reply(ctx, env.ReplyTo, ack); err != nil {