lpra/internal/server/fsstream.go

package server

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

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

// Files go between the browser and the agent as plain HTTP bodies, joined on
// the server by a pipe: the browser's upload is the body of the agent's GET,
// the agent's POST is the body of the browser's download. Nothing is held
// whole in memory and nothing is base64.

// fsStreamTimeout bounds one file transfer, however slow the links.
const fsStreamTimeout = 2 * time.Hour

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

// fsStream is one file on its way.
type fsStream struct {
	agentID string
	r       *io.PipeReader
	w       *io.PipeWriter
	size    int64         // bytes to come, -1 when not known
	ready   chan struct{} // download: closed when the agent starts sending
	once    sync.Once
	taken   atomic.Bool // the agent has come for it
}

// start marks the download as under way with its size; false when it already
// is.
func (st *fsStream) start(size int64) bool {
	ok := false
	st.once.Do(func() {
		st.size = size
		close(st.ready)
		ok = true
	})
	return ok
}

type fsStreams struct {
	mu sync.Mutex
	m  map[string]*fsStream
}

func (c *fsStreams) open(agentID string, size int64) (string, *fsStream) {
	st := &fsStream{agentID: agentID, size: size, ready: make(chan struct{})}
	st.r, st.w = io.Pipe()
	id := randToken()[:24]
	c.mu.Lock()
	if c.m == nil {
		c.m = map[string]*fsStream{}
	}
	c.m[id] = st
	c.mu.Unlock()
	return id, st
}

// get returns the stream if it belongs to the agent asking.
func (c *fsStreams) get(id, agentID string) *fsStream {
	c.mu.Lock()
	defer c.mu.Unlock()
	if st := c.m[id]; st != nil && st.agentID == agentID {
		return st
	}
	return nil
}

func (c *fsStreams) drop(id string) {
	c.mu.Lock()
	st := c.m[id]
	delete(c.m, id)
	c.mu.Unlock()
	if st != nil {
		st.r.CloseWithError(io.ErrClosedPipe)
		st.w.CloseWithError(io.ErrClosedPipe)
	}
}

// waitTask runs a task on the agent and waits for its result, or for the
// browser to leave.
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
}

// POST /api/panel/agents/{id}/fs/upload?path=... - write a file; the request
// body is the file, passed on to the agent as it arrives.
func (s *Server) handleFSUpload(w http.ResponseWriter, r *http.Request) {
	agentID, p := r.PathValue("id"), r.URL.Query().Get("path")
	if !s.fsAgent(w, agentID, true) {
		return
	}
	if p == "" {
		writeJSON(w, proto.TaskResult{Error: "не указан путь"})
		return
	}
	if r.ContentLength > proto.FSFileCap {
		writeJSON(w, proto.TaskResult{Error: "файл больше 512 МБ"})
		return
	}
	body := http.MaxBytesReader(w, r.Body, proto.FSFileCap)
	id, st := s.fsStreams.open(agentID, r.ContentLength)
	defer s.fsStreams.drop(id)
	copied := make(chan struct{})
	go func() {
		defer close(copied)
		n, err := io.Copy(st.w, body)
		if err != nil && !errors.Is(err, io.ErrClosedPipe) {
			log.Printf("fs upload %s to %s: browser stream broke after %d of %d bytes: %v", p, agentID, n, r.ContentLength, err)
		}
		st.w.CloseWithError(err) // nil: the reader sees the end of the file
	}()

	args := proto.FSStreamArgs{Path: p, Stream: id, Size: r.ContentLength}
	res, err := s.waitTask(r.Context(), agentID, proto.TaskFSUpload, args, fsStreamTimeout)
	if err == nil && unknownOp(res) && !st.taken.Load() {
		res = s.legacyUpload(r.Context(), agentID, p, st.r)
	}
	// Stop the copy and let it end before the handler does: the body is not
	// to be read after that.
	st.r.CloseWithError(io.ErrClosedPipe)
	<-copied
	if err != nil {
		writeJSON(w, proto.TaskResult{Error: err.Error()})
		return
	}
	writeJSON(w, res)
}

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

