From bc39ce1ddd3365af8137fe38fadf49ca194e247f Mon Sep 17 00:00:00 2001
From: Sakurasan <1173092237@qq.com>
Date: Sat, 10 Jun 2023 19:40:47 +0800
Subject: [PATCH] fix (#15)

---
 router/router.go | 21 +++++++++++++++++++--
 store/userdb.go  |  9 +++++++++
 2 files changed, 28 insertions(+), 2 deletions(-)

diff --git a/router/router.go b/router/router.go
index f2d76fd..839f9af 100644
--- a/router/router.go
+++ b/router/router.go
@@ -178,7 +178,8 @@ func Handleinit(c *gin.Context) {
 }
 
 func HandleMe(c *gin.Context) {
-	u, err := store.GetUserByID(1)
+	token := c.GetHeader("Authorization")
+	u, err := store.GetUserByToken(token[7:])
 	if err != nil {
 		c.JSON(http.StatusOK, gin.H{
 			"error": err.Error(),
@@ -190,7 +191,7 @@ func HandleMe(c *gin.Context) {
 		int(u.ID),
 		u.UpdatedAt.Format(time.RFC3339),
 		u.Name,
-		u.Token,
+		u.Token[:15] + "*****" + u.Token[len(u.Token)-15:],
 		u.CreatedAt.Format(time.RFC3339),
 	}
 	c.JSON(http.StatusOK, resJSON)
@@ -601,6 +602,22 @@ func Cost(model string, promptCount, completionCount int) float64 {
 func HandleUsage(c *gin.Context) {
 	fromStr := c.Query("from")
 	toStr := c.Query("to")
+	getMonthStartAndEnd := func() (start, end string) {
+		loc, _ := time.LoadLocation("Local")
+		now := time.Now().In(loc)
+
+		year, month, _ := now.Date()
+
+		startOfMonth := time.Date(year, month, 1, 0, 0, 0, 0, loc)
+		endOfMonth := startOfMonth.AddDate(0, 1, 0)
+
+		start = startOfMonth.Format("2006-01-02")
+		end = endOfMonth.Format("2006-01-02")
+		return
+	}
+	if fromStr == "" || toStr == "" {
+		fromStr, toStr = getMonthStartAndEnd()
+	}
 
 	usage, err := store.QueryUsage(fromStr, toStr)
 	if err != nil {
diff --git a/store/userdb.go b/store/userdb.go
index b621e7e..2e4ac78 100644
--- a/store/userdb.go
+++ b/store/userdb.go
@@ -73,6 +73,15 @@ func GetUserByName(name string) (*User, error) {
 	return &user, nil
 }
 
+func GetUserByToken(token string) (*User, error) {
+	var user User
+	result := db.Where("token = ?", token).First(&user)
+	if result.Error != nil {
+		return nil, result.Error
+	}
+	return &user, nil
+}
+
 func GetUserID(authkey string) (int, error) {
 	var user User
 	result := db.Where(&User{Token: authkey}).First(&user)