From aa25f19bf198e0ec86b30159500d3e5880323932 Mon Sep 17 00:00:00 2001
From: Armon Dadgar <armon.dadgar@gmail.com>
Date: Thu, 23 Jan 2014 11:15:53 -0800
Subject: [PATCH] Adding simple ruleset interface

---
 ruleset.go      | 42 ++++++++++++++++++++++++++++++++++++++++++
 ruleset_test.go | 21 +++++++++++++++++++++
 2 files changed, 63 insertions(+)
 create mode 100644 ruleset.go
 create mode 100644 ruleset_test.go

diff --git a/ruleset.go b/ruleset.go
new file mode 100644
index 0000000..cbdfc80
--- /dev/null
+++ b/ruleset.go
@@ -0,0 +1,42 @@
+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.IPAddr, dstPort int, srcIP *net.IPAddr, srcPort int) bool
+
+	// AllowBind is used to filter bind requests
+	AllowBind(dstIP *net.IPAddr, dstPort int, srcIP *net.IPAddr, srcPort int) bool
+
+	// AllowAssociate is used to filter associate requests
+	AllowAssociate(dstIP *net.IPAddr, dstPort int, srcIP *net.IPAddr, srcPort int) bool
+}
+
+// PermitAll is an returns a RuleSet which allows all types of connections
+func PermitAll() RuleSet {
+	return &PermitCommand{true, true, true}
+}
+
+// PermitCommand is an implementation of the RuleSet which
+// enables filtering supported commands
+type PermitCommand struct {
+	EnableConnect   bool
+	EnableBind      bool
+	EnableAssociate bool
+}
+
+func (p *PermitCommand) AllowConnect(*net.IPAddr, int, *net.IPAddr, int) bool {
+	return p.EnableConnect
+}
+
+func (p *PermitCommand) AllowBind(*net.IPAddr, int, *net.IPAddr, int) bool {
+	return p.EnableBind
+}
+
+func (p *PermitCommand) AllowAssociate(*net.IPAddr, int, *net.IPAddr, int) bool {
+	return p.EnableAssociate
+}
diff --git a/ruleset_test.go b/ruleset_test.go
new file mode 100644
index 0000000..180fea8
--- /dev/null
+++ b/ruleset_test.go
@@ -0,0 +1,21 @@
+package socks5
+
+import (
+	"testing"
+)
+
+func TestPermitCommand(t *testing.T) {
+	r := &PermitCommand{true, false, false}
+
+	if !r.AllowConnect(nil, 500, nil, 1000) {
+		t.Fatalf("expect connect")
+	}
+
+	if r.AllowBind(nil, 500, nil, 1000) {
+		t.Fatalf("do not expect bind")
+	}
+
+	if r.AllowAssociate(nil, 500, nil, 1000) {
+		t.Fatalf("do not expect associate")
+	}
+}