github.com/TeaOSLab/EdgeNode@v1.3.8/internal/nodes/ip_library_updater.go (about)

     1  // Copyright 2022 Liuxiangchao iwind.liu@gmail.com. All rights reserved. Official site: https://goedge.cn .
     2  
     3  package nodes
     4  
     5  import (
     6  	"errors"
     7  	"github.com/TeaOSLab/EdgeCommon/pkg/rpc/pb"
     8  	"github.com/TeaOSLab/EdgeNode/internal/remotelogs"
     9  	"github.com/TeaOSLab/EdgeNode/internal/rpc"
    10  	"github.com/iwind/TeaGo/Tea"
    11  	"github.com/iwind/TeaGo/types"
    12  	"io"
    13  	"os"
    14  )
    15  
    16  type IPLibraryUpdater struct {
    17  }
    18  
    19  func NewIPLibraryUpdater() *IPLibraryUpdater {
    20  	return &IPLibraryUpdater{}
    21  }
    22  
    23  // DataDir 文件目录
    24  func (this *IPLibraryUpdater) DataDir() string {
    25  	// data/
    26  	var dir = Tea.Root + "/data"
    27  	stat, err := os.Stat(dir)
    28  	if err == nil && stat.IsDir() {
    29  		return dir
    30  	}
    31  
    32  	err = os.Mkdir(dir, 0666)
    33  	if err == nil {
    34  		return dir
    35  	}
    36  
    37  	remotelogs.Error("IP_LIBRARY_UPDATER", "create directory '"+dir+"' failed: "+err.Error())
    38  
    39  	// 如果不能创建 data/ 目录，那么使用临时目录
    40  	return os.TempDir()
    41  }
    42  
    43  // FindLatestFile 检查最新的IP库文件
    44  func (this *IPLibraryUpdater) FindLatestFile() (code string, fileId int64, err error) {
    45  	rpcClient, err := rpc.SharedRPC()
    46  	if err != nil {
    47  		return "", 0, err
    48  	}
    49  	resp, err := rpcClient.IPLibraryArtifactRPC.FindPublicIPLibraryArtifact(rpcClient.Context(), &pb.FindPublicIPLibraryArtifactRequest{})
    50  	if err != nil {
    51  		return "", 0, err
    52  	}
    53  	var artifact = resp.IpLibraryArtifact
    54  	if artifact == nil {
    55  		return
    56  	}
    57  	return artifact.Code, artifact.FileId, nil
    58  }
    59  
    60  // DownloadFile 下载文件
    61  func (this *IPLibraryUpdater) DownloadFile(fileId int64, writer io.Writer) error {
    62  	if fileId <= 0 {
    63  		return errors.New("invalid fileId: " + types.String(fileId))
    64  	}
    65  
    66  	rpcClient, err := rpc.SharedRPC()
    67  	if err != nil {
    68  		return err
    69  	}
    70  
    71  	chunkIdsResp, err := rpcClient.FileChunkRPC.FindAllFileChunkIds(rpcClient.Context(), &pb.FindAllFileChunkIdsRequest{FileId: fileId})
    72  	if err != nil {
    73  		return err
    74  	}
    75  	for _, chunkId := range chunkIdsResp.FileChunkIds {
    76  		chunkResp, err := rpcClient.FileChunkRPC.DownloadFileChunk(rpcClient.Context(), &pb.DownloadFileChunkRequest{FileChunkId: chunkId})
    77  		if err != nil {
    78  			return err
    79  		}
    80  		var chunk = chunkResp.FileChunk
    81  		if chunk == nil {
    82  			return errors.New("can not find file chunk with chunk id '" + types.String(chunkId) + "'")
    83  		}
    84  		_, err = writer.Write(chunk.Data)
    85  		if err != nil {
    86  			return err
    87  		}
    88  	}
    89  	return nil
    90  }
    91  
    92  // LogInfo 普通日志
    93  func (this *IPLibraryUpdater) LogInfo(message string) {
    94  	remotelogs.Println("IP_LIBRARY_UPDATER", message)
    95  }
    96  
    97  // LogError 错误日志
    98  func (this *IPLibraryUpdater) LogError(err error) {
    99  	if err == nil {
   100  		return
   101  	}
   102  	remotelogs.Error("IP_LIBRARY_UPDATER", err.Error())
   103  }