lpra/internal/agent/session.go

package agent

import (
	"bytes"
	"context"
	"encoding/json"
	"io"
	"log"
	"net"
	"net/http"
	"strconv"
	"time"

	"github.com/x-grk/longpolling-remote-access/internal/proto"
)

// runTerminal spawns a PTY and bridges it to the session.
// user is the OS user to run the shell as ("" = the agent's own user).
func (a *Agent) runTerminal(parent context.Context, sid, user string) {
	term, err := StartPTY(user)
	if err != nil {
		log.Printf("term %s: start pty: %v", sid, err)
		// Show the reason in the browser terminal before closing it.
		a.postOut(context.Background(), sid, proto.OutMessage{Data: b64.EncodeToString([]byte("lpra: " + err.Error() + "\r\n"))})
		a.postOut(context.Background(), sid, proto.OutMessage{Exit: true})
		return
	}
	log.Printf("term %s: shell started (user %q)", sid, user)
	a.bridgeStream(parent, sid, term, term.Resize)
	log.Printf("term %s: ended", sid)
}

// runForward dials 127.0.0.1:port on the client and bridges it to the session,
// giving the host a reverse TCP tunnel over the same long-poll transport.
func (a *Agent) runForward(parent context.Context, sid string, port int) {
	d := net.Dialer{Timeout: 10 * time.Second}
	conn, err := d.DialContext(parent, "tcp", net.JoinHostPort("127.0.0.1", strconv.Itoa(port)))
	if err != nil {
		log.Printf("forward %s: dial :%d: %v", sid, port, err)
		a.postOut(context.Background(), sid, proto.OutMessage{Exit: true})
		return
	}
	// TCP_NODELAY (default on in Go) keeps latency low for small packets.
	if tc, ok := conn.(*net.TCPConn); ok {
		_ = tc.SetNoDelay(true)
	}
	log.Printf("forward %s: connected to :%d", sid, port)
	a.bridgeStream(parent, sid, conn, nil)
}

// bridgeStream wires a local byte stream (PTY or socket) to a hub session:
// server->agent frames are written to rw, and rw output is streamed back.
func (a *Agent) bridgeStream(parent context.Context, sid string, rw io.ReadWriteCloser, onResize func(cols, rows uint16)) {
	ctx, cancel := context.WithCancel(parent)
	defer cancel()
	defer rw.Close()
	defer func() {
		if l := a.link.Load(); l != nil {
			l.unsubscribe(sid)
		}
	}()
	// A session closed from the panel (a tab reloaded, a window shut) ends
	// the stream at once, not on its next output: an idle shell would live
	// on until the agent stopped.
	go func() { <-ctx.Done(); rw.Close() }()

	// Downstream: pull input frames and write them to the stream.
	go func() {
		defer cancel()
		for ctx.Err() == nil {
			if !a.pollIn(ctx, sid, rw, onResize) {
				return
			}
		}
	}()

	// Upstream: stream output until the stream (or session) ends.
	buf := make([]byte, 32*1024)
	for {
		n, err := rw.Read(buf)
		if n > 0 {
			if !a.postOut(ctx, sid, proto.OutMessage{Data: b64.EncodeToString(buf[:n])}) {
				return
			}
		}
		if err != nil {
			a.postOut(ctx, sid, proto.OutMessage{Exit: true})
			return
		}
	}
}

// pollIn performs one downstream long-poll and applies the frames.
// Returns false when the session is gone and the loop should stop.
func (a *Agent) pollIn(ctx context.Context, sid string, w io.Writer, onResize func(cols, rows uint16)) bool {
	// On the WebSocket the input comes pushed; if it drops, by a long poll.
	if l := a.link.Load(); l != nil {
		if ch := l.subscribe(sid); ch != nil {
			select {
			case in := <-ch:
				if in.gone {
					return false
				}
				return applyFrames(in.frames, w, onResize)
			case <-l.done:
			case <-ctx.Done():
				return false
			}
		}
	}
	req, _ := http.NewRequestWithContext(ctx, http.MethodGet, a.host()+"/api/agent/term/"+sid+"/in", nil)
	a.authHeaders(req)
	resp, err := a.http.Do(req)
	if err != nil {
		return ctx.Err() == nil
	}
	defer resp.Body.Close()
	if resp.StatusCode == http.StatusGone || resp.StatusCode == http.StatusUnauthorized {
		return false
	}
	if resp.StatusCode != http.StatusOK {
		return true
	}
	var batch proto.InBatch
	if err := json.NewDecoder(resp.Body).Decode(&batch); err != nil {
		return true
	}
	return applyFrames(batch.Frames, w, onResize)
}

// applyFrames writes a session's input to its stream; false once the panel
// closed it.
func applyFrames(frames []proto.InFrame, w io.Writer, onResize func(cols, rows uint16)) bool {
	for _, f := range frames {
		switch f.Type {
		case proto.FrameStdin:
			if data, err := b64.DecodeString(f.Data); err == nil {
				_, _ = w.Write(data)
			}
		case proto.FrameResize:
			if onResize != nil {
				onResize(f.Cols, f.Rows)
			}
		case proto.FrameClose:
			return false
		}
	}
	return true
}

// postOut sends one upstream message. Returns false if the session is gone.
func (a *Agent) postOut(ctx context.Context, sid string, msg proto.OutMessage) bool {
	if l := a.link.Load(); l != nil {
		if l.send(proto.WSMsg{T: "out", SID: sid, Out: &msg}) == nil {
			return true
		}
	}
	body, _ := json.Marshal(msg)
	req, _ := http.NewRequestWithContext(ctx, http.MethodPost, a.host()+"/api/agent/term/"+sid+"/out", bytes.NewReader(body))
	a.authHeaders(req)
	req.Header.Set("Content-Type", "application/json")
	resp, err := a.http.Do(req)
	if err != nil {
		return ctx.Err() == nil
	}
	defer resp.Body.Close()
	_, _ = io.Copy(io.Discard, resp.Body) // drain so the connection is reused
	return resp.StatusCode == http.StatusNoContent
}