nadir-agent/internal/modules/terminal/terminal.go

package terminal

import (
	"context"
	"encoding/json"
	"net/http"
	"os/exec"
	"strings"
	"sync/atomic"
	"time"

	"nadir/internal/auth"
	"nadir/internal/module"
	"nadir/internal/rbac"

	"github.com/coder/websocket"
	"github.com/creack/pty"
	"github.com/danielgtaylor/huma/v2"
	"github.com/danielgtaylor/huma/v2/adapters/humago"
)

// tagTerminal is the OpenAPI tag for this module (registered in server.go),
// keeping tags 1:1 with modules per the project convention.
const tagTerminal = "Terminal"

const (
	// maxTerminals caps concurrent shells so a buggy frontend or careless admin
	// can't pile up PTYs (guideline: limit everything). Raise if it's ever a real
	// limit in practice.
	maxTerminals = 10
	// idleTimeout closes a session after this long with no I/O in either
	// direction, reclaiming abandoned shells.
	idleTimeout = 15 * time.Minute
)

// terminalSem is the concurrency limiter; a slot is held for the life of a session.
var terminalSem = make(chan struct{}, maxTerminals)

type terminalModule struct {
	sessions *auth.SessionStore
}

// New creates a new Terminal module that allows interactive shell access.
func New(sessions *auth.SessionStore) module.Module {
	return &terminalModule{sessions: sessions}
}

func (m *terminalModule) ID() string   { return "terminal" }

func (m *terminalModule) Permissions() []rbac.Permission {
	return []rbac.Permission{rbac.Root}
}

type TerminalInput struct {
	Ctx huma.Context `json:"-"`
}

// Resolve extracts the huma.Context into the input struct.
func (i *TerminalInput) Resolve(ctx huma.Context) []error {
	i.Ctx = ctx
	return nil
}

type resizeMessage struct {
	Cols uint16 `json:"cols"`
	Rows uint16 `json:"rows"`
}

func (m *terminalModule) Register(api huma.API) {
	huma.Register(api, huma.Operation{
		OperationID: "terminal-connect",
		Method:      "GET",
		Path:        "/api/terminal",
		Summary:     "Connect to an interactive terminal",
		Description: "Upgrades the connection to a WebSocket and spawns a PTY shell as the logged-in user. Send JSON `{cols, rows}` text messages to resize, and raw binary/text messages for stdin. This is a raw WebSocket endpoint — it cannot be exercised from the API docs \"Try it\" panel; use a WebSocket client.",
		Tags:        []string{tagTerminal},
		Metadata:    map[string]any{"module": m.ID(), "permission": string(rbac.Root)},
		Errors:      []int{401, 403, 426, 500},
	}, func(ctx context.Context, in *TerminalInput) (*struct{}, error) {
		// The RBAC middleware already authenticated this request and enforced the
		// "root" permission before we got here. We re-read the session only to get
		// the username for `su`; the 401 below is the fallback for when the module
		// is mounted without the middleware (e.g. in unit tests).
		cookie, err := huma.ReadCookie(in.Ctx, "nadir_session_id")
		if err != nil || cookie == nil {
			return nil, huma.Error401Unauthorized("unauthorized")
		}

		sess, ok := m.sessions.GetByToken(cookie.Value)
		if !ok {
			return nil, huma.Error401Unauthorized("unauthorized")
		}

		req, res := humago.Unwrap(in.Ctx)
		if req == nil || res == nil {
			return nil, huma.Error500InternalServerError("missing http context")
		}

		// Reject plain GETs (e.g. the docs "Try it" button) with a clear 426 rather
		// than letting websocket.Accept emit a raw protocol-violation error.
		if !strings.EqualFold(req.Header.Get("Upgrade"), "websocket") {
			return nil, huma.NewError(http.StatusUpgradeRequired,
				"this endpoint requires a WebSocket connection; connect with a WebSocket client")
		}

		// InsecureSkipVerify is deliberately NOT set: coder/websocket then enforces
		// that the Origin host matches the request Host, rejecting cross-site upgrade
		// attempts. Defense-in-depth on top of the SameSite=Strict session cookie —
		// important because this endpoint hands out an interactive shell.
		conn, err := websocket.Accept(res, req, nil)
		if err != nil {
			// websocket.Accept already wrote the error response (e.g. 403 on an
			// Origin mismatch). Just stop; writing again would corrupt the response.
			return nil, nil
		}
		defer conn.CloseNow()

		// Bound concurrent shells: take a slot or reject (don't pile up PTYs).
		select {
		case terminalSem <- struct{}{}:
			defer func() { <-terminalSem }()
		default:
			conn.Close(websocket.StatusTryAgainLater, "too many terminal sessions")
			return nil, nil
		}

		// Launch the user's login shell via su.
		// "su - <username>" ensures we get their actual environment and shell.
		cmd := exec.CommandContext(req.Context(), "su", "-", "--", sess.Username)

		// Start the command with a PTY.
		ptmx, err := pty.Start(cmd)
		if err != nil {
			conn.Close(websocket.StatusInternalError, "failed to start pty")
			return nil, nil
		}
		defer ptmx.Close()

		// lastActive is bumped by both pumps; the watchdog uses it to close idle
		// sessions. Output activity (e.g. `top`) counts, so it isn't killed.
		var lastActive atomic.Int64
		lastActive.Store(time.Now().UnixNano())
		go func() {
			tick := time.NewTicker(idleTimeout / 4)
			defer tick.Stop()
			for {
				select {
				case <-req.Context().Done():
					return
				case <-tick.C:
					if time.Since(time.Unix(0, lastActive.Load())) > idleTimeout {
						conn.Close(websocket.StatusGoingAway, "idle timeout")
						return
					}
				}
			}
		}()

		// Pump stdout/stderr from PTY to WebSocket.
		go func() {
			buf := make([]byte, 8192)
			for {
				n, err := ptmx.Read(buf)
				if err != nil {
					break
				}
				lastActive.Store(time.Now().UnixNano())
				// We write PTY output as binary messages. The frontend (e.g., xterm.js)
				// can handle UTF-8 binary or text transparently.
				err = conn.Write(req.Context(), websocket.MessageBinary, buf[:n])
				if err != nil {
					break
				}
			}
			conn.Close(websocket.StatusNormalClosure, "")
		}()

		// Pump stdin and resize commands from WebSocket to PTY.
		for {
			typ, b, err := conn.Read(req.Context())
			if err != nil {
				break
			}
			lastActive.Store(time.Now().UnixNano())

			if typ == websocket.MessageText {
				var resize resizeMessage
				if err := json.Unmarshal(b, &resize); err == nil && resize.Cols > 0 && resize.Rows > 0 {
					// Handle resize
					_ = pty.Setsize(ptmx, &pty.Winsize{
						Cols: resize.Cols,
						Rows: resize.Rows,
					})
					continue
				}
			}

			// If not a valid resize message, or if it's MessageBinary, pass to PTY stdin.
			_, _ = ptmx.Write(b)
		}

		// The read loop has ended (client gone or PTY closed). Tear down the shell
		// and reap it: closing the PTY sends EOF, Kill covers shells that ignore it.
		_ = ptmx.Close()
		_ = cmd.Process.Kill()
		_ = cmd.Wait()
		return nil, nil
	})
}