added TransportFromProtocolSets funcs (#92)

np-guard · Jan 7, 2025 · da5d10b · da5d10b
1 parent cd12519
commit da5d10b
Show file tree

Hide file tree

Showing 2 changed files with 26 additions and 0 deletions.
diff --git a/pkg/netset/icmpset.go b/pkg/netset/icmpset.go
@@ -77,6 +77,18 @@ func NewICMPSet(minType, maxType, minCode, maxCode int64) *ICMPSet {
 	)
 }
 
+func ICMPSetFromICMP(icmp netp.ICMP) *ICMPSet {
+	if icmp.TypeCode == nil {
+		return AllICMPSet()
+	}
+	icmpType := int64(icmp.TypeCode.Type)
+	if icmp.TypeCode.Code == nil {
+		return NewICMPSet(icmpType, icmpType, int64(netp.MinICMPCode), int64(netp.MaxICMPCode))
+	}
+	icmpCode := int64(*icmp.TypeCode.Code)
+	return NewICMPSet(icmpType, icmpType, icmpCode, icmpCode)
+}
+
 func EmptyICMPSet() *ICMPSet {
 	return &ICMPSet{props: ds.NewProductLeft[*TypeSet, *CodeSet]()}
 }

diff --git a/pkg/netset/transportset.go b/pkg/netset/transportset.go
@@ -35,13 +35,27 @@ func NewUDPTransport(srcMinP, srcMaxP, dstMinP, dstMaxP int64) *TransportSet {
 	return NewTCPorUDPTransport(netp.ProtocolStringUDP, srcMinP, srcMaxP, dstMinP, dstMaxP)
 }
 
+func NewICMPTransportFromICMPSet(icmpSet *ICMPSet) *TransportSet {
+	return &TransportSet{ds.NewDisjoint(
+		EmptyTCPorUDPSet(),
+		icmpSet.Copy(),
+	)}
+}
+
 func NewICMPTransport(minType, maxType, minCode, maxCode int64) *TransportSet {
 	return &TransportSet{ds.NewDisjoint(
 		EmptyTCPorUDPSet(),
 		NewICMPSet(minType, maxType, minCode, maxCode),
 	)}
 }
 
+func NewTCPUDPTransportFromTCPUDPSet(tcpudpSet *TCPUDPSet) *TransportSet {
+	return &TransportSet{ds.NewDisjoint(
+		tcpudpSet.Copy(),
+		EmptyICMPSet(),
+	)}
+}
+
 func AllTCPorUDPTransport(protocol netp.ProtocolString) *TransportSet {
 	return NewTCPorUDPTransport(protocol, netp.MinPort, netp.MaxPort, netp.MinPort, netp.MaxPort)
 }