Widen match result type and emit safe str copies for mixed-length arms

The type checker computed unify_types between match arm types but discarded
the result and used the first arm's type. With arms returning strings of
different lengths (e.g. "high_priority": str(13) and "default_priority":
str(16)), the match result was sized to the first arm, then codegen did
__builtin_memcpy(&dst, &src, sizeof(src)) across the differently-laid-out
str_N_t types - corrupting the destination because src's .len field ended
up overwriting dst's .data[] bytes.

Fold unify_types over all arms to compute the actual join, consistent with
the str(N) <: str(M) subtyping the language already defines. Replace four
unsafe struct-memcpy sites in the codegen (IRAssign on temps, IRAssign on
named vars, IRMatch result assembly, and IRVariableDecl init - both the
hoisted-redeclaration and fresh-declaration paths) with a single
emit_str_copy helper that does length-respecting field copies, safe
regardless of layout.

Signed-off-by: Cong Wang <cwang@multikernel.io>

Author: congwang-mk

src/ebpf_c_codegen.ml

@@ -260,6 +260,17 @@ let initialize_context_generators () =
   Kernelscript_context.Fprobe_codegen.register ();
   Kernelscript_context.Perf_event_codegen.register ()
 
+(** Emit a safe str_N_t to str_M_t copy. The naive
+    `__builtin_memcpy(&dst, &src, sizeof(src))` is wrong whenever the two
+    str_N_t types have different sizes: layouts diverge, so the source's
+    `.len` field gets memcpy'd into the destination's `.data[]` array and the
+    destination's `.len` is left uninitialized (or filled with garbage past
+    the end of src). Field-level copy respects the length-prefixed semantics
+    regardless of either side's declared capacity. *)
+let emit_str_copy ctx ~dest ~src =
+  emit_line ctx (sprintf "%s.len = %s.len;" dest src);
+  emit_line ctx (sprintf "__builtin_memcpy(%s.data, %s.data, %s.len);" dest src src)
+
 (** Emit all pending string literal declarations *)
 let emit_pending_string_literals ctx =
   List.iter (fun (var_name, content, size) ->
@@ -1720,28 +1731,46 @@ let generate_c_expression ctx ir_expr =
           (* Generate regular if-else chain with temporary variable for complex cases *)
           let temp_var = fresh_var ctx "match_result" in
           let result_type = ebpf_type_from_ir_type ir_expr.expr_type in
-          
+          let result_str_size = match ir_expr.expr_type with IRStr n -> Some n | _ -> None in
+
           (* Generate temporary variable for the result *)
           emit_line ctx (sprintf "%s %s;" result_type temp_var);
-          
-          (* Defer string literals during match arm value generation *)
-          ctx.defer_string_literals <- true;
-          let arm_values = List.map (fun arm ->
-            (arm, generate_c_value ctx arm.ir_arm_value)
-          ) arms in
-          ctx.defer_string_literals <- false;
-          
-          (* Emit collected string literals before if-else chain *)
-          emit_pending_string_literals ctx;
-          
+
+          (* For str-typed results, the type checker has widened the match's type
+             to the LUB of all arm types (str(max N)). Emit each arm value directly
+             into the result-typed slot - literals as compound literals at the result
+             size, runtime values via field-level copy that respects .len. This
+             avoids the unsafe `memcpy(&dst, &src, sizeof(src))` across different
+             str_N_t layouts, which used to write src's .len bytes into dst's .data.
+             For non-str results, keep the existing deferred-literal flow. *)
+          let arm_values = match result_str_size with
+            | Some _ -> List.map (fun arm -> (arm, "")) arms
+            | None ->
+                ctx.defer_string_literals <- true;
+                let vals = List.map (fun arm ->
+                  (arm, generate_c_value ctx arm.ir_arm_value)
+                ) arms in
+                ctx.defer_string_literals <- false;
+                emit_pending_string_literals ctx;
+                vals
+          in
+
           (* Generate if-else chain *)
-          let needs_memcpy = match ir_expr.expr_type with IRStr _ -> true | _ -> false in
           let generate_match_arm is_first (arm, arm_val_str) =
             let emit_assignment () =
-              if needs_memcpy then
-                emit_line ctx (sprintf "__builtin_memcpy(&%s, &%s, sizeof(%s));" temp_var arm_val_str arm_val_str)
-              else
-                emit_line ctx (sprintf "%s = %s;" temp_var arm_val_str)
+              match result_str_size with
+              | Some result_size ->
+                  (match arm.ir_arm_value.value_desc with
+                   | IRLiteral (Ast.StringLit s) ->
+                       let len = min (String.length s) result_size in
+                       let truncated = if len < String.length s then String.sub s 0 len else s in
+                       emit_line ctx (sprintf "%s = (str_%d_t){ .data = \"%s\", .len = %d };"
+                         temp_var result_size (String.escaped truncated) len)
+                   | _ ->
+                       let val_str = generate_c_value ctx arm.ir_arm_value in
+                       emit_str_copy ctx ~dest:temp_var ~src:val_str)
+              | None ->
+                  emit_line ctx (sprintf "%s = %s;" temp_var arm_val_str)
             in
             match arm.ir_arm_pattern with
             | IRConstantPattern const_val ->
@@ -1759,14 +1788,14 @@ let generate_c_expression ctx ir_expr =
                 decrease_indent ctx;
                 emit_line ctx "}"
           in
-          
+
           (* Generate all arms *)
           (match arm_values with
            | [] -> () (* No arms - should not happen *)
            | first_arm :: rest_arms ->
                generate_match_arm true first_arm;
                List.iter (generate_match_arm false) rest_arms);
-          
+
           (* Return the temporary variable *)
           temp_var
 
@@ -1931,11 +1960,21 @@ and generate_c_instruction ctx ir_instr =
       (* Check if variable is already declared (e.g., in callback functions) *)
       let var_hash = Hashtbl.hash var_name in
       if Hashtbl.mem ctx.declared_registers var_hash then
-        (* Variable already declared, just generate assignment if there's an initializer *)
+        (* Variable already declared (typically hoisted to function top), emit
+           the initializer's assignment at the original position. Cross-size
+           str-to-str needs length-respecting field copy; see emit_str_copy. *)
         (match init_expr_opt with
          | Some init_expr ->
+             let cross_size_str = match typ, init_expr.expr_desc with
+               | IRStr d, IRValue src_val ->
+                   (match src_val.val_type with IRStr s -> s <> d | _ -> false)
+               | _ -> false
+             in
              let init_str = generate_c_expression ctx init_expr in
-             emit_line ctx (sprintf "%s = %s;" var_name init_str)
+             if cross_size_str then
+               emit_str_copy ctx ~dest:var_name ~src:init_str
+             else
+               emit_line ctx (sprintf "%s = %s;" var_name init_str)
          | None -> (* No initializer, no need to emit anything *) ())
       else
         (* Variable not declared yet, generate full declaration *)
@@ -1945,10 +1984,10 @@ and generate_c_instruction ctx ir_instr =
              (* Check if this is a string assignment that needs special handling *)
              (match typ, init_expr.expr_desc with
               | IRStr dest_size, IRValue src_val when (match src_val.val_type with IRStr src_size -> src_size <= dest_size | _ -> false) ->
-                  (* String to string assignment with compatible sizes - regenerate src with dest size *)
+                  (* String to string with compatible sizes. *)
                   (match src_val.value_desc with
                     | IRLiteral (StringLit s) ->
-                        (* Generate direct struct assignment for string literals *)
+                        (* Literal: emit at the destination type via designated initializer. *)
                         let len = String.length s in
                         let max_content_len = dest_size in
                         let actual_len = min len max_content_len in
@@ -1957,10 +1996,17 @@ and generate_c_instruction ctx ir_instr =
                         emit_line ctx (sprintf "    .data = \"%s\"," (String.escaped truncated_s));
                         emit_line ctx (sprintf "    .len = %d" actual_len);
                         emit_line ctx "};"
-                    | _ -> 
-                        (* For non-literal strings, use regular assignment *)
-                        let init_str = generate_c_expression ctx init_expr in
-                        emit_line ctx (sprintf "%s %s = %s;" type_str var_name init_str))
+                    | _ ->
+                        let src_size = match src_val.val_type with IRStr s -> s | _ -> dest_size in
+                        if src_size = dest_size then
+                          (* Same type: plain struct assignment. *)
+                          (let init_str = generate_c_expression ctx init_expr in
+                           emit_line ctx (sprintf "%s %s = %s;" type_str var_name init_str))
+                        else
+                          (* Cross-size: declare-then-field-copy (see emit_str_copy). *)
+                          (emit_line ctx (sprintf "%s %s;" type_str var_name);
+                           let src_str = generate_c_expression ctx init_expr in
+                           emit_str_copy ctx ~dest:var_name ~src:src_str))
               | IRStr _, _ ->
                   (* Other string expressions (concatenation, etc.) *)
                   let init_str = generate_c_expression ctx init_expr in
@@ -2570,13 +2616,14 @@ and generate_assignment ctx dest_val expr is_const =
                | None -> ())
           | _ -> ());
 
