mirror of
https://github.com/maride/afl-transmit.git
synced 2024-10-18 11:06:34 +00:00
Avoid reinventing the wheel, properly using TARs
This commit is contained in:
parent
f499b0563d
commit
0a4217b5b9
@ -4,23 +4,39 @@ import (
|
||||
"archive/tar"
|
||||
"bytes"
|
||||
"compress/flate"
|
||||
"encoding/base64"
|
||||
"fmt"
|
||||
"io/ioutil"
|
||||
"log"
|
||||
"os"
|
||||
"strings"
|
||||
)
|
||||
|
||||
// Packs a whole fuzzer directory - at least queue/, fuzz_bitmap, fuzzer_stats
|
||||
func PackFuzzer(fuzzerName string, directory string) ([]byte, error) {
|
||||
// Gather contents
|
||||
contentArray := [][]byte{
|
||||
[]byte(fuzzerName),
|
||||
packSingleFile(directory, "fuzz_bitmap"),
|
||||
packSingleFile(directory, "fuzzer_stats"),
|
||||
packQueueFiles(directory),
|
||||
// PackFuzzers packs all targeted fuzzers into a TAR - at least queue/, fuzz_bitmap, fuzzer_stats
|
||||
func PackFuzzers(fuzzers []string, fuzzerDirectory string) ([]byte, error) {
|
||||
// Create TAR archive
|
||||
var tarBuffer bytes.Buffer
|
||||
tarWriter := tar.NewWriter(&tarBuffer)
|
||||
|
||||
// Essentially we want to pack three things from each targeted fuzzer:
|
||||
// - the fuzz_bitmap file
|
||||
// - the fuzzer_stats file
|
||||
// - the is_main_fuzzer file if present
|
||||
// - the queue/ directory
|
||||
for _, fuzzer := range fuzzers {
|
||||
// We need full paths to read, but will write relative paths into the TAR archive
|
||||
absFuzzerPath := fuzzerDirectory
|
||||
relFuzzerPath := strings.TrimPrefix(fuzzer, fuzzerDirectory)
|
||||
|
||||
// Read-n-Pack™
|
||||
packSingleFile(tarWriter, absFuzzerPath, relFuzzerPath, "fuzz_bitmap", false)
|
||||
packSingleFile(tarWriter, absFuzzerPath, relFuzzerPath, "fuzzer_stats", false)
|
||||
packSingleFile(tarWriter, absFuzzerPath, relFuzzerPath, "is_main_fuzzer", true)
|
||||
packQueueFiles(tarWriter, absFuzzerPath, relFuzzerPath)
|
||||
}
|
||||
|
||||
// Close TAR archive
|
||||
tarWriter.Close()
|
||||
|
||||
// Prepare FLATE compression
|
||||
var flateBuffer bytes.Buffer
|
||||
flateWrite, flateErr := flate.NewWriter(&flateBuffer, flate.BestCompression)
|
||||
@ -28,52 +44,50 @@ func PackFuzzer(fuzzerName string, directory string) ([]byte, error) {
|
||||
return nil, fmt.Errorf("unable to prepare flate compressor: %s", flateErr)
|
||||
}
|
||||
|
||||
// Convert all parts to base64, and concat them to the packet
|
||||
firstRun := true
|
||||
for _, a := range contentArray {
|
||||
b64Buf := make([]byte, base64.StdEncoding.EncodedLen(len(a)))
|
||||
base64.StdEncoding.Encode(b64Buf, a)
|
||||
|
||||
// Add newline char as separator, avoiding it on the first run
|
||||
if firstRun {
|
||||
firstRun = false
|
||||
} else {
|
||||
flateWrite.Write([]byte("\n"))
|
||||
}
|
||||
|
||||
// Append base64 encoded content
|
||||
flateWrite.Write(b64Buf)
|
||||
}
|
||||
|
||||
// Apply FLATE compression
|
||||
flateWrite.Write(tarBuffer.Bytes())
|
||||
flateWrite.Close()
|
||||
|
||||
// Return result: a big byte array, representing concatted base64-encoded files
|
||||
// Return result: a DEFLATEd TAR archive
|
||||
return flateBuffer.Bytes(), nil
|
||||
}
|
||||
|
||||
// Reads a single file and returns it
|
||||
func packSingleFile(directory string, fileName string) []byte {
|
||||
path := fmt.Sprintf("%s%c%s", directory, os.PathSeparator, fileName)
|
||||
contents, readErr := ioutil.ReadFile(path)
|
||||
// packSingleFile packs a single file and writes it to the archive
|
||||
// fuzzerDirectory is the base directory, e.g. /project/fuzzers/
|
||||
// fuzzer is the name of the fuzzer itself, e.g. main-fuzzer-01
|
||||
// filename is the name of the file you want to pack, e.g. fuzzer_stats
|
||||
// ignoreNotFound is just used for files which may not be present in all fuzzer directories, like is_main_fuzzer
|
||||
func packSingleFile(tarWriter *tar.Writer, absPath string, relPath string, fileName string, ignoreNotFound bool) {
|
||||
// Read file
|
||||
readPath := fmt.Sprintf("%s%c%s%c%s", absPath, os.PathSeparator, relPath, os.PathSeparator, fileName)
|
||||
contents, readErr := ioutil.ReadFile(readPath)
|
||||
if readErr != nil {
|
||||
log.Printf("Failed to read file %s: %s", path, readErr)
|
||||
return nil
|
||||
if !ignoreNotFound {
|
||||
log.Printf("Failed to read file %s: %s", readPath, readErr)
|
||||
}
|
||||
return
|
||||
}
|
||||
|
||||
return contents
|
||||
// Create header for this file
|
||||
header := &tar.Header{
|
||||
Name: fmt.Sprintf("%s%c%s", relPath, os.PathSeparator, fileName),
|
||||
Mode: 0600,
|
||||
Size: int64(len(contents)),
|
||||
}
|
||||
|
||||
// Add header and contents to archive
|
||||
tarWriter.WriteHeader(header)
|
||||
tarWriter.Write(contents)
|
||||
}
|
||||
|
||||
// Packs the files in the given directory into a tar archive
|
||||
func packQueueFiles(directory string) []byte {
|
||||
var tarBuffer bytes.Buffer
|
||||
tarWriter := tar.NewWriter(&tarBuffer)
|
||||
|
||||
func packQueueFiles(tarWriter *tar.Writer, absPath string, relPath string) {
|
||||
// Get list of queue files
|
||||
queuePath := fmt.Sprintf("%s%cqueue", directory, os.PathSeparator)
|
||||
queuePath := fmt.Sprintf("%s%c%s%cqueue", absPath, os.PathSeparator, relPath, os.PathSeparator)
|
||||
filesInDir, readErr := ioutil.ReadDir(queuePath)
|
||||
if readErr != nil {
|
||||
log.Printf("Failed to list directory content of %s: %s", directory, readErr)
|
||||
return nil
|
||||
log.Printf("Failed to list directory content of %s: %s", queuePath, readErr)
|
||||
return
|
||||
}
|
||||
|
||||
// Walk over each file and add it to our archive
|
||||
@ -84,29 +98,7 @@ func packQueueFiles(directory string) []byte {
|
||||
continue
|
||||
}
|
||||
|
||||
// Create header for this file
|
||||
header := &tar.Header{
|
||||
Name: f.Name(),
|
||||
Mode: 0600,
|
||||
Size: f.Size(),
|
||||
}
|
||||
|
||||
// Read file
|
||||
path := fmt.Sprintf("%s%c%s", queuePath, os.PathSeparator, f.Name())
|
||||
contents, readErr := ioutil.ReadFile(path)
|
||||
if readErr != nil {
|
||||
log.Printf("Failed to read file %s: %s", path, readErr)
|
||||
continue
|
||||
}
|
||||
|
||||
// Add header and contents to archive
|
||||
tarWriter.WriteHeader(header)
|
||||
tarWriter.Write(contents)
|
||||
// Pack into the archive
|
||||
packSingleFile(tarWriter, absPath, relPath, fmt.Sprintf("queue%c%s", os.PathSeparator, f.Name()), false)
|
||||
}
|
||||
|
||||
// Close constructed tar archive
|
||||
tarWriter.Close()
|
||||
|
||||
// And return it
|
||||
return tarBuffer.Bytes()
|
||||
}
|
||||
|
||||
@ -4,16 +4,15 @@ import (
|
||||
"archive/tar"
|
||||
"bytes"
|
||||
"compress/flate"
|
||||
"encoding/base64"
|
||||
"fmt"
|
||||
"io"
|
||||
"io/ioutil"
|
||||
"log"
|
||||
"os"
|
||||
"strings"
|
||||
"path"
|
||||
)
|
||||
|
||||
// Unpacks a raw string, creates files and stores them in the target directory. May return an error if one occurrs
|
||||
// UnpackInto decrompesses the given bytes with DEFLATE, then unpacks the result as TAR archive into the targetDir
|
||||
func UnpackInto(raw []byte, targetDir string) error {
|
||||
// Prepare FLATE decompressor
|
||||
var flateBuffer bytes.Buffer
|
||||
@ -23,74 +22,11 @@ func UnpackInto(raw []byte, targetDir string) error {
|
||||
flateBuffer.Write(raw)
|
||||
raw, _ = ioutil.ReadAll(flateReader)
|
||||
|
||||
// Process raw bytes
|
||||
splitted := bytes.Split(raw, []byte("\n"))
|
||||
if len(splitted) != 4 {
|
||||
// We are currently packing four things in there (the fuzzer name, queue/, fuzz_bitmap, fuzzer_stats)
|
||||
// So if we don't get three parts, we have a malformed packet
|
||||
return fmt.Errorf("unable to unpack packet: Expected 4 parts, got %d", len(splitted))
|
||||
}
|
||||
|
||||
// base64 decode contents
|
||||
for i, s := range splitted {
|
||||
b64Buf := make([]byte, base64.StdEncoding.DecodedLen(len(s)))
|
||||
base64.StdEncoding.Decode(b64Buf, s)
|
||||
splitted[i] = b64Buf
|
||||
}
|
||||
|
||||
// Check filename, and process it
|
||||
fuzzerName := string(bytes.TrimRight(splitted[0], "\x00"))
|
||||
if strings.Contains(fuzzerName, "/") {
|
||||
return fmt.Errorf("received file name with a slash, discarding whole packet for fuzzer \"%s\"", fuzzerName)
|
||||
}
|
||||
|
||||
// Check if our target directory (this very fuzzers directory) already exists, or if we need to create it
|
||||
targetDir = fmt.Sprintf("%s%c%s", targetDir, os.PathSeparator, fuzzerName)
|
||||
_, folderErr := os.Stat(targetDir)
|
||||
if os.IsNotExist(folderErr) {
|
||||
// directory doesn't yet exist, create it
|
||||
mkdirErr := os.MkdirAll(targetDir, 0700)
|
||||
if mkdirErr != nil {
|
||||
// Creating the target directory failed, so we won't proceed unpacking into a non-existent directory
|
||||
return fmt.Errorf("unable to unpack packet: could not create directory at %s: %s", targetDir, mkdirErr)
|
||||
}
|
||||
}
|
||||
|
||||
// Process every single part
|
||||
unpackSingleFile(splitted[1], targetDir, "fuzz_bitmap")
|
||||
unpackSingleFile(splitted[2], targetDir, "fuzzer_stats")
|
||||
unpackQueueDir(splitted[3], targetDir)
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// Writes the contents to the target
|
||||
func unpackSingleFile(raw []byte, targetDirectory string, filename string) {
|
||||
path := fmt.Sprintf("%s%c%s", targetDirectory, os.PathSeparator, filename)
|
||||
|
||||
// Check if the file already exists - we won't overwrite it then
|
||||
_, fileInfoErr := os.Stat(path)
|
||||
if os.IsExist(fileInfoErr) {
|
||||
// File already exists, we don't need to write a thing
|
||||
return
|
||||
}
|
||||
|
||||
writeErr := ioutil.WriteFile(path, raw, 0644)
|
||||
if writeErr != nil {
|
||||
log.Printf("Unable to write to file %s: %s", path, writeErr)
|
||||
}
|
||||
}
|
||||
|
||||
// Writes all files in the raw byte array into the target directory
|
||||
func unpackQueueDir(raw []byte, targetDir string) {
|
||||
// Open TAR archive
|
||||
var tarBuffer bytes.Buffer
|
||||
tarBuffer.Write(raw)
|
||||
tarReader := tar.NewReader(&tarBuffer)
|
||||
|
||||
// Set correct path for files
|
||||
targetDir = fmt.Sprintf("%s%cqueue", targetDir, os.PathSeparator)
|
||||
|
||||
// Create queue directory if it doesn't exist yet
|
||||
_, folderErr := os.Stat(targetDir)
|
||||
if os.IsNotExist(folderErr) {
|
||||
@ -115,4 +51,36 @@ func unpackQueueDir(raw []byte, targetDir string) {
|
||||
io.Copy(&fileBuffer, tarReader)
|
||||
unpackSingleFile(fileBuffer.Bytes(), targetDir, header.Name)
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// Writes the contents to the target
|
||||
func unpackSingleFile(raw []byte, targetDirectory string, filename string) {
|
||||
destPath := fmt.Sprintf("%s%c%s", targetDirectory, os.PathSeparator, filename)
|
||||
|
||||
// Check if the file already exists - we won't overwrite it then
|
||||
_, fileInfoErr := os.Stat(destPath)
|
||||
if os.IsExist(fileInfoErr) {
|
||||
// File already exists, we don't need to write a thing
|
||||
return
|
||||
}
|
||||
|
||||
// Check if the target directory already exists - otherwise we create it
|
||||
dirOfFile := path.Dir(fmt.Sprintf("%s%c%s", targetDirectory, os.PathSeparator, filename))
|
||||
_, dirInfoErr := os.Stat(dirOfFile)
|
||||
if os.IsNotExist(dirInfoErr) {
|
||||
// Create directories as required
|
||||
mkdirErr := os.MkdirAll(dirOfFile, 0755)
|
||||
if mkdirErr != nil {
|
||||
log.Printf("Failed to create directory %s: %s", dirOfFile, mkdirErr)
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
// Write file
|
||||
writeErr := ioutil.WriteFile(destPath, raw, 0644)
|
||||
if writeErr != nil {
|
||||
log.Printf("Unable to write to file %s: %s", destPath, writeErr)
|
||||
}
|
||||
}
|
||||
|
||||
@ -8,7 +8,6 @@ import (
|
||||
"io/ioutil"
|
||||
"log"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"time"
|
||||
)
|
||||
|
||||
@ -27,20 +26,16 @@ func RegisterWatchdogFlags() {
|
||||
func WatchFuzzers(outputDirectory string) {
|
||||
// Loop forever
|
||||
for {
|
||||
// Loop over those fuzzer directories we want to share
|
||||
for _, localFuzzDir := range getTargetFuzzers(outputDirectory) {
|
||||
// Pack important parts of the fuzzer directory into a byte array
|
||||
fuzzerName := filepath.Base(localFuzzDir)
|
||||
packedFuzzer, packerErr := logistic.PackFuzzer(fuzzerName, localFuzzDir)
|
||||
if packerErr != nil {
|
||||
log.Printf("Failed to pack fuzzer: %s", packerErr)
|
||||
continue
|
||||
}
|
||||
|
||||
// and send it to our peers
|
||||
net.SendToPeers(packedFuzzer)
|
||||
// Pack important parts of the fuzzers into an archive
|
||||
packedFuzzers, packerErr := logistic.PackFuzzers(getTargetFuzzers(outputDirectory), outputDirectory)
|
||||
if packerErr != nil {
|
||||
log.Printf("Failed to pack fuzzer: %s", packerErr)
|
||||
continue
|
||||
}
|
||||
|
||||
// and send it to our peers
|
||||
net.SendToPeers(packedFuzzers)
|
||||
|
||||
// Sleep a bit
|
||||
time.Sleep(time.Duration(rescanSecs) * time.Second)
|
||||
}
|
||||
|
||||
Loading…
Reference in New Issue
Block a user