pg/cmd/pgcli/share/filemanager.go

package share

import (
	"bytes"
	"crypto/sha256"
	"encoding/binary"
	"fmt"
	"io"
	"log/slog"
	"net"
	"net/url"
	"os"
	"os/user"
	"path/filepath"
	"strings"
	"sync"
	"time"

	"github.com/schollz/progressbar/v3"
)

type FileManager struct {
	mutex sync.RWMutex
	index int
	files map[int]string
}

func (fm *FileManager) Add(file string) (int, error) {
	fm.mutex.Lock()
	defer fm.mutex.Unlock()
	absPath, err := filepath.Abs(file)
	if err != nil {
		return -1, err
	}
	if relPath, ok := strings.CutPrefix(file, "~"); ok {
		curUser, err := user.Current()
		if err != nil {
			return -1, err
		}
		absPath = filepath.Join(curUser.HomeDir, relPath)
	}
	fm.files[fm.index] = absPath
	fm.index++
	return fm.index - 1, nil
}

func (fm *FileManager) Open(index int) (*os.File, error) {
	fm.mutex.RLock()
	defer fm.mutex.RUnlock()
	if absPath, ok := fm.files[index]; ok {
		return os.Open(absPath)
	}
	return nil, os.ErrNotExist
}

func (fm *FileManager) HandleRequest(peerID string, conn net.Conn) {
	defer conn.Close()
	header := make([]byte, 4)
	_, err := io.ReadFull(conn, header)
	if err != nil || !bytes.Equal(header[:2], []byte{0, 0}) {
		conn.Write(buildErr(1))
		slog.Error("Magic not verified, closing connection...")
		return
	}

	index := binary.BigEndian.Uint16(header[2:])
	f, err := fm.Open(int(index))
	if err != nil {
		conn.Write(buildErr(2))
		slog.Error("Open file failed", "err", err)
		return
	}
	defer f.Close()

	stat, err := f.Stat()
	if err != nil {
		return
	}

	conn.Write(buildOK(stat.Size()))

	go func() {
		header := make([]byte, 1)
		io.ReadFull(conn, header)
		switch header[0] {
		case 1:
			conn.Close()
		}
	}()

	bar := progressbar.NewOptions64(
		stat.Size(),
		progressbar.OptionSetDescription(fmt.Sprintf("%s:%s", peerID, url.QueryEscape(stat.Name()))),
		progressbar.OptionSetWriter(os.Stderr),
		progressbar.OptionShowBytes(true),
		progressbar.OptionSetWidth(10),
		progressbar.OptionThrottle(200*time.Millisecond),
		progressbar.OptionShowCount(),
		progressbar.OptionShowElapsedTimeOnFinish(),
		progressbar.OptionOnCompletion(func() {
			fmt.Fprint(os.Stderr, "\n")
		}),
		progressbar.OptionSetTheme(progressbar.Theme{
			Saucer:        "=",
			SaucerHead:    ">",
			SaucerPadding: " ",
			BarStart:      "[",
			BarEnd:        "]",
		}),
		progressbar.OptionSpinnerType(14),
		progressbar.OptionFullWidth(),
		progressbar.OptionSetRenderBlankState(true),
	)

	sha256Checksum := sha256.New()
	_, err = io.Copy(io.MultiWriter(&sender{conn}, bar, sha256Checksum), f)
	if err != nil {
		slog.Info("Copy file failed", "err", err)
	}
	checksum := sha256Checksum.Sum(nil)
	slog.Debug("Checksum", "sum", checksum)
	conn.Write(checksum)
}

type sender struct {
	net.Conn
}

func (s *sender) Write(b []byte) (n int, err error) {
	s.Conn.SetWriteDeadline(time.Now().Add(5 * time.Second))
	return s.Conn.Write(b)
}

func buildOK(fileSize int64) []byte {
	pkt := make([]byte, 5)
	copy(pkt[1:], binary.BigEndian.AppendUint32(nil, uint32(fileSize)))
	return pkt
}

func buildErr(code byte) []byte {
	pkt := make([]byte, 5)
	pkt[0] = code
	return pkt
}