From 757124b2567b195a5e13db1bff6de8aec4533a95 Mon Sep 17 00:00:00 2001
From: Shion Ichikawa <shion1305@gmail.com>
Date: Fri, 23 Feb 2024 20:28:14 +0900
Subject: [PATCH] =?UTF-8?q?=E2=9A=A1=EF=B8=8F=20codeCache=E3=81=B8?=
 =?UTF-8?q?=E3=81=AERaceCondition=E5=AF=BE=E7=AD=96?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 pkg/token/issuer.go | 14 +++++++++++---
 1 file changed, 11 insertions(+), 3 deletions(-)

diff --git a/pkg/token/issuer.go b/pkg/token/issuer.go
index 1651f46..b2d5f3d 100644
--- a/pkg/token/issuer.go
+++ b/pkg/token/issuer.go
@@ -7,6 +7,7 @@ import (
 	"fmt"
 	"google.golang.org/appengine/v2/log"
 	"strconv"
+	"sync"
 	"time"
 	"ynufes-mypage-backend/pkg/jwt"
 )
@@ -16,6 +17,7 @@ type Issuer struct {
 	issuer    string
 	maxAge    time.Duration
 	codeCache map[string]credit
+	lock      sync.RWMutex // lock for codeCache, prevents race condition
 }
 
 type credit struct {
@@ -23,8 +25,8 @@ type credit struct {
 	time int64
 }
 
-func NewTokenIssuer(jwtSecret string, issuer string, maxAge time.Duration) Issuer {
-	return Issuer{
+func NewTokenIssuer(jwtSecret string, issuer string, maxAge time.Duration) *Issuer {
+	return &Issuer{
 		jwtSecret: jwtSecret,
 		issuer:    issuer,
 		maxAge:    maxAge,
@@ -32,7 +34,9 @@ func NewTokenIssuer(jwtSecret string, issuer string, maxAge time.Duration) Issue
 	}
 }
 
-func (v Issuer) IssueNewCode(id string) (string, error) {
+func (v *Issuer) IssueNewCode(id string) (string, error) {
+	v.lock.Lock()
+	defer v.lock.Unlock()
 	cnt := 0
 	newCode, err := v.generateSecret()
 	if err != nil {
@@ -53,6 +57,8 @@ func (v Issuer) IssueNewCode(id string) (string, error) {
 }
 
 func (v Issuer) IssueToken(code string) (string, error) {
+	v.lock.Lock()
+	defer v.lock.Unlock()
 	credit, ok := v.codeCache[code]
 	if !ok {
 		return "", fmt.Errorf("code not found")
@@ -72,6 +78,8 @@ func (v Issuer) IssueToken(code string) (string, error) {
 }
 
 func (v Issuer) RevokeOldCodes() {
+	v.lock.Lock()
+	defer v.lock.Unlock()
 	for s, t := range v.codeCache {
 		// If the code is older than 5 seconds, delete it
 		if time.Now().UnixMilli()-t.time > 5000 {