awakesecurity · ccycle · Feb 20, 2023 · Feb 8, 2023 · Apr 21, 2023 · Apr 27, 2023
diff --git a/src/Proto3/Suite/DotProto/Generate.hs b/src/Proto3/Suite/DotProto/Generate.hs
@@ -25,6 +25,7 @@ module Proto3.Suite.DotProto.Generate
   ( CompileError(..)
   , StringType(..)
   , RecordStyle (..)
+  , IsPrefixed(..)
   , parseStringType
   , TypeContext
   , CompileArgs(..)
@@ -83,6 +84,7 @@ data CompileArgs = CompileArgs
   , outputDir          :: FilePath
   , stringType         :: StringType
   , recordStyle        :: RecordStyle
+  , isPrefixed         :: IsPrefixed
   }
 
 data StringType = StringType String String
@@ -110,7 +112,7 @@ compileDotProtoFile CompileArgs{..} = runExceptT $ do
   Turtle.mktree (Turtle.directory modulePath)
 
   extraInstances <- foldMapM getExtraInstances extraInstanceFiles
-  haskellModule <- renderHsModuleForDotProto stringType recordStyle extraInstances dotProto importTypeContext
+  haskellModule <- renderHsModuleForDotProto stringType recordStyle isPrefixed extraInstances dotProto importTypeContext
 
   liftIO (writeFile (Turtle.encodeString modulePath) haskellModule)
   where
@@ -182,9 +184,10 @@ renderHsModuleForDotProto
     :: MonadError CompileError m
     => StringType
     -> RecordStyle
+    -> IsPrefixed
     -> ([HsImportDecl],[HsDecl]) -> DotProto -> TypeContext -> m String
-renderHsModuleForDotProto stringType recordStyle extraInstanceFiles dotProto importCtxt = do
-    haskellModule <- hsModuleForDotProto stringType recordStyle extraInstanceFiles dotProto importCtxt
+renderHsModuleForDotProto stringType recordStyle isPrefixed extraInstanceFiles dotProto importCtxt = do
+    haskellModule <- hsModuleForDotProto stringType recordStyle isPrefixed extraInstanceFiles dotProto importCtxt
 
     let languagePragmas = textUnlines $ map (\extn -> "{-# LANGUAGE " <> extn <> " #-}") $ sort extensions
         ghcOptionPragmas = textUnlines $ map (\opt -> "{-# OPTIONS_GHC " <> opt <> " #-}") $ sort options
@@ -207,6 +210,9 @@ renderHsModuleForDotProto stringType recordStyle extraInstanceFiles dotProto imp
                            , "TypeFamilies"
                            , "UndecidableInstances"
                            ]
