From 991b5836f61a45486505360a30ae3a240d244aa4 Mon Sep 17 00:00:00 2001
From: cheatfate <eugene.kabanov@status.im>
Date: Tue, 13 Apr 2021 11:56:21 +0300
Subject: [PATCH] Add `redirect` macro to allow path redirection. Add tests for
 `redirect` macro. Add check for unique patterns in path and tests.

---
 presto/btrees.nim     |  12 +++
 presto/route.nim      |  84 ++++++++++++++++---
 presto/segpath.nim    |  66 +++++++++++++++
 tests/testroute.nim   | 186 ++++++++++++++++++++++++++++++++++++++++++
 tests/testsegpath.nim |  11 +++
 5 files changed, 348 insertions(+), 11 deletions(-)

diff --git a/presto/btrees.nim b/presto/btrees.nim
index e8eed4c..cc4c93e 100644
--- a/presto/btrees.nim
+++ b/presto/btrees.nim
@@ -45,6 +45,18 @@ proc getOrDefault*[Key, Val](b: BTree[Key, Val], key: Key): Val =
   for j in 0..<x.entries:
     if eq(key, x.keys[j]): return x.vals[j]
 
+proc getOrDefault*[Key, Val](b: BTree[Key, Val], key: Key, default: Val): Val =
+  var x = b.root
+  while x.isInternal:
+    for j in 0..<x.entries:
+      if j+1 == x.entries or less(key, x.keys[j+1]):
+        x = x.links[j]
+        break
+  assert(not x.isInternal)
+  for j in 0..<x.entries:
+    if eq(key, x.keys[j]): return x.vals[j]
+  return default
+
 proc contains*[Key, Val](b: BTree[Key, Val], key: Key): bool =
   var x = b.root
   while x.isInternal:
diff --git a/presto/route.nim b/presto/route.nim
index 4fedd02..0ae7e0b 100644
--- a/presto/route.nim
+++ b/presto/route.nim
@@ -18,15 +18,25 @@ type
   RestApiCallback* = proc(request: HttpRequestRef, pathParams: HttpTable,
                           queryParams: HttpTable,
                           body: Option[ContentBody]): Future[RestApiResponse] {.
-                     raises: [Defect], gcsafe.}
+                       raises: [Defect], gcsafe.}
+
+  RestRouteKind* {.pure.} = enum
+    None, Handler, Redirect
+
   RestRoute* = object
     requestPath*: SegmentedPath
     routePath*: SegmentedPath
     callback*: RestApiCallback
 
   RestRouteItem* = object
+    case kind*: RestRouteKind
+    of RestRouteKind.None:
+      discard
+    of RestRouteKind.Handler:
+      callback: RestApiCallback
+    of RestRouteKind.Redirect:
+      redirectPath*: SegmentedPath
     path: SegmentedPath
-    callback: RestApiCallback
 
   RestRouter* = object
     patternCallback*: PatternCallback
