From 1582ecdae5f735cd47178ed7cfd808b56c29782c Mon Sep 17 00:00:00 2001
From: Marten Seemann <martenseemann@gmail.com>
Date: Sat, 24 Jul 2021 18:42:03 +0200
Subject: [PATCH] add functions to tell if an address contains a security
 protocol

---
 multiaddr.go      | 11 +++++++++++
 multiaddr_test.go | 30 ++++++++++++++++++++++++++++++
 2 files changed, 41 insertions(+)

diff --git a/multiaddr.go b/multiaddr.go
index 58fe8ce..c3fabc1 100644
--- a/multiaddr.go
+++ b/multiaddr.go
@@ -184,3 +184,14 @@ func (m *multiaddr) ValueForProtocol(code int) (value string, err error) {
 	})
 	return
 }
+
+func ContainsSecurityProtocol(addr Multiaddr) bool {
+	var contains bool
+	ForEach(addr, func(c Component) bool {
+		if code := c.Protocol().Code; code == P_TLS || code == P_NOISE {
+			contains = true
+		}
+		return true
+	})
+	return contains
+}
diff --git a/multiaddr_test.go b/multiaddr_test.go
index f617ff9..5359eaa 100644
--- a/multiaddr_test.go
+++ b/multiaddr_test.go
@@ -739,3 +739,33 @@ func TestComponentJSONMarshaler(t *testing.T) {
 		t.Error("expected equal components in circular marshaling test")
 	}
 }
+
+func TestContainsSecurity(t *testing.T) {
+	does := []string{
+		"/ip4/127.0.0.1/tcp/80/tls",
+		"/ip6/::1/tcp/4321/noise",
+	}
+	for _, s := range does {
+		addr, err := NewMultiaddr(s)
+		if err != nil {
+			t.Fatal(err)
+		}
+		if !ContainsSecurityProtocol(addr) {
+			t.Fatalf("expected %s to be recognized as a security protocol", s)
+		}
+	}
+	doesnt := []string{
+		"/ip4/127.0.0.1/tcp/80",
+		"/ip6/::1/tcp/4321",
+		"/ip4/127.0.0.1/udp/443/quic",
+	}
+	for _, s := range doesnt {
+		addr, err := NewMultiaddr(s)
+		if err != nil {
+			t.Fatal(err)
+		}
+		if ContainsSecurityProtocol(addr) {
+			t.Fatalf("didn't expect %s to be recognized as a security protocol", s)
+		}
+	}
+}