feat: add ConnectContext for context-aware connection

tcpclient.go

@@ -4,6 +4,7 @@ package gos7
 // This software may be modified and distributed under the terms
 // of the BSD license. See the LICENSE file for details.
 import (
+	"context"
 	"encoding/binary"
 	"fmt"
 	"io"
@@ -187,30 +188,15 @@ func (mb *tcpTransporter) Send(request []byte) (response []byte, err error) {
 // Connect establishes a new connection to the address in Address.
 // Connect and Close are exported so that multiple requests can be done with one session
 func (mb *tcpTransporter) Connect() error {
-	// mb.mu.Lock()
-	// defer mb.mu.Unlock()
-
-	return mb.connect()
-}
-func (mb *tcpTransporter) tcpConnect() error {
-	mb.mu.Lock()
-	defer mb.mu.Unlock()
-	if mb.conn == nil {
-		dialer := net.Dialer{Timeout: mb.Timeout}
-		conn, err := dialer.Dial("tcp", mb.Address)
-		if err != nil {
-			if conn != nil {
-				_ = conn.Close()
-			}
-			return err
-		}
-		mb.conn = conn
-	}
-	return nil
+	return mb.ConnectContext(context.Background())
 }
-func (mb *tcpTransporter) connect() error {
+
+// ConnectContext establishes a new connection to the address in Address.
+// It is the same as Connect but accepts a context that can be used to
+// cancel or set a deadline on the TCP dial and subsequent protocol handshake.
+func (mb *tcpTransporter) ConnectContext(ctx context.Context) error {
 	//first stage: TCP connection
-	err := mb.tcpConnect()
+	err := mb.tcpConnect(ctx)
 	if err != nil {
 		return err
 	}
@@ -224,7 +210,23 @@ func (mb *tcpTransporter) connect() error {
 	}
 	// Third stage : S7 protocol data unit negotiation
 	return mb.negotiatePduLength()
+}
 
+func (mb *tcpTransporter) tcpConnect(ctx context.Context) error {
+	mb.mu.Lock()
+	defer mb.mu.Unlock()
+	if mb.conn == nil {
+		dialer := net.Dialer{Timeout: mb.Timeout}
+		conn, err := dialer.DialContext(ctx, "tcp", mb.Address)
+		if err != nil {
+			if conn != nil {
+				_ = conn.Close()
+			}
+			return err
+		}
+		mb.conn = conn
+	}
+	return nil
 }
 
 func (mb *tcpTransporter) isoConnect() error {

tcpclient_test.go

@@ -5,6 +5,7 @@ package gos7
 // of the BSD license. See the LICENSE file for details.
 import (
 	"bytes"
+	"context"
 	"io"
 	"net"
 	"testing"
@@ -37,7 +38,7 @@ func TestTCPTransporter(t *testing.T) {
 	}
 	req := []byte{0, 1, 0, 17, 0, 2, 1, 2, 0, 1, 0, 17, 0, 2, 1, 2, 2} //lengh 17, > MinPduSize
 
-	client.tcpConnect() //assume tcp connect to test locally
+	client.tcpConnect(context.Background()) //assume tcp connect to test locally
 	rsp, err := client.Send(req)
 	if err != nil {
 		t.Fatal(err)