-         (* Use memcpy for cross-size string assignments *)
-         let use_memcpy = match dest_val.val_type, expr.expr_desc with
+         (* Cross-size string assignments need length-respecting field copy
+            (see emit_str_copy). Same-size stays as plain struct assignment. *)
+         let cross_size_str = match dest_val.val_type, expr.expr_desc with
            | IRStr d, IRValue src_val -> (match src_val.val_type with IRStr s -> s <> d | _ -> false)
            | _ -> false
          in
-         if use_memcpy then
-           emit_line ctx (sprintf "__builtin_memcpy(&%s, &%s, sizeof(%s));" dest_str expr_str expr_str)
+         if cross_size_str then
+           emit_str_copy ctx ~dest:dest_str ~src:expr_str
          else
            emit_line ctx (sprintf "%s%s = %s;" assignment_prefix dest_str expr_str)
        )
@@ -2596,12 +2643,14 @@ and generate_assignment ctx dest_val expr is_const =
        (* Check if this is a string assignment *)
        (match dest_val.val_type, expr.expr_desc with
         | IRStr dest_size, IRValue src_val when (match src_val.val_type with IRStr src_size -> src_size <= dest_size | _ -> false) ->
-            (* String to string assignment - use memcpy for cross-size compatibility *)
+            (* String to string assignment - length-respecting field copy for
+               cross-size cases (see emit_str_copy); plain struct assignment when
+               sizes match. *)
             let dest_str = generate_c_value ctx dest_val in
             let src_str = generate_c_value ctx src_val in
             (match src_val.val_type with
              | IRStr src_size when src_size <> dest_size ->
-                 emit_line ctx (sprintf "__builtin_memcpy(&%s, &%s, sizeof(%s));" dest_str src_str src_str)
+                 emit_str_copy ctx ~dest:dest_str ~src:src_str
              | _ ->
                  emit_line ctx (sprintf "%s = %s;" dest_str src_str))
         | IRStr _, _ ->
@@ -3119,13 +3168,14 @@ let generate_assignment ctx dest_val expr is_const =
                | None -> ())
           | _ -> ());
 
