Validate directives are in valid locations

2 files changed, 31 insertions, 4 deletions

diff --git a/src/Language/GraphQL/Validate/Rules.hs b/src/Language/GraphQL/Validate/Rules.hs
index 7cfa712..6c35f70 100644
--- a/src/Language/GraphQL/Validate/Rules.hs
+++ b/src/Language/GraphQL/Validate/Rules.hs
@@ -10,7 +10,8 @@
 
 -- | This module contains default rules defined in the GraphQL specification.
 module Language.GraphQL.Validate.Rules
-    ( executableDefinitionsRule
+    ( directivesInValidLocationsRule
+    , executableDefinitionsRule
     , fieldsOnCorrectTypeRule
     , fragmentsOnCompositeTypesRule
     , fragmentSpreadTargetDefinedRule
@@ -90,6 +91,7 @@ specifiedRules =
     , uniqueInputFieldNamesRule
     -- Directives.
     , knownDirectiveNamesRule
+    , directivesInValidLocationsRule
     , uniqueDirectiveNamesRule
     -- Variables.
     , uniqueVariableNamesRule
@@ -514,7 +516,7 @@ uniqueArgumentNamesRule = ArgumentsRule fieldRule directiveRule
 -- of each directive is allowed per location.
 uniqueDirectiveNamesRule :: forall m. Rule m
 uniqueDirectiveNamesRule = DirectivesRule
-    $ lift . filterDuplicates extract "directive"
+    $ const $ lift . filterDuplicates extract "directive"
   where
     extract (Directive directiveName _ location') = (directiveName, location')
 
@@ -818,7 +820,7 @@ knownArgumentNamesRule = ArgumentsRule fieldRule directiveRule
 -- | GraphQL servers define what directives they support. For each usage of a
 -- directive, the directive must be available on that server.
 knownDirectiveNamesRule :: Rule m
-knownDirectiveNamesRule = DirectivesRule $ \directives' -> do
+knownDirectiveNamesRule = DirectivesRule $ const $ \directives' -> do
     definitions' <- asks directives
     let directiveSet = HashSet.fromList $ fmap directiveName directives'
     let definitionSet = HashSet.fromList $ HashMap.keys definitions'
@@ -867,3 +869,27 @@ knownInputFieldNamesRule = ValueRule go constGo
         , Text.unpack typeName
         , "\"."
         ]
+
+-- | GraphQL servers define what directives they support and where they support
+-- them. For each usage of a directive, the directive must be used in a location
+-- that the server has declared support for.
+directivesInValidLocationsRule :: Rule m
+directivesInValidLocationsRule = DirectivesRule directivesRule
+  where
+    directivesRule directiveLocation directives' = do
+        Directive directiveName _ location <- lift $ Seq.fromList directives'
+        maybeDefinition <- asks $ HashMap.lookup directiveName . directives
+        case maybeDefinition of
+            Just (Schema.Directive _ allowedLocations _)
+                | directiveLocation `notElem` allowedLocations -> pure $ Error
+                    { message = errorMessage directiveName directiveLocation
+                    , locations = [location]
+                    }
+            _ -> lift mempty
+    errorMessage directiveName directiveLocation = concat
+        [ "Directive \"@"
+        , Text.unpack directiveName
+        , "\" may not be used on "
+        , show directiveLocation
+        , "."
+        ]
diff --git a/src/Language/GraphQL/Validate/Validation.hs b/src/Language/GraphQL/Validate/Validation.hs
index 0e9f1a8..32a454e 100644
--- a/src/Language/GraphQL/Validate/Validation.hs
+++ b/src/Language/GraphQL/Validate/Validation.hs
@@ -13,6 +13,7 @@ module Language.GraphQL.Validate.Validation
 import Control.Monad.Trans.Reader (ReaderT)
 import Data.HashMap.Strict (HashMap)
 import Data.Sequence (Seq)
+import Language.GraphQL.AST.DirectiveLocation (DirectiveLocation(..))
 import Language.GraphQL.AST.Document
 import qualified Language.GraphQL.Type.In as In
 import qualified Language.GraphQL.Type.Out as Out
@@ -45,7 +46,7 @@ data Rule m
     | FragmentSpreadRule (FragmentSpread -> RuleT m)
     | FieldRule (Maybe (Out.Type m) -> Field -> RuleT m)
     | ArgumentsRule (Maybe (Out.Type m) -> Field -> RuleT m) (Directive -> RuleT m)
-    | DirectivesRule ([Directive] -> RuleT m)
+    | DirectivesRule (DirectiveLocation -> [Directive] -> RuleT m)
     | VariablesRule ([VariableDefinition] -> RuleT m)
     | ValueRule (Maybe In.Type -> Value -> RuleT m) (Maybe In.Type -> ConstValue -> RuleT m)