Refactor logging

command/arp.go

@@ -1,7 +1,6 @@
 package command
 
 import (
-	"bufio"
 	"context"
 	"errors"
 	"net"
@@ -12,11 +11,11 @@ import (
 	"time"
 
 	"github.com/spf13/cobra"
+	"github.com/v-byte-cpu/sx/command/log"
 	"github.com/v-byte-cpu/sx/pkg/ip"
 	"github.com/v-byte-cpu/sx/pkg/packet/afpacket"
 	"github.com/v-byte-cpu/sx/pkg/scan"
 	"github.com/v-byte-cpu/sx/pkg/scan/arp"
-	"go.uber.org/zap"
 )
 
 var errSrcIP = errors.New("invalid source IP")
@@ -77,44 +76,22 @@ var arpCmd = &cobra.Command{
 			}
 		}
 
-		r := &scan.Range{Subnet: dstSubnet, Interface: iface, SrcIP: srcIP.To4(), SrcMAC: srcMAC}
-		return startEngine(r)
-	},
-}
-
-func logResults(logger *zap.Logger, results <-chan *arp.ScanResult) {
-	bw := bufio.NewWriter(os.Stdout)
-	defer bw.Flush()
-	for result := range results {
-		// TODO refactor it using logger facade interface
+		var logger log.Logger
 		if jsonFlag {
-			data, err := result.MarshalJSON()
-			if err != nil {
-				logger.Error("arp", zap.Error(err))
-			}
-			_, err = bw.Write(data)
-			if err != nil {
-				logger.Error("arp", zap.Error(err))
-			}
+			logger, err = log.NewJSONLogger(os.Stdout, "arp")
 		} else {
-			_, err := bw.WriteString(result.String())
-			if err != nil {
-				logger.Error("arp", zap.Error(err))
-			}
+			logger, err = log.NewPlainLogger(os.Stdout, "arp")
 		}
-		err := bw.WriteByte('\n')
 		if err != nil {
-			logger.Error("arp", zap.Error(err))
+			return err
 		}
-	}
-}
 
-func startEngine(r *scan.Range) error {
-	logger, err := zap.NewProduction()
-	if err != nil {
-		return err
-	}
+		r := &scan.Range{Subnet: dstSubnet, Interface: iface, SrcIP: srcIP.To4(), SrcMAC: srcMAC}
+		return startEngine(logger, r)
+	},
+}
 
