fix(stun): honour ctx and make timeout configurable (#47)

This diff ensures we immediately close the connection if the timeout is
done. It also allows to configure the timeout.

pkg/cli/stun/README.md

@@ -39,6 +39,12 @@ Writes structured logs to the given FILE. If FILE already exists, we
 append to it. If FILE does not exist, we create it. If FILE is a single
 dash (`-`), we write to the stdout.
 
+### `--max-time DURATION`
+
+Sets the maximum time that the STUN transaction operation is allowed to take
+in seconds (e.g., `--max-time 5`). If this flag is not specified, the
+default max time is 30 seconds.
+
 ### `--measure`
 
 Do not exit with `1` if communication with the endpoint fails. Only exit

pkg/cli/stun/stun.go

@@ -9,6 +9,7 @@ import (
 	"errors"
 	"fmt"
 	"io"
+	"time"
 
 	"github.com/rbmk-project/common/cliutils"
 	"github.com/rbmk-project/common/closepool"
@@ -51,6 +52,7 @@ func (cmd command) Main(ctx context.Context, env cliutils.Environment, argv ...s
 
 	// 4. add flags to the parser
 	logfile := clip.String("logs", "", "path where to write structured logs")
+	maxtime := clip.Int("maxtime", 30, "maximum time for transaction to complete (in seconds)")
 	measure := clip.Bool("measure", false, "do not exit 1 on measurement failure")
 
 	// 5. parse command line arguments
@@ -71,6 +73,7 @@ func (cmd command) Main(ctx context.Context, env cliutils.Environment, argv ...s
 
 	// 7. finish filling up the task
 	task.Endpoint = args[0]
+	task.MaxTime = time.Duration(*maxtime) * time.Second
 
 	// 8. handle --logs flag
 	var filepool closepool.Pool

pkg/cli/stun/task.go

@@ -25,14 +25,17 @@ type Task struct {
 	// LogsWriter is where we write structured logs
 	LogsWriter io.Writer
 
+	// MaxTime is the maximum time to wait for the operation to finish.
+	MaxTime time.Duration
+
 	// Output is where we write the results
 	Output io.Writer
 }
 
 // Run executes the STUN binding request task
 func (task *Task) Run(ctx context.Context) error {
 	// 1. Set up the overall operation timeout
-	ctx, cancel := context.WithTimeout(ctx, 30*time.Second)
+	ctx, cancel := context.WithTimeout(ctx, task.MaxTime)
 	defer cancel()
 
 	// 2. Set up the JSON logger for writing measurements
@@ -52,11 +55,20 @@ func (task *Task) Run(ctx context.Context) error {
 		return conn
 	}
 
-	// 5. Establish UDP connection to STUN server
+	// 5. Establish UDP connection to STUN server and make sure
+	// we have proper context deadline propagation. Also, make sure
+	// that we bail immediately if the context is done.
 	conn, err := netx.DialContext(ctx, "udp", task.Endpoint)
 	if err != nil {
 		return fmt.Errorf("cannot connect to STUN server: %w", err)
 	}
+	if d, ok := ctx.Deadline(); ok {
+		conn.SetDeadline(d)
+	}
+	go func() {
+		<-ctx.Done()
+		conn.Close()
+	}()
 
 	// 6. Build STUN binding request message
 	message := stun.MustBuild(stun.TransactionID, stun.BindingRequest)