reworking log, early check if db connects

README.md

@@ -46,7 +46,7 @@ go install github.com/Kyagara/hagelslag-go@latest
 
 ## TODO/Ideas
 
-- Convert the logging to a tui, maybe with some prompts ('e' to print the errors), so it's easier to read.
+- Improve logging.
 
 - At high rates, DB connection errors out.
 
@@ -56,6 +56,4 @@ go install github.com/Kyagara/hagelslag-go@latest
 
 - Maybe add bedrocks servers to the Minecraft scanner.
 
-- I think logging can be better.
-
 - Remove this weird virus that keeps adding Frieren in the code.

hagelslag.go

@@ -11,6 +11,8 @@ import (
 	"runtime"
 	"strings"
 	"sync"
+	"sync/atomic"
+	"syscall"
 	"time"
 
 	"github.com/rs/zerolog/log"
@@ -131,6 +133,7 @@ func (h Hagelslag) spawn(semaphore <-chan struct{}, ip string, port string, netw
 	address := net.JoinHostPort(ip, port)
 	conn, err := dialer.Dial(network, address)
 	if err != nil {
+		// Don't log timeouts
 		if errors.Is(err, os.ErrDeadlineExceeded) {
 			return
 		}
@@ -141,6 +144,7 @@ func (h Hagelslag) spawn(semaphore <-chan struct{}, ip string, port string, netw
 	defer conn.Close()
 
 	if h.OnlyConnect {
+		atomic.AddInt64(&successCount, 1)
 		return
 	}
 
@@ -157,20 +161,22 @@ func (h Hagelslag) spawn(semaphore <-chan struct{}, ip string, port string, netw
 	}
 
 	if err != nil {
-		// Don't log EOF and connection reset errors
-		if errors.Is(err, io.EOF) || errors.Is(err, net.ErrClosed) {
+		// Don't log timeouts, connection reset errors or  EOF
+		if errors.Is(err, os.ErrDeadlineExceeded) || errors.Is(err, syscall.ECONNRESET) || errors.Is(err, io.EOF) {
 			return
 		}
 
-		log.Log().Str("ip", ip).Err(err).Msg("SCAN")
+		fmt.Println()
+		log.Error().Str("ip", ip).Err(err).Msg("SCAN")
 		return
 	}
 
 	err = h.Scanner.Save(ip, latency, response, collection)
 	if err != nil {
-		log.Log().Str("ip", ip).Err(err).Msg("SAVE")
+		fmt.Println()
+		log.Error().Str("ip", ip).Err(err).Msg("SAVE")
 		return
 	}
 
-	log.Log().Str("ip", ip).Msg("SAVED")
+	atomic.AddInt64(&successCount, 1)
 }

http.go

@@ -2,7 +2,6 @@ package main
 
 import (
 	"context"
-	"fmt"
 	"io"
 	"net"
 	"strings"
@@ -35,14 +34,14 @@ func (s HTTP) Scan(ip string, conn net.Conn) ([]byte, int64, error) {
 	start := time.Now()
 	_, err := conn.Write(unsafe.Slice(unsafe.StringData(get), len(get)))
 	if err != nil {
-		return nil, 0, fmt.Errorf("sending GET request: %w", err)
+		return nil, 0, err
 	}
 
 	response := make([]byte, 17)
 
 	_, err = io.ReadFull(conn, response)
 	if err != nil {
-		return nil, 0, fmt.Errorf("reading status from GET response: %w", err)
+		return nil, 0, err
 	}
 
 	latency := time.Since(start).Milliseconds()
@@ -54,7 +53,7 @@ func (s HTTP) Scan(ip string, conn net.Conn) ([]byte, int64, error) {
 
 	response, err = io.ReadAll(conn)
 	if err != nil {
-		return nil, 0, fmt.Errorf("reading GET response: %w", err)
+		return nil, 0, err
 	}
 
 	return response, latency, nil

main.go

@@ -1,6 +1,7 @@
 package main
 
 import (
+	"context"
 	"fmt"
 	"os"
 	"os/signal"
@@ -12,41 +13,63 @@ import (
 
 	"github.com/rs/zerolog"
 	"github.com/rs/zerolog/log"
+	"go.mongodb.org/mongo-driver/mongo"
+	"go.mongodb.org/mongo-driver/mongo/options"
+)
+
+var (
+	// IP/s
+	ipRate = int64(0)
+	// Amount of times the scanner connect (if OnlyConnect is true) or saved to the database successfully
+	successCount = int64(0)
 )
 
 func main() {
 	zerolog.TimeFieldFormat = zerolog.TimeFormatUnix
+	log.Logger = log.Output(zerolog.ConsoleWriter{Out: os.Stderr})
 
 	hagelslag, err := NewHagelslag()
 	if err != nil {
-		panic(err)
+		fmt.Printf("Error initializing hagelslag: %s\n", err)
+		os.Exit(1)
 	}
 
-	ips := make(chan string, hagelslag.MaxTasks)
-
-	var wg sync.WaitGroup
-	for range hagelslag.NumWorkers {
-		wg.Add(1)
-		go hagelslag.worker(ips, &wg)
+	err = testDBConnection(hagelslag.URI)
+	if err != nil {
+		fmt.Printf("Error connecting to database: %s\n", err)
+		os.Exit(1)
 	}
 
 	segA, segB, segC, segD, err := parseIP(hagelslag.StartingIP)
 	if err != nil {
-		panic(err)
+		fmt.Printf("Error parsing starting IP: %s\n", err)
+		os.Exit(1)
 	}
 
+	ticker := time.NewTicker(1 * time.Second)
+
+	ips := make(chan string, hagelslag.MaxTasks)
 	signals := make(chan os.Signal, 1)
 	signal.Notify(signals, syscall.SIGINT, syscall.SIGTERM)
 
-	counter := int64(0)
-
-	go getIPPerSecond(&counter)
+	var wg sync.WaitGroup
+	for range hagelslag.NumWorkers {
+		wg.Add(1)
+		go hagelslag.worker(ips, &wg)
+	}
 
 	// Main loop
 	for {
 		select {
+		case <-ticker.C:
+			ipsPerSecond := atomic.LoadInt64(&ipRate)
+			success := atomic.LoadInt64(&successCount)
+			fmt.Printf("\r\033[KSuccess: %d | IP/s: %d", success, ipsPerSecond)
+			atomic.StoreInt64(&ipRate, 0)
+
 		case <-signals:
-			fmt.Println("Received signal, shutting down...")
+			fmt.Println()
+			fmt.Println("Shutting down...")
 			close(ips)
 			wg.Wait()
 			fmt.Printf("Last IP: %d.%d.%d.%d\n", segA, segB, segC, segD)
@@ -57,11 +80,12 @@ func main() {
 			ip := strconv.Itoa(segA) + "." + strconv.Itoa(segB) + "." + strconv.Itoa(segC) + "." + strconv.Itoa(segD)
 
 			if isReserved(&segA, &segB, &segC) {
-				log.Log().Str("ip", ip).Msg("Reserved range reached, skipping to next available range")
+				fmt.Println()
+				log.Info().Str("ip", ip).Msg("Reserved range reached, skipping to next available range")
 			}
 
 			ips <- ip
-			atomic.AddInt64(&counter, 1)
+			atomic.AddInt64(&ipRate, 1)
 
 			segD++
 			if segD >= 256 {
@@ -87,14 +111,24 @@ func main() {
 	}
 }
 
-func getIPPerSecond(counter *int64) {
-	secondTicker := time.NewTicker(1 * time.Second)
+// For checking earlier if the database is reachable
+func testDBConnection(uri string) error {
+	client, err := mongo.Connect(context.TODO(), options.Client().SetServerSelectionTimeout(3*time.Second).ApplyURI(uri))
+	if err != nil {
+		return err
+	}
+
+	err = client.Ping(context.TODO(), nil)
+	if err != nil {
+		return err
+	}
 
-	for range secondTicker.C {
-		ipsPerSecond := atomic.LoadInt64(counter)
-		log.Log().Int64("ips", ipsPerSecond).Msg("IPs per second")
-		atomic.StoreInt64(counter, 0)
+	err = client.Disconnect(context.TODO())
+	if err != nil {
+		return err
 	}
+
+	return nil
 }
 
 // ************#*#******####*########***#####################%%%%@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@%%@@@@%%%@@@@@@@@@@@@@@@@@@@