1
0
Fork 0

Refactor RuleSet api using Request struct

logger
ap4y 2016-01-11 20:36:53 +13:00
parent 1fcd99d996
commit aad92b1d46
2 changed files with 15 additions and 27 deletions

View File

@ -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
} }

View File

@ -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")
} }
} }