From c87857ed6c1620073e907d60bd21fc87b4e4d0bc Mon Sep 17 00:00:00 2001
From: taichi uchihara <contact@ucpr.dev>
Date: Tue, 25 Apr 2023 01:33:35 +0900
Subject: [PATCH] Fix invalid generation when using in-rule and not_in rule in
 map.keys.string (#847)

## Summary

I fixed invalid generation for Go and cc when using in-rule and not_in
rule in map.keys.string

## Issue

Closes https://github.com/bufbuild/protoc-gen-validate/issues/785

---------

Co-authored-by: Chris Roche <github@rodaine.com>
---
 templates/cc/msg.go             | 16 +++++++++++++++-
 templates/goshared/msg.go       | 14 ++++++++++++++
 tests/harness/cases/maps.proto  |  3 +++
 tests/harness/executor/cases.go |  4 ++++
 4 files changed, 36 insertions(+), 1 deletion(-)

diff --git a/templates/cc/msg.go b/templates/cc/msg.go
index 46cd290fe..c8944d801 100644
--- a/templates/cc/msg.go
+++ b/templates/cc/msg.go
@@ -91,7 +91,21 @@ const msgTpl = `
 	{{ end }}{{ end }}
         {{ end }}{{ end }}
 
-        {{ if has .Rules "Keys"}}{{ if .Rules.Keys }}
+	{{ if has .Rules "Keys"}}{{ if .Rules.Keys }}
+	{{ if has .Rules.Keys.GetString_ "In" }} {{ if .Rules.Keys.GetString_.In }}
+	const std::set<string> {{ lookup .Field "InLookup" }} = {
+			{{- range .Rules.Keys.GetString_.In }}
+				{{ inKey $f . }},
+			{{- end }}
+		};
+	{{ end }}{{ end }}
+	{{ if has .Rules.Keys.GetString_ "NotIn" }} {{ if .Rules.Keys.GetString_.NotIn }}
+	const std::set<string> {{ lookup .Field "NotInLookup" }} = {
+			{{- range .Rules.Keys.GetString_.NotIn }}
+				{{ inKey $f . }},
+			{{- end }}
+		};
+	{{ end }}{{ end }}
 	{{ if has .Rules.Keys.GetString_ "Pattern" }} {{ if .Rules.Keys.GetString_.Pattern }}
 		const re2::RE2 {{ lookup .Field "Pattern" }}(re2::StringPiece({{ lit .Rules.Keys.GetString_.GetPattern }},
                                               sizeof({{ lit .Rules.Keys.GetString_.GetPattern }}) - 1));
diff --git a/templates/goshared/msg.go b/templates/goshared/msg.go
index 8d9f4dcdc..030b94bce 100644
--- a/templates/goshared/msg.go
+++ b/templates/goshared/msg.go
@@ -234,6 +234,20 @@ var _ interface{
 	{{ end }}{{ end }}
 
 	{{ if has .Rules "Keys"}}{{ if .Rules.Keys }}
+	{{ if has .Rules.Keys.GetString_ "In" }} {{ if .Rules.Keys.GetString_.In }}
+		var {{ lookup .Field "InLookup" }} = map[{{ inType .Field .Rules.Keys.GetString_.In }}]struct{}{
+			{{- range .Rules.Keys.GetString_.In }}
+				{{ inKey $f . }}: {},
+			{{- end }}
+		}
+	{{ end }}{{ end }}
+	{{ if has .Rules.Keys.GetString_ "NotIn" }} {{ if .Rules.Keys.GetString_.NotIn }}
+		var {{ lookup .Field "NotInLookup" }} = map[{{ inType .Field .Rules.Keys.GetString_.NotIn }}]struct{}{
+			{{- range .Rules.Keys.GetString_.NotIn }}
+				{{ inKey $f . }}: {},
+			{{- end }}
+		}
+	{{ end }}{{ end }}
 	{{ if has .Rules.Keys.GetString_ "Pattern" }} {{ if .Rules.Keys.GetString_.Pattern }}
 		var {{ lookup .Field "Pattern" }} = regexp.MustCompile({{ lit .Rules.Keys.GetString_.GetPattern }})
 	{{ end }}{{ end }}
diff --git a/tests/harness/cases/maps.proto b/tests/harness/cases/maps.proto
index 2471f5315..1ddbbd570 100644
--- a/tests/harness/cases/maps.proto
+++ b/tests/harness/cases/maps.proto
@@ -36,3 +36,6 @@ message MultipleMaps {
   map <int32, bool> second = 2 [(validate.rules).map.keys.int32.lt = 0];
   map <int32, bool> third = 3 [(validate.rules).map.keys.int32.gt = 0];
 }
+
+message MapKeysIn {map<string, string> val = 1 [(validate.rules).map.keys.string = {in: ["foo", "bar"]}]; }
+message MapKeysNotIn {map<string, string> val = 1 [(validate.rules).map.keys.string = {not_in: ["foo", "bar"]}]; }
diff --git a/tests/harness/executor/cases.go b/tests/harness/executor/cases.go
index 0d4a8fa8f..13c37282b 100644
--- a/tests/harness/executor/cases.go
+++ b/tests/harness/executor/cases.go
@@ -1184,8 +1184,12 @@ var mapCases = []TestCase{
 	{"map - keys - valid", &cases.MapKeys{Val: map[int64]string{-1: "a", -2: "b"}}, 0},
 	{"map - keys - valid (empty)", &cases.MapKeys{Val: map[int64]string{}}, 0},
 	{"map - keys - valid (pattern)", &cases.MapKeysPattern{Val: map[string]string{"A": "a"}}, 0},
+	{"map - keys - valid (in)", &cases.MapKeysIn{Val: map[string]string{"foo": "value"}}, 0},
+	{"map - keys - valid (not_in)", &cases.MapKeysNotIn{Val: map[string]string{"baz": "value"}}, 0},
 	{"map - keys - invalid", &cases.MapKeys{Val: map[int64]string{1: "a"}}, 1},
 	{"map - keys - invalid (pattern)", &cases.MapKeysPattern{Val: map[string]string{"A": "a", "!@#$%^&*()": "b"}}, 1},
+	{"map - keys - invalid (in)", &cases.MapKeysIn{Val: map[string]string{"baz": "value"}}, 1},
+	{"map - keys - invalid (not_in)", &cases.MapKeysNotIn{Val: map[string]string{"foo": "value"}}, 1},
 
 	{"map - values - valid", &cases.MapValues{Val: map[string]string{"a": "Alpha", "b": "Beta"}}, 0},
 	{"map - values - valid (empty)", &cases.MapValues{Val: map[string]string{}}, 0},