+func startEngine(logger log.Logger, r *scan.Range) error {
 	ctx, cancel := signal.NotifyContext(context.Background(), os.Interrupt)
 	defer cancel()
 
@@ -136,7 +113,7 @@ func startEngine(r *scan.Range) error {
 	wg.Add(1)
 	go func() {
 		defer wg.Done()
-		logResults(logger, m.Results())
+		logger.LogResults(m.Results())
 	}()
 
 	// start scan
@@ -153,7 +130,7 @@ func startEngine(r *scan.Range) error {
 	go func() {
 		defer wg.Done()
 		for err := range errc {
-			logger.Error("arp", zap.Error(err))
+			logger.Error(err)
 		}
 	}()
 	wg.Wait()

command/log/json_logger.go

@@ -0,0 +1,45 @@
+package log
+
+import (
+	"bufio"
+	"io"
+
+	"github.com/v-byte-cpu/sx/pkg/scan/arp"
+	"go.uber.org/zap"
+)
+
+type JSONLogger struct {
+	logger *zap.Logger
+	writer io.Writer
+	label  string
+}
+
+func NewJSONLogger(w io.Writer, label string) (*JSONLogger, error) {
+	logger, err := zap.NewProduction()
+	if err != nil {
+		return nil, err
+	}
+	return &JSONLogger{logger, w, label}, nil
+}
+
+func (l *JSONLogger) Error(err error) {
+	l.logger.Error(l.label, zap.Error(err))
+}
+
+func (l *JSONLogger) LogResults(results <-chan *arp.ScanResult) {
+	bw := bufio.NewWriter(l.writer)
+	defer bw.Flush()
+	var err error
+	var data []byte
+	for result := range results {
+		if data, err = result.MarshalJSON(); err != nil {
+			l.Error(err)
+		}
+		if _, err = bw.Write(data); err != nil {
+			l.Error(err)
+		}
+		if err = bw.WriteByte('\n'); err != nil {
+			l.Error(err)
+		}
+	}
+}

command/log/json_logger_test.go

@@ -0,0 +1,71 @@
+package log
+
+import (
+	"bytes"
+	"net"
+	"strings"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+	"github.com/v-byte-cpu/sx/pkg/scan/arp"
+)
+
+func scanResultToJSON(t *testing.T, result *arp.ScanResult) string {
+	t.Helper()
+	data, err := result.MarshalJSON()
+	require.NoError(t, err)
+	return string(data)
+}
+
+func TestJSONLoggerResults(t *testing.T) {
+	t.Parallel()
+
+	tests := []struct {
+		name     string
+		expected []byte
+		results  []*arp.ScanResult
+	}{
+		{
+			name:     "emptyResults",
+			expected: nil,
+			results:  nil,
+		},
+		{
+			name:     "oneResult",
+			expected: []byte(scanResultToJSON(t, newScanResult(net.IPv4(192, 168, 0, 3).To4())) + "\n"),
+			results: []*arp.ScanResult{
+				newScanResult(net.IPv4(192, 168, 0, 3).To4()),
+			},
+		},
+		{
+			name: "twoResults",
+			expected: []byte(strings.Join([]string{
+				scanResultToJSON(t, newScanResult(net.IPv4(192, 168, 0, 3).To4())),
+				scanResultToJSON(t, newScanResult(net.IPv4(192, 168, 0, 5).To4())),
+			}, "\n") + "\n"),
+			results: []*arp.ScanResult{
+				newScanResult(net.IPv4(192, 168, 0, 3).To4()),
+				newScanResult(net.IPv4(192, 168, 0, 5).To4()),
+			},
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+
+			var buf bytes.Buffer
+			logger, err := NewJSONLogger(&buf, "arp")
+			require.NoError(t, err)
+
+			resultCh := make(chan *arp.ScanResult, len(tt.results))
+			for _, result := range tt.results {
+				resultCh <- result
+			}
+			close(resultCh)
+			logger.LogResults(resultCh)
+
+			assert.Equal(t, tt.expected, buf.Bytes())
+		})
+	}
+}

command/log/logger.go

@@ -0,0 +1,8 @@
+package log
+
+import "github.com/v-byte-cpu/sx/pkg/scan/arp"
+
+type Logger interface {
+	Error(err error)
+	LogResults(results <-chan *arp.ScanResult)
+}

command/log/plain_logger.go

@@ -0,0 +1,41 @@
+package log
+
+import (
+	"bufio"
+	"io"
+
+	"github.com/v-byte-cpu/sx/pkg/scan/arp"
+	"go.uber.org/zap"
+)
+
+type PlainLogger struct {
+	logger *zap.Logger
+	writer io.Writer
+	label  string
+}
+
+func NewPlainLogger(w io.Writer, label string) (*PlainLogger, error) {
+	logger, err := zap.NewProduction()
+	if err != nil {
+		return nil, err
+	}
+	return &PlainLogger{logger, w, label}, nil
+}
+
+func (l *PlainLogger) Error(err error) {
+	l.logger.Error(l.label, zap.Error(err))
+}
+
+func (l *PlainLogger) LogResults(results <-chan *arp.ScanResult) {
+	bw := bufio.NewWriter(l.writer)
+	defer bw.Flush()
+	var err error
+	for result := range results {
+		if _, err = bw.WriteString(result.String()); err != nil {
+			l.Error(err)
+		}
+		if err = bw.WriteByte('\n'); err != nil {
+			l.Error(err)
+		}
+	}
+}

command/log/plain_logger_test.go

@@ -0,0 +1,72 @@
+package log
+
+import (
+	"bytes"
+	"net"
+	"strings"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+	"github.com/v-byte-cpu/sx/pkg/scan/arp"
+)
+
+func newScanResult(ip net.IP) *arp.ScanResult {
+	return &arp.ScanResult{
+		IP:     ip.String(),
+		MAC:    net.HardwareAddr{0x11, 0x22, 0x33, 0x44, 0x55, 0x66}.String(),
+		Vendor: "Sunny Industries",
+	}
+}
+
+func TestPlainLoggerResults(t *testing.T) {
+	t.Parallel()
+
+	tests := []struct {
+		name     string
+		expected []byte
+		results  []*arp.ScanResult
+	}{
+		{
+			name:     "emptyResults",
+			expected: nil,
+			results:  nil,
+		},
+		{
+			name:     "oneResult",
+			expected: []byte(newScanResult(net.IPv4(192, 168, 0, 3).To4()).String() + "\n"),
+			results: []*arp.ScanResult{
+				newScanResult(net.IPv4(192, 168, 0, 3).To4()),
+			},
+		},
+		{
+			name: "twoResults",
+			expected: []byte(strings.Join([]string{
+				newScanResult(net.IPv4(192, 168, 0, 3).To4()).String(),
+				newScanResult(net.IPv4(192, 168, 0, 5).To4()).String(),
+			}, "\n") + "\n"),
+			results: []*arp.ScanResult{
+				newScanResult(net.IPv4(192, 168, 0, 3).To4()),
+				newScanResult(net.IPv4(192, 168, 0, 5).To4()),
+			},
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+
+			var buf bytes.Buffer
+			logger, err := NewPlainLogger(&buf, "arp")
+			require.NoError(t, err)
+
+			resultCh := make(chan *arp.ScanResult, len(tt.results))
+			for _, result := range tt.results {
+				resultCh <- result
+			}
+			close(resultCh)
+			logger.LogResults(resultCh)
+
+			assert.Equal(t, tt.expected, buf.Bytes())
+		})
+	}
+}