waf-platform/EdgeAPI/internal/db/models/server_dao_copy_plus.go

// Copyright 2023 Liuxiangchao iwind.liu@gmail.com. All rights reserved. Official site: https://goedge.cn .
//go:build plus

package models

import (
	"encoding/json"
	"github.com/TeaOSLab/EdgeCommon/pkg/nodeconfigs"
	"github.com/TeaOSLab/EdgeCommon/pkg/serverconfigs"
	"github.com/iwind/TeaGo/dbs"
	"github.com/iwind/TeaGo/types"
)

// 服务基本信息
type clusterServerList struct {
	clusterId int64
	serverIds []int64
}

// CopyServerConfigToServers 拷贝服务配置到一组服务
func (this *ServerDAO) CopyServerConfigToServers(tx *dbs.Tx, fromServerId int64, toServerIds []int64, configCode serverconfigs.ConfigCode, wafCopyRegions bool) error {
	if fromServerId <= 0 {
		return nil
	}
	if len(toServerIds) == 0 {
		return nil
	}

	fromWebId, err := SharedServerDAO.FindServerWebId(tx, fromServerId)
	if err != nil {
		return err
	}

	clusterServers, toWebIds, err := this.findServerClusterIdsAndWebIds(tx, toServerIds, fromWebId)
	if err != nil {
		return err
	}
	if len(clusterServers) == 0 {
		return nil
	}

	switch configCode {
	case serverconfigs.ConfigCodeStat: // 统计
		if fromWebId <= 0 || len(toWebIds) == 0 {
			return nil
		}

		err = SharedHTTPWebDAO.CopyWebConfigs(tx, fromWebId, toWebIds, "stat")
		if err != nil {
			return err
		}
	case serverconfigs.ConfigCodeCharset: // 字符编码
		if fromWebId <= 0 || len(toWebIds) == 0 {
			return nil
		}

		err = SharedHTTPWebDAO.CopyWebConfigs(tx, fromWebId, toWebIds, "charset")
		if err != nil {
			return err
		}
	case serverconfigs.ConfigCodeUserAgent: // UA名单
		if fromWebId <= 0 || len(toWebIds) == 0 {
			return nil
		}

		err = SharedHTTPWebDAO.CopyWebConfigs(tx, fromWebId, toWebIds, "userAgent")
		if err != nil {
			return err
		}
	case serverconfigs.ConfigCodeReferers: // 防盗链
		if fromWebId <= 0 || len(toWebIds) == 0 {
			return nil
		}

		err = SharedHTTPWebDAO.CopyWebConfigs(tx, fromWebId, toWebIds, "referers")
		if err != nil {
			return err
		}
	case serverconfigs.ConfigCodeWebp: // WebP配置
		if fromWebId <= 0 || len(toWebIds) == 0 {
			return nil
		}

		err = SharedHTTPWebDAO.CopyWebConfigs(tx, fromWebId, toWebIds, "webp")
		if err != nil {
			return err
		}
	case serverconfigs.ConfigCodeAccessLog: // 访问日志
		if fromWebId <= 0 || len(toWebIds) == 0 {
			return nil
		}

		err = SharedHTTPWebDAO.CopyWebConfigs(tx, fromWebId, toWebIds, "accessLog")
		if err != nil {
			return err
		}
	case serverconfigs.ConfigCodeCC: // CC
		if fromWebId <= 0 || len(toWebIds) == 0 {
			return nil
		}

		err = SharedHTTPWebDAO.CopyWebConfigs(tx, fromWebId, toWebIds, "cc")
		if err != nil {
			return err
		}
	case serverconfigs.ConfigCodeHostRedirects: // URL跳转
		if fromWebId <= 0 || len(toWebIds) == 0 {
			return nil
		}

		err = SharedHTTPWebDAO.CopyWebConfigs(tx, fromWebId, toWebIds, "hostRedirects")
		if err != nil {
			return err
		}
	case serverconfigs.ConfigCodeRoot: // ROOT配置
		if fromWebId <= 0 || len(toWebIds) == 0 {
			return nil
		}

		err = SharedHTTPWebDAO.CopyWebConfigs(tx, fromWebId, toWebIds, "root")
		if err != nil {
			return err
		}
	case serverconfigs.ConfigCodeRemoteAddr: // 访客地址
		if fromWebId <= 0 || len(toWebIds) == 0 {
			return nil
		}

		err = SharedHTTPWebDAO.CopyWebConfigs(tx, fromWebId, toWebIds, "remoteAddr")
		if err != nil {
			return err
		}
	case serverconfigs.ConfigCodeRequestLimit: // 请求限制
		if fromWebId <= 0 || len(toWebIds) == 0 {
			return nil
		}

		err = SharedHTTPWebDAO.CopyWebConfigs(tx, fromWebId, toWebIds, "requestLimit")
		if err != nil {
			return err
		}
	case serverconfigs.ConfigCodeCompression: // 压缩
		if fromWebId <= 0 || len(toWebIds) == 0 {
			return nil
		}

		// TODO 将来可能需要支持gzip、brotli等独立的配置
		err = SharedHTTPWebDAO.CopyWebConfigs(tx, fromWebId, toWebIds, "compression")
		if err != nil {
			return err
		}
	case serverconfigs.ConfigCodeOptimization: // 页面优化
		if fromWebId <= 0 || len(toWebIds) == 0 {
			return nil
		}

		err = SharedHTTPWebDAO.CopyWebConfigs(tx, fromWebId, toWebIds, "optimization")
		if err != nil {
			return err
		}
	case serverconfigs.ConfigCodeCache: //缓存
		if fromWebId <= 0 || len(toWebIds) == 0 {
			return nil
		}
		err = SharedHTTPWebDAO.CopyCacheConfigs(tx, fromWebId, toWebIds)
		if err != nil {
			return err
		}
	case serverconfigs.ConfigCodeUAM: // UAM
		err = this.CopyServerUAMConfigs(tx, fromServerId, toServerIds)
		if err != nil {
			return err
		}
	case serverconfigs.ConfigCodeWebsocket: // Websocket
		if fromWebId <= 0 || len(toWebIds) == 0 {
			return nil
		}

		err = SharedHTTPWebDAO.CopyWebsocketConfigs(tx, fromWebId, toWebIds)
		if err != nil {
			return err
		}
	case serverconfigs.ConfigCodePages: // 自定义页面
		if fromWebId <= 0 || len(toWebIds) == 0 {
			return nil
		}

		err = SharedHTTPWebDAO.CopyPageConfigs(tx, fromWebId, toWebIds)
		if err != nil {
			return err
		}
	case serverconfigs.ConfigCodeAuth: // 访问鉴权
		if fromWebId <= 0 || len(toWebIds) == 0 {
			return nil
		}

		err = SharedHTTPWebDAO.CopyAuthConfigs(tx, fromWebId, toWebIds)
		if err != nil {
			return err
		}
	case serverconfigs.ConfigCodeReverseProxy: // 反向代理
		err = this.CopyServerReverseProxyConfigs(tx, fromServerId, toServerIds)
		if err != nil {
			return err
		}
	case serverconfigs.ConfigCodeWAF: // WAF设置
		err = SharedHTTPWebDAO.CopyFirewallConfigs(tx, fromWebId, toWebIds, wafCopyRegions)
		if err != nil {
			return err
		}
	case serverconfigs.ConfigCodeMultimedia: // 音视频 Multimedia
		if fromWebId <= 0 || len(toWebIds) == 0 {
			return nil
		}

		err = SharedHTTPWebDAO.CopyWebConfigs(tx, fromWebId, toWebIds, HTTPWebField_Hls)
		if err != nil {
			return err
		}
	}

	// 通知更新
	for _, serverList := range clusterServers {
		err = SharedUpdatingServerListDAO.CreateList(tx, serverList.clusterId, serverList.serverIds)
		if err != nil {
			return err
		}

		err = SharedNodeTaskDAO.CreateClusterTask(tx, nodeconfigs.NodeRoleNode, serverList.clusterId, 0, 0, NodeTaskTypeUpdatingServers)
		if err != nil {
			return err
		}
	}

	return nil
}

