lpra/internal/agent/fsstream.go

package agent

import (
	"bytes"
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"net/http"
	"os"
	"path/filepath"
	"strconv"
	"strings"
	"sync"
	"syscall"
	"time"

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

// Files go to and from the browser in pieces, each a request of its own
// through the server: small enough for any proxy in front, and sent again
// alone when the link drops it. A few go at once, so a good link is kept as
// busy as by one long stream.

const (
	// xferWorkers is how many pieces are in flight at once.
	xferWorkers = 4
	// xferPieceTimeout bounds one piece's request: a link that slow, or
	// stuck, gets the piece sent again.
	xferPieceTimeout = 5 * time.Minute
	// xferTries bounds how often a piece is tried over a failing link.
	xferTries = 30
)

// errXferGone is the server saying the transfer is over: the browser left
// or gave up.
type errXferGone string

func (e errXferGone) Error() string { return string(e) }

func (a *Agent) xferURL(id string, i int) string {
	return a.host() + "/api/agent/fs/xfer/" + id + "/" + strconv.Itoa(i)
}

// xferRequest makes one request for a piece, trying again over a failing link.
// A 503 is the server waiting on the other side: asked again at once.
func (a *Agent) xferRequest(ctx context.Context, newReq func(context.Context) (*http.Request, error), keep func(*http.Response) error) error {
	fails := 0
	for {
		if err := ctx.Err(); err != nil {
			return err
		}
		rctx, cancel := context.WithTimeout(ctx, xferPieceTimeout)
		req, err := newReq(rctx)
		if err != nil {
			cancel()
			return err
		}
		a.authHeaders(req)
		resp, err := a.stream.Do(req)
		if err == nil {
			switch {
			case resp.StatusCode == http.StatusGone:
				msg, _ := io.ReadAll(io.LimitReader(resp.Body, 512))
				resp.Body.Close()
				cancel()
				return errXferGone(strings.TrimSpace(string(msg)))
			case resp.StatusCode == http.StatusServiceUnavailable:
				io.Copy(io.Discard, resp.Body)
				resp.Body.Close()
				cancel()
				continue
			case resp.StatusCode < 300:
				err = keep(resp)
				resp.Body.Close()
				cancel()
				if err == nil {
					return nil
				}
			default:
				io.Copy(io.Discard, resp.Body)
				resp.Body.Close()
				err = errStatus(resp.StatusCode)
			}
		}
		cancel()
		if fails++; fails >= xferTries {
			return fmt.Errorf("сеть не держит передачу: %w", err)
		}
		select {
		case <-time.After(min(time.Duration(fails)*time.Second, 10*time.Second)):
		case <-ctx.Done():
			return ctx.Err()
		}
	}
}

// pieces runs fn for every piece, xferWorkers at a time, until one fails.
func pieces(ctx context.Context, n int, fn func(ctx context.Context, i int) error) error {
	ctx, cancel := context.WithCancel(ctx)
	defer cancel()
	var (
		mu    sync.Mutex
		next  int
		first error
		wg    sync.WaitGroup
	)
	for w := 0; w < min(xferWorkers, n); w++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			for {
				mu.Lock()
				if first != nil || next >= n {
					mu.Unlock()
					return
				}
				i := next
				next++
				mu.Unlock()
				if err := fn(ctx, i); err != nil {
					mu.Lock()
					if first == nil {
						first = err
						cancel()
					}
					mu.Unlock()
					return
				}
			}
		}()
	}
	wg.Wait()
	return first
}

func pieceCount(size int64, chunk int) int {
	if size <= 0 {
		return 1
	}
	return int((size + int64(chunk) - 1) / int64(chunk))
}

