diff --git a/internal/pkg/util/utils.go b/internal/pkg/util/utils.go
index 36aab0b..9b22576 100644
--- a/internal/pkg/util/utils.go
+++ b/internal/pkg/util/utils.go
@@ -146,9 +146,22 @@
 }
 
 // GetUniPortFromFlow returns uni port from the flow data
-func GetUniPortFromFlow(nniPort string, flow *of.VoltSubFlow) uint32 {
+func GetUniPortFromFlow(uplinkPort string, nniPorts []string, flow *of.VoltSubFlow) uint32 {
 	var portNo uint32
-	if nniPort == strconv.Itoa(int(flow.Match.InPort)) {
+	var isDSFlow bool
+
+	if len(nniPorts) > 0 {
+		for _, nniPort := range nniPorts {
+			if nniPort == strconv.Itoa(int(flow.Match.InPort)) {
+				isDSFlow = true
+				break
+			}
+		}
+	} else if uplinkPort == strconv.Itoa(int(flow.Match.InPort)) {
+		isDSFlow = true
+	}
+
+	if isDSFlow {
 		if of.IPProtocolUDP == flow.Match.L4Protocol {
 			// For DHCP DS flow, uniport is not part of metadata. Hence retrieve it from cookie
 			portNo = GetUniFromDSDhcpFlow(flow.Cookie)
@@ -158,6 +171,7 @@
 	} else {
 		portNo = flow.Match.InPort
 	}
+
 	return portNo
 }