// 查找一组服务的集群和WebId信息
func (this *ServerDAO) findServerClusterIdsAndWebIds(tx *dbs.Tx, serverIds []int64, exceptWebId int64) (clusterServers []*clusterServerList, webIds []int64, err error) {
	if len(serverIds) == 0 {
		return
	}

	ones, err := this.Query(tx).
		Result("id", "webId", "clusterId").
		Pk(serverIds).
		Reuse(false).
		FindAll()
	if err != nil {
		return nil, nil, err
	}

	var clusterMap = map[int64]*clusterServerList{} // clusterId => servers

	for _, one := range ones {
		var server = one.(*Server)
		var clusterId = int64(server.ClusterId)
		if clusterId <= 0 {
			continue
		}

		serverList, ok := clusterMap[clusterId]
		if ok {
			serverList.serverIds = append(serverList.serverIds, int64(server.Id))
		} else {
			clusterMap[clusterId] = &clusterServerList{
				clusterId: clusterId,
				serverIds: []int64{int64(server.Id)},
			}
		}

		var webId = int64(server.WebId)
		if webId > 0 && webId != exceptWebId {
			webIds = append(webIds, webId)
		}
	}

	for _, serverList := range clusterMap {
		clusterServers = append(clusterServers, serverList)
	}

	return
}

