Refactor RuleSet api using Request struct
parent
1fcd99d996
commit
aad92b1d46
32
ruleset.go
32
ruleset.go
|
@ -1,19 +1,8 @@
|
||||||
package socks5
|
package socks5
|
||||||
|
|
||||||
import (
|
|
||||||
"net"
|
|
||||||
)
|
|
||||||
|
|
||||||
// RuleSet is used to provide custom rules to allow or prohibit actions
|
// RuleSet is used to provide custom rules to allow or prohibit actions
|
||||||
type RuleSet interface {
|
type RuleSet interface {
|
||||||
// AllowConnect is used to filter connect requests
|
Allow(req *Request) bool
|
||||||
AllowConnect(dstIP net.IP, dstPort int, srcIP net.IP, srcPort int) bool
|
|
||||||
|
|
||||||
// AllowBind is used to filter bind requests
|
|
||||||
AllowBind(dstIP net.IP, dstPort int, srcIP net.IP, srcPort int) bool
|
|
||||||
|
|
||||||
// AllowAssociate is used to filter associate requests
|
|
||||||
AllowAssociate(dstIP net.IP, dstPort int, srcIP net.IP, srcPort int) bool
|
|
||||||
}
|
}
|
||||||
|
|
||||||
// PermitAll returns a RuleSet which allows all types of connections
|
// PermitAll returns a RuleSet which allows all types of connections
|
||||||
|
@ -34,14 +23,15 @@ type PermitCommand struct {
|
||||||
EnableAssociate bool
|
EnableAssociate bool
|
||||||
}
|
}
|
||||||
|
|
||||||
func (p *PermitCommand) AllowConnect(net.IP, int, net.IP, int) bool {
|
func (p *PermitCommand) Allow(req *Request) bool {
|
||||||
return p.EnableConnect
|
switch req.Command {
|
||||||
}
|
case ConnectCommand:
|
||||||
|
return p.EnableConnect
|
||||||
|
case BindCommand:
|
||||||
|
return p.EnableBind
|
||||||
|
case AssociateCommand:
|
||||||
|
return p.EnableAssociate
|
||||||
|
}
|
||||||
|
|
||||||
func (p *PermitCommand) AllowBind(net.IP, int, net.IP, int) bool {
|
return false
|
||||||
return p.EnableBind
|
|
||||||
}
|
|
||||||
|
|
||||||
func (p *PermitCommand) AllowAssociate(net.IP, int, net.IP, int) bool {
|
|
||||||
return p.EnableAssociate
|
|
||||||
}
|
}
|
||||||
|
|
|
@ -1,21 +1,19 @@
|
||||||
package socks5
|
package socks5
|
||||||
|
|
||||||
import (
|
import "testing"
|
||||||
"testing"
|
|
||||||
)
|
|
||||||
|
|
||||||
func TestPermitCommand(t *testing.T) {
|
func TestPermitCommand(t *testing.T) {
|
||||||
r := &PermitCommand{true, false, false}
|
r := &PermitCommand{true, false, false}
|
||||||
|
|
||||||
if !r.AllowConnect(nil, 500, nil, 1000) {
|
if !r.Allow(&Request{Command: ConnectCommand}) {
|
||||||
t.Fatalf("expect connect")
|
t.Fatalf("expect connect")
|
||||||
}
|
}
|
||||||
|
|
||||||
if r.AllowBind(nil, 500, nil, 1000) {
|
if r.Allow(&Request{Command: BindCommand}) {
|
||||||
t.Fatalf("do not expect bind")
|
t.Fatalf("do not expect bind")
|
||||||
}
|
}
|
||||||
|
|
||||||
if r.AllowAssociate(nil, 500, nil, 1000) {
|
if r.Allow(&Request{Command: AssociateCommand}) {
|
||||||
t.Fatalf("do not expect associate")
|
t.Fatalf("do not expect associate")
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
Loading…
Reference in New Issue