diff --git a/README.md b/README.md index a380713..c4ee8d6 100644 --- a/README.md +++ b/README.md @@ -8,6 +8,7 @@ You need to specify your AFL output directory with `--fuzzer-directory`, and you Some other options exist to let you fine-tune your *afl-transmit* experience, have a look at them via `--help`. On default, *afl-transmit* opens port 1337/TCP to wait for incoming connections. If you are not on a private net, make sure to protect this port with a firewall, or anyone on the internet may send you files (although this might become interesting). +As a countermeasure, use the `--restrict-to-peers` flags to only allow connections from your known peers. ### Quickstart diff --git a/net/listener.go b/net/listener.go index f5acc96..da7c770 100644 --- a/net/listener.go +++ b/net/listener.go @@ -13,11 +13,13 @@ import ( var ( port int + restrictToPeers bool ) // Registers the flags required for the listener func RegisterListenFlags() { flag.IntVar(&port, "port", ServerPort, "Port to bind server component to") + flag.BoolVar(&restrictToPeers, "restrict-to-peers", false, "Only allow connections from peers") } // Sets up a listener and listens forever for packets on the given port, storing their contents in the outputDirectory @@ -40,8 +42,28 @@ func Listen(outputDirectory string) error { log.Printf("Encountered error while accepting from %s: %s", conn.RemoteAddr().String(), connErr) continue } - // Handle in a separate thread - go handle(conn, outputDirectory) + + // Check if we should restrict connections from peers + handleConnection := true + if restrictToPeers { + found := false + // Loop over peers + for _, p := range peers { + // Check if we found the remote address in our peers list + if p.Address == conn.RemoteAddr().String() { + found = true + break + } + } + + // Handle connection only if its a peer + handleConnection = found + } + + if handleConnection { + // Handle in a separate thread + go handle(conn, outputDirectory) + } } }