diff --git a/knockr.go b/knockr.go index 5428e00..538b897 100644 --- a/knockr.go +++ b/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 {