feat: support struct in Go as script parameters #705

backend/src/parser_go.rs

@@ -3,8 +3,8 @@
 use itertools::Itertools;
 
 use crate::error::to_anyhow;
-use crate::parser::{Arg, MainArgSignature, Typ};
-use crate::parser_go_ast::{self, Ident};
+use crate::parser::{Arg, MainArgSignature, ObjectProperty, Typ};
+use crate::parser_go_ast::{self, FieldList, Ident, StructType};
 use crate::parser_go_ast::{Decl, Expr};
 use crate::parser_go_scanner;
 use crate::parser_go_token::{Position, Token};
@@ -20,21 +20,8 @@ pub fn parse_go_sig(code: &str) -> crate::error::Result<MainArgSignature> {
             .list
             .iter()
             .map(|param| {
-                let (otyp, typ) = match &param.type_ {
-                    Some(typ) => parse_go_typ(typ),
-                    None => (None, Typ::Unknown),
-                };
-                Arg {
-                    name: param
-                        .names
-                        .as_ref()
-                        .and_then(|x| x.first().map(|y| y.name.to_string()))
-                        .unwrap_or_else(|| "".to_string()),
-                    otyp,
-                    typ,
-                    default: None,
-                    has_default: false,
-                }
+                let (otyp, typ) = get_type(param);
+                Arg { name: get_name(param), otyp, typ, default: None, has_default: false }
             })
             .collect_vec();
         Ok(MainArgSignature { star_args: false, star_kwargs: false, args })
@@ -45,6 +32,23 @@ pub fn parse_go_sig(code: &str) -> crate::error::Result<MainArgSignature> {
     }
 }
 
+fn get_type(param: &parser_go_ast::Field) -> (Option<String>, Typ) {
+    let (otyp, typ) = &param
+        .type_
+        .as_ref()
+        .map(|typ| parse_go_typ(typ))
+        .unwrap_or_else(|| (None, Typ::Unknown));
+    (otyp.clone(), typ.clone())
+}
+
+fn get_name(param: &parser_go_ast::Field) -> String {
+    param
+        .names
+        .as_ref()
+        .and_then(|x| x.first().map(|y| y.name.to_string()))
+        .unwrap_or_else(|| "".to_string())
+}
+
 fn parse_go_typ(typ: &parser_go_ast::Expr) -> (Option<String>, Typ) {
     match typ {
         Expr::Ident(Ident { name, .. }) => (
@@ -63,14 +67,47 @@ fn parse_go_typ(typ: &parser_go_ast::Expr) -> (Option<String>, Typ) {
                 Typ::List(Box::new(inner_typ)),
             )
         }
+        Expr::StructType(StructType { fields: Some(FieldList { list, .. }), .. }) => {
+            let (otyps, typs): (Vec<String>, Vec<ObjectProperty>) = list
+                .iter()
+                .map(|field| {
+                    let json_tag = field
+                        .tag
+                        .as_ref()
+                        .and_then(|x| x.value.strip_prefix("`json:\""))
+                        .and_then(|x| x.strip_suffix("\"`"))
+                        .and_then(|x| x.split(',').last().map(|x| x.to_string()));
+                    let (otyp, typ) = get_type(field);
+                    let name = get_name(field);
+                    let key = json_tag.unwrap_or_else(|| name.to_string());
+                    (
+                        format!("{name} {} `json:\"{key}\"`", otyp_to_string(otyp)),
+                        ObjectProperty { key, typ: Box::new(typ) },
+                    )
+                })
+                .collect::<Vec<_>>()
+                .into_iter()
+                .unzip();
+            (
+                Some(format!(
+                    "struct {{ {} }}",
+                    otyps.iter().join("; ").to_string()
+                )),
+                Typ::Object(typs),
+            )
+        }
         _ => (None, Typ::Unknown),
     }
 }
 
+pub fn otyp_to_string(otyp: Option<String>) -> String {
+    otyp.unwrap_or_else(|| "interface{}".to_string())
+}
+
 #[cfg(test)]
 mod tests {
 
-    use crate::parser::{Arg, MainArgSignature, Typ};
+    use crate::parser::{Arg, MainArgSignature, ObjectProperty, Typ};
 
     use super::*;
 
@@ -82,7 +119,7 @@ package main
 
 import "fmt"
 
-func main(x int, y string, z bool, l []string) {
+func main(x int, y string, z bool, l []string, o struct { Name string `json:"name"` }) {
     fmt.Println("hello world")
 }    
 
@@ -122,6 +159,16 @@ func main(x int, y string, z bool, l []string) {
                         default: None,
                         has_default: false
                     },
+                    Arg {
+                        otyp: Some("struct { Name string `json:\"name\"` }".to_string()),
+                        name: "o".to_string(),
+                        typ: Typ::Object(vec![ObjectProperty {
+                            key: "name".to_string(),
+                            typ: Box::new(Typ::Str(None))
+                        },]),
+                        default: None,
+                        has_default: false
+                    },
                 ]
             }
         );

backend/src/worker.rs

@@ -18,6 +18,7 @@ use crate::{
         JobKind, QueuedJob,
     },
     parser::Typ,
+    parser_go::otyp_to_string,
     parser_py,
     scripts::{ScriptHash, ScriptLang},
     users::{create_token_for_owner, get_email_from_username},
