From 434fac52b7f11aeb26094c25d0c0d3c859f77b3a Mon Sep 17 00:00:00 2001
From: Kristoffer Dalby <kradalby@kradalby.no>
Date: Sat, 30 Oct 2021 14:29:03 +0000
Subject: [PATCH] Fix lint error

---
 app.go                     | 4 +---
 cmd/headscale/cli/utils.go | 5 +----
 utils.go                   | 7 +++++++
 3 files changed, 9 insertions(+), 7 deletions(-)

diff --git a/app.go b/app.go
index 94517d54..707c111e 100644
--- a/app.go
+++ b/app.go
@@ -355,9 +355,7 @@ func (h *Headscale) Serve() error {
 		h.cfg.UnixSocket,
 		[]grpc.DialOption{
 			grpc.WithInsecure(),
-			grpc.WithDialer(func(addr string, timeout time.Duration) (net.Conn, error) {
-				return net.DialTimeout("unix", addr, timeout)
-			}),
+			grpc.WithContextDialer(GrpcSocketDialer),
 		}...,
 	)
 	if err != nil {
diff --git a/cmd/headscale/cli/utils.go b/cmd/headscale/cli/utils.go
index 1d6fdd61..8e044bf5 100644
--- a/cmd/headscale/cli/utils.go
+++ b/cmd/headscale/cli/utils.go
@@ -5,7 +5,6 @@ import (
 	"encoding/json"
 	"errors"
 	"fmt"
-	"net"
 	"net/url"
 	"os"
 	"path/filepath"
@@ -310,9 +309,7 @@ func getHeadscaleGRPCClient() (apiV1.HeadscaleServiceClient, *grpc.ClientConn) {
 		grpcOptions = append(
 			grpcOptions,
 			grpc.WithInsecure(),
-			grpc.WithDialer(func(addr string, timeout time.Duration) (net.Conn, error) {
-				return net.DialTimeout("unix", addr, timeout)
-			}),
+			grpc.WithContextDialer(headscale.GrpcSocketDialer),
 		)
 	} else {
 		// If we are not connecting to a local server, require an API key for authentication
diff --git a/utils.go b/utils.go
index cbe1d870..555227ec 100644
--- a/utils.go
+++ b/utils.go
@@ -6,10 +6,12 @@
 package headscale
 
 import (
+	"context"
 	"crypto/rand"
 	"encoding/json"
 	"fmt"
 	"io"
+	"net"
 	"strings"
 
 	"golang.org/x/crypto/nacl/box"
@@ -156,3 +158,8 @@ func tailNodesToString(nodes []*tailcfg.Node) string {
 func tailMapResponseToString(resp tailcfg.MapResponse) string {
 	return fmt.Sprintf("{ Node: %s, Peers: %s }", resp.Node.Name, tailNodesToString(resp.Peers))
 }
+
+func GrpcSocketDialer(ctx context.Context, addr string) (net.Conn, error) {
+	var d net.Dialer
+	return d.DialContext(ctx, "unix", addr)
+}