From eb5857b68aecbe27beee4703b7a39c8aece734b3 Mon Sep 17 00:00:00 2001
From: Victor Julien <victor@inliniac.net>
Date: Wed, 19 Apr 2017 15:57:08 +0200
Subject: [PATCH] unittests: add/improve helpers for stream/flow

---
 src/util-unittest-helper.c | 53 +++++++++++++++++++++++++++++++++++++-
 src/util-unittest-helper.h |  3 +++
 2 files changed, 55 insertions(+), 1 deletion(-)

diff --git a/src/util-unittest-helper.c b/src/util-unittest-helper.c
index f82da37bda..f678f129f1 100644
--- a/src/util-unittest-helper.c
+++ b/src/util-unittest-helper.c
@@ -1,4 +1,4 @@
-/* Copyright (C) 2007-2013 Open Information Security Foundation
+/* Copyright (C) 2007-2017 Open Information Security Foundation
  *
  * You can copy, redistribute or modify this Program under the terms of
  * the GNU General Public License version 2 as published by the Free
@@ -36,6 +36,9 @@
 #include "detect-engine.h"
 #include "detect-engine-sigorder.h"
 
+#include "stream-tcp.h"
+#include "stream-tcp-private.h"
+
 #include "util-debug.h"
 #include "util-time.h"
 #include "util-error.h"
@@ -490,6 +493,54 @@ void UTHFreeFlow(Flow *flow)
     }
 }
 
+int UTHAddStreamToFlow(Flow *f, int direction,
+    uint8_t *data, uint32_t data_len)
+{
+    FAIL_IF_NULL(f);
+    FAIL_IF_NOT(f->proto == IPPROTO_TCP);
+    FAIL_IF_NULL(f->protoctx);
+    TcpSession *ssn = f->protoctx;
+
+    StreamingBufferSegment seg;
+    TcpStream *stream = direction == 0 ? &ssn->client : &ssn->server;
+    int r = StreamingBufferAppend(&stream->sb, &seg, data, data_len);
+    FAIL_IF_NOT(r == 0);
+    stream->last_ack += data_len;
+    return 1;
+}
+
+int UTHAddSessionToFlow(Flow *f,
+    uint32_t ts_isn,
+    uint32_t tc_isn)
+{
+    FAIL_IF_NULL(f);
+
+    TcpSession *ssn = SCCalloc(1, sizeof(*ssn));
+    FAIL_IF_NULL(ssn);
+
+    StreamingBuffer x = STREAMING_BUFFER_INITIALIZER(&stream_config.sbcnf);
+    ssn->client.sb = x;
+    ssn->server.sb = x;
+
+    ssn->client.isn = ts_isn;
+    ssn->server.isn = tc_isn;
+
+    f->protoctx = ssn;
+    return 1;
+}
+
+int UTHRemoveSessionFromFlow(Flow *f)
+{
+    FAIL_IF_NULL(f);
+    FAIL_IF_NOT(f->proto == IPPROTO_TCP);
+    TcpSession *ssn = f->protoctx;
+    FAIL_IF_NULL(ssn);
+    StreamTcpSessionCleanup(ssn);
+    SCFree(ssn);
+    f->protoctx = NULL;
+    return 1;
+}
+
 /**
  * \brief UTHGenericTest: function that perfom a generic check taking care of
  *                      as maximum common unittest elements as possible.
diff --git a/src/util-unittest-helper.h b/src/util-unittest-helper.h
index 3b57a9e47a..7555e7f36c 100644
--- a/src/util-unittest-helper.h
+++ b/src/util-unittest-helper.h
@@ -43,6 +43,9 @@ void UTHFreePackets(Packet **, int);
 
 Flow *UTHBuildFlow(int family, char *src, char *dst, Port sp, Port dp);
 void UTHFreeFlow(Flow *flow);
+int UTHAddStreamToFlow(Flow *f, int direction, uint8_t *data, uint32_t data_len);
+int UTHAddSessionToFlow(Flow *f, uint32_t ts_isn, uint32_t tc_isn);
+int UTHRemoveSessionFromFlow(Flow *f);
 
 int UTHAppendSigs(DetectEngineCtx *, char **, int);
 int UTHMatchPackets(DetectEngineCtx *, Packet **, int);