github.com/TeaOSLab/EdgeNode@v1.3.8/internal/waf/action_post_307.go (about)

     1  package waf
     2  
     3  import (
     4  	"github.com/TeaOSLab/EdgeNode/internal/remotelogs"
     5  	"github.com/TeaOSLab/EdgeNode/internal/utils"
     6  	"github.com/TeaOSLab/EdgeNode/internal/waf/requests"
     7  	"github.com/iwind/TeaGo/types"
     8  	"io"
     9  	"net/http"
    10  	"time"
    11  )
    12  
    13  type Post307Action struct {
    14  	Life  int32  `yaml:"life" json:"life"`
    15  	Scope string `yaml:"scope" json:"scope"`
    16  
    17  	BaseAction
    18  }
    19  
    20  func (this *Post307Action) Init(waf *WAF) error {
    21  	return nil
    22  }
    23  
    24  func (this *Post307Action) Code() string {
    25  	return ActionPost307
    26  }
    27  
    28  func (this *Post307Action) IsAttack() bool {
    29  	return false
    30  }
    31  
    32  func (this *Post307Action) WillChange() bool {
    33  	return true
    34  }
    35  
    36  func (this *Post307Action) Perform(waf *WAF, group *RuleGroup, set *RuleSet, request requests.Request, writer http.ResponseWriter) PerformResult {
    37  	const cookieName = "WAF_VALIDATOR_ID"
    38  
    39  	// 仅限于POST
    40  	if request.WAFRaw().Method != http.MethodPost {
    41  		return PerformResult{
    42  			ContinueRequest: true,
    43  		}
    44  	}
    45  
    46  	// 是否已经在白名单中
    47  	if SharedIPWhiteList.Contains("set:"+types.String(set.Id), this.Scope, request.WAFServerId(), request.WAFRemoteIP()) {
    48  		return PerformResult{
    49  			ContinueRequest: true,
    50  		}
    51  	}
    52  
    53  	// 判断是否有Cookie
    54  	cookie, cookieErr := request.WAFRaw().Cookie(cookieName)
    55  	if cookieErr == nil && cookie != nil {
    56  		var remoteIP string
    57  		var life int64
    58  		var setId int64
    59  		var policyId int64
    60  		var groupId int64
    61  		var timestamp int64
    62  
    63  		var infoArg = &InfoArg{}
    64  		var success bool
    65  		decodeErr := infoArg.Decode(cookie.Value)
    66  		if decodeErr == nil && infoArg.IsValid() {
    67  			success = true
    68  
    69  			remoteIP = infoArg.RemoteIP
    70  			life = int64(infoArg.Life)
    71  			setId = infoArg.SetId
    72  			policyId = infoArg.PolicyId
    73  			groupId = infoArg.GroupId
    74  			timestamp = infoArg.Timestamp
    75  		} else {
    76  			// 兼容老版本
    77  			m, decodeMapErr := utils.SimpleDecryptMap(cookie.Value)
    78  			if decodeMapErr == nil {
    79  				success = true
    80  
    81  				remoteIP = m.GetString("remoteIP")
    82  				timestamp = m.GetInt64("timestamp")
    83  				life = m.GetInt64("life")
    84  				setId = m.GetInt64("setId")
    85  				groupId = m.GetInt64("groupId")
    86  				policyId = m.GetInt64("policyId")
    87  			}
    88  		}
    89  
    90  		if success && remoteIP == request.WAFRemoteIP() && time.Now().Unix() < timestamp+10 {
    91  			if life <= 0 {
    92  				life = 600 // 默认10分钟
    93  			}
    94  			SharedIPWhiteList.RecordIP("set:"+types.String(setId), this.Scope, request.WAFServerId(), request.WAFRemoteIP(), time.Now().Unix()+life, policyId, false, groupId, setId, "")
    95  			return PerformResult{
    96  				ContinueRequest: true,
    97  			}
    98  		}
    99  	}
   100  
   101  	var m = &InfoArg{
   102  		Timestamp:        time.Now().Unix(),
   103  		Life:             this.Life,
   104  		Scope:            this.Scope,
   105  		PolicyId:         waf.Id,
   106  		GroupId:          group.Id,
   107  		SetId:            set.Id,
   108  		RemoteIP:         request.WAFRemoteIP(),
   109  		UseLocalFirewall: false,
   110  	}
   111  	info, err := utils.SimpleEncryptObject(m)
   112  	if err != nil {
   113  		remotelogs.Error("WAF_POST_307_ACTION", "encode info failed: "+err.Error())
   114  		return PerformResult{
   115  			ContinueRequest: true,
   116  		}
   117  	}
   118  
   119  	// 清空请求内容
   120  	var req = request.WAFRaw()
   121  	if req.ContentLength > 0 && req.Body != nil {
   122  		var buf = utils.BytePool16k.Get()
   123  		_, _ = io.CopyBuffer(io.Discard, req.Body, buf.Bytes)
   124  		utils.BytePool16k.Put(buf)
   125  		_ = req.Body.Close()
   126  	}
   127  
   128  	// 设置Cookie
   129  	http.SetCookie(writer, &http.Cookie{
   130  		Name:   cookieName,
   131  		Path:   "/",
   132  		MaxAge: 10,
   133  		Value:  info,
   134  	})
   135  
   136  	request.DisableStat()
   137  	request.ProcessResponseHeaders(writer.Header(), http.StatusTemporaryRedirect)
   138  	http.Redirect(writer, request.WAFRaw(), request.WAFRaw().URL.String(), http.StatusTemporaryRedirect)
   139  
   140  	flusher, ok := writer.(http.Flusher)
   141  	if ok {
   142  		flusher.Flush()
   143  	}
   144  
   145  	return PerformResult{}
   146  }