From f93f7931edb46bcba91751062d52435210de3578 Mon Sep 17 00:00:00 2001
From: Artem Skriabin <artem.skriabin@distributedlab.com>
Date: Wed, 11 Sep 2024 14:42:26 +0300
Subject: [PATCH] Fix getting country code from pub signals

---
 internal/service/requests/verify_passport.go  |  9 +++--
 .../service/requests/verify_passport_test.go  | 40 +++++++++++++++++++
 2 files changed, 46 insertions(+), 3 deletions(-)
 create mode 100644 internal/service/requests/verify_passport_test.go

diff --git a/internal/service/requests/verify_passport.go b/internal/service/requests/verify_passport.go
index 08e2f66..780ceeb 100644
--- a/internal/service/requests/verify_passport.go
+++ b/internal/service/requests/verify_passport.go
@@ -64,16 +64,19 @@ func NewVerifyPassport(r *http.Request) (req resources.VerifyPassportRequest, er
 // ExtractCountry extracts country code from the proof, converting decimal UTF-8
 // code to ISO 3166-1 alpha-3 code.
 func ExtractCountry(proof zkptypes.ZKProof) (string, error) {
-	if len(proof.PubSignals) <= int(zk.Citizenship) {
+	if len(proof.PubSignals) <= zk.Indexes(zk.GlobalPassport)[zk.Citizenship] {
 		return "", val.Errors{"country_code": val.ErrLengthTooShort}.Filter()
 	}
 
-	b, ok := new(big.Int).SetString(proof.PubSignals[zk.Citizenship], 10)
+	getter := zk.PubSignalGetter{Signals: proof.PubSignals, ProofType: zk.GlobalPassport}
+	code := getter.Get(zk.Citizenship)
+
+	b, ok := new(big.Int).SetString(code, 10)
 	if !ok {
 		b = new(big.Int)
 	}
 
-	code := string(b.Bytes())
+	code = string(b.Bytes())
 
 	return code, val.Errors{"country_code": val.Validate(code, val.Required, is.CountryCode3)}.Filter()
 }
diff --git a/internal/service/requests/verify_passport_test.go b/internal/service/requests/verify_passport_test.go
new file mode 100644
index 0000000..ddfa984
--- /dev/null
+++ b/internal/service/requests/verify_passport_test.go
@@ -0,0 +1,40 @@
+package requests
+
+import (
+	zkptypes "github.com/iden3/go-rapidsnark/types"
+	"github.com/stretchr/testify/assert"
+	"testing"
+)
+
+func TestExtractCountry(t *testing.T) {
+	p := zkptypes.ZKProof{
+		PubSignals: []string{"9072791962122059526608165338435582217583998249448445715859811865975207897521",
+			"0",
+			"0",
+			"0",
+			"0",
+			"0",
+			"5589842",
+			"0",
+			"0",
+			"211985299740800702300256033401632392934377086534111448880928528431996790315",
+			"6334216257887790546056750444971451489366652820049478976520311646899632086040",
+			"6657866160187480897185968992500348980038797861763249572960971022605721484835",
+			"23073",
+			"55199728742705",
+			"0",
+			"1726052791",
+			"0",
+			"10",
+			"52983525027888",
+			"53009295421745",
+			"55199728742705",
+			"52983525027888",
+			"52983525027888",
+		},
+	}
+
+	code, err := ExtractCountry(p)
+	assert.NoError(t, err)
+	assert.Equal(t, "UKR", code)
+}