/
middleware.go
158 lines (142 loc) · 4.02 KB
/
middleware.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
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
package hocoosmiddleware
import (
"fmt"
"net/http"
"strconv"
"strings"
"time"
"github.com/caddyserver/caddy/v2"
"github.com/caddyserver/caddy/v2/caddyconfig/caddyfile"
"github.com/caddyserver/caddy/v2/caddyconfig/httpcaddyfile"
"github.com/caddyserver/caddy/v2/modules/caddyhttp"
"go.uber.org/zap"
)
func init() {
caddy.RegisterModule(&HocoosMiddleware{})
httpcaddyfile.RegisterHandlerDirective("hocoos_middleware", parseCaddyfile)
}
type HocoosMiddleware struct {
RedisURL string `json:"redis_url"`
PathPrefix string `json:"path_prefix"`
CacheTTL int `json:"cache_ttl"`
ExcludeHosts string `json:"exclude_hosts"`
logger *zap.SugaredLogger
client *redisClient
localCache *localCache
excludeHosts []string
}
// CaddyModule returns the Caddy module information.
func (m *HocoosMiddleware) CaddyModule() caddy.ModuleInfo {
return caddy.ModuleInfo{
ID: "http.handlers.hocoos_middleware",
New: func() caddy.Module { return new(HocoosMiddleware) },
}
}
// Provision implements caddy.Provisioner.
func (m *HocoosMiddleware) Provision(ctx caddy.Context) error {
if err := m.Validate(); err != nil {
return err
}
m.logger = ctx.Logger(m).Sugar()
if m.PathPrefix == "" {
m.PathPrefix = "Caddy/Config/Domains"
}
if m.CacheTTL <= 0 {
m.CacheTTL = 60
}
m.excludeHosts = strings.Split(m.ExcludeHosts, ",")
m.localCache = newLocalCache(time.Duration(m.CacheTTL) * time.Second)
m.client = newRedisClient()
if err := m.client.init(ctx, m.RedisURL, m.logger); err != nil {
return err
}
return nil
}
// Validate implements caddy.Validator.
func (m *HocoosMiddleware) Validate() error {
if m.RedisURL == "" {
return fmt.Errorf("redis_url is required")
}
return nil
}
// ServeHTTP implements caddyhttp.MiddlewareHandler.
func (m *HocoosMiddleware) ServeHTTP(w http.ResponseWriter, r *http.Request, next caddyhttp.Handler) error {
m.logger.Debugf("%s %s", r.Method, r.URL.Path)
for _, host := range m.excludeHosts {
if strings.Contains(r.Host, host) {
m.logger.Debugf("host %s excluded", r.Host)
return next.ServeHTTP(w, r)
}
}
key := fmt.Sprintf("%s/%s", m.PathPrefix, r.Host)
var data string
var found bool
data, found = m.localCache.Get(key)
if found {
m.logger.Debugf("%s %s found in local cache", r.Method, r.URL.Path)
} else {
var err error
data, err = m.client.get(r.Context(), key)
if err != nil {
m.logger.Debugf("get %s error: %v", key, err)
w.WriteHeader(http.StatusInternalServerError)
_, writeErr := w.Write([]byte(err.Error()))
return writeErr
}
m.logger.Debugf("get %s: %s from redis", key, data)
m.localCache.Set(key, data)
}
switch data {
case "0":
w.WriteHeader(http.StatusNotFound)
_, err := w.Write([]byte(fmt.Sprintf("%s host not allowed", r.Host)))
return err
case "1":
return next.ServeHTTP(w, r)
default:
w.WriteHeader(http.StatusInternalServerError)
_, writeErr := w.Write([]byte(fmt.Sprintf("permission data for host %s not found", r.Host)))
return writeErr
}
}
// UnmarshalCaddyfile implements caddyfile.Unmarshaler.
func (m *HocoosMiddleware) UnmarshalCaddyfile(d *caddyfile.Dispenser) error {
for d.Next() {
key := d.Val()
var value string
if !d.Args(&value) {
continue
}
switch key {
case "redis_url":
m.RedisURL = value
case "path_prefix":
m.PathPrefix = value
case "cache_ttl":
itoa, err := strconv.Atoi(value)
if err != nil {
m.CacheTTL = 60
} else {
m.CacheTTL = itoa
}
case "exclude_hosts":
m.ExcludeHosts = value
default:
return fmt.Errorf("unknown key %s", key)
}
}
return nil
}
// parseCaddyfile unmarshals tokens from h into a new HocoosMiddleware.
func parseCaddyfile(h httpcaddyfile.Helper) (caddyhttp.MiddlewareHandler, error) {
var m *HocoosMiddleware
err := m.UnmarshalCaddyfile(h.Dispenser)
return m, err
}
// Interface guards
var (
_ caddy.Provisioner = (*HocoosMiddleware)(nil)
_ caddy.Validator = (*HocoosMiddleware)(nil)
_ caddyhttp.MiddlewareHandler = (*HocoosMiddleware)(nil)
_ caddyfile.Unmarshaler = (*HocoosMiddleware)(nil)
)