feat: initial commit

internal/agent/session.go

@@ -0,0 +1,131 @@
+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.
+func (a *Agent) runTerminal(parent context.Context, sid string) {
+	term, err := StartPTY()
+	if err != nil {
+		log.Printf("term %s: start pty: %v", sid, err)
+		a.postOut(context.Background(), sid, proto.OutMessage{Exit: true})
+		return
+	}
+	log.Printf("term %s: shell started", sid)
+	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()
+
+	// 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 {
+	req, _ := http.NewRequestWithContext(ctx, http.MethodGet, a.cfg.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
+	}
+	for _, f := range batch.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 {
+	body, _ := json.Marshal(msg)
+	req, _ := http.NewRequestWithContext(ctx, http.MethodPost, a.cfg.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
+}