@@ -737,7 +738,6 @@ async fn handle_go_job(
     logs.push_str("\n\n--- GO CODE EXECUTION ---\n");
     set_logs(logs, job.id, db).await;
 
-    let sig = crate::parser_go::parse_go_sig(&inner_content)?;
     let token = create_token_for_owner(
         &db,
         &job.workspace_id,
@@ -748,22 +748,11 @@ async fn handle_go_job(
     )
     .await?;
     create_args_and_out_file(job, &token, base_internal_url, job_dir).await?;
+    {
+        let sig = crate::parser_go::parse_go_sig(&inner_content)?;
+        drop(inner_content);
 
-    let spread = sig
-        .args
-        .into_iter()
-        .map(|x| {
-            format!(
-                "json_arg[\"{}\"].({})",
-                x.name,
-                x.otyp.unwrap_or_else(|| "interface{}".to_string())
-            )
-        })
-        .join(", ");
-
-    let wrapper_content: String = format!(
-        r#"
-package main
+        const WRAPPER_CONTENT: &str = r#"package main
 
 import (
     "encoding/json"
@@ -773,20 +762,21 @@ import (
 )
 
 func main() {{
+
     dat, err := os.ReadFile("args.json")
     if err != nil {{
         fmt.Println(err)
         os.Exit(1)
     }}
 
-    var json_arg map[string]interface{{}}
+    var req inner.Req
 
-    if err := json.Unmarshal(dat, &json_arg); err != nil {{
+    if err := json.Unmarshal(dat, &req); err != nil {{
         fmt.Println(err)
         os.Exit(1)
     }}
 
-    res, err := inner.Inner_main({spread})
+    res, err := inner.Run(req)
     if err != nil {{
         fmt.Println(err)
         os.Exit(1)
@@ -806,11 +796,44 @@ func main() {{
         fmt.Println(err)
         os.Exit(1)
     }}
+}}"#;
+
+        write_file(job_dir, "main.go", WRAPPER_CONTENT).await?;
+
+        {
+            let spread = &sig
+                .args
+                .clone()
+                .into_iter()
+                .map(|x| format!("req.{}", capitalize(&x.name)))
+                .join(", ");
+            let req_body = &sig
+                .args
+                .into_iter()
+                .map(|x| {
+                    format!(
+                        "{} {} `json:\"{}\"`",
+                        capitalize(&x.name),
+                        otyp_to_string(x.otyp),
+                        x.name
+                    )
+                })
+                .join("\n");
+            let runner_content: String = format!(
+                r#"package inner
+type Req struct {{
+    {req_body}
+}}
+
+func Run(req Req) (interface{{}}, error){{
+    return main({spread})
 }}
 
 "#,
-    );
-    write_file(job_dir, "main.go", &wrapper_content).await?;
+            );
+            write_file(&format!("{job_dir}/inner"), "runner.go", &runner_content).await?;
+        }
+    }
     let mut reserved_variables = get_reserved_variables(job, &token, &base_url, db).await?;
     reserved_variables.insert("RUST_LOG".to_string(), "info".to_string());
 
@@ -862,6 +885,14 @@ func main() {{
     read_result(job_dir).await
 }
 
+fn capitalize(s: &str) -> String {
+    let mut c = s.chars();
+    match c.next() {
+        None => String::new(),
+        Some(f) => f.to_uppercase().collect::<String>() + c.as_str(),
+    }
+}
+
 async fn handle_deno_job(
     WorkerConfig { base_internal_url, base_url, disable_nuser, disable_nsjail, .. }: &WorkerConfig,
     Envs { nsjail_path, deno_path, path_env, .. }: &Envs,
@@ -1389,7 +1420,6 @@ async fn gen_go_mymod(code: &str, job_dir: &str) -> error::Result<()> {
     } else {
         format!("package inner; {code}")
     };
-    let code = code.replace("func main(", "func Inner_main(");
 
     let mymod_dir = format!("{job_dir}/inner");
     DirBuilder::new()

frontend/src/lib/components/DisplayResult.svelte

@@ -130,7 +130,7 @@
 				>
 			</div>
 		{:else if !forceJson && resultKind == 'error'}<div
-				><pre class="text-sm text-red-500">{result.error}</pre>
+				><pre class="text-sm text-red-500 whitespace-pre-wrap">{result.error}</pre>
 			</div>
 		{:else}<Highlight
 				language={json}

frontend/src/lib/components/ScriptEditor.svelte

@@ -167,19 +167,6 @@
 											Move the focus outside of the text editor to recompute the input schema from
 											main signature or press Ctrl/Cmd+S
 										</p>
-										<p class="mt-4">
-											{#if isValid}
-												<Icon data={faCheck} class="text-green-600 mr-1" scale={0.6} />
-												The current preview input matches requirements defined in arguments
-											{:else}
-												<Icon
-													data={faExclamationTriangle}
-													class="text-yellow-500 mr-1"
-													scale={0.6}
-												/>
-												The current preview input doesn't match requirements defined in arguments
-											{/if}
-										</p>
 									</div>
 									<SchemaForm {schema} bind:args bind:isValid />
 								</div>

nsjail/run.go.config.proto

@@ -84,6 +84,12 @@ mount {
     is_bind: true
 }
 
+mount {
+    src: "{JOB_DIR}/inner/runner.go"
+    dst: "/tmp/go/inner/runner.go"
+    is_bind: true
+}
+
 mount {
     src: "{JOB_DIR}/args.json"
     dst: "/tmp/go/args.json"