waf-platform/EdgeAPI/internal/rpc/services/service_http_web_plus.go

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

package services

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	teaconst "github.com/TeaOSLab/EdgeAPI/internal/const"
	"github.com/TeaOSLab/EdgeAPI/internal/db/models"
	"github.com/TeaOSLab/EdgeCommon/pkg/rpc/pb"
	"github.com/TeaOSLab/EdgeCommon/pkg/serverconfigs"
	"github.com/iwind/TeaGo/dbs"
	"github.com/iwind/TeaGo/types"
	stringutil "github.com/iwind/TeaGo/utils/string"
)

// UpdateHTTPWebUAM 修改UAM设置
func (this *HTTPWebService) UpdateHTTPWebUAM(ctx context.Context, req *pb.UpdateHTTPWebUAMRequest) (*pb.RPCSuccess, error) {
	if !teaconst.IsPlus {
		return nil, this.NotImplementedYet()
	}

	_, userId, err := this.ValidateAdminAndUser(ctx, true)
	if err != nil {
		return nil, err
	}

	var tx *dbs.Tx
	if userId > 0 {
		err = models.SharedHTTPWebDAO.CheckUserWeb(tx, userId, req.HttpWebId)
		if err != nil {
			return nil, err
		}
	}

	var config = &serverconfigs.UAMConfig{}
	err = json.Unmarshal(req.UamJSON, config)
	if err != nil {
		return nil, err
	}
	err = config.Init()
	if err != nil {
		return nil, errors.New("valid uam config failed: " + err.Error())
	}

	err = models.SharedHTTPWebDAO.UpdateWebUAM(tx, req.HttpWebId, config)
	if err != nil {
		return nil, err
	}

	return this.Success()
}

// FindHTTPWebUAM 查找UAM设置
func (this *HTTPWebService) FindHTTPWebUAM(ctx context.Context, req *pb.FindHTTPWebUAMRequest) (*pb.FindHTTPWebUAMResponse, error) {
	if !teaconst.IsPlus {
		return nil, this.NotImplementedYet()
	}

	_, userId, err := this.ValidateAdminAndUser(ctx, true)
	if err != nil {
		return nil, err
	}

	var tx *dbs.Tx
	if userId > 0 {
		err = models.SharedHTTPWebDAO.CheckUserWeb(tx, userId, req.HttpWebId)
		if err != nil {
			return nil, err
		}
	}

	uamJSON, err := models.SharedHTTPWebDAO.FindWebUAM(tx, req.HttpWebId)
	if err != nil {
		return nil, err
	}
	return &pb.FindHTTPWebUAMResponse{
		UamJSON: uamJSON,
	}, nil
}

// UpdateHTTPWebCC 修改CC设置
func (this *HTTPWebService) UpdateHTTPWebCC(ctx context.Context, req *pb.UpdateHTTPWebCCRequest) (*pb.RPCSuccess, error) {
	if !teaconst.IsPlus {
		return nil, this.NotImplementedYet()
	}

	_, userId, err := this.ValidateAdminAndUser(ctx, true)
	if err != nil {
		return nil, err
	}

	var tx *dbs.Tx
	if userId > 0 {
		err = models.SharedHTTPWebDAO.CheckUserWeb(tx, userId, req.HttpWebId)
		if err != nil {
			return nil, err
		}
	}

	var config = serverconfigs.DefaultHTTPCCConfig()
	err = json.Unmarshal(req.CcJSON, config)
	if err != nil {
		return nil, err
	}
	err = config.Init()
	if err != nil {
		return nil, errors.New("valid cc config failed: " + err.Error())
	}

	err = models.SharedHTTPWebDAO.UpdateWebCC(tx, req.HttpWebId, config)
	if err != nil {
		return nil, err
	}

	return this.Success()
}

// FindHTTPWebCC 查找CC设置
func (this *HTTPWebService) FindHTTPWebCC(ctx context.Context, req *pb.FindHTTPWebCCRequest) (*pb.FindHTTPWebCCResponse, error) {
	if !teaconst.IsPlus {
		return nil, this.NotImplementedYet()
	}

	_, userId, err := this.ValidateAdminAndUser(ctx, true)
	if err != nil {
		return nil, err
	}

	var tx *dbs.Tx
	if userId > 0 {
		err = models.SharedHTTPWebDAO.CheckUserWeb(tx, userId, req.HttpWebId)
		if err != nil {
			return nil, err
		}
	}

	ccJSON, err := models.SharedHTTPWebDAO.FindWebCC(tx, req.HttpWebId)
	if err != nil {
		return nil, err
	}
	return &pb.FindHTTPWebCCResponse{
		CcJSON: ccJSON,
	}, nil
}

