lxd-agent/exec: Use shared.NewExecWrapper and bring into line with co…

…ntainer exec

Fixes canonical#12210

Signed-off-by: Thomas Parrott <thomas.parrott@canonical.com>

lxd-agent/exec.go

@@ -282,12 +282,12 @@ func (s *execWs) Do(op *operations.Operation) error {
  stderr = ttys[execWSStderr]
  }
 
- attachedChildIsDead := make(chan struct{})
+ waitAttachedChildIsDead, markAttachedChildIsDead := context.WithCancel(context.Background())
  var wgEOF sync.WaitGroup
 
  finisher := func(cmdResult int, cmdErr error) error {
- // Close this before closing the control connection so control handler can detect command ending.
- close(attachedChildIsDead)
+ // Cancel this before closing the control connection so control handler can detect command ending.
+ markAttachedChildIsDead()
 
  for _, tty := range ttys {
  _ = tty.Close()
@@ -384,10 +384,8 @@ func (s *execWs) Do(op *operations.Operation) error {
  mt, r, err := conn.NextReader()
  if err != nil || mt == websocket.CloseMessage {
  // Check if command process has finished normally, if so, no need to kill it.
- select {
- case <-attachedChildIsDead:
+ if waitAttachedChildIsDead.Err() != nil {
  return
- default:
  }
 
  if mt == websocket.CloseMessage {
@@ -409,10 +407,8 @@ func (s *execWs) Do(op *operations.Operation) error {
  buf, err := io.ReadAll(r)
  if err != nil {
  // Check if command process has finished normally, if so, no need to kill it.
- select {
- case <-attachedChildIsDead:
+ if waitAttachedChildIsDead.Err() != nil {
  return
- default:
  }
 
  l.Warn("Failed reading control websocket message, killing command", logger.Ctx{"err": err})
@@ -469,7 +465,7 @@ func (s *execWs) Do(op *operations.Operation) error {
  conn := s.conns[0]
  s.connsLock.Unlock()
 
- readDone, writeDone := ws.Mirror(conn, ptys[0])
+ readDone, writeDone := ws.Mirror(conn, shared.NewExecWrapper(waitAttachedChildIsDead, ptys[0]))
 
  <-readDone
  <-writeDone