Merge pull request #2 from massix/feat/use-context

feat: use contexts instead of forwarding signals to terminate goroutines

.github/workflows/ci.yaml

@@ -4,15 +4,16 @@ on: [push]
 
 jobs:
   build:
-    name: Build
+    name: Build raw binary
     runs-on: ubuntu-latest
     steps:
       - uses: actions/checkout@v4
       - uses: actions/setup-go@v5
         with:
           go-version: ">=1.22"
       - run: go version
-      - run: go build -v ./cmd/protrans
+      - run: make all
+      - run: ldd ./protrans
   test:
     name: Test
     runs-on: ubuntu-latest
@@ -22,4 +23,14 @@ jobs:
         with:
           go-version: ">=1.22"
       - run: go version
-      - run: go test -v ./...
+      - run: make test
+  flake:
+    name: Build Flake
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+      - uses: cachix/install-nix-action@v27
+        with:
+          github_access_token: ${{ secrets.GITHUB_TOKEN }}
+      - run: nix build
+      - run: ldd ./result/bin/protrans

.gitignore

@@ -1 +1,2 @@
 .direnv/
+protrans

Makefile

@@ -0,0 +1,15 @@
+GO := $(shell which go)
+VERSION = 1.1
+
+.PHONY: clean
+
+all: protrans
+
+protrans: cmd/protrans/main.go
+	$(GO) build -ldflags "-X 'main.Version=${VERSION}'" -o $@ -v $^
+
+test:
+	$(GO) test -v ./...
+
+clean:
+	rm -fr protrans

cmd/protrans/main.go

