2020-06-19 13:31:20 +00:00
|
|
|
package net
|
|
|
|
|
|
|
|
import (
|
|
|
|
"flag"
|
|
|
|
"fmt"
|
|
|
|
"github.com/maride/afl-transmit/logistic"
|
|
|
|
"io"
|
2020-06-19 23:06:47 +00:00
|
|
|
"io/ioutil"
|
2020-06-19 13:31:20 +00:00
|
|
|
"log"
|
|
|
|
"net"
|
|
|
|
"strings"
|
|
|
|
)
|
|
|
|
|
|
|
|
var (
|
|
|
|
port int
|
|
|
|
)
|
|
|
|
|
|
|
|
// Registers the flags required for the listener
|
|
|
|
func RegisterListenFlags() {
|
|
|
|
flag.IntVar(&port, "port", ServerPort, "Port to bind server component to")
|
|
|
|
}
|
|
|
|
|
|
|
|
// Sets up a listener and listens forever for packets on the given port, storing their contents in the outputDirectory
|
|
|
|
func Listen(outputDirectory string) error {
|
|
|
|
// Create listener
|
|
|
|
addrStr := fmt.Sprintf(":%v", port)
|
|
|
|
listener, listenErr := net.Listen("tcp", addrStr)
|
|
|
|
if listenErr != nil {
|
|
|
|
return listenErr
|
|
|
|
}
|
|
|
|
|
|
|
|
// Prepare output directory path
|
|
|
|
outputDirectory = strings.TrimRight(outputDirectory, "/")
|
|
|
|
|
|
|
|
// Listen forever
|
|
|
|
for {
|
|
|
|
// Accept connection
|
|
|
|
conn, connErr := listener.Accept()
|
|
|
|
if connErr != nil {
|
|
|
|
log.Printf("Encountered error while accepting from %s: %s", conn.RemoteAddr().String(), connErr)
|
|
|
|
continue
|
|
|
|
}
|
|
|
|
// Handle in a separate thread
|
|
|
|
go handle(conn, outputDirectory)
|
|
|
|
}
|
|
|
|
}
|
|
|
|
|
|
|
|
// Handles a single connection, and unpacks the received data into outputDirectory
|
|
|
|
func handle(conn net.Conn, outputDirectory string) {
|
|
|
|
// Make sure to close connection on return
|
|
|
|
defer conn.Close()
|
|
|
|
|
2020-06-19 23:06:47 +00:00
|
|
|
// Read raw content
|
|
|
|
cont, contErr := ioutil.ReadAll(conn) // bufio.NewReader(conn).ReadString('\x00')
|
2020-06-19 13:31:20 +00:00
|
|
|
|
2020-06-19 23:06:47 +00:00
|
|
|
if contErr == nil || contErr == io.EOF {
|
|
|
|
// We received the whole content, time to process it
|
|
|
|
unpackErr := logistic.UnpackInto([]byte(cont), outputDirectory)
|
|
|
|
if unpackErr != nil {
|
|
|
|
log.Printf("Encountered error processing packet from %s: %s", conn.RemoteAddr().String(), unpackErr)
|
2020-06-19 13:31:20 +00:00
|
|
|
}
|
|
|
|
|
2020-06-19 23:06:47 +00:00
|
|
|
return
|
|
|
|
} else {
|
|
|
|
// We encountered an error on that connection
|
|
|
|
log.Printf("Encountered error while reading from %s: %s", conn.RemoteAddr().String(), contErr)
|
|
|
|
return
|
2020-06-19 13:31:20 +00:00
|
|
|
}
|
|
|
|
}
|