[waf]如果状态码是444，则直接关闭连接

Author: iwind

teaproxy/response.go

@@ -1,10 +1,12 @@
 package teaproxy
 
 import (
+	"bufio"
 	"bytes"
 	"compress/gzip"
 	"github.com/TeaWeb/code/teaconfigs"
 	"github.com/iwind/TeaGo/logs"
+	"net"
 	"net/http"
 )
 
@@ -225,3 +227,12 @@ func (this *ResponseWriter) Close() {
 		this.gzipWriter = nil
 	}
 }
+
+// Hijack
+func (this *ResponseWriter) Hijack() (conn net.Conn, buf *bufio.ReadWriter, err error) {
+	hijack, ok := this.writer.(http.Hijacker)
+	if ok {
+		return hijack.Hijack()
+	}
+	return
+}

teawaf/action_block.go

@@ -27,6 +27,19 @@ type BlockAction struct {
 
 func (this *BlockAction) Perform(waf *WAF, request *requests.Request, writer http.ResponseWriter) (allow bool) {
 	if writer != nil {
+		// if status code eq 444, we close the connection
+		if this.StatusCode == 444 {
+			hijack, ok := writer.(http.Hijacker)
+			if ok {
+				conn, _, _ := hijack.Hijack()
+				if conn != nil {
+					_ = conn.Close()
+					return
+				}
+			}
+		}
+
+		// output response
 		if this.StatusCode > 0 {
 			writer.WriteHeader(this.StatusCode)
 		} else {