-
Notifications
You must be signed in to change notification settings - Fork 34
/
action_block.go
77 lines (70 loc) · 1.66 KB
/
action_block.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
package actions
import (
"github.com/TeaWeb/code/teautils"
"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 = teautils.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(writer http.ResponseWriter) (allow bool) {
if writer != nil {
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 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 TeaWAF"))
}
}
return false
}