1 files changed, 37 insertions, 24 deletions

diff --git a/udp_flow.c b/udp_flow.c
index 8907f2f..0023837 100644
--- a/udp_flow.c
+++ b/udp_flow.c
@@ -17,8 +17,6 @@
 #include "udp_internal.h"
 #include "epoll_ctl.h"
 
-#define UDP_CONN_TIMEOUT	180 /* s, timeout for ephemeral or local bind */
-
 /**
  * udp_at_sidx() - Get UDP specific flow at given sidx
  * @sidx:    Flow and side to retrieve
@@ -74,11 +72,6 @@ static int udp_flow_sock(const struct ctx *c,
 {
 	const struct flowside *side = &uflow->f.side[sidei];
 	uint8_t pif = uflow->f.pif[sidei];
-	union {
-		flow_sidx_t sidx;
-		uint32_t data;
-	} fref = { .sidx = FLOW_SIDX(uflow, sidei) };
-	union epoll_ref ref;
 	int rc;
 	int s;
 
@@ -88,14 +81,9 @@ static int udp_flow_sock(const struct ctx *c,
 		return s;
 	}
 
-	ref.type = EPOLL_TYPE_UDP;
-	ref.data = fref.data;
-	ref.fd = s;
-
 	flow_epollid_set(&uflow->f, EPOLLFD_ID_DEFAULT);
-
-	rc = epoll_add(flow_epollfd(&uflow->f), EPOLLIN, ref);
-	if (rc < 0) {
+	if (flow_epoll_set(&uflow->f, EPOLL_CTL_ADD, EPOLLIN, s, sidei) < 0) {
+		rc = -errno;
 		close(s);
 		return rc;
 	}
@@ -109,6 +97,7 @@ static int udp_flow_sock(const struct ctx *c,
 		flow_dbg_perror(uflow, "Couldn't connect flow socket");
 		return rc;
 	}
+	uflow->s[sidei] = s;
 
 	/* It's possible, if unlikely, that we could receive some packets in
 	 * between the bind() and connect() which may or may not be for this
@@ -139,6 +128,7 @@ static int udp_flow_sock(const struct ctx *c,
  * udp_flow_new() - Common setup for a new UDP flow
  * @c:		Execution context
  * @flow:	Initiated flow
+ * @rule_hint:	Index of forwarding rule, or -1 if unknown
  * @now:	Timestamp
  *
  * Return: sidx for the target side of the new UDP flow, or FLOW_SIDX_NONE
@@ -147,23 +137,25 @@ static int udp_flow_sock(const struct ctx *c,
  * #syscalls getsockname
  */
 static flow_sidx_t udp_flow_new(const struct ctx *c, union flow *flow,
-				const struct timespec *now)
+				int rule_hint, const struct timespec *now)
 {
 	struct udp_flow *uflow = NULL;
 	const struct flowside *tgt;
 	unsigned sidei;
 
-	if (!(tgt = flow_target(c, flow, IPPROTO_UDP)))
+	if (!(tgt = flow_target(c, flow, rule_hint, IPPROTO_UDP)))
 		goto cancel;
 
 	uflow = FLOW_SET_TYPE(flow, FLOW_UDP, udp);
 	uflow->ts = now->tv_sec;
 	uflow->s[INISIDE] = uflow->s[TGTSIDE] = -1;
 	uflow->ttl[INISIDE] = uflow->ttl[TGTSIDE] = 0;
+	uflow->activity[INISIDE] = 1;
+	uflow->activity[TGTSIDE] = 0;
 
 	flow_foreach_sidei(sidei) {
 		if (pif_is_socket(uflow->f.pif[sidei]))
-			if ((uflow->s[sidei] = udp_flow_sock(c, uflow, sidei)) < 0)
+			if (udp_flow_sock(c, uflow, sidei) < 0)
 				goto cancel;
 	}
 
@@ -216,6 +208,7 @@ cancel:
  * @dst:	Our (local) address to which the datagram is arriving
  * @port:	Our (local) port number to which the datagram is arriving
  * @s_in:	Source socket address, filled in by recvmmsg()
+ * @rule_hint:	Index of forwarding rule, or -1 if unknown
  * @now:	Timestamp
  *
  * #syscalls fcntl arm:fcntl64 ppc64:fcntl64|fcntl i686:fcntl64
@@ -226,7 +219,7 @@ cancel:
 flow_sidx_t udp_flow_from_sock(const struct ctx *c, uint8_t pif,
 			       const union inany_addr *dst, in_port_t port,
 			       const union sockaddr_inany *s_in,
-			       const struct timespec *now)
+			       int rule_hint, const struct timespec *now)
 {
 	const struct flowside *ini;
 	struct udp_flow *uflow;
@@ -235,7 +228,7 @@ flow_sidx_t udp_flow_from_sock(const struct ctx *c, uint8_t pif,
 
 	sidx = flow_lookup_sa(c, IPPROTO_UDP, pif, s_in, dst, port);
 	if ((uflow = udp_at_sidx(sidx))) {
-		uflow->ts = now->tv_sec;
+		udp_flow_activity(uflow, sidx.sidei, now);
 		return flow_sidx_opposite(sidx);
 	}
 
@@ -260,7 +253,7 @@ flow_sidx_t udp_flow_from_sock(const struct ctx *c, uint8_t pif,
 		return FLOW_SIDX_NONE;
 	}
 
-	return udp_flow_new(c, flow, now);
+	return udp_flow_new(c, flow, rule_hint, now);
 }
 
 /**
@@ -292,7 +285,7 @@ flow_sidx_t udp_flow_from_tap(const struct ctx *c,
 	sidx = flow_lookup_af(c, IPPROTO_UDP, pif, af, saddr, daddr,
 			      srcport, dstport);
 	if ((uflow = udp_at_sidx(sidx))) {
-		uflow->ts = now->tv_sec;
+		udp_flow_activity(uflow, sidx.sidei, now);
 		return flow_sidx_opposite(sidx);
 	}
 
@@ -316,7 +309,7 @@ flow_sidx_t udp_flow_from_tap(const struct ctx *c,
 		return FLOW_SIDX_NONE;
 	}
 
-	return udp_flow_new(c, flow, now);
+	return udp_flow_new(c, flow, FWD_NO_HINT, now);
 }
 
 /**
@@ -332,7 +325,7 @@ static void udp_flush_flow(const struct ctx *c,
 {
 	/* We don't know exactly where the datagrams will come from, but we know
 	 * they'll have an interface and oport matching this flow */
-	udp_sock_fwd(c, uflow->s[sidei], uflow->f.pif[sidei],
+	udp_sock_fwd(c, uflow->s[sidei], -1, uflow->f.pif[sidei],
 		     uflow->f.side[sidei].oport, now);
 }
 
@@ -369,9 +362,29 @@ bool udp_flow_defer(const struct ctx *c, struct udp_flow *uflow,
 bool udp_flow_timer(const struct ctx *c, struct udp_flow *uflow,
 		    const struct timespec *now)
 {
-	if (now->tv_sec - uflow->ts <= UDP_CONN_TIMEOUT)
+	int timeout = c->udp.timeout;
+
+	if (uflow->activity[TGTSIDE] &&
+	    (uflow->activity[INISIDE] > 1 || uflow->activity[TGTSIDE] > 1))
+		timeout = c->udp.stream_timeout;
+
+	if (now->tv_sec - uflow->ts <= timeout)
 		return false;
 
 	udp_flow_close(c, uflow);
 	return true;
 }
+
+/**
+ * udp_flow_activity() - Track activity of a UDP flow
+ * @uflow:	UDP flow
+ * @sidei:	Side index of the flow (INISIDE or TGTSIDE)
+ * @now:	Current timestamp
+ */
+void udp_flow_activity(struct udp_flow *uflow, unsigned int sidei,
+		       const struct timespec *now)
+{
+	uflow->ts = now->tv_sec;
+	if (uflow->activity[sidei] < UINT8_MAX)
+		uflow->activity[sidei]++;
+}