EdgeNode/internal/waf/action_block.go

package waf

import (
	teaconst "github.com/TeaOSLab/EdgeNode/internal/const"
	"github.com/TeaOSLab/EdgeNode/internal/utils"
	"github.com/TeaOSLab/EdgeNode/internal/waf/requests"
	"github.com/iwind/TeaGo/Tea"
	"github.com/iwind/TeaGo/logs"
	"io"
	"io/ioutil"
	"net/http"
	"os"
	"path/filepath"
	"regexp"
	"time"
)

// url client configure
var urlPrefixReg = regexp.MustCompile("^(?i)(http|https)://")
var httpClient = utils.SharedHttpClient(5 * time.Second)

type BlockAction struct {
	StatusCode int    `yaml:"statusCode" json:"statusCode"`
	Body       string `yaml:"body" json:"body"` // supports HTML
	URL        string `yaml:"url" json:"url"`
}

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 {
			writer.WriteHeader(http.StatusForbidden)
		}
		if len(this.URL) > 0 {
			if urlPrefixReg.MatchString(this.URL) {
				req, err := http.NewRequest(http.MethodGet, this.URL, nil)
				if err != nil {
					logs.Error(err)
					return false
				}
				resp, err := httpClient.Do(req)
				if err != nil {
					logs.Error(err)
					return false
				}
				defer func() {
					_ = resp.Body.Close()
				}()

				for k, v := range resp.Header {
					for _, v1 := range v {
						writer.Header().Add(k, v1)
					}
				}

				buf := make([]byte, 1024)
				_, _ = io.CopyBuffer(writer, resp.Body, buf)
			} else {
				path := this.URL
				if !filepath.IsAbs(this.URL) {
					path = Tea.Root + string(os.PathSeparator) + path
				}

				data, err := ioutil.ReadFile(path)
				if err != nil {
					logs.Error(err)
					return false
				}
				_, _ = writer.Write(data)
			}
			return false
		}
		if len(this.Body) > 0 {
			_, _ = writer.Write([]byte(this.Body))
		} else {
			_, _ = writer.Write([]byte("The request is blocked by " + teaconst.ProductName))
		}
	}
	return false
}