refactor: Refactor main to improve testability

README.md

@@ -14,12 +14,13 @@ GoBSIPS is a simple, single-file, static go binary that installs itself into Lin
 To install GoBSIPS, run the following command:
 
 ```
-sudo ./gobsips
+sudo mv gobsips /usr/local/bin
+sudo /usr/local/bin/gobsips
 ```
 
 This will:
 1. Create a default configuration file at `/etc/sysconfig/gobsips`
-2. Install a systemd service file
+2. Install a systemd service that runs to `/usr/local/bin/gobsips`
 3. Enable and start the service
 
 ## Configuration

go.mod

@@ -4,4 +4,13 @@ go 1.22.6
 
 require github.com/armon/go-socks5 v0.0.0-20160902184237-e75332964ef5
 
-require golang.org/x/net v0.28.0 // indirect
+require (
+	github.com/davecgh/go-spew v1.1.1 // indirect
+	github.com/pmezard/go-difflib v1.0.0 // indirect
+	gopkg.in/yaml.v3 v3.0.1 // indirect
+)
+
+require (
+	github.com/stretchr/testify v1.9.0
+	golang.org/x/net v0.28.0 // indirect
+)

go.sum

@@ -1,4 +1,13 @@
 github.com/armon/go-socks5 v0.0.0-20160902184237-e75332964ef5 h1:0CwZNZbxp69SHPdPJAN/hZIm0C4OItdklCFmMRWYpio=
 github.com/armon/go-socks5 v0.0.0-20160902184237-e75332964ef5/go.mod h1:wHh0iHkYZB8zMSxRWpUBQtwG5a7fFgvEO+odwuTv2gs=
+github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c=
+github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
+github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
+github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
+github.com/stretchr/testify v1.9.0 h1:HtqpIVDClZ4nwg75+f6Lvsy/wHu+3BoSGCbBAcpTsTg=
+github.com/stretchr/testify v1.9.0/go.mod h1:r2ic/lqez/lEtzL7wO/rwa5dbSLXVDPFyf8C91i36aY=
 golang.org/x/net v0.28.0 h1:a9JDOJc5GMUJ0+UDqmLT86WiEy7iWyIhz8gz8E4e5hE=
 golang.org/x/net v0.28.0/go.mod h1:yqtgsTWOOnlGLG9GFRrK3++bGOUEkNBoHZc8MEDWPNg=
+gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
+gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA=
+gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=

main.go

@@ -15,26 +15,32 @@ var (
 )
 
 func main() {
+	if err := do(); err != nil {
+		log.Fatal(err)
+	}
+}
+
+func do() error {
 	flag.Parse()
 	if *flagDaemon {
 		*flagInstall = false
 	}
-	if *flagInstall && os.Getuid() != 0 {
-		log.Fatal("You must install this program as root.")
+	if *flagInstall && getuid() != 0 {
+		return fmt.Errorf("you must install this program as root")
 	}
 
 	switch {
 
 	case *flagInstall:
 
 		if err := installDefaultConfig(); err != nil {
-			log.Fatalf("Failed to write default config to %s: %v", configPath, err)
+			return fmt.Errorf("failed to write default config to %s: %v", configPath, err)
 		}
 		if err := installSystemdService(); err != nil {
-			log.Fatal("Failed to install systemd service:", err)
+			return fmt.Errorf("failed to install systemd service: %v", err)
 		}
 		fmt.Println("gobsips installed and started as a systemd service.")
-		return
+		return nil
 
 	case *flagDaemon:
 
@@ -50,15 +56,25 @@ func main() {
 		// Create SOCKS5 server
 		server, err := socks5.New(conf)
 		if err != nil {
-			log.Fatal(err)
+			return fmt.Errorf("failed to create SOCKS5 server: %v", err)
 		}
 
 		// Start server
 		addr := fmt.Sprintf("%s:%s", config.ListenHost, config.ListenPort)
 		log.Printf("Starting gobsips on %s", addr)
-		if err := server.ListenAndServe("tcp", addr); err != nil {
-			log.Fatal(err)
+		if err := listenAndServe(server, "tcp", addr); err != nil {
+			return fmt.Errorf("failed to start server: %v", err)
 		}
 
 	}
+
+	return nil
 }
+
+var (
+	getuid = os.Getuid
+
+	listenAndServe = func(server *socks5.Server, network, addr string) error {
+		return server.ListenAndServe(network, addr)
+	}
+)

main_test.go

@@ -0,0 +1,82 @@
+package main
+
+import (
+	"flag"
+	"fmt"
+	"os"
+	"os/exec"
+	"testing"
+
+	"github.com/armon/go-socks5"
+	"github.com/stretchr/testify/assert"
+)
+
+func TestRun(t *testing.T) {
+	// Save original flag values and restore them after the test
+	origDaemonFlag := *flagDaemon
+	origInstallFlag := *flagInstall
+	defer func() {
+		*flagDaemon = origDaemonFlag
+		*flagInstall = origInstallFlag
+	}()
+
+	// Test installation path
+	*flagInstall = true
+	*flagDaemon = false
+
+	// Redirect stdout to capture output
+	oldStdout := os.Stdout
+	r, w, _ := os.Pipe()
+	os.Stdout = w
+
+	err := do()
+	if os.Getuid() != 0 {
+		assert.Error(t, err, "you must install this program as root")
+	}
+
+	// Restore stdout
+	w.Close()
+	os.Stdout = oldStdout
+
+	// Read captured output
+	var output []byte
+	r.Read(output)
+
+	if os.Getuid() != 0 {
+		assert.Equal(t, "", string(output))
+	}
+
+	getuid = func() int {
+		return 0
+	}
+	err = do()
+	assert.Error(t, err, "failed to write default config to %s: %v", configPath, err)
+	configPath = "/tmp/test-config"
+	err = do()
+	assert.Error(t, err, "failed to install systemd service: %v", err)
+	run = func(cmd *exec.Cmd) error {
+		return nil
+	}
+	systemdPath = "/tmp/test-systemd-service"
+	err = do()
+	assert.NoError(t, err)
+
+	*flagInstall = false
+	*flagDaemon = true
+	listenAndServe = func(_ *socks5.Server, network, addr string) error {
+		return fmt.Errorf("testing without starting server (%v:%v)", network, addr)
+	}
+	err = do()
+	assert.Error(t, err, "failed to start server: testing without starting server (tcp:0.0.0.0:1080)")
+	listenAndServe = func(_ *socks5.Server, _, _ string) error {
+		return nil
+	}
+	err = do()
+	assert.NoError(t, err)
+}
+
+func TestMain(m *testing.M) {
+	// This is needed to properly initialize flags for tests
+	flag.Parse()
+	os.Exit(m.Run())
+}