// doFSPut 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) doFSPut(ctx context.Context, env proto.TaskEnvelope) proto.TaskResult {
	var args proto.FSXferArgs
	if err := json.Unmarshal(env.Args, &args); err != nil {
		return proto.TaskResult{Error: err.Error()}
	}
	if args.Size < 0 || args.Size > proto.FSFileCap || args.Chunk <= 0 {
		return proto.TaskResult{Error: "файл больше 512 МБ"}
	}
	dir, base := filepath.Split(args.Path)
	if dir == "" {
		dir = "."
	}
	if info, err := os.Stat(args.Path); err == nil && info.IsDir() {
		return proto.TaskResult{Error: "это каталог"}
	}
	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())
		}
	}()

	// Pieces written so far, to tell the server which it can let go.
	n := pieceCount(args.Size, args.Chunk)
	var (
		mu     sync.Mutex
		got    = make([]bool, n)
		contig int
	)
	err = pieces(ctx, n, func(ctx context.Context, i int) error {
		want := min(int64(args.Chunk), args.Size-int64(i)*int64(args.Chunk))
		mu.Lock()
		ack := contig
		mu.Unlock()
		var data []byte
		err := a.xferRequest(ctx, func(ctx context.Context) (*http.Request, error) {
			return http.NewRequestWithContext(ctx, http.MethodGet, a.xferURL(args.Xfer, i)+"?ack="+strconv.Itoa(ack), nil)
		}, func(resp *http.Response) error {
			b, err := io.ReadAll(io.LimitReader(resp.Body, int64(args.Chunk)+1))
			if err == nil && int64(len(b)) != want {
				err = errors.New("кусок не той длины")
			}
			data = b
			return err
		})
		if err != nil {
			return err
		}
		if _, err := tmp.WriteAt(data, int64(i)*int64(args.Chunk)); err != nil {
			return err
		}
		mu.Lock()
		got[i] = true
		for contig < n && got[contig] {
			contig++
		}
		mu.Unlock()
		return nil
	})
	if err != nil {
		var gone errXferGone
		if errors.As(err, &gone) {
			return proto.TaskResult{Error: "загрузка прервана: " + string(gone)}
		}
		return proto.TaskResult{Error: "загрузка прервана: " + err.Error()}
	}
	if err := tmp.Truncate(args.Size); err != nil {
		return proto.TaskResult{Error: err.Error()}
	}
	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 {
		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": args.Size})
	return proto.TaskResult{OK: true, Data: data}
}

// doFSSend sends a file to the browser that asked for it.
func (a *Agent) doFSSend(ctx context.Context, env proto.TaskEnvelope) proto.TaskResult {
	var args proto.FSXferArgs
	if err := json.Unmarshal(env.Args, &args); err != nil {
		return proto.TaskResult{Error: err.Error()}
	}
	if args.Chunk <= 0 {
		return proto.TaskResult{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: "это каталог"}
	}
	size := info.Size()
	if size > proto.FSFileCap {
		return proto.TaskResult{Error: "файл больше 512 МБ"}
	}
	sizeHdr := strconv.FormatInt(size, 10)
	err = pieces(ctx, pieceCount(size, args.Chunk), func(ctx context.Context, i int) error {
		off := int64(i) * int64(args.Chunk)
		buf := make([]byte, min(int64(args.Chunk), size-off))
		if _, err := f.ReadAt(buf, off); err != nil && err != io.EOF {
			return err
		}
		return a.xferRequest(ctx, func(ctx context.Context) (*http.Request, error) {
			req, err := http.NewRequestWithContext(ctx, http.MethodPost, a.xferURL(args.Xfer, i), bytes.NewReader(buf))
			if err == nil {
				req.Header.Set("X-Size", sizeHdr)
				req.Header.Set("Content-Type", "application/octet-stream")
			}
			return req, err
		}, func(*http.Response) error { return nil })
	})
	if err != nil {
		return proto.TaskResult{Error: "передача прервана: " + err.Error()}
	}
	data, _ := json.Marshal(map[string]any{"path": args.Path, "size": size})
	return proto.TaskResult{OK: true, Data: data}
}