// CopyServerConfigToGroups 拷贝服务配置到分组
func (this *ServerDAO) CopyServerConfigToGroups(tx *dbs.Tx, fromServerId int64, groupIds []int64, configCode string, wafCopyRegions bool) error {
	if len(groupIds) == 0 {
		return nil
	}

	var serverIds = []int64{}
	for _, groupId := range groupIds {
		var query = this.Query(tx)

		// 使用用户ID查询提升性能
		userId, err := SharedServerGroupDAO.FindGroupUserId(tx, groupId)
		if err != nil {
			return err
		}
		if userId > 0 {
			query.Attr("userId", userId)
		}

		ones, err := query.
			ResultPk().
			State(ServerStateEnabled).
			Where("JSON_CONTAINS(groupIds, :groupId)").
			Param("groupId", types.String(groupId)).
			FindAll()
		if err != nil {
			return err
		}
		for _, one := range ones {
			serverIds = append(serverIds, int64(one.(*Server).Id))
		}
	}

	return this.CopyServerConfigToServers(tx, fromServerId, serverIds, configCode, wafCopyRegions)
}

// CopyServerConfigToCluster 拷贝服务配置到集群
func (this *ServerDAO) CopyServerConfigToCluster(tx *dbs.Tx, fromServerId int64, clusterId int64, configCode string, wafCopyRegions bool) error {
	ones, err := this.Query(tx).
		ResultPk().
		State(ServerStateEnabled).
		Attr("clusterId", clusterId).
		UseIndex("clusterId").
		FindAll()
	if err != nil {
		return err
	}
	var serverIds = []int64{}
	for _, one := range ones {
		serverIds = append(serverIds, int64(one.(*Server).Id))
	}
	return this.CopyServerConfigToServers(tx, fromServerId, serverIds, configCode, wafCopyRegions)
}

// CopyServerConfigToUser 拷贝服务配置到用户
func (this *ServerDAO) CopyServerConfigToUser(tx *dbs.Tx, fromServerId int64, userId int64, configCode string, wafCopyRegions bool) error {
	ones, err := this.Query(tx).
		ResultPk().
		State(ServerStateEnabled).
		Attr("userId", userId).
		UseIndex("userId").
		FindAll()
	if err != nil {
		return err
	}
	var serverIds = []int64{}
	for _, one := range ones {
		serverIds = append(serverIds, int64(one.(*Server).Id))
	}
	return this.CopyServerConfigToServers(tx, fromServerId, serverIds, configCode, wafCopyRegions)
}

// CopyServerUAMConfigs 复制UAM设置
func (this *ServerDAO) CopyServerUAMConfigs(tx *dbs.Tx, fromServerId int64, toServerIds []int64) error {
	if fromServerId <= 0 || len(toServerIds) == 0 {
		return nil
	}

	var configField = "uam"

	uamJSON, err := this.Query(tx).
		Pk(fromServerId).
		Result(configField).
		FindJSONCol()
	if err != nil {
		return err
	}

	if IsNull(uamJSON) {
		// 暂时不支持设置为空
		return nil
	}

	return this.Query(tx).
		Pk(toServerIds).
		Neq("id", fromServerId).
		Set(configField, uamJSON).
		Reuse(false).
		UpdateQuickly()
}

// CopyServerReverseProxyConfigs 复制反向代理
func (this *ServerDAO) CopyServerReverseProxyConfigs(tx *dbs.Tx, fromServerId int64, toServerIds []int64) error {
	if fromServerId <= 0 || len(toServerIds) == 0 {
		return nil
	}

	reverseProxyJSON, err := this.Query(tx).
		Pk(fromServerId).
		Result("reverseProxy").
		FindJSONCol()
	if err != nil {
		return err
	}
	if IsNull(reverseProxyJSON) {
		return nil
	}

	var reverseProxyRef = &serverconfigs.ReverseProxyRef{}
	err = json.Unmarshal(reverseProxyJSON, reverseProxyRef)
	if err != nil {
		return err
	}
	if reverseProxyRef.ReverseProxyId <= 0 {
		return this.Query(tx).
			Pk(toServerIds).
			Neq("id", fromServerId).
			Set("reverseProxy", reverseProxyJSON).
			Reuse(false).
			UpdateQuickly()
	}

	var oldReverseProxyId = reverseProxyRef.ReverseProxyId
	for _, toServerId := range toServerIds {
		if fromServerId == toServerId {
			continue
		}
		newReverseProxyId, err := SharedReverseProxyDAO.CloneReverseProxy(tx, oldReverseProxyId)
		if err != nil {
			return err
		}
		if newReverseProxyId > 0 {
			reverseProxyRef.ReverseProxyId = newReverseProxyId
			newReverseProxyRefJSON, err := json.Marshal(reverseProxyRef)
			if err != nil {
				return err
			}
			err = this.Query(tx).
				Pk(toServerId).
				Set("reverseProxy", newReverseProxyRefJSON).
				UpdateQuickly()
			if err != nil {
				return err
			}
		}
	}

	return nil
}