From aad92b1d46ffc82851d173d69a00331e55d60e55 Mon Sep 17 00:00:00 2001
From: ap4y <lod@pisem.net>
Date: Mon, 11 Jan 2016 20:36:53 +1300
Subject: [PATCH] Refactor RuleSet api using Request struct

---
 ruleset.go      | 32 +++++++++++---------------------
 ruleset_test.go | 10 ++++------
 2 files changed, 15 insertions(+), 27 deletions(-)

diff --git a/ruleset.go b/ruleset.go
index 1b19a9a..1f67878 100644
--- a/ruleset.go
+++ b/ruleset.go
@@ -1,19 +1,8 @@
 package socks5
 
-import (
-	"net"
-)
-
 // RuleSet is used to provide custom rules to allow or prohibit actions
 type RuleSet interface {
-	// AllowConnect is used to filter connect requests
-	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
+	Allow(req *Request) bool
 }
 
 // PermitAll returns a RuleSet which allows all types of connections
@@ -34,14 +23,15 @@ type PermitCommand struct {
 	EnableAssociate bool
 }
 
-func (p *PermitCommand) AllowConnect(net.IP, int, net.IP, int) bool {
-	return p.EnableConnect
-}
+func (p *PermitCommand) Allow(req *Request) bool {
+	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 p.EnableBind
-}
-
-func (p *PermitCommand) AllowAssociate(net.IP, int, net.IP, int) bool {
-	return p.EnableAssociate
+	return false
 }
diff --git a/ruleset_test.go b/ruleset_test.go
index 180fea8..abdaed4 100644
--- a/ruleset_test.go
+++ b/ruleset_test.go
@@ -1,21 +1,19 @@
 package socks5
 
-import (
-	"testing"
-)
+import "testing"
 
 func TestPermitCommand(t *testing.T) {
 	r := &PermitCommand{true, false, false}
 
-	if !r.AllowConnect(nil, 500, nil, 1000) {
+	if !r.Allow(&Request{Command: ConnectCommand}) {
 		t.Fatalf("expect connect")
 	}
 
-	if r.AllowBind(nil, 500, nil, 1000) {
+	if r.Allow(&Request{Command: BindCommand}) {
 		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")
 	}
 }