Initial commit (code only without large binaries)

EdgeAdmin/internal/web/actions/default/settings/ip-library/upload.go

@@ -0,0 +1,213 @@
+// Copyright 2022 Liuxiangchao iwind.liu@gmail.com. All rights reserved. Official site: https://goedge.cn .
+
+package iplibrary
+
+import (
+	"bytes"
+	"encoding/json"
+	"fmt"
+	"github.com/TeaOSLab/EdgeAdmin/internal/utils/sizes"
+	"github.com/TeaOSLab/EdgeAdmin/internal/web/actions/actionutils"
+	"github.com/TeaOSLab/EdgeAdmin/internal/web/actions/default/settings/ip-library/iplibraryutils"
+	iplib "github.com/TeaOSLab/EdgeCommon/pkg/iplibrary"
+	"github.com/TeaOSLab/EdgeCommon/pkg/rpc/pb"
+	"github.com/iwind/TeaGo/Tea"
+	"github.com/iwind/TeaGo/actions"
+	"github.com/iwind/TeaGo/logs"
+	"io"
+	"os"
+	"path/filepath"
+	"strings"
+)
+
+type UploadAction struct {
+	actionutils.ParentAction
+}
+
+func (this *UploadAction) Init() {
+	this.Nav("", "", "upload")
+}
+
+func (this *UploadAction) RunGet(params struct{}) {
+	this.Data["canAccess"] = iplibraryutils.CanAccess()
+
+	this.Show()
+}
+
+func (this *UploadAction) RunPost(params struct {
+	Name string
+	File *actions.File
+
+	Must *actions.Must
+	CSRF *actionutils.CSRF
+}) {
+	if len(params.Name) == 0 {
+		this.FailField("name", "请输入IP库名称")
+		return
+	}
+
+	if params.File == nil {
+		this.Fail("请选择要上传的IP库文件")
+		return
+	}
+
+	fp, err := params.File.OriginFile.Open()
+	if err != nil {
+		this.Fail("读取IP库文件失败：" + err.Error())
+		return
+	}
+	defer func() {
+		_ = fp.Close()
+	}()
+
+	data, err := io.ReadAll(io.LimitReader(fp, 64*sizes.M)) // 最大不超过64M
+	if err != nil {
+		this.Fail("读取IP库文件失败：" + err.Error())
+		return
+	}
+
+	// 只支持 MaxMind 格式文件（.mmdb）
+	filename := strings.ToLower(params.File.Filename)
+	if !strings.HasSuffix(filename, ".mmdb") {
+		this.Fail("只支持 MaxMind 格式文件（.mmdb），请上传 GeoLite2-City.mmdb 或 GeoLite2-ASN.mmdb 文件")
+		return
+	}
+
+	// MaxMind 格式文件，保存到 data/iplibrary/ 目录
+	iplibDir := Tea.Root + "/data/iplibrary"
+	err = os.MkdirAll(iplibDir, 0755)
+	if err != nil {
+		this.Fail("创建IP库目录失败：" + err.Error())
+		return
+	}
+
+	// 根据文件名判断是 City 还是 ASN
+	var targetPath string
+	if strings.Contains(filename, "city") {
+		targetPath = filepath.Join(iplibDir, "maxmind-city.mmdb")
+	} else if strings.Contains(filename, "asn") {
+		targetPath = filepath.Join(iplibDir, "maxmind-asn.mmdb")
+	} else {
+		this.Fail("MaxMind 文件名必须包含 'city' 或 'asn'")
+		return
+	}
+
+	// 保存文件（使用临时文件原子替换）
+	tmpPath := targetPath + ".tmp"
+	err = os.WriteFile(tmpPath, data, 0644)
+	if err != nil {
+		this.Fail("保存IP库文件失败：" + err.Error())
+		return
+	}
+
+	// 原子替换
+	err = os.Rename(tmpPath, targetPath)
+	if err != nil {
+		os.Remove(tmpPath)
+		this.Fail("替换IP库文件失败：" + err.Error())
+		return
+	}
+
+	// 通过 RPC 将文件上传到 EdgeAPI
+	_, err = this.RPC().IPLibraryRPC().UploadMaxMindFile(this.AdminContext(), &pb.UploadMaxMindFileRequest{
+		Filename: params.File.Filename,
+		Data:     data,
+	})
+	if err != nil {
+		logs.Println("[IP_LIBRARY]upload MaxMind file to EdgeAPI failed: " + err.Error())
+		// 继续执行，不影响本地保存
+	}
+
+	// 创建简单的 Meta
+	meta := &iplib.Meta{
+		Version: 3,
+		Code:    "maxmind",
+		Author:  "MaxMind",
+	}
+	meta.Init()
+
+	// TODO 检查是否要自动创建省市区
+
+	// 上传IP库文件到数据库
+	fileResp, err := this.RPC().FileRPC().CreateFile(this.AdminContext(), &pb.CreateFileRequest{
+		Filename: params.File.Filename,
+		Size:     int64(len(data)),
+		IsPublic: false,
+		MimeType: "",
+		Type:     "ipLibraryArtifact",
+	})
+	if err != nil {
+		this.ErrorPage(err)
+		return
+	}
+	var fileId = fileResp.FileId
+
+	var buf = make([]byte, 256*1024)
+	var dataReader = bytes.NewReader(data)
+	for {
+		n, err := dataReader.Read(buf)
+		if n > 0 {
+			_, chunkErr := this.RPC().FileChunkRPC().CreateFileChunk(this.AdminContext(), &pb.CreateFileChunkRequest{
+				FileId: fileId,
+				Data:   buf[:n],
+			})
+			if chunkErr != nil {
+				this.Fail("上传文件到数据库失败：" + chunkErr.Error())
+				return
+			}
+		}
+		if err != nil {
+			break
+		}
+	}
+
+	// 创建IP库信息
+	metaJSON, err := json.Marshal(meta)
+	if err != nil {
+		this.Fail("元数据编码失败：" + err.Error())
+		return
+	}
+
+	createResp, err := this.RPC().IPLibraryArtifactRPC().CreateIPLibraryArtifact(this.AdminContext(), &pb.CreateIPLibraryArtifactRequest{
+		FileId:   fileId,
+		MetaJSON: metaJSON,
+		Name:     params.Name,
+	})
+	if err != nil {
+		this.Fail("创建IP库失败：" + err.Error())
+		return
+	}
+
+	// 获取所有IP库记录，将其他记录的 isPublic 设为 false，新上传的设为 true
+	artifactsResp, err := this.RPC().IPLibraryArtifactRPC().FindAllIPLibraryArtifacts(this.AdminContext(), &pb.FindAllIPLibraryArtifactsRequest{})
+	if err != nil {
+		// 如果获取列表失败，不影响上传成功，只记录日志
+		logs.Println("[IP_LIBRARY]failed to get all artifacts: " + err.Error())
+	} else {
+		// 将所有其他记录设为未使用
+		for _, artifact := range artifactsResp.IpLibraryArtifacts {
+			if artifact.Id != createResp.IpLibraryArtifactId {
+				if artifact.IsPublic {
+					// 将其他使用中的记录设为未使用
+					_, err = this.RPC().IPLibraryArtifactRPC().UpdateIPLibraryArtifactIsPublic(this.AdminContext(), &pb.UpdateIPLibraryArtifactIsPublicRequest{
+						IpLibraryArtifactId: artifact.Id,
+						IsPublic:            false,
+					})
+					if err != nil {
+						logs.Println("[IP_LIBRARY]failed to update artifact " + fmt.Sprintf("%d", artifact.Id) + " isPublic: " + err.Error())
+					}
+				}
+			}
+		}
+		// 将新上传的记录设为使用中
+		_, err = this.RPC().IPLibraryArtifactRPC().UpdateIPLibraryArtifactIsPublic(this.AdminContext(), &pb.UpdateIPLibraryArtifactIsPublicRequest{
+			IpLibraryArtifactId: createResp.IpLibraryArtifactId,
+			IsPublic:            true,
+		})
+		if err != nil {
+			logs.Println("[IP_LIBRARY]failed to set new artifact as public: " + err.Error())
+		}
+	}
+
+	this.Success()
+}