+          ++ case isPrefixed of
+            IsPrefixed True -> []
+            IsPrefixed False -> ["DuplicateRecordFields"]
 
         options :: [T.Text]
         options = [ "-fno-warn-unused-imports"
@@ -254,6 +260,8 @@ hsModuleForDotProto
     -- ^ the module and the type for string
     -> RecordStyle
     -- ^ kind of records to generate
+    -> IsPrefixed
+    -- ^ flag for prefix of field names
     -> ([HsImportDecl], [HsDecl])
     -- ^ Extra user-define instances that override default generated instances
     -> DotProto
@@ -264,6 +272,7 @@ hsModuleForDotProto
 hsModuleForDotProto
     stringType
     recordStyle
+    isPrefixed
     (extraImports, extraInstances)
     dotProto@DotProto{ protoMeta = DotProtoMeta { metaModulePath = modulePath }
                      , protoPackage
@@ -289,7 +298,7 @@ hsModuleForDotProto
        typeContext <- dotProtoTypeContext dotProto
 
        let toDotProtoDeclaration =
-             dotProtoDefinitionD stringType recordStyle protoPackage (typeContext <> importTypeContext)
+             dotProtoDefinitionD stringType recordStyle isPrefixed protoPackage (typeContext <> importTypeContext)
 
        let extraInstances' = instancesForModule moduleName extraInstances
 
@@ -733,19 +742,20 @@ validMapKey = (`elem` [ Int32, Int64, SInt32, SInt64, UInt32, UInt64
 dotProtoDefinitionD :: MonadError CompileError m
                     => StringType
                     -> RecordStyle
+                    -> IsPrefixed
                     -> DotProtoPackageSpec
                     -> TypeContext
                     -> DotProtoDefinition
                     -> m [HsDecl]
-dotProtoDefinitionD stringType recordStyle pkgSpec ctxt = \case
+dotProtoDefinitionD stringType recordStyle isPrefixed pkgSpec ctxt = \case
   DotProtoMessage _ messageName messageParts ->
-    dotProtoMessageD stringType recordStyle ctxt Anonymous messageName messageParts
+    dotProtoMessageD stringType recordStyle isPrefixed ctxt Anonymous messageName messageParts
 
   DotProtoEnum _ enumName enumParts ->
     dotProtoEnumD Anonymous enumName enumParts
 
   DotProtoService _ serviceName serviceParts ->
-    dotProtoServiceD stringType pkgSpec ctxt serviceName serviceParts
+    dotProtoServiceD stringType isPrefixed pkgSpec ctxt serviceName serviceParts
 
 -- | Generate 'Named' instance for a type in this package
 namedInstD :: String -> HsDecl
@@ -774,12 +784,13 @@ dotProtoMessageD
      . MonadError CompileError m
     => StringType
     -> RecordStyle
+    -> IsPrefixed
     -> TypeContext
     -> DotProtoIdentifier
     -> DotProtoIdentifier
     -> [DotProtoMessagePart]
     -> m [HsDecl]
-dotProtoMessageD stringType recordStyle ctxt parentIdent messageIdent messageParts = do
+dotProtoMessageD stringType recordStyle isPrefixed ctxt parentIdent messageIdent messageParts = do
     messageName <- qualifiedMessageName parentIdent messageIdent
 
     let mkDataDecl flds =
@@ -800,18 +811,18 @@ dotProtoMessageD stringType recordStyle ctxt parentIdent messageIdent messagePar
           , pure (nfDataInstD messageDataDecl messageName)
           , pure (namedInstD messageName)
           , pure (hasDefaultInstD messageName)
-          , messageInstD stringType ctxt' parentIdent messageIdent messageParts
+          , messageInstD stringType isPrefixed ctxt' parentIdent messageIdent messageParts
 
-          , toJSONPBMessageInstD stringType ctxt' parentIdent messageIdent messageParts
-          , fromJSONPBMessageInstD stringType ctxt' parentIdent messageIdent messageParts
+          , toJSONPBMessageInstD stringType isPrefixed ctxt' parentIdent messageIdent messageParts
+          , fromJSONPBMessageInstD stringType isPrefixed ctxt' parentIdent messageIdent messageParts
 
             -- Generate Aeson instances in terms of JSONPB instances
           , pure (toJSONInstDecl messageName)
           , pure (fromJSONInstDecl messageName)
 
 #ifdef SWAGGER
           -- And the Swagger ToSchema instance corresponding to JSONPB encodings
-          , toSchemaInstanceDeclaration stringType ctxt' messageName Nothing
+          , toSchemaInstanceDeclaration stringType isPrefixed ctxt' messageName Nothing
               =<< foldMapM getName messageParts
 #endif
 
@@ -843,12 +854,12 @@ dotProtoMessageD stringType recordStyle ctxt parentIdent messageIdent messagePar
 
     messagePartFieldD :: String -> DotProtoMessagePart -> m [([HsName], HsBangType)]
     messagePartFieldD messageName (DotProtoMessageField DotProtoField{..}) = do
-      fullName <- prefixedFieldName messageName =<< dpIdentUnqualName dotProtoFieldName
+      fullName <- prefixedFieldNameWithFlag isPrefixed messageName =<< dpIdentUnqualName dotProtoFieldName
       fullTy <- dptToHsType WithinMessage stringType ctxt' dotProtoFieldType
       pure [ ([HsIdent fullName], HsUnBangedTy fullTy ) ]
 
     messagePartFieldD messageName (DotProtoMessageOneOf fieldName _) = do
-      fullName <- prefixedFieldName messageName =<< dpIdentUnqualName fieldName
+      fullName <- prefixedFieldNameWithFlag isPrefixed messageName =<< dpIdentUnqualName fieldName
       qualTyName <- prefixedConName messageName =<< dpIdentUnqualName fieldName
       let fullTy = HsTyApp (HsTyCon (haskellName "Maybe")) . type_ $ qualTyName
       pure [ ([HsIdent fullName], HsUnBangedTy fullTy) ]
@@ -858,7 +869,7 @@ dotProtoMessageD stringType recordStyle ctxt parentIdent messageIdent messagePar
     nestedDecls :: DotProtoDefinition -> m [HsDecl]
     nestedDecls (DotProtoMessage _ subMsgName subMessageDef) = do
       parentIdent' <- concatDotProtoIdentifier parentIdent messageIdent
-      dotProtoMessageD stringType recordStyle ctxt' parentIdent' subMsgName subMessageDef
+      dotProtoMessageD stringType recordStyle isPrefixed ctxt' parentIdent' subMsgName subMessageDef
 
     nestedDecls (DotProtoEnum _ subEnumName subEnumDef) = do
       parentIdent' <- concatDotProtoIdentifier parentIdent messageIdent
@@ -873,7 +884,7 @@ dotProtoMessageD stringType recordStyle ctxt parentIdent messageIdent messagePar
       (cons, idents) <- fmap unzip (mapM (oneOfCons fullName) fields)
 
 #ifdef SWAGGER
-      toSchemaInstance <- toSchemaInstanceDeclaration stringType ctxt' fullName (Just idents)
+      toSchemaInstance <- toSchemaInstanceDeclaration stringType isPrefixed ctxt' fullName (Just idents)
                             =<< mapM getFieldNameForSchemaInstanceDeclaration fields
 #endif
 
@@ -906,14 +917,15 @@ messageInstD
     :: forall m
      . MonadError CompileError m
     => StringType
+    -> IsPrefixed
     -> TypeContext
     -> DotProtoIdentifier
     -> DotProtoIdentifier
     -> [DotProtoMessagePart]
     -> m HsDecl
-messageInstD stringType ctxt parentIdent msgIdent messageParts = do
+messageInstD stringType isPrefixed ctxt parentIdent msgIdent messageParts = do
      msgName         <- qualifiedMessageName parentIdent msgIdent
-     qualifiedFields <- getQualifiedFields msgName messageParts
+     qualifiedFields <- getQualifiedFields isPrefixed msgName messageParts
 
      encodedFields   <- mapM encodeMessageField qualifiedFields
      decodedFields   <- mapM decodeMessageField qualifiedFields
@@ -1041,14 +1053,15 @@ toJSONPBMessageInstD
     :: forall m
      . MonadError CompileError m
     => StringType
+    -> IsPrefixed
     -> TypeContext
     -> DotProtoIdentifier
     -> DotProtoIdentifier
     -> [DotProtoMessagePart]
     -> m HsDecl
-toJSONPBMessageInstD stringType ctxt parentIdent msgIdent messageParts = do
+toJSONPBMessageInstD stringType isPrefixed ctxt parentIdent msgIdent messageParts = do
     msgName    <- qualifiedMessageName parentIdent msgIdent
-    qualFields <- getQualifiedFields msgName messageParts
+    qualFields <- getQualifiedFields isPrefixed msgName messageParts
 
     let applyE nm oneofNm = do
           fs <- traverse (encodeMessageField oneofNm) qualFields
@@ -1162,14 +1175,15 @@ fromJSONPBMessageInstD
     :: forall m
      . MonadError CompileError m
     => StringType
+    -> IsPrefixed
     -> TypeContext
     -> DotProtoIdentifier
     -> DotProtoIdentifier
     -> [DotProtoMessagePart]
     -> m HsDecl
-fromJSONPBMessageInstD stringType ctxt parentIdent msgIdent messageParts = do
+fromJSONPBMessageInstD stringType isPrefixed ctxt parentIdent msgIdent messageParts = do
     msgName    <- qualifiedMessageName parentIdent msgIdent
-    qualFields <- getQualifiedFields msgName messageParts
+    qualFields <- getQualifiedFields isPrefixed msgName messageParts
 
     fieldParsers <- traverse parseField qualFields
 
@@ -1306,6 +1320,7 @@ getFieldNameForSchemaInstanceDeclaration fld = do
 toSchemaInstanceDeclaration
     :: MonadError CompileError m
     => StringType
+    -> IsPrefixed
     -> TypeContext
     -> String
     -- ^ Name of the message type to create an instance for
@@ -1315,10 +1330,10 @@ toSchemaInstanceDeclaration
     -- ^ Field names, with every field that is not actually a oneof
     -- combining fields paired with its options and protobuf type
     -> m HsDecl
-toSchemaInstanceDeclaration stringType ctxt messageName maybeConstructors fieldNamesEtc = do
+toSchemaInstanceDeclaration stringType isPrefixed ctxt messageName maybeConstructors fieldNamesEtc = do
   let fieldNames = map snd fieldNamesEtc
 
-  qualifiedFieldNames <- mapM (prefixedFieldName messageName) fieldNames
+  qualifiedFieldNames <- mapM (prefixedFieldNameWithFlag isPrefixed messageName) fieldNames
 
   let messageConstructor = HsCon (UnQual (HsIdent messageName))
 
@@ -1651,12 +1666,13 @@ dotProtoEnumD parentIdent enumIdent enumParts = do
 dotProtoServiceD
     :: MonadError CompileError m
     => StringType
+    -> IsPrefixed
     -> DotProtoPackageSpec
     -> TypeContext
     -> DotProtoIdentifier
     -> [DotProtoServicePart]
     -> m [HsDecl]
-dotProtoServiceD stringType pkgSpec ctxt serviceIdent service = do
+dotProtoServiceD stringType isPrefixed pkgSpec ctxt serviceIdent service = do
      serviceName <- typeLikeName =<< dpIdentUnqualName serviceIdent
 
      endpointPrefix <-
@@ -1667,7 +1683,7 @@ dotProtoServiceD stringType pkgSpec ctxt serviceIdent service = do
          DotProtoNoPackage -> pure $ "/" ++ serviceName ++ "/"
 
      let serviceFieldD (DotProtoServiceRPCMethod RPCMethod{..}) = do
-           fullName <- prefixedMethodName serviceName =<< dpIdentUnqualName rpcMethodName
+           fullName <- prefixedMethodNameWithFlag isPrefixed serviceName =<< dpIdentUnqualName rpcMethodName
 
            methodName <- case rpcMethodName of
                            Single nm -> pure nm

diff --git a/src/Proto3/Suite/DotProto/Internal.hs b/src/Proto3/Suite/DotProto/Internal.hs
@@ -33,6 +33,7 @@ import           Data.List                 (intercalate)
 import qualified Data.List.NonEmpty        as NE
 import qualified Data.Map                  as M
 import           Data.Maybe                (fromMaybe)
+import qualified Data.Set                  as S
 import qualified Data.Text                 as T
 import           Data.Tuple                (swap)
 import qualified NeatInterpolation         as Neat
@@ -499,6 +500,10 @@ prefixedConName msgName conName = do
   constructor <- typeLikeName conName
   return (msgName ++ constructor)
 
+newtype IsPrefixed = IsPrefixed Bool
+instance Show IsPrefixed where
+  show (IsPrefixed b) = show b
+
 -- | @'prefixedMethodName' service method@ produces a Haskell record selector name for the service method @method@ by
 -- joining the names @service@, @method@ under concatenation on a camel-casing transformation.
 prefixedMethodName :: MonadError CompileError m => String -> String -> m String
@@ -509,13 +514,38 @@ prefixedMethodName serviceName (x : xs)
       method <- typeLikeName (x : xs)
       return (fieldLikeName serviceName ++ method)
 
+prefixedMethodNameWithFlag :: MonadError CompileError m => IsPrefixed -> String -> String -> m String
+prefixedMethodNameWithFlag _ _ "" = invalidTypeNameError "<empty name>"
+prefixedMethodNameWithFlag (IsPrefixed flag) serviceName (x : xs)
+  | flag = prefixedMethodName serviceName (x : xs)
+  | name `S.member` haskellKeywords = return (name ++ "_")
+  | otherwise = return name
+  where
+    name = (toCamelCase . fieldLikeName) (x : xs)
+
 -- | @'prefixedFieldName' prefix field@ constructs a Haskell record selector name by prepending @prefix@ in camel-case
 -- to the message field/service method name @field@.
 prefixedFieldName :: MonadError CompileError m => String -> String -> m String
 prefixedFieldName msgName fieldName = do
   field <- typeLikeName fieldName
   return (fieldLikeName msgName ++ field)
 
+prefixedFieldNameWithFlag :: MonadError CompileError m => IsPrefixed -> String -> String -> m String
+prefixedFieldNameWithFlag (IsPrefixed flag) msgName fieldName = do
+  if flag then prefixedFieldName msgName fieldName else return $
+    -- Avoid parse error occurring when the field name matches any of Haskell keywords
+    if name `S.member` haskellKeywords then name ++ "_" else name
+  where
+    name = (toCamelCase . fieldLikeName) fieldName
+
+haskellKeywords :: S.Set String
+haskellKeywords = S.fromList
+  ["case","class","data","default","deriving","do","else"
+  ,"if","import","in","infix","infixl","infixr","instance","let","module"
+  ,"newtype","of","then","type","where","_"
+  ,"foreign"
+  ]
+
 dpIdentUnqualName :: MonadError CompileError m => DotProtoIdentifier -> m String
 dpIdentUnqualName (Single name)       = pure name
 dpIdentUnqualName (Dots (Path names)) = pure (NE.last names)
@@ -594,11 +624,11 @@ data OneofSubfield = OneofSubfield
   } deriving Show
 
 getQualifiedFields :: MonadError CompileError m
-                   => String -> [DotProtoMessagePart] -> m [QualifiedField]
-getQualifiedFields msgName msgParts = flip foldMapM msgParts $ \case
+                   => IsPrefixed -> String -> [DotProtoMessagePart] -> m [QualifiedField]
+getQualifiedFields isPrefixed msgName msgParts = flip foldMapM msgParts $ \case
   DotProtoMessageField DotProtoField{..} -> do
     fieldName <- dpIdentUnqualName dotProtoFieldName
-    qualName <- prefixedFieldName msgName fieldName
+    qualName <- prefixedFieldNameWithFlag isPrefixed msgName fieldName
     pure . (:[]) $ QualifiedField { recordFieldName = coerce qualName
                                   , fieldInfo = FieldNormal (coerce fieldName)
                                                             dotProtoFieldNumber
@@ -611,7 +641,7 @@ getQualifiedFields msgName msgParts = flip foldMapM msgParts $ \case
 
   DotProtoMessageOneOf oneofIdent fields -> do
     ident <- dpIdentUnqualName oneofIdent
-    oneofName <- prefixedFieldName msgName ident
+    oneofName <- prefixedFieldNameWithFlag isPrefixed msgName ident
     oneofTypeName <- prefixedConName msgName ident
 
     let mkSubfield DotProtoField{..} = do