From eb7e7de463b10489867bd4365d20f752fe25f573 Mon Sep 17 00:00:00 2001
From: Alexander <zavorotynskiy@pm.me>
Date: Fri, 15 Nov 2024 10:56:58 +0100
Subject: [PATCH] feat(backend): fix to save the latest message tracker just
 after the token has been expired

---
 backend/internal/http/router/handlers-web.go | 27 ++++++++++++++++----
 backend/pkg/token/http.go                    |  2 +-
 backend/pkg/token/tokenizer.go               | 14 +++++++---
 3 files changed, 34 insertions(+), 9 deletions(-)

diff --git a/backend/internal/http/router/handlers-web.go b/backend/internal/http/router/handlers-web.go
index e5518053e..01f47947c 100644
--- a/backend/internal/http/router/handlers-web.go
+++ b/backend/internal/http/router/handlers-web.go
@@ -302,9 +302,14 @@ func (e *Router) pushMessagesHandlerWeb(w http.ResponseWriter, r *http.Request)
 	if sessionData != nil {
 		r = r.WithContext(context.WithValue(r.Context(), "sessionID", fmt.Sprintf("%d", sessionData.ID)))
 	}
+	tokenJustExpired := false
 	if err != nil {
-		e.ResponseWithError(r.Context(), w, http.StatusUnauthorized, err, startTime, r.URL.Path, bodySize)
-		return
+		if errors.Is(err, token.JUST_EXPIRED) {
+			tokenJustExpired = true
+		} else {
+			e.ResponseWithError(r.Context(), w, http.StatusUnauthorized, err, startTime, r.URL.Path, bodySize)
+			return
+		}
 	}
 
 	// Add sessionID and projectID to context
@@ -314,13 +319,21 @@ func (e *Router) pushMessagesHandlerWeb(w http.ResponseWriter, r *http.Request)
 
 	// Check request body
 	if r.Body == nil {
-		e.ResponseWithError(r.Context(), w, http.StatusBadRequest, errors.New("request body is empty"), startTime, r.URL.Path, bodySize)
+		errCode := http.StatusBadRequest
+		if tokenJustExpired {
+			errCode = http.StatusUnauthorized
+		}
+		e.ResponseWithError(r.Context(), w, errCode, errors.New("request body is empty"), startTime, r.URL.Path, bodySize)
 		return
 	}
 
 	bodyBytes, err := e.readBody(w, r, e.getBeaconSize(sessionData.ID))
 	if err != nil {
-		e.ResponseWithError(r.Context(), w, http.StatusRequestEntityTooLarge, err, startTime, r.URL.Path, bodySize)
+		errCode := http.StatusRequestEntityTooLarge
+		if tokenJustExpired {
+			errCode = http.StatusUnauthorized
+		}
+		e.ResponseWithError(r.Context(), w, errCode, err, startTime, r.URL.Path, bodySize)
 		return
 	}
 	bodySize = len(bodyBytes)
@@ -329,7 +342,11 @@ func (e *Router) pushMessagesHandlerWeb(w http.ResponseWriter, r *http.Request)
 	err = e.services.Producer.Produce(e.cfg.TopicRawWeb, sessionData.ID, bodyBytes)
 	if err != nil {
 		e.log.Error(r.Context(), "can't send messages batch to queue: %s", err)
-		e.ResponseWithError(r.Context(), w, http.StatusInternalServerError, errors.New("can't save message, try again"), startTime, r.URL.Path, bodySize)
+		errCode := http.StatusInternalServerError
+		if tokenJustExpired {
+			errCode = http.StatusUnauthorized
+		}
+		e.ResponseWithError(r.Context(), w, errCode, errors.New("can't save message, try again"), startTime, r.URL.Path, bodySize)
 		return
 	}
 
diff --git a/backend/pkg/token/http.go b/backend/pkg/token/http.go
index 898bad559..146904aae 100644
--- a/backend/pkg/token/http.go
+++ b/backend/pkg/token/http.go
@@ -11,7 +11,7 @@ const BEARER_SCHEMA = "Bearer "
 func (tokenizer *Tokenizer) ParseFromHTTPRequest(r *http.Request) (*TokenData, error) {
 	header := r.Header.Get("Authorization")
 	if !strings.HasPrefix(header, BEARER_SCHEMA) {
-		return nil, errors.New("Missing token")
+		return nil, errors.New("missing token")
 	}
 	token := header[len(BEARER_SCHEMA):]
 	return tokenizer.Parse(token)
diff --git a/backend/pkg/token/tokenizer.go b/backend/pkg/token/tokenizer.go
index dd45907a8..efed1ec7f 100644
--- a/backend/pkg/token/tokenizer.go
+++ b/backend/pkg/token/tokenizer.go
@@ -11,7 +11,10 @@ import (
 	"github.com/btcsuite/btcutil/base58"
 )
 
-var EXPIRED = errors.New("token expired")
+var (
+	EXPIRED      = errors.New("token expired")
+	JUST_EXPIRED = errors.New("token just expired")
+)
 
 type Tokenizer struct {
 	secret []byte
@@ -64,8 +67,13 @@ func (tokenizer *Tokenizer) Parse(token string) (*TokenData, error) {
 	if err != nil {
 		return nil, err
 	}
+	res := &TokenData{id, delay, expTime}
 	if expTime <= time.Now().UnixMilli() {
-		return &TokenData{id, delay, expTime}, EXPIRED
+		// If token is expired less than 30 seconds ago, we still consider it semi-valid
+		if expTime+30000 > time.Now().UnixMilli() {
+			return res, JUST_EXPIRED
+		}
+		return res, EXPIRED
 	}
-	return &TokenData{id, delay, expTime}, nil
+	return res, nil
 }