From b2bf9cbc413d26a2a3601af90ad8e487dfc1a5ee Mon Sep 17 00:00:00 2001
From: David Adrian <davidcadrian@gmail.com>
Date: Tue, 18 Feb 2020 15:46:36 -0500
Subject: [PATCH] Remove dependency on config in output.go (#250)

Refactor the output handlers to take the necessary writers as arguments.

https://github.com/zmap/zgrab2/pull/250
---
 config.go |  3 ++-
 output.go | 22 +++++++++++++++-------
 2 files changed, 17 insertions(+), 8 deletions(-)

diff --git a/config.go b/config.go
index cd0e08a..f70cdc7 100644
--- a/config.go
+++ b/config.go
@@ -89,7 +89,8 @@ func validateFrameworkConfiguration() {
 			log.Fatal(err)
 		}
 	}
-	SetOutputFunc(OutputResultsFile)
+	outputFunc := OutputResultsWriterFunc(config.outputFile)
+	SetOutputFunc(outputFunc)
 
 	if config.MetaFileName == "-" {
 		config.metaFile = os.Stderr
diff --git a/output.go b/output.go
index 5b6e0bc..5d234fc 100644
--- a/output.go
+++ b/output.go
@@ -3,6 +3,7 @@ package zgrab2
 import (
 	"bufio"
 	"fmt"
+	"io"
 )
 
 // FlagMap is a function that maps a single-bit bitmask (i.e. a number of the
@@ -133,16 +134,23 @@ func WidenMapKeys(input map[int]string) map[uint64]string {
 // results or error.
 type OutputResultsFunc func(results <-chan []byte) error
 
-// OutputResultsFile is an OutputResultsFunc that write results to
-// a filename provided on the command line.
-func OutputResultsFile(results <-chan []byte) error {
-	out := bufio.NewWriter(config.outputFile)
-	defer out.Flush()
+// OutputResultsWriterFunc returns an OutputResultsFunc that wraps an io.Writer
+// in a buffered writer, and uses OutputResults.
+func OutputResultsWriterFunc(w io.Writer) OutputResultsFunc {
+	buf := bufio.NewWriter(w)
+	return func(result <-chan []byte) error {
+		defer buf.Flush()
+		return OutputResults(buf, result)
+	}
+}
+
+// OutputResults writes results to a buffered Writer from a channel.
+func OutputResults(w *bufio.Writer, results <-chan []byte) error {
 	for result := range results {
-		if _, err := out.Write(result); err != nil {
+		if _, err := w.Write(result); err != nil {
 			return err
 		}
-		if err := out.WriteByte('\n'); err != nil {
+		if err := w.WriteByte('\n'); err != nil {
 			return err
 		}
 	}