@@ -1,6 +1,7 @@
 package main
 
 import (
+	"context"
 	"fmt"
 	"os"
 	"os/signal"
@@ -23,9 +24,14 @@ func dumpTransmissionConfiguration(conf *config.ProtransConfiguration) string {
 	return fmt.Sprintf("\tHost: %s\n\tPort: %d\n\tUsername: %s\n", conf.Transmission.Host, conf.Transmission.Port, conf.Transmission.Username)
 }
 
+var Version string
+
 func main() {
 	var configurationPath string
 
+	var wg sync.WaitGroup
+	defer wg.Wait()
+
 	if len(os.Args) > 1 {
 		configurationPath = os.Args[1]
 		logrus.Infof("Parsing configuration from '%s'", configurationPath)
@@ -35,7 +41,9 @@ func main() {
 
 	conf := config.NewConfiguration(configurationPath, true)
 	logrus.SetLevel(conf.LogrusLogLevel())
-	logrus.Infof("Log level: %s\n", conf.LogrusLogLevel().String())
+
+	logrus.Infof("Protrans version: %s", Version)
+	logrus.Infof("Log level: %s", conf.LogrusLogLevel().String())
 	logrus.Infof("NAT Configuration:\n%s", dumpNatConfiguration(conf))
 	logrus.Infof("Transmission Configuration:\n%s", dumpTransmissionConfiguration(conf))
 
@@ -47,12 +55,11 @@ func main() {
 		logrus.Panic(err)
 	}
 
+	ctx, cancel := context.WithCancel(context.Background())
+
 	// Register to some signals
-	done := make(chan os.Signal, 126)
-	signal.Notify(done, syscall.SIGTERM)
-	signal.Notify(done, syscall.SIGINT)
-	signal.Notify(done, syscall.SIGABRT)
-	signal.Notify(done, syscall.SIGHUP)
+	sigChan := make(chan os.Signal, 1)
+	signal.Notify(sigChan, syscall.SIGTERM, syscall.SIGINT, syscall.SIGABRT, syscall.SIGHUP)
 
 	// Buffered channel to make sure we're refreshing it constantly
 	ipChan := make(chan string, 30)
@@ -62,20 +69,27 @@ func main() {
 	defer func() {
 		close(ipChan)
 		close(portChan)
-		close(done)
+		close(sigChan)
 	}()
 
-	var wg sync.WaitGroup
 	wg.Add(3)
 
 	// This goroutine will constantly check the external IP address and send it to a channel
-	go flow.FetchExternalIP(natClient, &wg, ipChan, done)
+	go flow.FetchExternalIP(ctx, natClient, &wg, ipChan)
 
 	// This goroutine will receive the IP address and create a port mapping which will be sent to another channel
-	go flow.MapPorts(natClient, int(conf.Nat.PortLifetime), &wg, ipChan, portChan, done)
+	go flow.MapPorts(ctx, natClient, int(conf.Nat.PortLifetime), &wg, ipChan, portChan)
 
 	// This goroutine will receive the mapped port and send it to Transmission if connected
-	go flow.TransmissionArgSetter(transmissionClient, &wg, portChan, done)
+	go flow.TransmissionArgSetter(ctx, transmissionClient, &wg, portChan)
+
+	select {
+	case <-ctx.Done():
+		logrus.Infof("Context closed, leaving")
+	case s := <-sigChan:
+		logrus.Infof("Received signal %q, leaving", s)
+		cancel()
+	}
 
-	wg.Wait()
+	logrus.Info("Waiting for goroutines to finish")
 }

flake.nix

@@ -5,20 +5,26 @@
     let
       system = "x86_64-linux";
       pkgs = import nixpkgs { inherit system; };
+      hardeningDisable = [ "fortify" ];
+      version = "1.1";
     in
     {
       devShells.${system}.default = pkgs.mkShell {
+        inherit hardeningDisable;
         packages = with pkgs; [ go ];
-        hardeningDisable = [ "fortify" ];
       };
 
       packages.${system}.default = pkgs.buildGoModule {
+        inherit hardeningDisable version;
         pname = "protrans";
-        version = "1.0";
-        hardeningDisable = [ "fortify" ];
 
-        src = ./.;
+        src =
+          let
+            noSrcs = [ ".vscode" ".git" ".github" ".gitignore" ".envrc" ];
+          in
+          builtins.filterSource (path: _: ! builtins.elem (baseNameOf path) noSrcs) ./.;
 
+        ldflags = [ "-X 'main.Version=${version}'" ];
         vendorHash = "sha256-H79018dCud68fYT0l3IGZXQvD22byhnw/GchsiYJc68=";
       };
 

pkg/config/config.go

@@ -33,6 +33,7 @@ type (
 func overrideWithEnvironment(iface, concrete any) {
 	types := reflect.TypeOf(iface)
 	values := reflect.ValueOf(concrete).Elem()
+
 	for i := 0; i < types.NumField(); i++ {
 		typeField := types.Field(i)
 		valueField := values.FieldByName(typeField.Name)

pkg/flow/flow.go

@@ -1,7 +1,7 @@
 package flow
 
 import (
-	"os"
+	"context"
 	"sync"
 	"time"
 
@@ -10,9 +10,12 @@ import (
 	"github.com/sirupsen/logrus"
 )
 
-func FetchExternalIP(natClient nat.NatClientI, wg *sync.WaitGroup, ipChan chan<- string, done chan os.Signal) {
+func FetchExternalIP(ctx context.Context, natClient nat.NatClientI, wg *sync.WaitGroup, ipChan chan<- string) {
 	running := true
 
+	timer := time.NewTimer(30 * time.Second)
+	defer timer.Stop()
+
 	for running {
 		ip, err := nat.GetExternalIP(natClient)
 		if err != nil {
@@ -23,26 +26,25 @@ func FetchExternalIP(natClient nat.NatClientI, wg *sync.WaitGroup, ipChan chan<-
 		}
 
 		select {
-		case s := <-done:
+		case <-ctx.Done():
 			logrus.Info("Gracefully stopping gateway detector")
 			running = false
-			done <- s // Make sure everyone is leaving
-		case <-time.After(30 * time.Second):
+		case <-timer.C:
 			logrus.Debug("No signals received in 30 seconds, refreshing IP...")
+			timer.Reset(30 * time.Second)
 		}
 	}
 
 	wg.Done()
 }
 
-func MapPorts(natClient nat.NatClientI, portLifetime int, wg *sync.WaitGroup, ipChan <-chan string, portChan chan<- int, done chan os.Signal) {
+func MapPorts(ctx context.Context, natClient nat.NatClientI, portLifetime int, wg *sync.WaitGroup, ipChan <-chan string, portChan chan<- int) {
 	running := true
 
 	for running {
 		select {
 		case ip := <-ipChan:
-			var mappedTcpPort int
-			var mappedUdpPort int
+			var mappedTcpPort, mappedUdpPort int
 			var err error
 			logrus.Debugf("Mapping port for external IP: %s", ip)
 
@@ -65,17 +67,16 @@ func MapPorts(natClient nat.NatClientI, portLifetime int, wg *sync.WaitGroup, ip
 
 			logrus.Debugf("Sending port %d to channel", mappedTcpPort)
 			portChan <- mappedTcpPort
-		case s := <-done:
+		case <-ctx.Done():
 			logrus.Info("Gracefully stopping port mapper")
 			running = false
-			done <- s
 		}
 	}
 
 	wg.Done()
 }
 
-func TransmissionArgSetter(transmissionClient transmission.TransmissionClient, wg *sync.WaitGroup, portChan <-chan int, done chan os.Signal) {
+func TransmissionArgSetter(ctx context.Context, transmissionClient transmission.TransmissionClient, wg *sync.WaitGroup, portChan <-chan int) {
 	running := true
 
 	for running {
@@ -104,10 +105,9 @@ func TransmissionArgSetter(transmissionClient transmission.TransmissionClient, w
 			} else {
 				logrus.Warnf("Should set port: %d but Transmission is not connected", mappedPort)
 			}
-		case s := <-done:
+		case <-ctx.Done():
 			logrus.Info("Gracefully stopping Transmission Client")
 			running = false
-			done <- s
 		}
 	}