@@ -42,18 +52,47 @@ proc init*(t: typedesc[RestRouter],
 proc addRoute*(rr: var RestRouter, request: HttpMethod, path: string,
                handler: RestApiCallback) {.raises: [Defect].} =
   let spath = SegmentedPath.init(request, path, rr.patternCallback)
-  let route = rr.routes.getOrDefault(spath)
-  doAssert(isNil(route.callback), "The route is already in the routing table")
-  rr.routes.add(spath, RestRouteItem(path: spath, callback: handler))
+  let route = rr.routes.getOrDefault(spath,
+                                     RestRouteItem(kind: RestRouteKind.None))
+  case route.kind
+  of RestRouteKind.None:
+    let item = RestRouteItem(kind: RestRouteKind.Handler,
+                           path: spath, callback: handler)
+    rr.routes.add(spath, item)
+  else:
+    raiseAssert("The route is already in the routing table")
+
+proc addRedirect*(rr: var RestRouter, request: HttpMethod, srcPath: string,
+                  dstPath: string) {.raises: [Defect].} =
+  let spath = SegmentedPath.init(request, srcPath, rr.patternCallback)
+  let dpath = SegmentedPath.init(request, dstPath, rr.patternCallback)
+  let route = rr.routes.getOrDefault(spath,
+                                     RestRouteItem(kind: RestRouteKind.None))
+  case route.kind
+  of RestRouteKind.None:
+    let item = RestRouteItem(kind: RestRouteKind.Redirect,
+                             path: spath, redirectPath: dpath)
+    rr.routes.add(spath, item)
+  else:
+    raiseAssert("The route is already in the routing table")
 
 proc getRoute*(rr: RestRouter,
                spath: SegmentedPath): Option[RestRoute] {.raises: [Defect].} =
-  let route = rr.routes.getOrDefault(spath)
-  if isNil(route.callback):
-    none[RestRoute]()
-  else:
-    some[RestRoute](RestRoute(requestPath: spath, routePath: route.path,
-                              callback: route.callback))
+  var path = spath
+  while true:
+    let route = rr.routes.getOrDefault(path,
+                                       RestRouteItem(kind: RestRouteKind.None))
+    case route.kind
+    of RestRouteKind.None:
+      return none[RestRoute]()
+    of RestRouteKind.Handler:
+      # Route handler was found
+      let item = RestRoute(requestPath: path, routePath: route.path,
+                           callback: route.callback)
+      return some(item)
+    of RestRouteKind.Redirect:
+      # Route redirection was found, so we perform path transformation
+      path = rewritePath(route.path, route.redirectPath, path)
 
 iterator params*(route: RestRoute): string {.raises: [Defect].} =
   var pats = route.routePath.patterns
@@ -168,6 +207,29 @@ proc isPathArg(typeNode: NimNode): bool =
   isBytesArg(typeNode) or (not(isOptionalArg(typeNode)) and
                            not(isSequenceArg(typeNode)))
 
+macro redirect*(router: RestRouter, meth: static[HttpMethod],
+                fromPath: static[string], toPath: static[string]): untyped =
+  ## Define REST API endpoint which redirects request to different compatible
+  ## endpoint ("/somecall" will be redirected to "/api/somecall").
+  let
+    srcPathStr = $fromPath
+    dstPathStr = $toPath
+    srcSegPath = SegmentedPath.init(meth, srcPathStr, nil)
+    dstSegPath = SegmentedPath.init(meth, dstPathStr, nil)
+    # Not sure about this, it creates HttpMethod(int).
+    methIdent = newLit(meth)
+
+  if not(isEqual(srcSegPath, dstSegPath)):
+    error("Source and destination path patterns should be equal", router)
+
+  var res = newStmtList()
+  res.add quote do:
+    `router`.addRedirect(`methIdent`, `fromPath`, `toPath`)
+
+  when defined(nimDumpRestAPI):
+    echo "\n", path, ": ", repr(res)
+  return res
+
 macro api*(router: RestRouter, meth: static[HttpMethod],
            path: static[string], body: untyped): untyped =
   ## Define REST API endpoint and implementation.
diff --git a/presto/segpath.nim b/presto/segpath.nim
index 72172db..806af39 100644
--- a/presto/segpath.nim
+++ b/presto/segpath.nim
@@ -23,6 +23,10 @@ type
     patternCb: PatternCallback
     patterns*: uint64
 
+  KeyValueTuple* = tuple
+    key: string
+    value: string
+
 template isPattern*(spath: SegmentedPath, pos: int): bool =
   (spath.patterns and (1'u64 shl pos)) != 0'u64
 
@@ -43,6 +47,16 @@ iterator keys*(spath: SegmentedPath): string =
     yield spath.data[index][1 .. ^2]
     pats = pats and not(1'u64 shl index)
 
+iterator indexes*(spath: SegmentedPath): int =
+  ## Iterate over all patterns indexes in path ``spath``.
+  var pats = spath.patterns
+  while pats != 0'u64:
+    let index = firstOne(pats) - 1
+    doAssert(index < len(spath.data))
+    if index < len(spath.data):
+      yield index
+    pats = pats and not(1'u64 shl index)
+
 proc getPatterns*(spath: SegmentedPath): seq[string] =
   ## Returns all the patterns in path ``spath``.
   var res = newSeq[string]()
@@ -50,6 +64,53 @@ proc getPatterns*(spath: SegmentedPath): seq[string] =
     res.add(item)
   res
 
+proc isEqual*(spath1, spath2: SegmentedPath): bool =
+  ## Returns ``true`` if both path has equal patterns (number of patterns are
+  ## equal, and pattern names are equal).
+  let pats1 = spath1.getPatterns()
+  let pats2 = spath2.getPatterns()
+  if len(pats1) != len(pats2):
+    false
+  else:
+    for item in pats1:
+      if item notin pats2:
+        return false
+    true
+
+iterator pairs*(spath: SegmentedPath, vpath: SegmentedPath): KeyValueTuple =
+  doAssert(len(spath.data) == len(vpath.data))
+  for index in spath.indexes():
+    yield (spath.data[index], vpath.data[index])
+
+proc getPairs*(spath: SegmentedPath, vpath: SegmentedPath): seq[KeyValueTuple] =
+  var res: seq[tuple[key: string, value: string]]
+  for item in pairs(spath, vpath):
+    res.add(item)
+  res
+
+proc getValue(data: seq[KeyValueTuple], key: string): Option[string] =
+  for item in data:
+    if item.key == key:
+      return some(item.value)
+  return none[string]()
+
+proc rewritePath*(spath: SegmentedPath, dpath: SegmentedPath,
+                  vpath: SegmentedPath): SegmentedPath =
+  doAssert(spath.patternsCount() == dpath.patternsCount())
+  let values = getPairs(spath, vpath)
+  if len(values) == 0:
+    SegmentedPath(data: dpath.data)
+  else:
+    var res = SegmentedPath(data: dpath.data)
+    var k = 0
+    for i in 0 ..< len(res.data):
+      if dpath.isPattern(i):
+        let vres = values.getValue(dpath.data[i])
+        doAssert(vres.isSome())
+        res.data[i] = vres.get()
+        inc(k)
+    res
+
 proc `==`*(s1, s2: SegmentedPath): bool =
   if len(s1.data) == len(s2.data):
     if not(s1.hasPatterns()) and not(s2.hasPatterns()):
@@ -134,12 +195,17 @@ proc init*(st: typedesc[SegmentedPath],
            upath: string, patternCb: PatternCallback): SegmentedPath =
   var res = SegmentedPath(patternCb: patternCb, patterns: 0'u64)
   var counter = 0
+  var patterns: seq[string]
   for item in upath.split("/"):
     doAssert(counter < 64, "Path has too many segments (more then 64)")
     if len(item) >= 2:
       if item[0] == '{' and item[^1] == '}':
         doAssert(len(item) > 2, "Patterns with empty names are not allowed")
         res.patterns = res.patterns or (1'u64 shl counter)
+        if item in patterns:
+          raiseAssert "Only unique patterns allowed in path"
+        else:
+          patterns.add(item)
         res.data.add(item)
         inc(counter)
         continue
diff --git a/tests/testroute.nim b/tests/testroute.nim
index a2d4d74..fe995f1 100644
--- a/tests/testroute.nim
+++ b/tests/testroute.nim
@@ -345,6 +345,22 @@ suite "REST API router & macro tests":
     router.addRoute(MethodPost, "/unique/path/{pattern1}", apiCallback)
     router.addRoute(MethodPost, "/unique/path/{pattern2}", apiCallback)
 
+    # Use HTTP method GET and redirect
+    router.addRedirect(MethodGet, "/redirect/path/1", "/unique/path/1")
+    router.addRedirect(MethodGet, "/redirect/path/2", "/unique/path/2")
+    router.addRedirect(MethodGet, "/redirect/path/{pattern1}",
+                                  "/unique/path/{pattern1}")
+    router.addRedirect(MethodGet, "/redirect/path/{pattern2}",
+                                  "/unique/path/{pattern2")
+
+    # Use HTTP method POST and redirect
+    router.addRedirect(MethodPost, "/redirect/path/1", "/unique/path/1")
+    router.addRedirect(MethodPost, "/redirect/path/2", "/unique/path/2")
+    router.addRedirect(MethodPost, "/redirect/path/{pattern1}",
+                                   "/unique/path/{pattern1}")
+    router.addRedirect(MethodPost, "/redirect/path/{pattern2}",
+                                   "/unique/path/{pattern2")
+
     expect AssertionError:
       router.addRoute(MethodGet, "/unique/path/1", apiCallback)
     expect AssertionError:
@@ -354,6 +370,15 @@ suite "REST API router & macro tests":
     expect AssertionError:
       router.addRoute(MethodGet, "/unique/path/{pattern2}", apiCallback)
 
+    expect AssertionError:
+      router.addRoute(MethodGet, "/redirect/path/1", apiCallback)
+    expect AssertionError:
+      router.addRoute(MethodGet, "/redirect/path/2", apiCallback)
+    expect AssertionError:
+      router.addRoute(MethodGet, "/redirect/path/{pattern1}", apiCallback)
+    expect AssertionError:
+      router.addRoute(MethodGet, "/redirect/path/{pattern2}", apiCallback)
+
     expect AssertionError:
       router.addRoute(MethodPost, "/unique/path/1", apiCallback)
     expect AssertionError:
@@ -362,3 +387,164 @@ suite "REST API router & macro tests":
       router.addRoute(MethodPost, "/unique/path/{pattern1}", apiCallback)
     expect AssertionError:
       router.addRoute(MethodPost, "/unique/path/{pattern2}", apiCallback)
+
+    expect AssertionError:
+      router.addRoute(MethodPost, "/redirect/path/1", apiCallback)
+    expect AssertionError:
+      router.addRoute(MethodPost, "/redirect/path/2", apiCallback)
+    expect AssertionError:
+      router.addRoute(MethodPost, "/redirect/path/{pattern1}", apiCallback)
+    expect AssertionError:
+      router.addRoute(MethodPost, "/redirect/path/{pattern2}", apiCallback)
+
+    expect AssertionError:
+      router.addRedirect(MethodGet, "/unique/path/1", "/unique/1")
+    expect AssertionError:
+      router.addRedirect(MethodGet, "/unique/path/2", "/unique/2")
+    expect AssertionError:
+      router.addRedirect(MethodGet, "/unique/path/{pattern1}",
+                                    "/unique/{pattern1}")
+    expect AssertionError:
+      router.addRedirect(MethodGet, "/unique/path/{pattern2}",
+                                    "/unique/{pattern2}")
+
+    expect AssertionError:
+      router.addRedirect(MethodGet, "/redirect/path/1", "/another/1")
+    expect AssertionError:
+      router.addRedirect(MethodGet, "/redirect/path/2", "/another/2")
+    expect AssertionError:
+      router.addRedirect(MethodGet, "/redirect/path/{pattern1}",
+                                    "/another/{pattern1}")
+    expect AssertionError:
+      router.addRedirect(MethodGet, "/redirect/path/{pattern2}",
+                                    "/another/{pattern2}")
+
+    expect AssertionError:
+      router.addRedirect(MethodPost, "/unique/path/1", "/unique/1")
+    expect AssertionError:
+      router.addRedirect(MethodPost, "/unique/path/2", "/unique/2")
+    expect AssertionError:
+      router.addRedirect(MethodPost, "/unique/path/{pattern1}",
+                                     "/unique/{pattern1}")
+    expect AssertionError:
+      router.addRedirect(MethodPost, "/unique/path/{pattern2}",
+                                     "/unique/{pattern2}")
+
+    expect AssertionError:
+      router.addRedirect(MethodPost, "/redirect/path/1", "/another/1")
+    expect AssertionError:
+      router.addRedirect(MethodPost, "/redirect/path/2", "/another/2")
+    expect AssertionError:
+      router.addRedirect(MethodPost, "/redirect/path/{pattern1}",
+                                     "/another/{pattern1}")
+    expect AssertionError:
+      router.addRedirect(MethodPost, "/redirect/path/{pattern2}",
+                                     "/another/{pattern2}")
+
+  test "Redirection test":
+    var router = RestRouter.init(testValidate)
+
+    router.api(MethodGet, "/test/empty/1") do (
+      ) -> RestApiResponse:
+      return ok(ContentBody(contentType: "test/test", data: "ok-1".toBytes()))
+    router.api(MethodGet, "/test/empty/p1/p2/p3/1") do (
+      ) -> RestApiResponse:
+      return ok(ContentBody(contentType: "test/test", data: "ok-2".toBytes()))
+    router.api(MethodGet, "/test/empty/p1/p2/p3/p5/p6/p7/p8/p9/1") do (
+      ) -> RestApiResponse:
+      return ok(ContentBody(contentType: "test/test", data: "ok-3".toBytes()))
+    router.api(MethodGet, "/test/basic/path1/{smp1}") do (
+      smp1: int) -> RestApiResponse:
+      let s1 = smp1.get()
+      if s1 == 999999:
+        return ok(ContentBody(contentType: "test/test",
+                              data: "ok-1".toBytes()))
+    router.api(MethodGet, "/test/basic/path1/path2/{smp1}/{smp2}") do (
+      smp1: int, smp2: string) -> RestApiResponse:
+      let s1 = smp1.get()
+      let s2 = smp2.get()
+      if (s1 == 999999) and (s2 == "string1"):
+        return ok(ContentBody(contentType: "test/test",
+                              data: "ok-2".toBytes()))
+    router.api(MethodGet,
+               "/test/basic/path1/path2/path3/{smp1}/{smp2}/{smp3}") do (
+      smp1: int, smp2: string, smp3: seq[byte]) -> RestApiResponse:
+      let s1 = smp1.get()
+      let s2 = smp2.get()
+      let s3 = smp3.get()
+      if (s1 == 999999) and (s2 == "string1") and
+         (bytesToString(s3) == "bytes1"):
+        return ok(ContentBody(contentType: "test/test",
+                              data: "ok-3".toBytes()))
+
+    router.redirect(MethodGet, "/api/1", "/test/empty/1")
+    router.redirect(MethodGet, "/api/2", "/test/empty/p1/p2/p3/1")
+    router.redirect(MethodGet, "/api/3",
+                    "/test/empty/p1/p2/p3/p5/p6/p7/p8/p9/1")
+    router.redirect(MethodGet, "/api/basic/{smp1}",
+                    "/test/basic/path1/{smp1}")
+    router.redirect(MethodGet, "/api/basic/{smp1}/{smp2}",
+                    "/test/basic/path1/path2/{smp1}/{smp2}")
+    router.redirect(MethodGet, "/api/basic/{smp1}/{smp2}/{smp3}",
+                    "/test/basic/path1/path2/path3/{smp1}/{smp2}/{smp3}")
+    # Patterns with mixed order
+    router.redirect(MethodGet, "/api/basic/{smp3}/{smp1}/{smp2}",
+                    "/test/basic/path1/path2/path3/{smp1}/{smp2}/{smp3}")
+    router.redirect(MethodGet, "/api/basic/{smp2}/{smp3}/{smp1}",
+                    "/test/basic/path1/path2/path3/{smp1}/{smp2}/{smp3}")
+    router.redirect(MethodGet, "/api/basic/{smp2}/{smp1}/{smp3}",
+                    "/test/basic/path1/path2/path3/{smp1}/{smp2}/{smp3}")
+    router.redirect(MethodGet, "/api/basic/{smp2}/p1/p2/p3/{smp1}",
+                    "/test/basic/path1/path2/{smp1}/{smp2}")
+    router.redirect(MethodGet, "/api/basic/{smp2}/p1/p2/p3/p4/p5/p6/{smp1}",
+                    "/test/basic/path1/path2/{smp1}/{smp2}")
+    router.redirect(MethodGet,
+                    "/api/basic/{smp2}/p1/p2/p3/p4/p5/p6/p7/{smp1}/p8",
+                    "/test/basic/path1/path2/{smp1}/{smp2}")
+
+    let r1 = router.sendMockRequest(MethodGet, "http://l.to/api/1")
+    let r2 = router.sendMockRequest(MethodGet, "http://l.to/api/2")
+    let r3 = router.sendMockRequest(MethodGet, "http://l.to/api/3")
+    let r4 = router.sendMockRequest(MethodGet,
+                                    "http://l.to/api/basic/999999")
+    let r5 = router.sendMockRequest(MethodGet,
+                                    "http://l.to/api/basic/999999/string1")
+    let r6 = router.sendMockRequest(MethodGet,
+                          "http://l.to/api/basic/999999/string1/0x627974657331")
+    let r7 = router.sendMockRequest(MethodGet,
+                          "http://l.to/api/basic/0x627974657331/999999/string1")
+    let r8 = router.sendMockRequest(MethodGet,
+                          "http://l.to/api/basic/string1/0x627974657331/999999")
+    let r9 = router.sendMockRequest(MethodGet,
+                          "http://l.to/api/basic/string1/999999/0x627974657331")
+    let r10 = router.sendMockRequest(MethodGet,
+                                "http://l.to/api/basic/string1/p1/p2/p3/999999")
+    let r11 = router.sendMockRequest(MethodGet,
+                       "http://l.to/api/basic/string1/p1/p2/p3/p4/p5/p6/999999")
+    let r12 = router.sendMockRequest(MethodGet,
+                 "http://l.to/api/basic/string1/p1/p2/p3/p4/p5/p6/p7/999999/p8")
+    check:
+      r1.isOk()
+      r2.isOk()
+      r3.isOk()
+      r4.isOk()
+      r5.isOk()
+      r6.isOk()
+      r7.isOk()
+      r8.isOk()
+      r9.isOk()
+      r10.isOk()
+      r11.isOk()
+      r12.isOk()
+      bytesToString(r1.get().data) == "ok-1"
+      bytesToString(r2.get().data) == "ok-2"
+      bytesToString(r3.get().data) == "ok-3"
+      bytesToString(r4.get().data) == "ok-1"
+      bytesToString(r5.get().data) == "ok-2"
+      bytesToString(r6.get().data) == "ok-3"
+      bytesToString(r7.get().data) == "ok-3"
+      bytesToString(r8.get().data) == "ok-3"
+      bytesToString(r9.get().data) == "ok-3"
+      bytesToString(r10.get().data) == "ok-2"
+      bytesToString(r11.get().data) == "ok-2"
+      bytesToString(r12.get().data) == "ok-2"
diff --git a/tests/testsegpath.nim b/tests/testsegpath.nim
index 3f05496..fdb8a7f 100644
--- a/tests/testsegpath.nim
+++ b/tests/testsegpath.nim
@@ -53,3 +53,14 @@ suite "SegmentedPath test suite":
       check:
         len(path.data) == (i + 2)
         path.getPatterns() == names
+
+  test "Non-unique patterns test":
+    const NonUniqueVectors = [
+      "/{item1}/{item2}/{item1}",
+      "/{i1}/{i1}",
+      "/{a}/{b}/{a}"
+    ]
+    for item in NonUniqueVectors:
+      expect AssertionError:
+        let path {.used.} = SegmentedPath.init(HttpMethod.MethodGet, item,
+                                               validate)