o/devicemgmtstate: add task to queue response messages (canonical#17254)

* o/devicemgmtstate: add task to queue response messages

* o/devicemgmtstate: test that a task chain error does not affect other messages

* o/devicemgmtstate: set error message for unknown message kind in setMessageResponseFromChange

* o/devicemgmtstate: do not advance Applied on non-success response

* o/devicemgmtstate: add reuseable noopTask in tests

* o/devicemgmtstate: improve error messages

* o/devicemgmtstate: fixes after review

Author: st3v3nmw

overlord/devicemgmtstate/devicemgmtmgr.go

@@ -25,6 +25,7 @@
 package devicemgmtstate
 
 import (
+	"encoding/json"
 	"errors"
 	"fmt"
 	"sort"
@@ -58,6 +59,8 @@ var (
 	maxSequences                  = 256
 	maxBlockedMessagesPerSequence = 8
 
+	awaitSubsystemRetryInterval = 30 * time.Second
+
 	deviceMgmtExchangeChangeKind = swfeats.RegisterChangeKind("device-management-exchange")
 )
 
@@ -180,6 +183,23 @@ func (ms *deviceMgmtState) getRequestMessage(id string) (*RequestMessage, error)
 	return nil, fmt.Errorf("cannot find message %q", id)
 }
 
+// removeRequestMessage removes a processed request message from its sequence,
+// leaving the sequence entry in place so its Applied progress is preserved for
+// later messages in the same sequence.
+func (ms *deviceMgmtState) removeRequestMessage(msg *RequestMessage) {
+	seq := ms.Sequences[msg.BaseID]
+	if seq == nil {
+		return
+	}
+
+	for i, m := range seq.Messages {
+		if m.SeqNum == msg.SeqNum {
+			seq.Messages = append(seq.Messages[:i], seq.Messages[i+1:]...)
+			return
+		}
+	}
+}
+
 // enqueueRequestMessages queues incoming request messages for processing
 // and updates polling state accordingly.
 func (ms *deviceMgmtState) enqueueRequestMessages(pollResp *store.MessageExchangeResponse) {
@@ -479,8 +499,6 @@ func (m *DeviceMgmtManager) dispatchSequence(dispatchTask *state.Task, seq *sequ
 // the final task so callers can chain subsequent messages after it.
 func (m *DeviceMgmtManager) dispatchMessage(prevTask *state.Task, msg *RequestMessage) *state.Task {
 	chg := prevTask.Change()
-	// TODO: add tests verifying that a failure in one message's task chain does not
-	// affect other messages (lanes provide this isolation, but it needs test coverage).
 	lane := m.state.NewLane()
 
 	addTask := func(kind, summary string) {
@@ -586,9 +604,99 @@ func (m *DeviceMgmtManager) doApplyMessage(t *state.Task, _ *tomb.Tomb) error {
 }
 
 // doQueueResponse builds a response, signs it, and queues it for transmission on the next exchange.
-// Retries until subsystem change completes.
+// Retries until the subsystem change (if any) completes.
 func (m *DeviceMgmtManager) doQueueResponse(t *state.Task, _ *tomb.Tomb) error {
-	// TODO: implement this task, no-op for now.
+	m.state.Lock()
+	defer m.state.Unlock()
+
+	ms, err := m.getState()
+	if err != nil {
+		return err
+	}
+
+	var msgID string
+	err = t.Get("message-id", &msgID)
+	if err != nil {
+		return err
+	}
+
+	msg, err := ms.getRequestMessage(msgID)
+	if err != nil {
+		// Message already processed on a prior run.
+		return nil
+	}
+
+	err = m.setMessageResponseFromChange(msg)
+	if err != nil {
+		return err
+	}
+
+	bodyBytes, err := json.Marshal(msg.ResponseBody)
+	if err != nil {
+		return fmt.Errorf("cannot marshal response body: %w", err)
+	}
+
+	// TODO: determine reasonable behavior for internal errors (e.g., signing or marshal failures).
+	// Since tasks are idempotent, a failed message will not be re-dispatched or re-applied on the
+	// next change, but the request message remains in state until doQueueResponse completes,
+	// so such failures leave it hanging indefinitely.
+
+	resAs, err := m.signer.SignResponseMessage(msg.AccountID, msg.ID(), msg.ResponseStatus, bodyBytes)
+	if err != nil {
+		return fmt.Errorf("cannot sign response message: %w", err)
+	}
+
+	ms.ReadyResponses[msg.ID()] = store.Message{
+		Format: "assertion",
+		Data:   string(asserts.Encode(resAs)),
+	}
+
+	// TODO: rejecting sequences currently happens in 2 ways:
+	// 1. doDispatchMessage can evict the sequence immediately if it's rejected early.
+	// 2. If it errors elsewhere (in validate, apply, or queue-response), we end
+	//    up not advancing Applied, which means we accumulate messages until we
+	//    hit the sequence cap.
+	// Refactor sequence rejection to always evict immediately.
+	if msg.SeqNum > 0 && msg.ResponseStatus == asserts.MessageStatusSuccess {
+		ms.Sequences[msg.BaseID].Applied = msg.SeqNum
+	}
+	ms.removeRequestMessage(msg)
+
+	m.setState(ms)
+
+	return nil
+}
+
+// setMessageResponseFromChange populates msg's response fields from the completed apply change.
+func (m *DeviceMgmtManager) setMessageResponseFromChange(msg *RequestMessage) error {
+	if msg.ResponseStatus != "" {
+		return nil
+	}
+
+	handler, ok := m.handlers[msg.Kind]
+	if !ok {
+		msg.ResponseStatus = asserts.MessageStatusError
+		msg.ResponseBody = map[string]any{"message": fmt.Sprintf("cannot find handler for message kind %q", msg.Kind)}
+		return nil
+	}
+
+	change := m.state.Change(msg.ApplyChangeID)
+	if change == nil {
+		return fmt.Errorf("internal error: cannot find subsystem change %q", msg.ApplyChangeID)
+	}
+	if !change.Status().Ready() {
+		return &state.Retry{After: awaitSubsystemRetryInterval}
+	}
+
+	body, err := handler.ResultFromChange(change)
+	if err != nil {
+		msg.ResponseStatus = asserts.MessageStatusError
+		msg.ResponseBody = map[string]any{"message": fmt.Sprintf("cannot process message: %v", err)}
+	} else {
+		msg.ResponseStatus = asserts.MessageStatusSuccess
+		msg.ResponseBody = body
+	}
+
 	return nil
 }