Skip to content

Commit f99ebbf

Browse files
mrhaoxxmrhaoxx
committed
chore: use generic type for BufferedLookup
1 parent 4d8b79a commit f99ebbf

12 files changed

Lines changed: 75 additions & 58 deletions

File tree

‎auth/policybase.go‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -69,7 +69,7 @@ func (p *policy) check(username string, path string) uint8 {
6969

7070
type policyBaseAuth struct {
7171
policies []*policy
72-
policyLookupBuf *utils.BufferedLookup
72+
policyLookupBuf *utils.BufferedLookup[[]*policy]
7373

7474
backends backendGroup
7575

@@ -124,7 +124,7 @@ func NewPBAuth() *policyBaseAuth {
124124
sessions: map[string]*session{},
125125
}
126126

127-
po.policyLookupBuf = utils.NewBufferedLookup(func(s string) interface{} {
127+
po.policyLookupBuf = utils.NewBufferedLookup(func(s string) []*policy {
128128
var r []*policy = nil
129129
for _, p := range po.policies {
130130
if p.hosts != nil && p.hosts.MatchString(s) {
@@ -499,7 +499,7 @@ func (LGM *policyBaseAuth) AddPolicy(name string, allow bool, users []string, ho
499499
}
500500

501501
func (mgr *policyBaseAuth) determine(host, path, user string) (v uint8) {
502-
pls := mgr.policyLookupBuf.Lookup(host).([]*policy)
502+
pls := mgr.policyLookupBuf.Lookup(host)
503503
if len(pls) == 0 {
504504
return 0
505505
}

‎dns/server.go‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -28,7 +28,7 @@ type filter struct {
2828
type server struct {
2929
records []*record
3030
filters []*filter
31-
bufferedLookupForFilters *utils.BufferedLookup
31+
bufferedLookupForFilters *utils.BufferedLookup[bool]
3232
// bufferedLookupForRecords *utils.BufferedLookup
3333

3434
domain string
@@ -72,7 +72,7 @@ func (s *server) ServeDNS(w dns.ResponseWriter, req *dns.Msg) {
7272
}()
7373

7474
for _, q := range req.Question {
75-
if s.bufferedLookupForFilters.Lookup(strings.ToLower(q.Name)).(bool) {
75+
if s.bufferedLookupForFilters.Lookup(strings.ToLower(q.Name)) {
7676
goto allowed
7777
} else {
7878
m.Rcode = dns.RcodeRefused
@@ -163,7 +163,7 @@ func NewServer() (ret *server) {
163163
filters: []*filter{},
164164
count: 0,
165165
}
166-
ret.bufferedLookupForFilters = utils.NewBufferedLookup(func(s string) interface{} {
166+
ret.bufferedLookupForFilters = utils.NewBufferedLookup(func(s string) bool {
167167
for _, r := range ret.filters {
168168
if ok, _ := r.name.MatchString(s); ok {
169169
if r.allowance {

‎http/cgi.go‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -10,7 +10,7 @@ func (mid *Midware) ngCgi(RequestCtx *HttpCtx, RequestPath *string) {
1010

1111
path := strings.TrimPrefix(RequestCtx.Req.URL.Path, PrefixNg)
1212

13-
s := mid.bufferedLookupForCgi.Lookup(path).([]*CgiStruct)
13+
s := mid.bufferedLookupForCgi.Lookup(path)
1414

1515
if len(s) == 0 {
1616
RequestCtx.Resp.ErrorPage(StatusNotFound, "The requested URL "+RequestCtx.Req.RequestURI+" was not found on this server.")

‎http/forward.go‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,7 @@ package http
22

33
func (h *Midware) ngForwardProxy(ctx *HttpCtx, RequestPath *string) {
44

5-
ServicesToExecute := h.bufferedLookupForForward.Lookup(ctx.Req.Host).([]*ServiceStruct)
5+
ServicesToExecute := h.bufferedLookupForForward.Lookup(ctx.Req.Host)
66
for i := 0; i < len(ServicesToExecute); i++ {
77

88
*RequestPath += ServicesToExecute[i].Id + " "

‎http/http.go‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -154,7 +154,7 @@ func newReqID() string {
154154
return hex.EncodeToString(b)
155155
}
156156

157-
func (h *Midware) preparetls(rw http.ResponseWriter, r *http.Request, conn *tcp.Conn) {
157+
func (h *Midware) preparetls(_ http.ResponseWriter, r *http.Request, conn *tcp.Conn) {
158158
if r.TLS == nil {
159159
head := conn.Head()
160160
if head > 1 {

‎http/midware.go‎

Lines changed: 10 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -21,16 +21,16 @@ import (
2121
//ng:generate def obj Midware
2222
type Midware struct {
2323
sni utils.GroupRegexp
24-
bufferedLookupForSNI *utils.BufferedLookup
24+
bufferedLookupForSNI *utils.BufferedLookup[bool]
2525

2626
current []*ServiceStruct
27-
bufferedLookupForHost *utils.BufferedLookup
27+
bufferedLookupForHost *utils.BufferedLookup[[]*ServiceStruct]
2828

2929
currentCgi []*CgiStruct
30-
bufferedLookupForCgi *utils.BufferedLookup
30+
bufferedLookupForCgi *utils.BufferedLookup[[]*CgiStruct]
3131

3232
currentForward []*ServiceStruct
33-
bufferedLookupForForward *utils.BufferedLookup
33+
bufferedLookupForForward *utils.BufferedLookup[[]*ServiceStruct]
3434

3535
muActiveRequest sync.RWMutex
3636
activeRequests map[string]*HttpCtx
@@ -85,7 +85,7 @@ var h2s = &http2.Server{}
8585
func (h *Midware) Handle(c *tcp.Conn) tcp.SerRet {
8686
top := c.TopProtocol()
8787
sni, ok := c.Load(tcp.KeyTlsSni)
88-
if ok && !h.bufferedLookupForSNI.Lookup(sni.(string)).(bool) {
88+
if ok && !h.bufferedLookupForSNI.Lookup(sni.(string)) {
8989
return tcp.Continue
9090
}
9191
switch top {
@@ -189,7 +189,7 @@ func (h *Midware) Process(RequestCtx *HttpCtx) {
189189
}
190190

191191
{
192-
ServicesToExecute := h.bufferedLookupForHost.Lookup(RequestCtx.Req.Host).([]*ServiceStruct)
192+
ServicesToExecute := h.bufferedLookupForHost.Lookup(RequestCtx.Req.Host)
193193
for i := 0; i < len(ServicesToExecute); i++ {
194194

195195
RequestPath += ServicesToExecute[i].Id + " " // record the executed service
@@ -228,7 +228,7 @@ func NewHttpMidware(sni []string) *Midware {
228228
},
229229
}
230230

231-
hmw.bufferedLookupForHost = utils.NewBufferedLookup(func(s string) interface{} {
231+
hmw.bufferedLookupForHost = utils.NewBufferedLookup(func(s string) []*ServiceStruct {
232232
ret := make([]*ServiceStruct, 0)
233233
for _, r := range hmw.current {
234234
if r.Hosts.MatchString(s) {
@@ -238,7 +238,7 @@ func NewHttpMidware(sni []string) *Midware {
238238
return ret
239239
})
240240

241-
hmw.bufferedLookupForCgi = utils.NewBufferedLookup(func(s string) interface{} {
241+
hmw.bufferedLookupForCgi = utils.NewBufferedLookup(func(s string) []*CgiStruct {
242242
var m []*CgiStruct = nil
243243
for _, t := range hmw.currentCgi {
244244
for _, r := range t.CgiPaths {
@@ -250,7 +250,7 @@ func NewHttpMidware(sni []string) *Midware {
250250
return m
251251
})
252252

253-
hmw.bufferedLookupForForward = utils.NewBufferedLookup(func(s string) interface{} {
253+
hmw.bufferedLookupForForward = utils.NewBufferedLookup(func(s string) []*ServiceStruct {
254254
ret := make([]*ServiceStruct, 0)
255255
for _, r := range hmw.currentForward {
256256
if r.Hosts.MatchString(s) {
@@ -260,7 +260,7 @@ func NewHttpMidware(sni []string) *Midware {
260260
return ret
261261
})
262262

263-
hmw.bufferedLookupForSNI = utils.NewBufferedLookup(func(s string) interface{} {
263+
hmw.bufferedLookupForSNI = utils.NewBufferedLookup(func(s string) bool {
264264
return hmw.sni == nil || hmw.sni.MatchString(s)
265265
})
266266

‎http/proxy.go‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -120,7 +120,7 @@ func (h *HttpHost) Init() {
120120
type ReverseProxy struct {
121121
hosts []*HttpHost
122122

123-
buf *utils.BufferedLookup
123+
buf *utils.BufferedLookup[*HttpHost]
124124

125125
allowhosts utils.GroupRegexp
126126
}
@@ -129,7 +129,7 @@ func (h *ReverseProxy) HandleHTTPCgi(ctx *HttpCtx, path string) Ret {
129129
_host := h.buf.Lookup(ctx.Req.Host)
130130
var id string
131131
if _host != nil {
132-
id = _host.(*HttpHost).Id
132+
id = _host.Id
133133
} else {
134134
id = "nohit"
135135
}
@@ -147,7 +147,7 @@ func NewHTTPProxier(allowedhosts []string) *ReverseProxy {
147147
allowhosts: utils.MustCompileRegexp(dns.Dnsnames2Regexps(allowedhosts)),
148148
}
149149

150-
hpx.buf = utils.NewBufferedLookup(func(host string) interface{} {
150+
hpx.buf = utils.NewBufferedLookup(func(host string) *HttpHost {
151151
for _, t := range hpx.hosts {
152152
if t.ServerName.MatchString(host) {
153153
// fmt.Println(t.ServerName.String(), host, "success")
@@ -166,7 +166,7 @@ func (h *ReverseProxy) HandleHTTP(ctx *HttpCtx) Ret {
166166
return Continue
167167
}
168168

169-
host := _host.(*HttpHost)
169+
host := _host
170170

171171
defer func() {
172172
recover()

‎ssh/midware.go‎

Lines changed: 18 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -59,7 +59,7 @@ func (ctx *Ctx) initUserAlt() {
5959
func (ctx *Ctx) Error(err_msg string) {
6060
for ch := range ctx.nc {
6161
n, _, _ := ch.Accept()
62-
n.Stderr().Write([]byte(err_msg))
62+
n.Stderr().Write([]byte(err_msg + "\r\n"))
6363
break
6464
}
6565
ctx.sshconn.Close()
@@ -86,17 +86,15 @@ type Midware struct {
8686
PublicKeyCallback PublicKeyCbFn
8787

8888
current []srv
89-
bufferedLookup *utils.BufferedLookup
89+
bufferedLookup *utils.BufferedLookup[ConnHandler]
90+
91+
basecfg ssh.ServerConfig
9092
}
9193

9294
var cur uint64
9395

9496
func (ctl *Midware) Handle(c *tcp.Conn) tcp.SerRet {
95-
serv := ssh.ServerConfig{}
96-
serv.ServerVersion = "SSH-2.0-OpenNG"
97-
for _, v := range ctl.private_keys {
98-
serv.AddHostKey(v)
99-
}
97+
serv := ctl.basecfg
10098

10199
ctx := Ctx{
102100
Id: atomic.AddUint64(&cur, 1),
@@ -235,13 +233,14 @@ func (ctl *Midware) Handle(c *tcp.Conn) tcp.SerRet {
235233
f := ctl.bufferedLookup.Lookup(ctx.Alt)
236234

237235
if f == nil {
236+
ctx.Error("SSH/2.0 418 I'm a teapot")
238237
path += "#"
239238
return tcp.Close
240239
}
241240

242241
path += "."
243242

244-
f.(ConnHandler).HandleConn(&ctx)
243+
f.HandleConn(&ctx)
245244

246245
return tcp.Close
247246
}
@@ -259,7 +258,17 @@ func NewSSHController(private_keys []ssh.Signer, banner string, quotes []string,
259258
rnd_quotes: quotes,
260259
}
261260

262-
Midware.bufferedLookup = utils.NewBufferedLookup(func(s string) interface{} {
261+
basecfg := ssh.ServerConfig{
262+
ServerVersion: "SSH-2.0-OpenNG",
263+
}
264+
265+
for _, v := range private_keys {
266+
basecfg.AddHostKey(v)
267+
}
268+
269+
Midware.basecfg = basecfg
270+
271+
Midware.bufferedLookup = utils.NewBufferedLookup(func(s string) ConnHandler {
263272
for _, t := range Midware.current {
264273
if t.matchalt.MatchString(s) {
265274
return t.hdr

‎tls/certificate.go‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -20,7 +20,7 @@ type Cert struct {
2020

2121
type TlsMgr struct {
2222
certs map[string]Cert
23-
lookup *utils.BufferedLookup
23+
lookup *utils.BufferedLookup[*tls.Certificate]
2424

2525
muCerts sync.RWMutex
2626
}
@@ -31,7 +31,7 @@ func NewTlsMgr() *TlsMgr {
3131
certs: make(map[string]Cert),
3232
}
3333

34-
mgr.lookup = utils.NewBufferedLookup(func(s string) interface{} {
34+
mgr.lookup = utils.NewBufferedLookup(func(s string) *tls.Certificate {
3535
mgr.muCerts.RLock()
3636
defer mgr.muCerts.RUnlock()
3737

@@ -48,7 +48,7 @@ func NewTlsMgr() *TlsMgr {
4848

4949
func (m *TlsMgr) getCertificate(dnsname string) *tls.Certificate {
5050
if cert := m.lookup.Lookup(dnsname); cert != nil {
51-
return cert.(*tls.Certificate)
51+
return cert
5252
} else {
5353
panic(errors.New("no certificate for " + dnsname))
5454
}

‎ui/builtin.go‎

Lines changed: 11 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -677,6 +677,14 @@ var _builtin_refs_assertions = map[string]Assert{
677677
Type: "ptr",
678678
Required: true,
679679
},
680+
"serv": {
681+
Type: "list",
682+
Sub: AssertMap{
683+
"_": {Type: "string"},
684+
},
685+
Desc: "matching services by regex pattern",
686+
Default: []*ArgNode{{Type: "string", Value: ".*$"}},
687+
},
680688
},
681689
},
682690
},
@@ -1550,16 +1558,17 @@ var _builtin_refs = map[string]Inst{
15501558
for _, srv := range services {
15511559
name := srv.MustGet("name").ToString()
15521560
logi := srv.MustGet("logi")
1561+
serv := srv.MustGet("serv").ToStringList()
15531562

15541563
service, ok := logi.Value.(ssh.ConnHandler)
15551564
if !ok {
15561565
return nil, errors.New("ptr " + name + " is not a ssh.ConnHandler")
15571566
}
15581567

1559-
midware.AddHandler(service, utils.MustCompileRegexp([]string{"^.*$"}))
1568+
midware.AddHandler(service, utils.MustCompileRegexp(serv))
15601569

15611570
// log.Verboseln(fmt.Sprintf("new ssh service %#v: logi=%T", name, logi.Value))
1562-
zlog.Debug().Str("name", name).Type("logi", logi.Value).Msg("new ssh service")
1571+
zlog.Debug().Str("name", name).Type("logi", logi.Value).Strs("serv", serv).Msg("new ssh service")
15631572
}
15641573
return midware, nil
15651574
},

0 commit comments

Comments
 (0)