feat: Add body to TCP, UDP, and TLS endpoints and templating (#1134)

* feat(endpoints): Add body to TCP, UDP, and TLS endpoints and templating * Changed the template to be more consistent with the rest of the application and added additional substritutions. * Changed getModifiedBody to getParsedBody and fixed connected response * Apply suggestion from @TwiN * Apply suggestion from @TwiN * Apply suggestion from @TwiN * Apply suggestion from @TwiN * Apply suggestion from @TwiN * Apply suggestion from @TwiN * Apply suggestion from @TwiN * Apply suggestion from @TwiN * Apply suggestion from @TwiN * Apply suggestion from @TwiN * Apply suggestion from @TwiN * Apply suggestion from @TwiN * Update client/client.go --------- Co-authored-by: TwiN <twin@linux.com>
2025-07-19 19:24:03 -06:00
parent f4a667549e
commit bdaffbca77
5 changed files with 111 additions and 33 deletions
--- a/client/client.go
+++ b/client/client.go
@@ -76,24 +76,37 @@ func GetDomainExpiration(hostname string) (domainExpiration time.Duration, err e
 	return domainExpiration, nil
 }

-// CanCreateTCPConnection checks whether a connection can be established with a TCP endpoint
-func CanCreateTCPConnection(address string, config *Config) bool {
-	conn, err := net.DialTimeout("tcp", address, config.Timeout)
-	if err != nil {
-		return false
-	}
-	_ = conn.Close()
-	return true
+// parseLocalAddressPlaceholder returns a string with the local address replaced
+func parseLocalAddressPlaceholder(item string, localAddr net.Addr) string {
+	item = strings.ReplaceAll(item, "[LOCAL_ADDRESS]", localAddr.String())
+	return item
 }

-// CanCreateUDPConnection checks whether a connection can be established with a UDP endpoint
-func CanCreateUDPConnection(address string, config *Config) bool {
-	conn, err := net.DialTimeout("udp", address, config.Timeout)
+// CanCreateNetworkConnection checks whether a connection can be established with a TCP or UDP endpoint
+func CanCreateNetworkConnection(netType string, address string, body string, config *Config) (bool, []byte) {
+	const (
+		MaximumMessageSize = 1024 // in bytes
+	)
+	connection, err := net.DialTimeout(netType, address, config.Timeout)
 	if err != nil {
-		return false
+		return false, nil
 	}
-	_ = conn.Close()
-	return true
+	defer connection.Close()
+	if body != "" {
+		body = parseLocalAddressPlaceholder(body, connection.LocalAddr())
+		connection.SetDeadline(time.Now().Add(config.Timeout))
+		_, err = connection.Write([]byte(body))
+		if err != nil {
+			return false, nil
+		}
+		buf := make([]byte, MaximumMessageSize)
+		n, err := connection.Read(buf)
+		if err != nil {
+			return false, nil
+		}
+		return true, buf[:n]
+	}
+	return true, nil
 }

 // CanCreateSCTPConnection checks whether a connection can be established with a SCTP endpoint
@@ -152,7 +165,10 @@ func CanPerformStartTLS(address string, config *Config) (connected bool, certifi
 }

 // CanPerformTLS checks whether a connection can be established to an address using the TLS protocol
-func CanPerformTLS(address string, config *Config) (connected bool, certificate *x509.Certificate, err error) {
+func CanPerformTLS(address string, body string, config *Config) (connected bool, response []byte, certificate *x509.Certificate, err error) {
+	const (
+		MaximumMessageSize = 1024 // in bytes
+	)
 	connection, err := tls.DialWithDialer(&net.Dialer{Timeout: config.Timeout}, "tcp", address, &tls.Config{
 		InsecureSkipVerify: config.Insecure,
 	})
@@ -166,9 +182,27 @@ func CanPerformTLS(address string, config *Config) (connected bool, certificate
 	// Reference: https://pkg.go.dev/crypto/tls#PeerCertificates
 	if len(verifiedChains) == 0 || len(verifiedChains[0]) == 0 {
 		peerCertificates := connection.ConnectionState().PeerCertificates
-		return true, peerCertificates[0], nil
+		certificate = peerCertificates[0]
+	} else {
+		certificate = verifiedChains[0][0]
 	}
-	return true, verifiedChains[0][0], nil
+	connected = true
+	if body != "" {
+		body = parseLocalAddressPlaceholder(body, connection.LocalAddr())
+		connection.SetDeadline(time.Now().Add(config.Timeout))
+		_, err = connection.Write([]byte(body))
+		if err != nil {
+			return
+		}
+		buf := make([]byte, MaximumMessageSize)
+		var n int
+		n, err = connection.Read(buf)
+		if err != nil {
+			return
+		}
+		response = buf[:n]
+	}
+	return
 }

 // CanCreateSSHConnection checks whether a connection can be established and a command can be executed to an address
@@ -234,6 +268,7 @@ func ExecuteSSHCommand(sshClient *ssh.Client, body string, config *Config) (bool
 	}
 	defer sshClient.Close()
 	var b Body
+	body = parseLocalAddressPlaceholder(body, sshClient.Conn.LocalAddr())
 	if err := json.Unmarshal([]byte(body), &b); err != nil {
 		return false, 0, err
 	}
@@ -304,6 +339,7 @@ func QueryWebSocket(address, body string, config *Config) (bool, []byte, error)
 		return false, nil, fmt.Errorf("error dialing websocket: %w", err)
 	}
 	defer ws.Close()
+	body = parseLocalAddressPlaceholder(body, ws.LocalAddr())
 	// Write message
 	if _, err := ws.Write([]byte(body)); err != nil {
 		return false, nil, fmt.Errorf("error writing websocket body: %w", err)
--- a/client/client_test.go
+++ b/client/client_test.go
@@ -223,7 +223,7 @@ func TestCanPerformTLS(t *testing.T) {
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
 			t.Parallel()
-			connected, _, err := CanPerformTLS(tt.args.address, &Config{Insecure: tt.args.insecure, Timeout: 5 * time.Second})
+			connected, _, _, err := CanPerformTLS(tt.args.address, "", &Config{Insecure: tt.args.insecure, Timeout: 5 * time.Second})
 			if (err != nil) != tt.wantErr {
 				t.Errorf("CanPerformTLS() err=%v, wantErr=%v", err, tt.wantErr)
 				return
@@ -235,11 +235,13 @@ func TestCanPerformTLS(t *testing.T) {
 	}
 }

-func TestCanCreateTCPConnection(t *testing.T) {
-	if CanCreateTCPConnection("127.0.0.1", &Config{Timeout: 5 * time.Second}) {
+func TestCanCreateConnection(t *testing.T) {
+	connected, _ := CanCreateNetworkConnection("tcp", "127.0.0.1", "", &Config{Timeout: 5 * time.Second})
+	if connected {
 		t.Error("should've failed, because there's no port in the address")
 	}
-	if !CanCreateTCPConnection("1.1.1.1:53", &Config{Timeout: 5 * time.Second}) {
+	connected, _ = CanCreateNetworkConnection("tcp", "1.1.1.1:53", "", &Config{Timeout: 5 * time.Second})
+	if !connected {
 		t.Error("should've succeeded, because that IP should always™ be up")
 	}
 }