-         (* Use memcpy for cross-size string assignments *)
-         let use_memcpy = match dest_val.val_type, expr.expr_desc with
+         (* Cross-size string assignments need length-respecting field copy
+            (see emit_str_copy). Same-size stays as plain struct assignment. *)
+         let cross_size_str = match dest_val.val_type, expr.expr_desc with
            | IRStr d, IRValue src_val -> (match src_val.val_type with IRStr s -> s <> d | _ -> false)
            | _ -> false
          in
-         if use_memcpy then
-           emit_line ctx (sprintf "__builtin_memcpy(&%s, &%s, sizeof(%s));" dest_str expr_str expr_str)
+         if cross_size_str then
+           emit_str_copy ctx ~dest:dest_str ~src:expr_str
          else
            emit_line ctx (sprintf "%s%s = %s;" assignment_prefix dest_str expr_str)
        )

src/type_checker.ml

@@ -1484,29 +1484,28 @@ and type_check_expression ctx expr =
         { tarm_pattern = arm.arm_pattern; tarm_body = typed_arm_body; tarm_pos = arm.arm_pos }
       ) arms in
 
-      (* Determine the result type - all arms must have compatible types *)
+      (* Result type is the least upper bound of all arm types under the language's
+         subtyping relation (e.g. str(N) ⊔ str(M) = str(max N M)). unify_types
+         already implements this LUB; the match checker must use the unified type
+         it returns rather than discard it - otherwise a narrower first arm forces
+         the whole expression to be too narrow for later arms. *)
+      let extract_arm_type arm =
+        match arm.tarm_body with
+        | TSingleExpr expr -> expr.texpr_type
+        | TBlock stmts -> extract_block_return_type stmts arm.tarm_pos
+      in
       let result_type = match typed_arms with
         | [] -> type_error "Match expression must have at least one arm" expr.expr_pos
         | first_arm :: rest_arms ->