// UpdateHTTPWebRequestScripts 修改请求脚本
func (this *HTTPWebService) UpdateHTTPWebRequestScripts(ctx context.Context, req *pb.UpdateHTTPWebRequestScriptsRequest) (*pb.RPCSuccess, error) {
	if !teaconst.IsPlus {
		return nil, this.NotImplementedYet()
	}

	_, userId, err := this.ValidateAdminAndUser(ctx, true)
	if err != nil {
		return nil, err
	}

	var tx = this.NullTx()

	if userId > 0 {
		err = models.SharedHTTPWebDAO.CheckUserWeb(tx, userId, req.HttpWebId)
		if err != nil {
			return nil, err
		}
	}

	var config = &serverconfigs.HTTPRequestScriptsConfig{}
	err = json.Unmarshal(req.RequestScriptsJSON, config)
	if err != nil {
		return nil, err
	}

	err = config.Init()
	if err != nil {
		return nil, fmt.Errorf("validate config failed: %w", err)
	}

	// 代码最大长度
	// TODO 需要可以在管理员系统中配置
	const codeMaxLength = 8192

	// 检查是否需要审核
	if userId > 0 {
		for _, group := range config.AllGroups() {
			for _, script := range group.Scripts {
				if len(script.Code) > codeMaxLength {
					return nil, errors.New("code length should not more than '" + types.String(codeMaxLength) + "'")
				}

				var realCode = script.TrimCode()
				if len(realCode) == 0 {
					continue
				}

				// 是否已审核通过
				var codeMD5 = stringutil.Md5(realCode)
				isPassed, existErr := models.SharedUserScriptDAO.ExistsPassedCodeMD5(tx, codeMD5)
				if existErr != nil {
					return nil, existErr
				}
				if isPassed {
					script.AuditingCodeMD5 = ""
					script.AuditingCode = ""
					continue
				}

				// 是否已存在
				scriptId, findErr := models.SharedUserScriptDAO.FindUserScriptIdWithCodeMD5(tx, userId, codeMD5)
				if findErr != nil {
					return nil, findErr
				}
				if scriptId <= 0 {
					scriptId, err = models.SharedUserScriptDAO.CreateUserScript(tx, userId, realCode)
					if err != nil {
						return nil, err
					}
				}

				// 保存WebId用于以后的更新
				err = models.SharedUserScriptDAO.AddWebIdToUserScript(tx, scriptId, req.HttpWebId)
				if err != nil {
					return nil, err
				}

				// 清空代码，等待审核
				// TODO 将 script.Code 还原为老版本的代码
				script.AuditingCode = script.Code // not realCode, to keep raw format
				script.Code = ""
				script.AuditingCodeMD5 = codeMD5
			}
		}
	}

	err = models.SharedHTTPWebDAO.UpdateWebRequestScripts(tx, req.HttpWebId, config)
	if err != nil {
		return nil, err
	}

	return this.Success()
}

// UpdateHTTPWebHLS 修改HLS设置
func (this *HTTPWebService) UpdateHTTPWebHLS(ctx context.Context, req *pb.UpdateHTTPWebHLSRequest) (*pb.RPCSuccess, error) {
	if !teaconst.IsPlus {
		return nil, this.NotImplementedYet()
	}

	_, userId, err := this.ValidateAdminAndUser(ctx, true)
	if err != nil {
		return nil, err
	}

	var tx = this.NullTx()

	if userId > 0 {
		err = models.SharedHTTPWebDAO.CheckUserWeb(tx, userId, req.HttpWebId)
		if err != nil {
			return nil, err
		}
	}

	if len(req.HlsJSON) == 0 {
		return nil, errors.New("require 'hlsJSON'")
	}

	var hlsConfig = &serverconfigs.HLSConfig{}
	err = json.Unmarshal(req.HlsJSON, hlsConfig)
	if err != nil {
		return nil, err
	}

	err = hlsConfig.Init()
	if err != nil {
		return nil, fmt.Errorf("validate config failed: %w", err)
	}

	err = models.SharedHTTPWebDAO.UpdateWebHLS(tx, req.HttpWebId, hlsConfig)
	if err != nil {
		return nil, err
	}

	return this.Success()
}

// FindHTTPWebHLS 查找HLS设置
func (this *HTTPWebService) FindHTTPWebHLS(ctx context.Context, req *pb.FindHTTPWebHLSRequest) (*pb.FindHTTPWebHLSResponse, error) {
	if !teaconst.IsPlus {
		return nil, this.NotImplementedYet()
	}

	_, userId, err := this.ValidateAdminAndUser(ctx, true)
	if err != nil {
		return nil, err
	}

	var tx = this.NullTx()

	if userId > 0 {
		err = models.SharedHTTPWebDAO.CheckUserWeb(tx, userId, req.HttpWebId)
		if err != nil {
			return nil, err
		}
	}

	hlsJSON, err := models.SharedHTTPWebDAO.FindWebHLS(tx, req.HttpWebId)
	if err != nil {
		return nil, err
	}

	return &pb.FindHTTPWebHLSResponse{
		HlsJSON: hlsJSON,
	}, nil
}