PoC: Lambda dev server

renatoalencar · renatoalencar · commit 5d089a952dd0 · 2021-11-17T07:55:17.000-03:00
diff --git a/dune-project b/dune-project
@@ -2,6 +2,4 @@
 
 (using fmt 1.1)
 
-(implicit_transitive_deps false)
-
 (name lambda-runtime)
diff --git a/lib/dune b/lib/dune
@@ -4,4 +4,4 @@
  (preprocess
   (pps ppx_deriving_yojson))
  (libraries yojson ppx_deriving_yojson.runtime bigstringaf bigarray-compat
-   uri logs.lwt lwt lwt.unix piaf result))
+   uri logs.lwt lwt lwt.unix piaf result uuidm))
diff --git a/lib/http.ml b/lib/http.ml
@@ -50,7 +50,7 @@ type api_gateway_request_identity =
   ; user_agent : string option [@key "userAgent"]
   ; user : string option
   }
-[@@deriving of_yojson]
+[@@deriving yojson]
 
 (* APIGatewayProxyRequestContext contains the information to identify the AWS
  * account and resources invoking the Lambda function. It also includes Cognito
@@ -69,7 +69,7 @@ type api_gateway_proxy_request_context =
   ; path : string option [@default None]
   ; api_id : string [@key "apiId"] (* The API Gateway REST API ID *)
   }
-[@@deriving of_yojson { strict = false }]
+[@@deriving yojson { strict = false }]
 
 type api_gateway_proxy_request =
   { resource : string
@@ -86,7 +86,7 @@ type api_gateway_proxy_request =
   ; body : string option
   ; is_base64_encoded : bool [@key "isBase64Encoded"]
   }
