jwebfs/scalib/main.go

package main

import (
	"bufio"
	"flag"
	"fmt"
	"net/rpc"
	"net/rpc/jsonrpc"
	"os"
	"path/filepath"

	"scalib/util"
	"strings"

	"go.uber.org/zap"
)

// FileInfo 包含文件的元信息
type FileInfo struct {
	FileName string
	FileSize int64
}

// 文件块
type FileChunk struct {
	Data     []byte
	FileName string
	DirPath  string
	Offset   int64
	IsLast   bool
}

// 客户端
type UpFileClient struct {
	rpcClient *rpc.Client
}

// todo 后续增加安全验证
// 使用方法 eg:  ./scalib -h 192.168.66.92:9098 -f /www/gfs/sync_zips/BIU_20251009_153640.zip /www/afs/logs/
// 主入口
func main() {
	// 启用日志
	logger := util.NewProductionLogger()
	defer logger.Sync()

	// 使用flag解析命令行参数
	remote := flag.String("h", "", "The remote server address (ip:port)")
	curPath := flag.String("f", "", "The local file path to upload")
	uploadPath := flag.String("u", "", "The remote upload path")
	flag.Parse()
	// for debug
	fmt.Printf("remote: %s, curPath: %s, uploadPath: %s\n", *remote, *curPath, *uploadPath)
	// 将服务器的信息拆分为ip和端口
	harr := strings.Split(*remote, ":")
	dstip := harr[0]
	dport := harr[1]

	// 连接到RPC服务器
	service := fmt.Sprintf("%v:%v", dstip, dport)
	client, err := NewUpFileClient(service)
	if err != nil {
		logger.Error("NewUpFileClient failed", zap.Error(err))
	}
	defer client.rpcClient.Close()
	// 调用 transferFile
	err = transferFile(client, *curPath, *uploadPath)
	if err != nil {
		logger.Error("TransferFile failed", zap.Error(err))
	}
	// fmt.Printf("TransferFile success\n")
}

// 创建新的客户端
func NewUpFileClient(addr string) (*UpFileClient, error) {
	client, err := jsonrpc.Dial("tcp", addr)
	if err != nil {
		return nil, err
	}
	return &UpFileClient{rpcClient: client}, nil
}

// 传输文件
func transferFile(c *UpFileClient, curPath string, uploadPath string) error {
	// 启用日志
	logger := util.NewProductionLogger()
	defer logger.Sync()
	// 发送文件信息
	// fmt.Printf("TransferFile filePath: %s, uploadPath: %s\n", curPath, uploadPath)
	// 提取文件名
	fileName := filepath.Base(curPath)
	// 获取文件信息
	fileInfo, err := os.Stat(curPath)
	if err != nil {
		// panic(err)
		// fmt.Printf("获取文件信息失败: %v\n", err)
		return err
	}

	// 异步
	// go func() {
	// 发送文件信息
	var reply string
	c.rpcClient.Call("UpFileService.SendFileInfo", FileInfo{
		FileName: uploadPath,
		FileSize: fileInfo.Size(),
	}, &reply)
	// 输出执行的结果
	if reply != "true" {
		logger.Error("SendFileInfo failed", zap.String("reply", reply))
	}
	fmt.Printf("SendFileInfo success, reply: %s\n", reply)

	// 打开文件准备读取
	file, err := os.Open(curPath)
	if err != nil {
		panic(err)
	}
	defer file.Close()

	// 读取并发送文件块
	buffer := make([]byte, 4096) // 4KB块大小
	reader := bufio.NewReader(file)
	var offset int64 = 0

	// 目标文件目录
	dirPath := filepath.Dir(uploadPath)

	fmt.Printf("开始发送文件: %s (大小: %d bytes)\n", fileName, fileInfo.Size())

	var scReply bool
	for {
		n, err := reader.Read(buffer)
		if err != nil {
			break
		}

		// 判断是否是最后一块
		isLast := offset+int64(n) >= fileInfo.Size()

		// 发送文件块
		err = c.rpcClient.Call("UpFileService.SendFileChunk", FileChunk{
			Data:     buffer[:n],
			FileName: fileName,
			DirPath:  dirPath,
			Offset:   offset,
			IsLast:   isLast,
		}, &scReply)

		if err != nil || !scReply {
			// panic("发送文件块失败: " + err.Error())
			logger.Error("SendFileChunk failed", zap.Error(err))
		}

		offset += int64(n)

		// 显示进度
		progress := float64(offset) / float64(fileInfo.Size()) * 100
		fmt.Printf("\r发送进度: %.2f%%", progress)
	}

	fmt.Println("\n文件发送完成!")

	// }()

	return nil
}