lpra/internal/agent/fsstream.go

package agent

import (
	"context"
	"encoding/json"
	"fmt"
	"io"
	"log"
	"net/http"
	"os"
	"path/filepath"
	"strconv"
	"syscall"
	"time"

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

// File transfers that do not fit a task's JSON: the bytes go as a plain HTTP
// body through a stream the server holds open for the browser on the other
// end, so a file is never held whole in memory.

func (a *Agent) streamURL(id string) string { return a.cfg.Host + "/api/agent/fs/stream/" + id }

// doFSUpload writes the file the browser is uploading. It lands in a temporary
// file next to the target first, so a broken upload leaves the old file as it
// was.
func (a *Agent) doFSUpload(ctx context.Context, env proto.TaskEnvelope) proto.TaskResult {
	var args proto.FSStreamArgs
	if err := json.Unmarshal(env.Args, &args); err != nil {
		return proto.TaskResult{Error: err.Error()}
	}
	if args.Size > proto.FSFileCap {
		return proto.TaskResult{Error: "файл больше 512 МБ"}
	}
	dir, base := filepath.Split(args.Path)
	if dir == "" {
		dir = "."
	}
	tmp, err := os.CreateTemp(dir, "."+base+".lpra-*")
	if err != nil {
		return proto.TaskResult{Error: err.Error()}
	}
	done := false
	defer func() {
		if !done {
			tmp.Close()
			os.Remove(tmp.Name())
		}
	}()

	req, err := http.NewRequestWithContext(ctx, http.MethodGet, a.streamURL(args.Stream), nil)
	if err != nil {
		return proto.TaskResult{Error: err.Error()}
	}
	a.authHeaders(req)
	resp, err := a.stream.Do(req)
	if err != nil {
		return proto.TaskResult{Error: err.Error()}
	}
	defer resp.Body.Close()
	if resp.StatusCode != http.StatusOK {
		return proto.TaskResult{Error: errStatus(resp.StatusCode).Error()}
	}
	n, err := io.Copy(tmp, io.LimitReader(resp.Body, proto.FSFileCap+1))
	if err != nil {
		log.Printf("fs receive %s: broke after %d of %d bytes: %v", args.Path, n, args.Size, err)
		return proto.TaskResult{Error: "загрузка прервана: " + err.Error()}
	}
	if n > proto.FSFileCap {
		return proto.TaskResult{Error: "файл больше 512 МБ"}
	}
	if args.Size >= 0 && n != args.Size {
		return proto.TaskResult{Error: fmt.Sprintf("загрузка прервана: пришло %d из %d байт", n, args.Size)}
	}
	if err := tmp.Close(); err != nil {
		return proto.TaskResult{Error: err.Error()}
	}
	// A file that is being replaced keeps its mode and owner.
	mode := os.FileMode(0o644)
	if info, err := os.Stat(args.Path); err == nil {
		if info.IsDir() {
			return proto.TaskResult{Error: "это каталог"}
		}
		mode = info.Mode().Perm()
		if st, ok := info.Sys().(*syscall.Stat_t); ok {
			_ = os.Chown(tmp.Name(), int(st.Uid), int(st.Gid))
		}
	}
	if err := os.Chmod(tmp.Name(), mode); err != nil {
		return proto.TaskResult{Error: err.Error()}
	}
	if err := os.Rename(tmp.Name(), args.Path); err != nil {
		return proto.TaskResult{Error: err.Error()}
	}
	done = true
	data, _ := json.Marshal(map[string]any{"path": args.Path, "size": n})
	return proto.TaskResult{OK: true, Data: data}
}

// doFSGet sends a file to the browser that asked for it.
func (a *Agent) doFSGet(ctx context.Context, env proto.TaskEnvelope) proto.TaskResult {
	var args proto.FSStreamArgs
	if err := json.Unmarshal(env.Args, &args); err != nil {
		return proto.TaskResult{Error: err.Error()}
	}
	f, err := os.Open(args.Path)
	if err != nil {
		return proto.TaskResult{Error: err.Error()}
	}
	defer f.Close()
	info, err := f.Stat()
	if err != nil {
		return proto.TaskResult{Error: err.Error()}
	}
	if info.IsDir() {
		return proto.TaskResult{Error: "это каталог"}
	}
	if info.Size() > proto.FSFileCap {
		return proto.TaskResult{Error: "файл больше 512 МБ"}
	}
	req, err := http.NewRequestWithContext(ctx, http.MethodPost, a.streamURL(args.Stream), io.LimitReader(f, info.Size()))
	if err != nil {
		return proto.TaskResult{Error: err.Error()}
	}
	req.ContentLength = info.Size()
	req.Header.Set("X-Size", strconv.FormatInt(info.Size(), 10))
	a.authHeaders(req)
	start := time.Now()
	resp, err := a.stream.Do(req)
	if err != nil {
		log.Printf("fs send %s (%d bytes): failed after %s: %v", args.Path, info.Size(), time.Since(start).Round(time.Millisecond), err)
		return proto.TaskResult{Error: err.Error()}
	}
	io.Copy(io.Discard, resp.Body)
	resp.Body.Close()
	if resp.StatusCode != http.StatusNoContent {
		log.Printf("fs send %s (%d bytes): server answered %d after %s", args.Path, info.Size(), resp.StatusCode, time.Since(start).Round(time.Millisecond))
		return proto.TaskResult{Error: "браузер не дочитал файл"}
	}
	data, _ := json.Marshal(map[string]any{"path": args.Path, "size": info.Size()})
	return proto.TaskResult{OK: true, Data: data}
}