diff --git a/src/Routes/RouteRegister.c b/src/Routes/RouteRegister.c
index 6425fce..9fe2e5b 100644
--- a/src/Routes/RouteRegister.c
+++ b/src/Routes/RouteRegister.c
@@ -55,7 +55,7 @@ ROUTE_IMPL(RouteRegister, args)
 
     User *user = NULL;
 
-    Array *uiaFlows;
+    Array *uiaFlows = NULL;
     int uiaResult;
 
     if (MATRIX_PATH_PARTS(args->path) == 0)
@@ -241,6 +241,7 @@ ROUTE_IMPL(RouteRegister, args)
 
         UserUnlock(user);
 finish:
+        UiaFlowsFree(uiaFlows);
         Free(username);
         Free(password);
         Free(deviceId);
diff --git a/src/Uia.c b/src/Uia.c
index e63d886..2ec3d89 100644
--- a/src/Uia.c
+++ b/src/Uia.c
@@ -135,6 +135,8 @@ BuildResponse(Array * flows, char *session, Db * db, HashMap ** response)
         DbUnlock(db, ref);
 
         HashMapSet(*response, "completed", JsonValueArray(ArrayCreate()));
+        HashMapSet(*response, "session", JsonValueString(session));
+        Free(session);
     }
     else
     {
@@ -167,11 +169,11 @@ BuildResponse(Array * flows, char *session, Db * db, HashMap ** response)
         }
 
         HashMapSet(*response, "completed", JsonValueArray(completed));
+        HashMapSet(*response, "session", JsonValueString(session));
 
         DbUnlock(db, ref);
     }
 
-    HashMapSet(*response, "session", JsonValueString(session));
     return 0;
 }
 
@@ -200,7 +202,7 @@ UiaBuildStage(char *type, HashMap * params)
         return NULL;
     }
 
-    stage->type = type;
+    stage->type = StrDuplicate(type);
     stage->params = params;
 
     return stage;
@@ -218,9 +220,6 @@ UiaComplete(Array * flows, HttpServerContext * context, Db * db,
     DbRef *dbRef;
     HashMap *dbJson;
 
-    size_t i, j;
-    int ret = 0;
-
     if (!flows)
     {
         return -1;
@@ -228,8 +227,7 @@ UiaComplete(Array * flows, HttpServerContext * context, Db * db,
 
     if (!context || !db || !request || !response)
     {
-        ret = -1;
-        goto finish;
+        return -1;
     }
 
     val = HashMapGet(request, "auth");
@@ -237,16 +235,14 @@ UiaComplete(Array * flows, HttpServerContext * context, Db * db,
     if (!val)
     {
         HttpResponseStatus(context, HTTP_UNAUTHORIZED);
-        ret = BuildResponse(flows, NULL, db, response);
-        goto finish;
+        return BuildResponse(flows, NULL, db, response);
     }
 
     if (JsonValueType(val) != JSON_OBJECT)
     {
         HttpResponseStatus(context, HTTP_BAD_REQUEST);
         *response = MatrixErrorCreate(M_BAD_JSON);
-        ret = 0;
-        goto finish;
+        return 0;
     }
 
     auth = JsonValueAsObject(val);
@@ -256,8 +252,7 @@ UiaComplete(Array * flows, HttpServerContext * context, Db * db,
     {
         HttpResponseStatus(context, HTTP_BAD_REQUEST);
         *response = MatrixErrorCreate(M_BAD_JSON);
-        ret = 0;
-        goto finish;
+        return 0;
     }
 
     session = JsonValueAsString(val);
@@ -267,8 +262,7 @@ UiaComplete(Array * flows, HttpServerContext * context, Db * db,
     {
         HttpResponseStatus(context, HTTP_BAD_REQUEST);
         *response = MatrixErrorCreate(M_BAD_JSON);
-        ret = 0;
-        goto finish;
+        return 0;
     }
 
     authType = JsonValueAsString(val);
@@ -277,17 +271,26 @@ UiaComplete(Array * flows, HttpServerContext * context, Db * db,
     if (!dbRef)
     {
         HttpResponseStatus(context, HTTP_UNAUTHORIZED);
-        ret = BuildResponse(flows, StrDuplicate(session), db, response);
-        goto finish;
+        return BuildResponse(flows, StrDuplicate(session), db, response);
     }
 
     dbJson = DbJson(dbRef);
 
     DbUnlock(db, dbRef);
 
-    ret = 1;
+    return 1;
+}
+
+void
+UiaFlowsFree(Array *flows)
+{
+    size_t i, j;
+
+    if (!flows)
+    {
+        return;
+    }
 
-finish:
     for (i = 0; i < ArraySize(flows); i++)
     {
         Array *stages = ArrayGet(flows, i);
@@ -304,7 +307,6 @@ finish:
         ArrayFree(stages);
     }
     ArrayFree(flows);
-    return ret;
 }
 
 void
diff --git a/src/include/Uia.h b/src/include/Uia.h
index d672efc..2ef575b 100644
--- a/src/include/Uia.h
+++ b/src/include/Uia.h
@@ -41,6 +41,9 @@ extern void
  UiaCleanup(MatrixHttpHandlerArgs *);
 
 extern int
- UiaComplete(Array * stages, HttpServerContext *, Db *, HashMap *, HashMap **);
+ UiaComplete(Array *, HttpServerContext *, Db *, HashMap *, HashMap **);
+
+extern void
+UiaFlowsFree(Array *);
 
 #endif