lpra/internal/server/fsstream.go

package server

import (
	"context"
	"encoding/base64"
	"encoding/json"
	"errors"
	"io"
	"net/http"
	"path"
	"strconv"
	"strings"
	"sync"
	"time"

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

// Files go between the browser and the agent in pieces. Each piece is a
// request of its own, small enough for any proxy in front (Cloudflare takes
// 100 MB a request at most), and one that fails is sent again instead of the
// whole file. Several pieces are in flight at once, so a good link stays as
// busy as it would with one long stream. The server holds only a few pieces
// of a transfer: whoever sends waits while they wait to be taken.
//
// A download keeps one plain response towards the browser, which a proxy
// passes on whatever its size: links, previews and the hex editor read it.

const (
	// xferChunk is the size of a piece.
	xferChunk = 4 << 20
	// xferBuffer is how many pieces of a transfer the server holds.
	xferBuffer = 8
	// xferWait is how long a request waits for its piece, or for room for it,
	// before it is told to ask again: well under a proxy's timeout.
	xferWait = 50 * time.Second
	// xferIdle ends a transfer nobody has touched for this long.
	xferIdle = 10 * time.Minute
	// xferTimeout bounds one transfer, however slow the links.
	xferTimeout = 2 * time.Hour
)

// fsLegacyCap is the most an agent that predates the pieces can take in or
// hand over: the file goes base64 inside a task then.
const fsLegacyCap = 32 << 20

// xfer is one file on its way.
type xfer struct {
	id, agentID string

	mu      sync.Mutex
	changed chan struct{} // closed and replaced on every change
	size    int64         // -1 until known
	n       int           // pieces in all, once the size is known
	chunks  map[int][]byte
	// low is the first piece still to go: the ones below it were written by
	// the agent (upload) or to the browser (download).
	low    int
	sized  chan struct{} // closed once the size is known
	err    string        // why the transfer failed
	done   bool
	result proto.TaskResult
	last   time.Time
}

func pieces(size int64) int {
	if size <= 0 {
		return 1 // an empty file is one empty piece
	}
	return int((size + xferChunk - 1) / xferChunk)
}

// pieceLen is how long piece i of the file is.
func (x *xfer) pieceLen(i int) int64 {
	return min(int64(xferChunk), x.size-int64(i)*xferChunk)
}

// notify wakes whoever waits on the transfer; x.mu is held.
func (x *xfer) notify() {
	x.last = time.Now()
	close(x.changed)
	x.changed = make(chan struct{})
}

// await waits until ready() holds and returns true with x.mu held; false,
// with it free, when the wait ran out or ctx ended.
func (x *xfer) await(ctx context.Context, d time.Duration, ready func() bool) bool {
	t := time.NewTimer(d)
	defer t.Stop()
	for {
		x.mu.Lock()
		if ready() {
			return true
		}
		ch := x.changed
		x.mu.Unlock()
		select {
		case <-ch:
		case <-t.C:
			return false
		case <-ctx.Done():
			return false
		}
	}
}

func (x *xfer) fail(msg string) {
	x.mu.Lock()
	if x.err == "" && !x.done {
		x.err = msg
		x.chunks = nil
		x.notify()
	}
	x.mu.Unlock()
}

func (x *xfer) finish(res proto.TaskResult) {
	x.mu.Lock()
	x.done, x.result, x.chunks = true, res, nil
	x.notify()
	x.mu.Unlock()
}

// setSize records the file's size; false when it differs from the one known.
func (x *xfer) setSize(size int64) bool {
	x.mu.Lock()
	defer x.mu.Unlock()
	if x.size >= 0 {
		return x.size == size
	}
	x.size, x.n = size, pieces(size)
	close(x.sized)
	x.notify()
	return true
}

// put stores piece i once there is room for it, and answers with the HTTP
// status for its sender. A piece already taken or held is accepted again: it
// is a retry.
func (x *xfer) put(ctx context.Context, i int, data []byte) (int, string) {
	if !x.await(ctx, xferWait, func() bool { return x.err != "" || x.done || i < x.low+xferBuffer }) {
		return http.StatusServiceUnavailable, "нет места, повторите"
	}
	defer x.mu.Unlock()
	if x.err != "" || x.done {
		return http.StatusGone, x.gone()
	}
	if i < 0 || i >= x.n || int64(len(data)) != x.pieceLen(i) {
		return http.StatusBadRequest, "не тот кусок"
	}
	if i >= x.low && x.chunks[i] == nil {
		x.chunks[i] = data
		x.notify()
	}
	return http.StatusNoContent, ""
}

// take hands over piece i, the next one in order, and lets it go.
func (x *xfer) take(ctx context.Context, i int, d time.Duration) ([]byte, error) {
	if !x.await(ctx, d, func() bool { return x.err != "" || x.chunks[i] != nil }) {
		return nil, errors.New("кусок не пришёл вовремя")
	}
	defer x.mu.Unlock()
	if x.err != "" {
		return nil, errors.New(x.err)
	}
	data := x.chunks[i]
	delete(x.chunks, i)
	x.low = i + 1
	x.notify()
	return data, nil
}

func (x *xfer) gone() string {
	if x.err != "" {
		return x.err
	}
	return "передача закончена"
}

// Read makes an upload a plain reader, for an agent that takes a file whole.
type xferReader struct {
	ctx context.Context
	x   *xfer
	i   int
	buf []byte
}

func (r *xferReader) Read(p []byte) (int, error) {
	for len(r.buf) == 0 {
		if r.i >= r.x.n {
			return 0, io.EOF
		}
		data, err := r.x.take(r.ctx, r.i, xferIdle)
		if err != nil {
			return 0, err
		}
		r.buf, r.i = data, r.i+1
	}
	n := copy(p, r.buf)
	r.buf = r.buf[n:]
	return n, nil
}

type xfers struct {
	mu sync.Mutex
	m  map[string]*xfer
}

func (c *xfers) open(agentID string, size int64) *xfer {
	x := &xfer{
		id: randToken()[:24], agentID: agentID, size: -1,
		changed: make(chan struct{}), chunks: map[int][]byte{},
		sized: make(chan struct{}), last: time.Now(),
	}
	if size >= 0 {
		x.size, x.n = size, pieces(size)
		close(x.sized)
	}
	c.mu.Lock()
	if c.m == nil {
		c.m = map[string]*xfer{}
	}
	// Forget the ones left behind: a closed tab, a client gone for good.
	for id, old := range c.m {
		old.mu.Lock()
		stale := time.Since(old.last) > xferIdle
		old.mu.Unlock()
		if stale {
			old.fail("передача брошена")
			delete(c.m, id)
		}
	}
	c.m[x.id] = x
	c.mu.Unlock()
	return x
}

func (c *xfers) get(id string) *xfer {
	c.mu.Lock()
	defer c.mu.Unlock()
	return c.m[id]
}

// forAgent returns the transfer if it is the asking agent's.
func (c *xfers) forAgent(id, agentID string) *xfer {
	if x := c.get(id); x != nil && x.agentID == agentID {
		return x
	}
	return nil
}

func (c *xfers) drop(id string) {
	c.mu.Lock()
	delete(c.m, id)
	c.mu.Unlock()
}

// waitTask runs a task on the agent and waits for its result, or for ctx to
// end.
func (s *Server) waitTask(ctx context.Context, agentID, typ string, args any, timeout time.Duration) (proto.TaskResult, error) {
	raw, _ := json.Marshal(args)
	taskID, ch := s.hub.CreateTask(agentID, typ, raw)
	defer s.hub.DiscardTask(taskID)
	t := time.NewTimer(timeout)
	defer t.Stop()
	select {
	case res := <-ch:
		return res, nil
	case <-ctx.Done():
		return proto.TaskResult{}, ctx.Err()
	case <-t.C:
		return proto.TaskResult{}, errors.New("таймаут ожидания агента")
	}
}

// unknownOp tells an agent that predates the task.
func unknownOp(res proto.TaskResult) bool {
	return !res.OK && strings.Contains(res.Error, "неизвестная операция")
}

// fsAgent answers the request itself and returns false when the agent is
// unknown or offline.
func (s *Server) fsAgent(w http.ResponseWriter, agentID string, asJSON bool) bool {
	a, err := s.db.GetAgent(agentID)
	if err != nil || a == nil {
		http.Error(w, "not found", http.StatusNotFound)
		return false
	}
	if !s.online(a) {
		if asJSON {
			writeJSON(w, proto.TaskResult{Error: "машина офлайн"})
		} else {
			http.Error(w, "машина офлайн", http.StatusServiceUnavailable)
		}
		return false
	}
	return true
}

// readPiece reads a piece's body into a buffer of its own size, so pieces
// leave no trail of grown buffers behind.
func readPiece(w http.ResponseWriter, r *http.Request) ([]byte, error) {
	if r.ContentLength >= 0 && r.ContentLength <= xferChunk {
		data := make([]byte, r.ContentLength)
		_, err := io.ReadFull(r.Body, data)
		return data, err
	}
	return io.ReadAll(http.MaxBytesReader(w, r.Body, xferChunk))
}

// ---------- upload: browser -> agent ----------

// POST /api/panel/agents/{id}/fs/upload?path=...&size=N - start writing a
// file. The answer names the transfer and the piece size; the browser then
// PUTs the pieces and asks for the result.
func (s *Server) handleFSUploadStart(w http.ResponseWriter, r *http.Request) {
	agentID, p := r.PathValue("id"), r.URL.Query().Get("path")
	if !s.fsAgent(w, agentID, true) {
		return
	}
	size, err := strconv.ParseInt(r.URL.Query().Get("size"), 10, 64)
	if p == "" || err != nil || size < 0 {
		writeJSON(w, proto.TaskResult{Error: "не указан путь или размер"})
		return
	}
	if size > proto.FSFileCap {
		writeJSON(w, proto.TaskResult{Error: "файл больше 512 МБ"})
		return
	}
	x := s.xfers.open(agentID, size)
	go s.runUpload(x, p)
	data, _ := json.Marshal(map[string]any{"id": x.id, "chunk": xferChunk})
	writeJSON(w, proto.TaskResult{OK: true, Data: data})
}

// runUpload has the agent write the file and keeps the result for the
// browser to collect.
func (s *Server) runUpload(x *xfer, p string) {
	ctx, cancel := context.WithTimeout(context.Background(), xferTimeout)
	defer cancel()
	args := proto.FSXferArgs{Path: p, Xfer: x.id, Size: x.size, Chunk: xferChunk}
	res, err := s.waitTask(ctx, x.agentID, proto.TaskFSPut, args, xferTimeout)
	if err == nil && unknownOp(res) {
		res = s.legacyUpload(ctx, x, p)
	}
	if err != nil {
		res = proto.TaskResult{Error: err.Error()}
	}
	if !res.OK {
		x.fail(res.Error)
	}
	x.finish(res)
	// The result waits a while for the browser to come for it.
	time.AfterFunc(xferIdle, func() { s.xfers.drop(x.id) })
}

// legacyUpload writes a file through an agent that predates the pieces.
func (s *Server) legacyUpload(ctx context.Context, x *xfer, p string) proto.TaskResult {
	if x.size > fsLegacyCap {
		return proto.TaskResult{Error: "клиент принимает файлы до 32 МБ - обновите его"}
	}
	data, err := io.ReadAll(&xferReader{ctx: ctx, x: x})
	if err != nil {
		return proto.TaskResult{Error: err.Error()}
	}
	args := proto.FSWriteArgs{Path: p, Data: base64.StdEncoding.EncodeToString(data)}
	res, err := s.waitTask(ctx, x.agentID, proto.TaskFSWrite, args, fsTaskTimeout)
	if err != nil {
		return proto.TaskResult{Error: err.Error()}
	}
	return res
}

// PUT /api/panel/fs/xfer/{xid}/{i} - piece i of a file being uploaded. 503:
// no room yet, send it again; 410: the transfer is over.
func (s *Server) handleXferPut(w http.ResponseWriter, r *http.Request) {
	x := s.xfers.get(r.PathValue("xid"))
	i, err := strconv.Atoi(r.PathValue("i"))
	if x == nil || err != nil {
		http.Error(w, "передача не найдена", http.StatusGone)
		return
	}
	data, err := readPiece(w, r)
	if err != nil {
		http.Error(w, "кусок не дошёл", http.StatusBadRequest)
		return
	}
	code, msg := x.put(r.Context(), i, data)
	if code == http.StatusNoContent {
		w.WriteHeader(code)
		return
	}
	http.Error(w, msg, code)
}

// POST /api/panel/fs/xfer/{xid}/cancel - the browser gives the upload up.
func (s *Server) handleXferCancel(w http.ResponseWriter, r *http.Request) {
	if x := s.xfers.get(r.PathValue("xid")); x != nil {
		x.fail("отменено")
	}
	w.WriteHeader(http.StatusNoContent)
}

// GET /api/panel/fs/xfer/{xid}/result - the upload's result, once the agent
// has written the file; {"done":false} means ask again.
func (s *Server) handleXferResult(w http.ResponseWriter, r *http.Request) {
	x := s.xfers.get(r.PathValue("xid"))
	if x == nil {
		writeJSON(w, map[string]any{"done": true, "result": proto.TaskResult{Error: "передача не найдена"}})
		return
	}
	if !x.await(r.Context(), xferWait, func() bool { return x.done }) {
		writeJSON(w, map[string]any{"done": false})
		return
	}
	res := x.result
	x.mu.Unlock()
	writeJSON(w, map[string]any{"done": true, "result": res})
}

// GET /api/agent/fs/xfer/{xid}/{i}?ack=K - piece i of the upload, for the
// agent to write. K says the agent has written every piece below it, so
// those are let go. 503: not here yet, ask again; 410: the transfer is over.
func (s *Server) handleAgentXferGet(w http.ResponseWriter, r *http.Request) {
	a := s.authenticateAgent(r)
	if a == nil {
		http.Error(w, "unauthorized", http.StatusUnauthorized)
		return
	}
	x := s.xfers.forAgent(r.PathValue("xid"), a.ID)
	i, err := strconv.Atoi(r.PathValue("i"))
	if x == nil || err != nil {
		http.Error(w, "передача не найдена", http.StatusGone)
		return
	}
	if ack, err := strconv.Atoi(r.URL.Query().Get("ack")); err == nil {
		x.mu.Lock()
		if ack > x.low {
			for k := x.low; k < ack; k++ {
				delete(x.chunks, k)
			}
			x.low = ack
			x.notify()
		}
		x.mu.Unlock()
	}
	if !x.await(r.Context(), xferWait, func() bool { return x.err != "" || x.chunks[i] != nil }) {
		http.Error(w, "кусок ещё не пришёл", http.StatusServiceUnavailable)
		return
	}
	if x.err != "" {
		msg := x.err
		x.mu.Unlock()
		http.Error(w, msg, http.StatusGone)
		return
	}
	data := x.chunks[i]
	x.mu.Unlock()
	w.Header().Set("Content-Type", "application/octet-stream")
	w.Header().Set("Content-Length", strconv.Itoa(len(data)))
	_, _ = w.Write(data)
}

// ---------- download: agent -> browser ----------

// inlineTypes are the files the panel shows in place rather than downloads.
var inlineTypes = map[string]string{
	".png": "image/png", ".jpg": "image/jpeg", ".jpeg": "image/jpeg", ".gif": "image/gif",
	".webp": "image/webp", ".bmp": "image/bmp", ".ico": "image/x-icon", ".avif": "image/avif",
	".svg": "image/svg+xml",
	".mp4": "video/mp4", ".m4v": "video/mp4", ".webm": "video/webm", ".ogv": "video/ogg", ".mov": "video/quicktime",
	".mp3": "audio/mpeg", ".wav": "audio/wav", ".ogg": "audio/ogg", ".oga": "audio/ogg", ".m4a": "audio/mp4",
	".aac": "audio/aac", ".flac": "audio/flac", ".opus": "audio/ogg",
	".pdf": "application/pdf",
}

// fileHeaders sets what the browser needs to download or show the file. A
// file shown in place runs nothing: no scripts in an SVG, whatever it holds.
// A PDF is the exception, as the browser's viewer will not open in a sandbox;
// it runs in that viewer, not in the panel's page.
func fileHeaders(w http.ResponseWriter, p string, inline bool) {
	name := path.Base(p)
	ctype := "application/octet-stream"
	if t, ok := inlineTypes[strings.ToLower(path.Ext(name))]; ok && inline {
		ctype = t
	}
	disp := "attachment"
	if inline {
		disp = "inline"
	}
	w.Header().Set("Content-Type", ctype)
	w.Header().Set("Content-Disposition", disp+"; filename=\""+strings.ReplaceAll(name, "\"", "")+"\"")
	w.Header().Set("X-Content-Type-Options", "nosniff")
	if ctype != "application/pdf" {
		w.Header().Set("Content-Security-Policy", "sandbox; default-src 'none'; img-src data:; style-src 'unsafe-inline'")
	}
	w.Header().Set("Cache-Control", "no-store")
}

// GET /api/panel/agents/{id}/fs/download?path=...[&inline=1] - the file, as
// one response fed by the pieces the agent sends. inline=1 serves pictures,
// video and sound for the page to show.
func (s *Server) handleFSDownload(w http.ResponseWriter, r *http.Request) {
	agentID, p := r.PathValue("id"), r.URL.Query().Get("path")
	inline := r.URL.Query().Get("inline") != ""
	if !s.fsAgent(w, agentID, false) {
		return
	}
	x := s.xfers.open(agentID, -1)
	defer s.xfers.drop(x.id)
	defer x.fail("браузер ушёл") // stops the agent if the response ends early

	raw, _ := json.Marshal(proto.FSXferArgs{Path: p, Xfer: x.id, Size: -1, Chunk: xferChunk})
	taskID, ch := s.hub.CreateTask(agentID, proto.TaskFSSend, raw)
	defer s.hub.DiscardTask(taskID)
	// The agent's result: an error ends the transfer wherever it stands.
	results := make(chan proto.TaskResult, 1)
	go func() {
		select {
		case res := <-ch:
			if !res.OK && !unknownOp(res) {
				x.fail(res.Error)
			}
			results <- res
		case <-r.Context().Done():
		}
	}()

	select {
	case <-x.sized:
	case res := <-results:
		if unknownOp(res) {
			s.legacyDownload(w, r, agentID, p, inline)
			return
		}
		msg := res.Error
		if msg == "" {
			msg = "агент не отдал файл"
		}
		http.Error(w, msg, http.StatusBadGateway)
		return
	case <-r.Context().Done():
		return
	case <-time.After(fsTaskTimeout):
		http.Error(w, "таймаут ожидания агента", http.StatusGatewayTimeout)
		return
	}

	fileHeaders(w, p, inline)
	w.Header().Set("Content-Length", strconv.FormatInt(x.size, 10))
	rc := http.NewResponseController(w)
	for i := 0; i < x.n; i++ {
		data, err := x.take(r.Context(), i, xferIdle)
		if err == nil {
			_, err = w.Write(data)
		}
		if err != nil {
			// Cut the response off, so the browser does not take a part of
			// the file for all of it.
			x.fail(err.Error())
			panic(http.ErrAbortHandler)
		}
		_ = rc.Flush()
	}
}

// legacyDownload gets a file from an agent that predates the pieces.
func (s *Server) legacyDownload(w http.ResponseWriter, r *http.Request, agentID, p string, inline bool) {
	res, err := s.waitTask(r.Context(), agentID, proto.TaskFSDownload, proto.FSPathArgs{Path: p}, fsTaskTimeout)
	if err != nil {
		http.Error(w, err.Error(), http.StatusGatewayTimeout)
		return
	}
	if !res.OK {
		http.Error(w, res.Error, http.StatusBadGateway)
		return
	}
	var rr proto.FSReadResult
	if err := json.Unmarshal(res.Data, &rr); err != nil {
		http.Error(w, "bad result", http.StatusBadGateway)
		return
	}
	data, _ := base64.StdEncoding.DecodeString(rr.Content)
	fileHeaders(w, p, inline)
	if rr.Truncated {
		w.Header().Set("X-Truncated", "1")
	}
	_, _ = w.Write(data)
}

// POST /api/agent/fs/xfer/{xid}/{i} - piece i of the file the browser asked
// for, X-Size its whole size. 503: no room yet, send it again; 410: the
// browser is gone.
func (s *Server) handleAgentXferPost(w http.ResponseWriter, r *http.Request) {
	a := s.authenticateAgent(r)
	if a == nil {
		http.Error(w, "unauthorized", http.StatusUnauthorized)
		return
	}
	x := s.xfers.forAgent(r.PathValue("xid"), a.ID)
	i, err := strconv.Atoi(r.PathValue("i"))
	if x == nil || err != nil {
		http.Error(w, "передача не найдена", http.StatusGone)
		return
	}
	size, err := strconv.ParseInt(r.Header.Get("X-Size"), 10, 64)
	if err != nil || size < 0 || size > proto.FSFileCap || !x.setSize(size) {
		http.Error(w, "не тот размер", http.StatusBadRequest)
		return
	}
	data, err := readPiece(w, r)
	if err != nil {
		http.Error(w, "кусок не дошёл", http.StatusBadRequest)
		return
	}
	code, msg := x.put(r.Context(), i, data)
	if code == http.StatusNoContent {
		w.WriteHeader(code)
		return
	}
	http.Error(w, msg, code)
}