termd/internal/attach/client.go

// Package attach implements the interactive attach client: it proxies the
// local terminal to a session over the daemon's websocket endpoint.
package attach

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"os"
	"os/signal"
	"strings"
	"syscall"

	"github.com/coder/websocket"
	"golang.org/x/term"

	"termd/internal/api"
	"termd/internal/client"
)

// control mirrors the server's attach control frames.
type control struct {
	Hello  *api.SessionInfo   `json:"hello,omitempty"`
	Resize *api.ResizeRequest `json:"resize,omitempty"`
	Exited *struct {
		Code int `json:"code"`
	} `json:"exited,omitempty"`
}

// Options configures an attach run.
type Options struct {
	DetachByte byte // 0x1d = C-], the default
	Resize     bool // push local terminal size to the session
	Write      bool // forward keystrokes; default is a read-only view
	Render     bool // emulator-rendered frames instead of the raw stream
	// (for viewers larger than the session)

	// I/O; in real use these are os.Stdin/os.Stdout with RawMode set.
	Stdin   io.Reader
	Stdout  io.Writer
	RawMode bool // put Stdin's fd into raw mode (requires *os.File)
}

// errDetached signals a clean user detach.
var errDetached = errors.New("detached")

// Run attaches to the session until the user detaches, the session ends, or
// ctx is canceled. Returns nil on clean detach/session end.
func Run(ctx context.Context, c *client.Client, id string, opts Options) error {
	if opts.DetachByte == 0 {
		opts.DetachByte = 0x1d // C-]
	}
	if opts.Stdin == nil {
		opts.Stdin = os.Stdin
	}
	if opts.Stdout == nil {
		opts.Stdout = os.Stdout
	}

	var params []string
	if !opts.Write {
		// The server enforces read-only; the client not sending is courtesy.
		params = append(params, "readonly=1")
	}
	if opts.Render {
		params = append(params, "view=render")
	}
	url := c.AttachURL(id)
	if len(params) > 0 {
		url += "?" + strings.Join(params, "&")
	}
	conn, _, err := websocket.Dial(ctx, url, &websocket.DialOptions{
		HTTPClient: c.HTTPClient(),
		HTTPHeader: authHeader(c),
	})
	if err != nil {
		return fmt.Errorf("dialing attach websocket: %w", err)
	}
	defer conn.CloseNow()
	// The repaint snapshot is one big frame; raise the limit generously.
	conn.SetReadLimit(16 << 20)

	ctx, cancel := context.WithCancel(ctx)
	defer cancel()

	if opts.RawMode {
		f, ok := opts.Stdin.(*os.File)
		if !ok {
			return fmt.Errorf("raw mode requires stdin to be a real terminal")
		}
		oldState, err := term.MakeRaw(int(f.Fd()))
		if err != nil {
			return fmt.Errorf("setting raw mode: %w", err)
		}
		defer term.Restore(int(f.Fd()), oldState) //nolint:errcheck
	}

	sendResize := func() {
		f, ok := opts.Stdout.(*os.File)
		if !ok {
			return
		}
		cols, rows, err := term.GetSize(int(f.Fd()))
		if err != nil {
			return
		}
		b, _ := json.Marshal(control{Resize: &api.ResizeRequest{Cols: cols, Rows: rows}})
		_ = conn.Write(ctx, websocket.MessageText, b)
	}
	if opts.Resize {
		sendResize()
		winch := make(chan os.Signal, 1)
		signal.Notify(winch, syscall.SIGWINCH)
		defer signal.Stop(winch)
		go func() {
			for {
				select {
				case <-winch:
					sendResize()
				case <-ctx.Done():
					return
				}
			}
		}()
	}

	// stdin → ws, scanning for the detach byte. In read-only mode stdin is
	// still consumed (the detach byte must work) but nothing is forwarded.
	inErr := make(chan error, 1)
	go func() {
		buf := make([]byte, 4096)
		for {
			n, err := opts.Stdin.Read(buf)
			if n > 0 {
				chunk := buf[:n]
				for i, b := range chunk {
					if b == opts.DetachByte {
						if opts.Write && i > 0 {
							if werr := conn.Write(ctx, websocket.MessageBinary, chunk[:i]); werr != nil {
								inErr <- werr
								return
							}
						}
						inErr <- errDetached
						return
					}
				}
				if opts.Write {
					if werr := conn.Write(ctx, websocket.MessageBinary, chunk); werr != nil {
						inErr <- werr
						return
					}
				}
			}
			if err != nil {
				inErr <- err
				return
			}
		}
	}()

	// ws → stdout.
	outErr := make(chan error, 1)
	exited := make(chan int, 1)
	go func() {
		for {
			typ, data, err := conn.Read(ctx)
			if err != nil {
				outErr <- err
				return
			}
			switch typ {
			case websocket.MessageBinary:
				if _, err := opts.Stdout.Write(data); err != nil {
					outErr <- err
					return
				}
			case websocket.MessageText:
				var c control
				if err := json.Unmarshal(data, &c); err == nil && c.Exited != nil {
					exited <- c.Exited.Code
					return
				}
			}
		}
	}()

	select {
	case err := <-inErr:
		if errors.Is(err, errDetached) {
			return nil
		}
		return err
	case code := <-exited:
		// Leave raw mode land before printing (deferred Restore runs after,
		// but a trailing CRLF keeps the shell prompt clean either way).
		fmt.Fprintf(opts.Stdout, "\r\n[session exited with code %d]\r\n", code)
		return nil
	case err := <-outErr:
		var closeErr websocket.CloseError
		if errors.As(err, &closeErr) && closeErr.Code == websocket.StatusNormalClosure {
			return nil
		}
		if errors.Is(err, context.Canceled) {
			return nil
		}
		return fmt.Errorf("connection lost: %w", err)
	case <-ctx.Done():
		return nil
	}
}

func authHeader(c *client.Client) map[string][]string {
	if c.Token() == "" {
		return nil
	}
	return map[string][]string{"Authorization": {"Bearer " + c.Token()}}
}