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

// 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()
}