add timeout functionality

This commit is contained in:
maride 2017-09-28 11:35:05 +02:00
parent bdcc67e653
commit 259e6f49d7

View File

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