From dcc8ba6ec52b0d2fcfa9af41755f6b3597a67c04 Mon Sep 17 00:00:00 2001
From: Tricia Bogen <tricia-github@cernu.us>
Date: Sat, 19 Aug 2023 08:01:20 -0700
Subject: [PATCH] add --timeout

---
 cmd/root.go    |  7 +++++--
 cmd/unifi.go   |  6 ++++--
 types/unifi.go | 46 +++++++++++++++++++++++++---------------------
 3 files changed, 34 insertions(+), 25 deletions(-)

diff --git a/cmd/root.go b/cmd/root.go
index 7f074fe..33e6df6 100644
--- a/cmd/root.go
+++ b/cmd/root.go
@@ -1,11 +1,12 @@
 package cmd
 
 import (
-	"github.com/rs/zerolog"
-	"golang.org/x/crypto/ssh/terminal"
 	"os"
 	"time"
 
+	"github.com/rs/zerolog"
+	"golang.org/x/crypto/ssh/terminal"
+
 	"github.com/spf13/cobra"
 )
 
@@ -23,6 +24,8 @@ func init() {
 	Root.Flags().String("host", "localhost", "hostname of the UniFi "+
 		"controller")
 	Root.Flags().Int("port", 8443, "UniFi controller port")
+	Root.Flags().Duration("timeout", 5*time.Second, "unifi server connection "+
+		"timeout")
 	Root.Flags().Bool("verify-tls", false, "if true, verify the TLS "+
 		"certificate of the UniFi controller")
 	Root.Flags().String("username", "", "login user on UniFi controller "+
diff --git a/cmd/unifi.go b/cmd/unifi.go
index e12111c..fb2db1d 100644
--- a/cmd/unifi.go
+++ b/cmd/unifi.go
@@ -2,9 +2,10 @@ package cmd
 
 import (
 	"fmt"
-	"github.com/rs/zerolog"
 	"os"
 
+	"github.com/rs/zerolog"
+
 	"github.com/pdbogen/unifi2mqtt/types"
 	"github.com/spf13/cobra"
 )
@@ -26,8 +27,9 @@ func startUnifi(cmd *cobra.Command, log zerolog.Logger) (<-chan types.Client, er
 	port, _ := cmd.Flags().GetInt("port")
 	verifyTls, _ := cmd.Flags().GetBool("verify-tls")
 	deviceTimeout, _ := cmd.Flags().GetDuration("unifi-timeout")
+	connectTimeout, _ := cmd.Flags().GetDuration("timeout")
 
-	u, err := types.NewUnifi(username, password, host, port, verifyTls, deviceTimeout, matchers, log)
+	u, err := types.NewUnifi(username, password, host, port, verifyTls, deviceTimeout, connectTimeout, matchers, log)
 	if err != nil {
 		return nil, fmt.Errorf("starting unifi client: %w", err)
 	}
diff --git a/types/unifi.go b/types/unifi.go
index 649c5e4..2ef59d6 100644
--- a/types/unifi.go
+++ b/types/unifi.go
@@ -15,17 +15,18 @@ type Client struct {
 }
 
 type Unifi struct {
-	Username      string
-	Password      string
-	Host          string
-	Port          int
-	VerifyTls     bool
-	DeviceTimeout time.Duration
-	Matchers      []*Matcher
-	seenClients   map[string]string
-	unifiClient   *unifi.Unifi
-	ch            chan Client
-	loginBackoff  time.Duration
+	Username       string
+	Password       string
+	Host           string
+	Port           int
+	VerifyTls      bool
+	DeviceTimeout  time.Duration
+	ConnectTimeout time.Duration
+	Matchers       []*Matcher
+	seenClients    map[string]string
+	unifiClient    *unifi.Unifi
+	ch             chan Client
+	loginBackoff   time.Duration
 }
 
 func (u *Unifi) Login(log zerolog.Logger) error {
@@ -43,6 +44,7 @@ func (u *Unifi) Login(log zerolog.Logger) error {
 		DebugLog: func(msg string, fmt ...interface{}) {
 			log.Debug().Msgf(msg, fmt...)
 		},
+		Timeout: u.ConnectTimeout,
 	})
 
 	if err == nil {
@@ -93,18 +95,20 @@ func (u *Unifi) Start(log zerolog.Logger) error {
 }
 
 func NewUnifi(username, password, host string, port int, verifyTls bool,
-	deviceTimeout time.Duration, matchers []*Matcher, log zerolog.Logger,
+	deviceTimeout time.Duration, connectTimeout time.Duration,
+	matchers []*Matcher, log zerolog.Logger,
 ) (*Unifi, error) {
 	ret := &Unifi{
-		Username:      username,
-		Password:      password,
-		Host:          host,
-		Port:          port,
-		VerifyTls:     verifyTls,
-		DeviceTimeout: deviceTimeout,
-		Matchers:      matchers,
-		seenClients:   map[string]string{},
-		ch:            make(chan Client),
+		Username:       username,
+		Password:       password,
+		Host:           host,
+		Port:           port,
+		VerifyTls:      verifyTls,
+		DeviceTimeout:  deviceTimeout,
+		ConnectTimeout: connectTimeout,
+		Matchers:       matchers,
+		seenClients:    map[string]string{},
+		ch:             make(chan Client),
 	}
 
 	if err := ret.Start(log); err != nil {