grpc-pb: Check required fields ins submessages

Signed-off-by: Johannes Zottele <[email protected]>

Author: Jozott00

grpc/grpc-core/src/commonMain/kotlin/kotlinx/rpc/grpc/pb/InternalMessage.kt

@@ -17,29 +17,30 @@ public abstract class InternalMessage(fieldsWithPresence: Int) {
     public abstract val _size: Int
 }
 
+@InternalRpcApi
 public class MsgFieldDelegate<T : Any>(
     private val presenceIdx: Int? = null,
     private val defaultProvider: (() -> T)? = null
 ) : ReadWriteProperty<InternalMessage, T> {
 
     private var valueSet = false
-    private var _value: T? = null
+    private var value: T? = null
 
     override operator fun getValue(thisRef: InternalMessage, property: KProperty<*>): T {
         if (!valueSet) {
             if (defaultProvider != null) {
-                _value = defaultProvider.invoke()
+                value = defaultProvider.invoke()
                 valueSet = true
             } else {
                 error("Property ${property.name} not initialized")
             }
         }
-        return _value as T
+        return value as T
     }
 
     override operator fun setValue(thisRef: InternalMessage, property: KProperty<*>, value: T) {
         presenceIdx?.let { thisRef.presenceMask[it] = true }
-        _value = value
+        this@MsgFieldDelegate.value = value
         valueSet = true
     }
 }

grpc/grpc-core/src/commonTest/kotlin/kotlinx/rpc/grpc/pb/ProtosTest.kt

@@ -6,6 +6,7 @@ package kotlinx.rpc.grpc.pb
 
 import OneOfMsg
 import OneOfMsgInternal
+import OneOfWithRequired
 import Outer
 import OuterInternal
 import asInternal
@@ -87,10 +88,20 @@ class ProtosTest {
     }
 
     @Test
-    fun testPresenceCheckProto() {
+    fun testRepeatedWithRequiredSubField() {
+        assertFailsWith<IllegalStateException> {
+            RepeatedWithRequired {
+                // we construct the message using the internal class,
+                // so it is not invoking the checkRequired method on construction
+                msgList = listOf(PresenceCheck { RequiredPresence = 2 }, PresenceCheckInternal())
+            }
+        }
+    }
 
+    @Test
+    fun testPresenceCheckProto() {
         // Check a missing required field in a user-constructed message
-        assertFailsWith<IllegalStateException>("PresenceCheck is missing required field: RequiredPresence") {
+        assertFailsWith<IllegalStateException> {
             PresenceCheck {}
         }
 
@@ -100,7 +111,7 @@ class ProtosTest {
         encoder.writeFloat(2, 1f)
         encoder.flush()
 
-        assertFailsWith<IllegalStateException>("PresenceCheck is missing required field: RequiredPresence") {
+        assertFailsWith<IllegalStateException> {
             PresenceCheckInternal.CODEC.decode(buffer)
         }
     }
@@ -217,6 +228,17 @@ class ProtosTest {
         assertEquals(OneOfMsg.Field.Fixed(123u), decoded.field)
     }
 
+    @Test
+    fun testOneOfRequiredSubField() {
+        assertFailsWith<IllegalStateException> {
+            OneOfWithRequired {
+                // we construct the message using the internal class,
+                // so it is not invoking the checkRequired method on construction
+                field = OneOfWithRequired.Field.Msg(PresenceCheckInternal())
+            }
+        }
+    }
+
     @Test
     fun testOneOfNull() {
         // write two values on the oneOf field.
@@ -239,7 +261,7 @@ class ProtosTest {
 
     @Test
     fun testRecursiveReqNotSet() {
-        assertFailsWith<IllegalStateException>("RecursiveReq is missing required field: rec") {
+        assertFailsWith<IllegalStateException> {
             val msg = RecursiveReq {
                 rec = RecursiveReq {
                     rec = RecursiveReq {

grpc/grpc-core/src/commonTest/proto/oneof.proto

@@ -1,5 +1,6 @@
 import "sub_message.proto";
 import "enum.proto";
+import "presence_check.proto";
 
 message OneOfMsg {
   oneof field {
@@ -8,4 +9,11 @@ message OneOfMsg {
     test.submsg.Other other = 4;
     kotlinx.rpc.grpc.test.MyEnum enum = 5;
   }
+}
+
+message OneOfWithRequired {
+  oneof field {
+    int32 sint = 1;
+    kotlinx.rpc.grpc.test.common.PresenceCheck msg = 2;
+  }
 }

grpc/grpc-core/src/commonTest/proto/repeated.proto

@@ -1,5 +1,7 @@
 syntax = "proto3";
 
+import "presence_check.proto";
+
 package kotlinx.rpc.grpc.test.common;
 
 message Repeated {
@@ -13,4 +15,8 @@ message Repeated {
   message Other {
     int32 a = 1;
   }
+}
+
+message RepeatedWithRequired {
+  repeated PresenceCheck msgList = 1;
 }

protoc-gen/src/main/kotlin/kotlinx/rpc/protobuf/ModelToKotlinCommonGenerator.kt

@@ -530,6 +530,35 @@ class ModelToKotlinCommonGenerator(
                 code("${field.name}.asInternal().checkRequiredFields()")
             })
         }
+
+        // check submessages in oneofs
+        declaration.fields().filter { it.second.type is FieldType.OneOf }.forEach { (_, field) ->
+            val oneOfType = field.type as FieldType.OneOf
+            val messageVariants = oneOfType.dec.variants.filter { it.type is FieldType.Message }
+            if (messageVariants.isEmpty()) return@forEach
+
+            scope("${field.name}?.also") {
+                whenBlock {
+                    messageVariants.forEach { variant ->
+                        val variantClassName = "${field.type.dec.name.safeFullName()}.${variant.name}"
+                        whenCase("it is $variantClassName") {
+                            code("it.value.asInternal().checkRequiredFields()")
+                        }
+                    }
+                }
+            }
+        }
+
+        // check submessages in lists
+        declaration.fields().filter { it.second.type is FieldType.List }.forEach { (_, field) ->
+            val listType = field.type as FieldType.List
+            if (listType.value !is FieldType.Message) return@forEach
+
+            scope("${field.name}.forEach") {
+                code("it.asInternal().checkRequiredFields()")
+            }
+        }
+
     }
 
     private fun CodeGenerator.generateInternalComputeSize(declaration: MessageDeclaration) {