-            let first_type = match first_arm.tarm_body with
-              | TSingleExpr expr -> expr.texpr_type
-              | TBlock stmts -> 
-                  extract_block_return_type stmts first_arm.tarm_pos
-            in
-            List.iter (fun arm ->
-              let arm_type = match arm.tarm_body with
-                | TSingleExpr expr -> expr.texpr_type
-                | TBlock stmts -> 
-                    extract_block_return_type stmts arm.tarm_pos
-              in
-              match unify_types first_type arm_type with
-              | Some _ -> () (* Compatible *)
-              | None -> 
-                  type_error ("All match arms must return compatible types. Expected " ^ 
-                             string_of_bpf_type first_type ^ " but got " ^ 
+            List.fold_left (fun acc arm ->
+              let arm_type = extract_arm_type arm in
+              match unify_types acc arm_type with
+              | Some unified -> unified
+              | None ->
+                  type_error ("All match arms must return compatible types. Expected " ^
+                             string_of_bpf_type acc ^ " but got " ^
                              string_of_bpf_type arm_type) arm.tarm_pos
-            ) rest_arms;
-            first_type
+            ) (extract_arm_type first_arm) rest_arms
       in
 
       { texpr_desc = TMatch (typed_matched_expr, typed_arms); texpr_type = result_type; texpr_pos = expr.expr_pos }

tests/dune

@@ -284,7 +284,7 @@
 (executable
  (name test_match)
  (modules test_match)
- (libraries kernelscript alcotest test_utils))
+ (libraries kernelscript alcotest test_utils str))
 
 (executable
  (name test_pinned_globals)

tests/test_match.ml

@@ -634,6 +634,47 @@ let test_enum_constant_resolution_in_match () =
      should be resolved to their actual values, not hardcoded to 0 *)
   ()
 
+(** Mixed-length string match arms: the result type is the least upper bound
+    str(max N_i) of all arm types, and codegen does length-respecting field
+    copies. The unsafe `memcpy(&dst, &src, sizeof(src))` pattern (which writes
+    src.len into dst.data on cross-size str_N_t copies) must not appear. *)
+let test_match_mixed_length_string_arms () =
+  let input = {|
+    @helper fn use_label(s: str(20)) -> u32 { return 0 }
+
+    @xdp fn classify(ctx: *xdp_md) -> xdp_action {
+      var n: u32 = 1
+      var label = match (n) {
+        1: "short",
+        2: "medium_str",
+        default: "longer_default_str"
+      }
+      var _ = use_label(label)
+      return 2
+    }
+
+    fn main() -> i32 { return 0 }
+  |} in
+  let ast = Parse.parse_string input in
+  let symbol_table = Symbol_table.build_symbol_table ast in
+  let (typed_ast, _) = Type_checker.type_check_and_annotate_ast ast in
+  let ir = Ir_generator.generate_ir typed_ast symbol_table "test" in
+  let (generated_code, _) = Ebpf_c_codegen.compile_multi_to_c_with_analysis ir in
+  let contains substr =
+    try ignore (Str.search_forward (Str.regexp_string substr) generated_code 0); true
+    with Not_found -> false
+  in
+  (* Widest arm is "longer_default_str" (18). With type-checker LUB widening,
+     the result temp must be at least str_18_t (or wider if propagated from
+     the use_label parameter type). *)
+  check bool "match result widened to a wide-enough str type" true
+    (contains "str_18_t" || contains "str_20_t");
+  (* Cross-size struct memcpy with sizeof(src) corrupts the destination - the
+     match codegen must not emit it. Field copies (`.len = ...` + memcpy of
+     `.data` with `.len` bytes) are safe. *)
+  check bool "no unsafe struct-memcpy for str match arms" false
+    (contains "memcpy(&__match_result")
+
 let suite = [
   "test_basic_match_parsing", `Quick, test_basic_match_parsing;
   "test_match_with_enums", `Quick, test_match_with_enums;
@@ -646,6 +687,7 @@ let suite = [
   "test_nested_match_structures", `Quick, test_nested_match_structures;
   "test_match_block_implicit_returns", `Quick, test_match_block_implicit_returns;
   "test_enum_constant_resolution_in_match", `Quick, test_enum_constant_resolution_in_match;
+  "test_match_mixed_length_string_arms", `Quick, test_match_mixed_length_string_arms;
 ]
 
 let () = run "Match Construct Tests" [