-[@@deriving of_yojson { strict = false }]
+[@@deriving yojson { strict = false }]
 
 type api_gateway_proxy_response =
   { status_code : int [@key "statusCode"]
@@ -97,13 +97,103 @@ type api_gateway_proxy_response =
 [@@deriving to_yojson]
 
 module API_gateway_request = struct
-  type t = api_gateway_proxy_request [@@deriving of_yojson]
+  type t = api_gateway_proxy_request [@@deriving yojson]
+
+  let mock (ctx: 'a Piaf.Server.ctx) =
+    let piaf_headers_to_string_map headers =
+      headers
+      |> Piaf.Headers.to_list
+      |> List.to_seq
+      |> StringMap.of_seq
+    in
+    let identity = { cognito_identity_pool_id = None
+                   ; account_id = None
+                   ; cognito_identity_id = None
+                   ; caller = None
+                   ; access_key = None
+                   ; api_key = None
+                     (* TODO: Use the actual IP address from Unix.sockaddr *)
+                   ; source_ip = "127.0.0.1"
+                   ; cognito_authentication_type = None
+                   ; cognito_authentication_provider = None
+                   ; user_arn = None
+                   ; user_agent = Piaf.Headers.get ctx.request.headers "User-Agent"
+                   ; user = None }
+    in
+    let uri = Piaf.Request.uri ctx.request in
+    let request_id = 
+      (* TODO: Not sure if this is equivalent to aws_request_id, but probably. Need to confirm. *)
+      Random.self_init ();
+      Uuidm.to_string @@ Uuidm.v4_gen (Random.get_state ()) ()
+    in
+    let request_context = { account_id = "123456789012"
+                          ; resource_id = "123456"
+                          ; stage = "dev"
+                          ; request_id
+                          ; identity
+                          ; resource_path = "/{proxy+}"
+                          ; authorizer = None
+                          ; http_method = Piaf.Method.to_string ctx.request.meth
+                          ; protocol = Some (Piaf.Versions.HTTP.to_string ctx.request.version)
+                          ; path = Some (Uri.path uri)
+                          ; api_id = "1234567890" }
+    in
+    let body =
+      Lwt_result.map_err Piaf.Error.to_string (Piaf.Body.to_string ctx.request.body)
+    in
+    let query_string_parameters =
+      Uri.query uri
+      |> List.map (fun (key, values) ->
+             match List.length values with
+               | 0 -> (key, "")
+               | 1 -> (key, List.hd values)
+                   (* TODO: Property handle this one *)
+               | _ -> failwith "Multiple values not supported for query strings for now")
+      |> List.to_seq
+      |> StringMap.of_seq
+    in
+    Lwt_result.map (fun body ->
+        (to_yojson { resource = (Uri.path uri)
+                   ; path = "/{proxy+}"
+                   ; http_method = request_context.http_method
+                   ; headers = piaf_headers_to_string_map ctx.request.headers
+                   ; query_string_parameters
+                   ; path_parameters = StringMap.empty
+                   ; stage_variables = StringMap.empty
+                   ; request_context
+                   ; body = if body = String.empty then None else Some body
+                   ; is_base64_encoded = false }))
+      body
+
+  let response _ =
+    Error "Not implemented"
 end
 
 module API_gateway_response = struct
   type t = api_gateway_proxy_response [@@deriving to_yojson]
 
   let to_yojson t = Lwt.return (to_yojson t)
+
+  let mock _ =
+    Lwt_result.fail "Not implemented"
+
+  let response response =
+    let string_map_to_piaf_headers headers =
+      headers
+      |> StringMap.to_seq
+      |> List.of_seq
+      |> Piaf.Headers.of_list
+    in
+    (* TODO: Support base64? *)
+    assert (not response.is_base64_encoded);
+
+    let response = 
+      Piaf.Response.of_string
+        ~headers:(string_map_to_piaf_headers response.headers)
+        ~body:(response.body)
+        (Piaf.Status.of_code response.status_code)
+    in
+    Ok response
 end
 
 include Runtime.Make (API_gateway_request) (API_gateway_response)
diff --git a/lib/json.ml b/lib/json.ml
@@ -34,6 +34,26 @@ module Id = struct
   type t = Yojson.Safe.t [@@deriving yojson]
 
   let to_yojson t = Lwt.return (to_yojson t)
+
+  let mock (ctx: 'a Piaf.Server.ctx) =
+    let event =
+      let open Lwt_result.Infix in
+      Piaf.Body.to_string ctx.request.body
+      >|= (fun body -> if body = String.empty then None else Some body)
+      >|= Option.map Yojson.Safe.from_string
+      >|= Option.value ~default:`Null
+    in
+    Lwt_result.map_err Piaf.Error.to_string event
+
+  let response response =
+    let body = Yojson.Safe.to_string response in
+    let response = Piaf.Response.of_string
+                     ~headers:(Piaf.Headers.of_list [ ("Content-Type", "application/json") ])
+                     ~body:body
+                     `OK
+    in
+    Ok response
+
 end
 
 include Runtime.Make (Id) (Id)
diff --git a/lib/runtime.ml b/lib/runtime.ml
@@ -32,6 +32,8 @@
 
 open Lwt.Infix
 
+module StringMap = Map.Make(String)
+
 module Make
     (Event : Runtime_intf.LambdaEvent)
     (Response : Runtime_intf.LambdaResponse) =
@@ -151,19 +153,119 @@ struct
       failwith
         (Format.asprintf "Could not start HTTP client: %a" Piaf.Error.pp_hum e)
 
+  let make_mocked_context () =
+    let function_settings = Config.{ function_name = "unknown"
+                                   ; memory_size = 128 (* typical value for a lambda *)
+                                   ; version = "$LATEST"
+                                     (* TODO: Find values for those *)
+                                   ; log_stream = ""
+                                   ; log_group = "" }
+    in
+    let request_id =
+      Random.self_init ();
+      Uuidm.to_string @@ Uuidm.v4_gen (Random.get_state ()) ()
+    in
+    Context.make
+      ~invoked_function_arn:("arn:aws:lambda:local:123456789012:function:" ^ function_settings.function_name)
+      ~aws_request_id:request_id
+      ~xray_trace_id:None
+      ~client_context:None
+      ~identity:None
+      ~deadline:0L
+      function_settings
+
+  (* TODO: Maybe local invocation functions can be in another module. *)
+  let invoke_locally ~lift handler event =
+    let context = make_mocked_context () in
+    let open Lwt_result.Infix in
+    Lwt_result.lift @@ Event.of_yojson event
+    >>= fun event -> lift (handler event context)
+
+  let start_locally ~lift handler =
+    let make_response res =
+      match Result.bind res Response.response with
+      | Ok response ->
+         response
+      | Error msg ->
+         let body = Printf.sprintf "Error: %s" msg in
+         Piaf.Response.of_string ~body `Internal_server_error
+    in
+    let server_handler (_ctx: Unix.sockaddr Piaf.Server.ctx) =
+      let event =
+        Event.mock _ctx
+        >|= (function
+             | Ok event -> event
+             | Error msg -> failwith msg)
+      in
+      let response =
+        event
+        >>= invoke_locally ~lift handler
+      in
+        response >|= make_response
+    in
+    Lwt.async (fun () -> 
+        let address = Unix.(ADDR_INET (inet_addr_loopback, 5000)) in
+        Lwt_io.establish_server_with_client_socket
+               address
+               (Piaf.Server.create server_handler)
+        >|= fun _server ->
+        print_endline "Server started at: http://127.0.0.1:5000"
+    );
+    let forever, _ = Lwt.wait () in
+    forever
+
+  let handle_local_invoke ~lift handler event_file =
+    let event =
+      match event_file with
+      | Some filename ->
+         Lwt_io.(open_file ~mode:Input filename)
+         >>= Lwt_io.read
+         >|= Yojson.Safe.from_string
+      | None ->
+         Lwt.return `Null
+    in
+    event 
+    >>= invoke_locally ~lift handler
+    >>= (function
+         | Ok response ->
+            Response.to_yojson response
+            >|= Yojson.Safe.to_string
+         | Error msg ->
+            Lwt.return ("Error: " ^ msg))
+    >|= print_endline
+
   let start_lambda ~lift handler =
-    match Config.get_runtime_api_endpoint () with
-    | Ok endpoint ->
-      (match Config.get_function_settings () with
-      | Ok function_config ->
-        let p =
-          start_with_runtime_endpoint ~lift handler function_config endpoint
+    let p = 
+      if Array.length Sys.argv = 1 then
+        let endpoint =
+          match Config.get_runtime_api_endpoint () with
+          | Ok endpoint -> endpoint
+          | Error msg -> failwith msg
         in
-        Lwt_main.run p
-      | Error msg ->
-        failwith msg)
-    | Error msg ->
-      failwith msg
+        let function_config = 
+          match Config.get_function_settings () with
+          | Ok function_config -> function_config
+          | Error msg -> failwith msg
+        in
+        start_with_runtime_endpoint ~lift handler function_config endpoint
+      else
+        (* TODO: Handle commands and options with a more robust solution
+           like Cmdliner or Arg. *)
+        match Sys.argv.(1) with
+        | "invoke" ->
+           let event_file =
+             if Array.length Sys.argv > 2 then
+               Some Sys.argv.(2)
+             else
+               None
+           in
+           handle_local_invoke ~lift handler event_file
+        | "start-api" ->
+           start_locally ~lift handler
+        | command ->
+           failwith ("Invalid command: " ^ command)
+    in
+    Lwt_main.run p
 
   let lambda handler = start_lambda ~lift:Lwt.return handler
 
diff --git a/lib/runtime_intf.ml b/lib/runtime_intf.ml
@@ -34,12 +34,20 @@ module type LambdaEvent = sig
   type t
 
   val of_yojson : Yojson.Safe.t -> (t, string) result
+
+  val mock : 'a Piaf.Server.ctx -> (Yojson.Safe.t, string) Lwt_result.t
+
+  val response : t -> (Piaf.Response.t, string) result
 end
 
 module type LambdaResponse = sig
   type t
 
   val to_yojson : t -> Yojson.Safe.t Lwt.t
+
+  val mock : 'a Piaf.Server.ctx -> (Yojson.Safe.t, string) Lwt_result.t
+
+  val response : t -> (Piaf.Response.t, string) result
 end
 
 module type LambdaRuntime = sig