// 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,
// streamed from the agent as it reads it. 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
	}
	id, st := s.fsStreams.open(agentID, -1)
	defer s.fsStreams.drop(id)
	raw, _ := json.Marshal(proto.FSStreamArgs{Path: p, Stream: id, Size: -1})
	taskID, ch := s.hub.CreateTask(agentID, proto.TaskFSGet, raw)
	defer s.hub.DiscardTask(taskID)

	select {
	case <-st.ready:
		fileHeaders(w, p, inline)
		if st.size >= 0 {
			w.Header().Set("Content-Length", strconv.FormatInt(st.size, 10))
		}
		if n, err := io.Copy(w, st.r); err != nil {
			// Cut the response off, so the browser does not take a part of
			// the file for all of it.
			log.Printf("fs download %s from %s: cut off after %d of %d bytes: %v", p, agentID, n, st.size, err)
			st.r.CloseWithError(err)
			panic(http.ErrAbortHandler)
		}
	case res := <-ch:
		if unknownOp(res) {
			s.legacyDownload(w, r, agentID, p, inline)
			return
		}
		msg := res.Error
		if msg == "" {
			msg = "агент не отдал файл"
		}
		http.Error(w, msg, http.StatusBadGateway)
	case <-r.Context().Done():
	case <-time.After(fsTaskTimeout):
		http.Error(w, "таймаут ожидания агента", http.StatusGatewayTimeout)
	}
}

// legacyDownload gets a file from an agent that predates the streams.
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)
}

// GET /api/agent/fs/stream/{sid} - the file the browser is uploading, for the
// agent to write.
func (s *Server) handleAgentFSStreamGet(w http.ResponseWriter, r *http.Request) {
	a := s.authenticateAgent(r)
	if a == nil {
		http.Error(w, "unauthorized", http.StatusUnauthorized)
		return
	}
	st := s.fsStreams.get(r.PathValue("sid"), a.ID)
	if st == nil || !st.taken.CompareAndSwap(false, true) {
		http.Error(w, "gone", http.StatusGone)
		return
	}
	w.Header().Set("Content-Type", "application/octet-stream")
	w.Header().Set("X-Accel-Buffering", "no")
	if st.size >= 0 {
		w.Header().Set("Content-Length", strconv.FormatInt(st.size, 10))
	}
	if n, err := io.Copy(w, st.r); err != nil {
		// The browser broke off: cut the response, so the agent does not
		// take a part of the file for all of it.
		log.Printf("fs upload to %s: cut off after %d of %d bytes: %v", a.ID, n, st.size, err)
		panic(http.ErrAbortHandler)
	}
}

// POST /api/agent/fs/stream/{sid} - the file the browser asked for, sent by
// the agent. 204 once the browser has it all.
func (s *Server) handleAgentFSStreamPost(w http.ResponseWriter, r *http.Request) {
	a := s.authenticateAgent(r)
	if a == nil {
		http.Error(w, "unauthorized", http.StatusUnauthorized)
		return
	}
	st := s.fsStreams.get(r.PathValue("sid"), a.ID)
	if st == nil {
		http.Error(w, "gone", http.StatusGone)
		return
	}
	size, err := strconv.ParseInt(r.Header.Get("X-Size"), 10, 64)
	if err != nil {
		size = -1
	}
	if !st.taken.CompareAndSwap(false, true) || !st.start(size) {
		http.Error(w, "gone", http.StatusGone)
		return
	}
	n, err := io.Copy(st.w, http.MaxBytesReader(w, r.Body, proto.FSFileCap))
	st.w.CloseWithError(err)
	if err != nil {
		log.Printf("fs download from %s: agent stream broke after %d of %d bytes: %v", a.ID, n, size, err)
		http.Error(w, "gone", http.StatusGone)
		return
	}
	w.WriteHeader(http.StatusNoContent)
}