1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
use crate::request_parts::RequestParts;
use crate::response_options::ResponseOptions;
use crate::{
    request::SpinRequest,
    response::{SpinBody, SpinResponse},
};
use dashmap::DashMap;
use futures::{SinkExt, StreamExt};
use http::Method as HttpMethod;
use leptos::server_fn::middleware::Service;
/// Leptos Spin Integration for server functions
use leptos::server_fn::{codec::Encoding, initialize_server_fn_map, ServerFn, ServerFnTraitObj};
use leptos::{create_runtime, provide_context};
use multimap::MultiMap;
use once_cell::sync::Lazy;
use spin_sdk::http::{Headers, IncomingRequest, OutgoingResponse, ResponseOutparam};
use url::Url;

#[allow(unused)] // used by server integrations
type LazyServerFnMap<Req, Res> = Lazy<DashMap<&'static str, ServerFnTraitObj<Req, Res>>>;

static REGISTERED_SERVER_FUNCTIONS: LazyServerFnMap<SpinRequest, SpinResponse> =
    initialize_server_fn_map!(SpinRequest, SpinResponse);

/// Explicitly register a server function. This is only necessary if you are
/// running the server in a WASM environment (or a rare environment that the
/// `inventory`) crate doesn't support. Spin is one of those environments
pub fn register_explicit<T>()
where
    T: ServerFn<ServerRequest = SpinRequest, ServerResponse = SpinResponse> + 'static,
{
    REGISTERED_SERVER_FUNCTIONS.insert(
        T::PATH,
        ServerFnTraitObj::new(
            T::PATH,
            T::InputEncoding::METHOD,
            |req| Box::pin(T::run_on_server(req)),
            T::middlewares,
        ),
    );
}

/// The set of all registered server function paths.
pub fn server_fn_paths() -> impl Iterator<Item = (&'static str, HttpMethod)> {
    REGISTERED_SERVER_FUNCTIONS
        .iter()
        .map(|item| (item.path(), item.method()))
}

pub async fn handle_server_fns(req: IncomingRequest, resp_out: ResponseOutparam) {
handle_server_fns_with_context(req, resp_out, ||{}).await;
}
pub async fn handle_server_fns_with_context(req: IncomingRequest, resp_out: ResponseOutparam, additional_context: impl Fn() + 'static + Clone + Send) {
    let pq = req.path_with_query().unwrap_or_default();

    let (spin_res, req_parts, res_options, runtime) = 
        match crate::server_fn::get_server_fn_by_path(&pq) {
            Some(lepfn) => {
            let runtime = create_runtime();
            let req_parts = RequestParts::new_from_req(&req);
            provide_context(req_parts.clone());
            let res_options = ResponseOptions::default_without_headers();
            provide_context(res_options.clone());
            additional_context();
            let spin_req = SpinRequest::new_from_req(req);
            (lepfn.clone().run(spin_req).await, req_parts, res_options, runtime)
        },
            None => panic!("Server FN path {} not found", &pq)
        
    };
        // If the Accept header contains text/html, than this is a request from 
        // a regular html form, so we should setup a redirect to either the referrer
        // or the user specified location

        let req_headers = Headers::from_list(req_parts.headers()).expect("Failed to construct Headers from Request Input for a Server Fn.");
        let accepts_html = &req_headers.get(&"Accept".to_string())[0];
        let accepts_html_bool = String::from_utf8_lossy(accepts_html).contains("text/html");

        if accepts_html_bool {
            
            let referrer = &req_headers.get(&"Referer".to_string())[0];
            let location = &req_headers.get(&"Location".to_string());
            if location.is_empty(){
                res_options.insert_header("location", referrer.to_owned());
            }
            // Set status and header for redirect
            if !res_options.status_is_set(){
                res_options.set_status(302);
            }

        } 

    let headers = merge_headers(spin_res.0.headers, res_options.headers());
    let status = res_options.status().unwrap_or(spin_res.0.status_code);
    match spin_res.0.body {
        SpinBody::Plain(r) => {
            let og = OutgoingResponse::new(headers);
            og.set_status_code(status).expect("Failed to set Status");
            let mut ogbod = og.take_body();
            resp_out.set(og);
            ogbod.send(r).await.unwrap();
        }
        SpinBody::Streaming(mut s) => {
            let og = OutgoingResponse::new(headers);
            og.set_status_code(status).expect("Failed to set Status");
            let mut res_body = og.take_body();
            resp_out.set(og);

            while let Some(Ok(chunk)) = s.next().await {
                let _ = res_body.send(chunk.to_vec()).await;
            }
        }
    }
    runtime.dispose();
}

/// Returns the server function at the given path
pub fn get_server_fn_by_path(path: &str) -> Option<ServerFnTraitObj<SpinRequest, SpinResponse>> {
    // Sanitize Url to prevent query string or ids causing issues. To do that Url wants a full url,
    // so we give it a fake one. We're only using the path anyway!
    let full_url =format!("http://leptos.dev{}", path);
    let Ok(url) = Url::parse(&full_url) else{
        println!("Failed to parse: {full_url:?}");
    return None;
    };
    REGISTERED_SERVER_FUNCTIONS.get_mut(url.path()).map(|f| f.clone())
}

/// Merge together two sets of headers, deleting any in the first set of Headers that have a key in
/// the second set of headers.
pub fn merge_headers(h1: Headers, h2: Headers) -> Headers {
    //1. Get all keys in H1 and H2
    let entries1 = h1.entries();
    let entries2 = h2.entries();

    let mut mmap1 = MultiMap::new();
    entries1.iter().for_each(|(k, v)| {
        mmap1.insert(k, v);
    });
    let mut mmap2 = MultiMap::new();
    entries2.iter().for_each(|(k, v)| {
        mmap2.insert(k, v);
    });

    //2. Delete any keys in H1 that are present in H2
    mmap1.retain(|&k, &_v| mmap2.get(k).is_none());

    //3. Iterate through H2, adding them to H1
    mmap1.extend(mmap2);

    //4. Profit
    let mut merged_vec: Vec<(String, Vec<u8>)> = vec![];
    mmap1.iter_all().for_each(|(k, v)| {
        for v in v.iter() {
            merged_vec.push((k.to_string(), v.to_vec()))
        }
    });
    Headers::from_list(&merged_vec).expect("Failed to create headers from merged list")
}