add timeout functionality

knockr.go

@@ -4,6 +4,7 @@ import (
 	"fmt"
 	"io"
 	"net"
+	"time"
 	"github.com/mkideal/cli"
 )
 
@@ -12,9 +13,10 @@ type knockArguments struct {
 	WhitelistPort int `cli:'wp' usage:'The port to launch the whitelist server on'`
 	GatewayPort int `cli:'gp' usage:'The port to protect'`
 	Destination string `cli:'d' usage:'The destination to relay traffic to'`
+	Timeout int64 `cli:'t' usage:'Time in seconds after which a whitelist entry will be removed'`
 }
 
-var whitelist []string
+var whitelist = make(map[string]int64)
 var arguments *knockArguments
 
 func main() {
@@ -68,20 +70,35 @@ func gateway_handler(c net.Conn) {
 
 func add_to_whitelist(addr string) {
 	if ! is_whitelisted(addr) {
-		whitelist = append(whitelist, addr)
+		update_whitelist_time(addr)
 	}
 }
 
+func remove_from_whitelist(addr string) {
+	delete(whitelist, addr)
+}
+
 func is_whitelisted(addr string) bool {
-	for i:=0; i < len(whitelist); i++ {
-		if whitelist[i] == addr {
+	if _, present := whitelist[addr]; present {
+		// Key is present in whitelist map
+		if (whitelist[addr] + arguments.Timeout) >= time.Now().Unix() {
+			// AND we are still in the timing window
+			update_whitelist_time(addr)
 			return true
+		} else {
+			// But we're outside of the timing window
+			remove_from_whitelist(addr)
+			return false
 		}
 	}
-
+	// Entry is not present.
 	return false
 }
 
+func update_whitelist_time(addr string) {
+	whitelist[addr] = time.Now().Unix()
+}
+
 func proxy(c net.Conn) {
 	ln, err := net.Dial("tcp", arguments.Destination)
 	if err != nil {